From: compilade <redacted>
Date: Tue, 8 Jul 2025 15:37:47 +0000 (-0400)
Subject: memory : fix broken batch splits for recurrent cache (#14575)
X-Git-Tag: upstream/0.0.5882~36
X-Git-Url: https://git.djapps.eu/?a=commitdiff_plain;h=bb4f7a9e4eec171fecf0f640b1337a1c24485560;p=pkg%2Fggml%2Fsources%2Fllama.cpp

memory : fix broken batch splits for recurrent cache (#14575)

Splits producing more than one ubatch per batch for recurrent models
were broken with #14512.

This fixes it by moving the completeness check after the ubatch split loop.
---

diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index 4b90dac7..a1b5b1a2 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -377,14 +377,18 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 ubatch = balloc.split_equal(n_ubatch, false);
             }
 
-            if (balloc.get_n_used() < balloc.get_n_tokens()) {
-                // failed to find a suitable split
+            if (ubatch.n_tokens == 0) {
                 break;
             }
 
             ubatches.push_back(std::move(ubatch)); // NOLINT
         }
 
+        if (balloc.get_n_used() < balloc.get_n_tokens()) {
+            // failed to find a suitable split
+            break;
+        }
+
         if (!prepare(ubatches)) {
             break;
         }