Commit cbe28bb3c for llama.cpp

commit cbe28bb3cb0758f0c566a9a3561766821aeb0748
Author: Xuan-Son Nguyen <son@huggingface.co>
Date:   Tue Oct 13 00:23:01 2026 +0200

    model: fix vllm d3 multi seq case (#30408)

    * model: fix vllm d3 multi seq case

    * add TODO

diff --git a/src/llama-batch.h b/src/llama-batch.h
index 8882c9941..e1e43979b 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -186,6 +186,7 @@ public:
     llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);

     // sequence-set-wise split - each ubatch contains a single sequence-set
+    // TODO: extend to support multiple _full_ sequences in the ubatch
     llama_ubatch split_seq(uint32_t n_ubatch);

     // a helper method for creating a well-defined ubatch of tokens
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 23170d8a3..231b34dc2 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -1838,6 +1838,8 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {

     // when computing embeddings, all tokens are output
     const bool output_all   = cparams.embeddings;
+    // non-causal attention needs the whole sequence in one ubatch
+    const bool embd_all     = output_all || !cparams.causal_attn;
     const bool has_samplers = !sampling.samplers.empty();

     const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
@@ -1916,7 +1918,7 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
     llama_memory_context_ptr mctx;

     while (true) {
-        mctx = memory->init_batch(*balloc, cparams.n_ubatch, output_all);
+        mctx = memory->init_batch(*balloc, cparams.n_ubatch, embd_all);
         if (!mctx) {
             return -2;
         }
diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp
index e91866469..a6ce547e6 100644
--- a/src/llama-kv-cache-iswa.cpp
+++ b/src/llama-kv-cache-iswa.cpp
@@ -157,8 +157,6 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_iswa::memory_breakdo
 }

 llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
-    GGML_UNUSED(embd_all);
-
     // first try simple split
     do {
         if (!unified) {
@@ -206,7 +204,7 @@ llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & ba

         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = balloc.split_equal(n_ubatch, !unified, 0);
+            auto ubatch = embd_all ? balloc.split_seq(n_ubatch) : balloc.split_equal(n_ubatch, !unified, 0);

             if (ubatch.n_tokens == 0) {
                 break;
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 463510caa..804c90532 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -721,14 +721,12 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
             llama_batch_allocr & balloc,
             uint32_t n_ubatch,
             bool embd_all) {
-    GGML_UNUSED(embd_all);
-
     do {
         balloc.split_reset();

         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
+            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : embd_all ? balloc.split_seq(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);

             if (ubatch.n_tokens == 0) {
                 break;
diff --git a/src/llama-memory.h b/src/llama-memory.h
index db8253966..b33c7aee4 100644
--- a/src/llama-memory.h
+++ b/src/llama-memory.h
@@ -85,6 +85,7 @@ struct llama_memory_i {
     // split the input batch into a set of ubatches and verify that they can fit into the cache
     // return a context object containing the ubatches and memory state required to process them
     // check the llama_memory_context_i::get_status() for the result
+    // embd_all: keep all the tokens of a sequence in one ubatch (all tokens are output, or the attention is non-causal)
     virtual llama_memory_context_ptr init_batch(
             llama_batch_allocr & balloc,
             uint32_t n_ubatch,