From 0dfa5e1104685cbccbced5d5b283c264a4d30f5d Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 15 Jun 2026 23:09:49 +0200 Subject: [PATCH 01/23] (wip) add llama_batch_ext --- include/llama.h | 61 +++++++++++++++++++++++++ src/llama-batch.cpp | 28 ++++++++++++ src/llama-batch.h | 108 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 197 insertions(+) diff --git a/include/llama.h b/include/llama.h index 27e480674282..61f10450c085 100644 --- a/include/llama.h +++ b/include/llama.h @@ -270,6 +270,11 @@ extern "C" { LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA, }; + enum llama_process_type { + LLAMA_PROCESS_TYPE_ENCODE, + LLAMA_PROCESS_TYPE_DECODE, + }; + struct llama_model_kv_override { enum llama_model_kv_override_type tag; @@ -959,6 +964,62 @@ extern "C" { struct llama_context * ctx, struct llama_batch batch); + // + // Extended batch API + // + + struct llama_batch_ext; + + LLAMA_API struct llama_batch_ext * llama_batch_ext_init (struct llama_context * ctx); + LLAMA_API void llama_batch_ext_free (struct llama_batch_ext * batch); + LLAMA_API void llama_batch_ext_clear(struct llama_batch_ext * batch); + + struct llama_batch_token { + llama_token id; // if id != LLAMA_TOKEN_NULL, embd must be nullptr + float * embd; // if embd != nullptr, token id must be LLAMA_TOKEN_NULL + float * embd_nextn; // used by nextn layers + llama_pos * pos; // if nullptr, the position will be automatically assigned + // for M-RoPE models, embedding tokens must have multiple positions per token; text token only requires one single position per token + llama_seq_id seq_id; + bool output; + }; + + // Add an input token to the batch + // Returns the batch index (>= 0) + // On error: + // -1: batch is full + // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) + // -3: invalid sequence id + LLAMA_API int32_t llama_batch_ext_add_token( + struct llama_batch_ext * batch, + llama_batch_token token); + + // Set output = true for the last added token in the batch + // Returns the batch index (>= 0) + LLAMA_API bool llama_batch_ext_set_output( + struct llama_batch_ext * batch, + int32_t idx, + bool output_last); + + // Get the logits for the output token with the specified batch index + LLAMA_API float * llama_batch_ext_get_logits( + struct llama_context * ctx, + struct llama_batch_ext * batch, + int32_t idx); + + // Get the embeddings for the output token with the specified batch index + // The length of each embedding is equal to llama_model_n_embd_out() + LLAMA_API float * llama_batch_ext_get_embeddings( + struct llama_context * ctx, + struct llama_batch_ext * batch, + int32_t idx); + + // Return values are the same as llama_decode() + LLAMA_API int32_t llama_process( + struct llama_context * ctx, + enum llama_process_type type, + struct llama_batch_ext * batch); + // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) // n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 6bf76939cddc..3d0adc4a5cf0 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -917,3 +917,31 @@ void llama_batch_free(struct llama_batch batch) { } if (batch.logits) free(batch.logits); } + + +// llama_batch_ext + +llama_batch_ext * llama_batch_ext_init(llama_context * ctx) { + return new llama_batch_ext(ctx); +} + +void llama_batch_ext_free(llama_batch_ext * batch) { + delete batch; +} + +void llama_batch_ext_clear(llama_batch_ext * batch) { + batch->clear(); +} + +int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_batch_token token) { + return batch->add_token(&token); +} + +bool llama_batch_ext_set_output(llama_batch_ext * batch, int32_t idx, bool output_last) { + return batch->set_output(idx, output_last); +} + +int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { + // TODO: implement llama_process + return 0; +} diff --git a/src/llama-batch.h b/src/llama-batch.h index f77520e86c3c..a78f5352d5d3 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -68,6 +68,114 @@ struct llama_ubatch { std::shared_ptr data; }; +struct llama_batch_ext { + size_t n_tokens_max; // max number of tokens that can be stored in the batch + size_t n_embd_inp; // number of embedding dimensions per token + llama_seq_id n_seq_max; // max number of sequences + std::vector pos_max; // keep track of the current position + llama_memory_i * memory; // memory for position inference + llama_token n_vocab; // max token ID that we accept + + struct token { + llama_token id = LLAMA_TOKEN_NULL; + size_t embd_off = 0; // index offset in the embd array + llama_pos pos = 0; + llama_seq_id seq_id = 0; + bool output = false; + }; + std::vector tokens; + std::vector embd; + + llama_batch_ext(llama_context * ctx) : + n_tokens_max(llama_n_batch(ctx)), + n_embd_inp(llama_model_n_embd_inp(llama_get_model(ctx))), + n_seq_max(llama_n_seq_max(ctx)), + memory(llama_get_memory(ctx)), + n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))) { + clear(); // initialize pos_max + } + + void clear() { + tokens.clear(); + embd .clear(); + for (llama_seq_id i = 0; i < n_seq_max; ++i) { + pos_max[i] = llama_memory_seq_pos_max(memory, i); + } + } + + llama_pos next_pos(llama_seq_id seq_id) const { + GGML_ASSERT(seq_id >= 0 && seq_id < n_seq_max); + return pos_max[seq_id] + 1; + } + + bool set_output(int32_t idx, bool output_last) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + tokens[idx].output = output_last; + return true; + } + + int32_t add_token(llama_batch_token * token_in) { + if ((token_in->id == LLAMA_TOKEN_NULL && token_in->embd == nullptr) + || (token_in->id != LLAMA_TOKEN_NULL && token_in->embd != nullptr)) { + return -2; // invalid input + } + if (token_in->embd) { + return add_token_embd(token_in); + } else { + return add_token_id(token_in); + } + } + + int32_t add_token_id(llama_batch_token * token_in) { + if (tokens.size() >= n_tokens_max) { + return -1; // size limit reached + } + if (token_in->id < 0 || token_in->id >= n_vocab) { + return -2; // invalid token id + } + if (token_in->seq_id < 0 || token_in->seq_id >= n_seq_max) { + return -3; // invalid sequence id + } + + token t; + t.id = token_in->id; + t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); + t.seq_id = token_in->seq_id; + t.output = token_in->output; + + tokens.push_back(t); + + return (int32_t)(tokens.size() - 1); + } + + int32_t add_token_embd(llama_batch_token * token_in) { + if (tokens.size() >= n_tokens_max) { + return -1; // size limit reached + } + if (!token_in->embd) { + return -2; // invalid embedding + } + if (token_in->seq_id < 0 || token_in->seq_id >= n_seq_max) { + return -3; // invalid sequence id + } + + token t; + t.embd_off = embd.size(); + t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); + t.seq_id = token_in->seq_id; + t.output = token_in->output; + + tokens.push_back(t); + embd.insert(embd.end(), + token_in->embd, + token_in->embd + n_embd_inp); + + return (int32_t)(tokens.size() - 1); + } +}; + // a helper for sanitizing, fulfilling and splitting a batch class llama_batch_allocr { public: From c7c5468eb23c324798738b9d85f0a01f5b1ad911 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 15 Jun 2026 23:15:24 +0200 Subject: [PATCH 02/23] wip --- include/llama.h | 3 --- src/llama-batch.h | 2 -- 2 files changed, 5 deletions(-) diff --git a/include/llama.h b/include/llama.h index 61f10450c085..c9749a28bf5b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -981,7 +981,6 @@ extern "C" { llama_pos * pos; // if nullptr, the position will be automatically assigned // for M-RoPE models, embedding tokens must have multiple positions per token; text token only requires one single position per token llama_seq_id seq_id; - bool output; }; // Add an input token to the batch @@ -1003,14 +1002,12 @@ extern "C" { // Get the logits for the output token with the specified batch index LLAMA_API float * llama_batch_ext_get_logits( - struct llama_context * ctx, struct llama_batch_ext * batch, int32_t idx); // Get the embeddings for the output token with the specified batch index // The length of each embedding is equal to llama_model_n_embd_out() LLAMA_API float * llama_batch_ext_get_embeddings( - struct llama_context * ctx, struct llama_batch_ext * batch, int32_t idx); diff --git a/src/llama-batch.h b/src/llama-batch.h index a78f5352d5d3..4bf6e851155d 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -143,7 +143,6 @@ struct llama_batch_ext { t.id = token_in->id; t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); t.seq_id = token_in->seq_id; - t.output = token_in->output; tokens.push_back(t); @@ -165,7 +164,6 @@ struct llama_batch_ext { t.embd_off = embd.size(); t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); t.seq_id = token_in->seq_id; - t.output = token_in->output; tokens.push_back(t); embd.insert(embd.end(), From c0fb071232902bd61504c9ce027f1a7760114ce2 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Tue, 16 Jun 2026 22:29:48 +0200 Subject: [PATCH 03/23] updated design --- include/llama.h | 90 +++++++++++++---------- src/llama-batch.cpp | 161 ++++++++++++++++++++++++++++++++++++++++-- src/llama-batch.h | 103 +++++---------------------- src/llama-hparams.cpp | 3 +- 4 files changed, 227 insertions(+), 130 deletions(-) diff --git a/include/llama.h b/include/llama.h index c9749a28bf5b..1fa72a1cf9e0 100644 --- a/include/llama.h +++ b/include/llama.h @@ -974,48 +974,64 @@ extern "C" { LLAMA_API void llama_batch_ext_free (struct llama_batch_ext * batch); LLAMA_API void llama_batch_ext_clear(struct llama_batch_ext * batch); - struct llama_batch_token { - llama_token id; // if id != LLAMA_TOKEN_NULL, embd must be nullptr - float * embd; // if embd != nullptr, token id must be LLAMA_TOKEN_NULL - float * embd_nextn; // used by nextn layers - llama_pos * pos; // if nullptr, the position will be automatically assigned - // for M-RoPE models, embedding tokens must have multiple positions per token; text token only requires one single position per token - llama_seq_id seq_id; - }; - - // Add an input token to the batch + // Add an input token to the batch, with default values: + // id = LLAMA_TOKEN_NULL + // embd = nullptr + // pos = next available position for the sequence id // Returns the batch index (>= 0) // On error: - // -1: batch is full - // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) - // -3: invalid sequence id - LLAMA_API int32_t llama_batch_ext_add_token( - struct llama_batch_ext * batch, - llama_batch_token token); - - // Set output = true for the last added token in the batch - // Returns the batch index (>= 0) - LLAMA_API bool llama_batch_ext_set_output( - struct llama_batch_ext * batch, - int32_t idx, - bool output_last); - - // Get the logits for the output token with the specified batch index - LLAMA_API float * llama_batch_ext_get_logits( - struct llama_batch_ext * batch, - int32_t idx); - - // Get the embeddings for the output token with the specified batch index - // The length of each embedding is equal to llama_model_n_embd_out() - LLAMA_API float * llama_batch_ext_get_embeddings( - struct llama_batch_ext * batch, - int32_t idx); + // -1: batch is full + // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) + // -3: invalid sequence id + LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); + + // Add an input token to the batch, with a specified token ID or token embedding + LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_token id, llama_seq_id seq_id); + LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, float * embd, llama_seq_id seq_id); + + // Add the token at index idx in the batch to another sequence id. The position will stays the same. + // Note: this should be called before other _set() functions + LLAMA_API bool llama_batch_ext_add_seq( + struct llama_batch_ext * batch, + int32_t idx, + llama_seq_id seq_id); + + /* + // Set nextn embedding for the token at index idx in the batch + LLAMA_API bool llama_batch_ext_set_embd_nextn( + struct llama_batch_ext * batch, + int32_t idx, + float * embd_nextn); + + // Set output embedding for the token at index idx in the batch + LLAMA_API bool llama_batch_ext_set_output_embd( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + + // Set output logits for the token at index idx in the batch + LLAMA_API bool llama_batch_ext_set_output_logits( + struct llama_batch_ext * batch, + int32_t idx, + bool value); + */ + + // Set custom position for the token at index idx in the batch + // For M-RoPE models: + // - Embedding tokens must have multiple positions per token + // - Text token only requires one single position per token + LLAMA_API bool llama_batch_ext_set_pos( + struct llama_batch_ext * batch, + int32_t idx, + llama_pos * pos); + + // TODO: implement get_embeddings() and get_logits() for llama_batch_ext // Return values are the same as llama_decode() LLAMA_API int32_t llama_process( - struct llama_context * ctx, - enum llama_process_type type, - struct llama_batch_ext * batch); + struct llama_context * ctx, + enum llama_process_type type, + struct llama_batch_ext * batch); // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 3d0adc4a5cf0..2ad9e541a739 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -3,6 +3,8 @@ #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h" +#include "llama-hparams.h" +#include "llama-model.h" #include #include @@ -921,6 +923,123 @@ void llama_batch_free(struct llama_batch batch) { // llama_batch_ext +llama_batch_ext::llama_batch_ext(llama_context * ctx) : + n_tokens_max(llama_n_batch(ctx)), + n_embd_inp(llama_model_n_embd_inp(llama_get_model(ctx))), + n_seq_max(llama_n_seq_max(ctx)), + memory(llama_get_memory(ctx)), + n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), + n_pos_per_embd(llama_get_model(ctx)->hparams.n_pos_per_embd()) { + clear(); // initialize pos_max +} + +void llama_batch_ext::clear() { + tokens.clear(); + embd .clear(); + for (llama_seq_id i = 0; i < n_seq_max; ++i) { + pos_max[i] = llama_memory_seq_pos_max(memory, i); + } +} + +// advance the position and return the post-incremented value +llama_pos llama_batch_ext::advance_pos(llama_seq_id seq_id) { + GGML_ASSERT(seq_id >= 0 && seq_id < n_seq_max); + return pos_max[seq_id]++; +} + +int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { + if (tokens.size() >= n_tokens_max) { + return -1; // size limit reached + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return -3; // invalid sequence id + } + + token t; + t.seq_ids.insert(seq_id); + t.pos = { advance_pos(seq_id), 0, 0, 0 }; + + tokens.push_back(t); + + return (int32_t)(tokens.size() - 1); +} + +bool llama_batch_ext::add_seq(int32_t idx, llama_seq_id seq_id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (seq_id < 0 || seq_id >= n_seq_max) { + return false; + } + + token & t = tokens[idx]; + + t.seq_ids.insert(seq_id); + + return true; +} + +bool llama_batch_ext::set_token_id(int32_t idx, llama_token id) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (id < 0 || id >= n_vocab) { + return false; + } + tokens[idx].id = id; + return true; +} + +bool llama_batch_ext::set_token_embd(int32_t idx, float * embd_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!embd_in) { + return false; + } + + token & t = tokens[idx]; + + t.embd_off = embd.size(); + embd.insert(embd.end(), embd_in, embd_in + n_embd_inp); + + return true; +} + +bool llama_batch_ext::set_token_pos(int32_t idx, llama_pos * pos_in) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + if (!pos_in) { + return false; + } + + token & t = tokens[idx]; + + size_t n_pos = t.id != LLAMA_TOKEN_NULL ? 1 : n_pos_per_embd; + for (size_t i = 0; i < n_pos; ++i) { + t.pos[i] = pos_in[i]; + } + + // also update seq pos_max + auto new_temporal_pos = pos_in[0]; + for (llama_seq_id seq : t.seq_ids) { + pos_max[seq] = std::max(pos_max[seq], new_temporal_pos); + } + + return true; +} + +bool llama_batch_ext::set_output(int32_t idx, bool output_last) { + if (idx < 0 || idx >= (int32_t) tokens.size()) { + return false; + } + tokens[idx].output = output_last; + return true; +} + +// llama_batch_ext C API + llama_batch_ext * llama_batch_ext_init(llama_context * ctx) { return new llama_batch_ext(ctx); } @@ -933,15 +1052,45 @@ void llama_batch_ext_clear(llama_batch_ext * batch) { batch->clear(); } -int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_batch_token token) { - return batch->add_token(&token); +int32_t llama_batch_ext_add(llama_batch_ext * batch, llama_seq_id seq_id) { + return batch->add_token(seq_id); +} + +int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_token id, llama_seq_id seq_id) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_id(idx, id)) { + return -2; + } + return idx; +} + +int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, float * embd, llama_seq_id seq_id) { + int32_t idx = batch->add_token(seq_id); + if (idx < 0) { + return idx; + } + if (!batch->set_token_embd(idx, embd)) { + return -2; + } + return idx; +} + +bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id seq_id) { + return batch->add_seq(idx, seq_id); +} + +bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); } -bool llama_batch_ext_set_output(llama_batch_ext * batch, int32_t idx, bool output_last) { - return batch->set_output(idx, output_last); +bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * pos) { + return batch->set_token_pos(idx, pos); } int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { - // TODO: implement llama_process - return 0; + // for now, we simply translate the llama_batch_ext into a llama_batch_allocr + return -1; } diff --git a/src/llama-batch.h b/src/llama-batch.h index 4bf6e851155d..93809767d209 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -10,6 +10,7 @@ #include #include #include +#include // keep this struct lightweight struct llama_ubatch { @@ -72,106 +73,36 @@ struct llama_batch_ext { size_t n_tokens_max; // max number of tokens that can be stored in the batch size_t n_embd_inp; // number of embedding dimensions per token llama_seq_id n_seq_max; // max number of sequences - std::vector pos_max; // keep track of the current position llama_memory_i * memory; // memory for position inference llama_token n_vocab; // max token ID that we accept + size_t n_pos_per_embd; + + std::vector pos_max; // keep track of the current position struct token { llama_token id = LLAMA_TOKEN_NULL; size_t embd_off = 0; // index offset in the embd array - llama_pos pos = 0; - llama_seq_id seq_id = 0; - bool output = false; + bool output = false; // TODO: have dedicated output flags + std::unordered_set seq_ids; + std::array pos = {0, 0, 0, 0}; }; std::vector tokens; std::vector embd; - llama_batch_ext(llama_context * ctx) : - n_tokens_max(llama_n_batch(ctx)), - n_embd_inp(llama_model_n_embd_inp(llama_get_model(ctx))), - n_seq_max(llama_n_seq_max(ctx)), - memory(llama_get_memory(ctx)), - n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))) { - clear(); // initialize pos_max - } + llama_batch_ext(llama_context * ctx); - void clear() { - tokens.clear(); - embd .clear(); - for (llama_seq_id i = 0; i < n_seq_max; ++i) { - pos_max[i] = llama_memory_seq_pos_max(memory, i); - } - } - - llama_pos next_pos(llama_seq_id seq_id) const { - GGML_ASSERT(seq_id >= 0 && seq_id < n_seq_max); - return pos_max[seq_id] + 1; - } - - bool set_output(int32_t idx, bool output_last) { - if (idx < 0 || idx >= (int32_t) tokens.size()) { - return false; - } - tokens[idx].output = output_last; - return true; - } + void clear(); - int32_t add_token(llama_batch_token * token_in) { - if ((token_in->id == LLAMA_TOKEN_NULL && token_in->embd == nullptr) - || (token_in->id != LLAMA_TOKEN_NULL && token_in->embd != nullptr)) { - return -2; // invalid input - } - if (token_in->embd) { - return add_token_embd(token_in); - } else { - return add_token_id(token_in); - } - } + // advance the position and return the post-incremented value + llama_pos advance_pos(llama_seq_id seq_id); - int32_t add_token_id(llama_batch_token * token_in) { - if (tokens.size() >= n_tokens_max) { - return -1; // size limit reached - } - if (token_in->id < 0 || token_in->id >= n_vocab) { - return -2; // invalid token id - } - if (token_in->seq_id < 0 || token_in->seq_id >= n_seq_max) { - return -3; // invalid sequence id - } - - token t; - t.id = token_in->id; - t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); - t.seq_id = token_in->seq_id; - - tokens.push_back(t); - - return (int32_t)(tokens.size() - 1); - } + int32_t add_token(llama_seq_id seq_id); - int32_t add_token_embd(llama_batch_token * token_in) { - if (tokens.size() >= n_tokens_max) { - return -1; // size limit reached - } - if (!token_in->embd) { - return -2; // invalid embedding - } - if (token_in->seq_id < 0 || token_in->seq_id >= n_seq_max) { - return -3; // invalid sequence id - } - - token t; - t.embd_off = embd.size(); - t.pos = token_in->pos ? *token_in->pos : next_pos(token_in->seq_id); - t.seq_id = token_in->seq_id; - - tokens.push_back(t); - embd.insert(embd.end(), - token_in->embd, - token_in->embd + n_embd_inp); - - return (int32_t)(tokens.size() - 1); - } + bool add_seq(int32_t idx, llama_seq_id seq_id); + bool set_token_id(int32_t idx, llama_token id); + bool set_token_embd(int32_t idx, float * embd_in); + bool set_token_pos(int32_t idx, llama_pos * pos_in); + bool set_output(int32_t idx, bool output_last); }; // a helper for sanitizing, fulfilling and splitting a batch diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 2bf576873824..7f4457eaf37b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -226,7 +226,8 @@ bool llama_hparams::is_recr(uint32_t il) const { } uint32_t llama_hparams::n_pos_per_embd() const { - return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1; + return (rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE) + ? GGML_MROPE_SECTIONS : 1; } bool llama_hparams::is_swa(uint32_t il) const { From 31323711421226f2a28208ec8e1bc25eec8d23a5 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 17 Jun 2026 11:30:59 +0200 Subject: [PATCH 04/23] updated impl --- src/llama-batch.cpp | 250 ++++++++++++++++++++++++++---------------- src/llama-batch.h | 19 ++-- src/llama-context.cpp | 67 ++++++----- src/llama-context.h | 4 + 4 files changed, 208 insertions(+), 132 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2ad9e541a739..3d48719f1773 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -25,135 +25,141 @@ llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd } bool llama_batch_allocr::init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all) { clear(); - batch = batch_inp; + this->vocab = &vocab; + this->n_embd = batch_inp.n_embd_inp; + this->n_seq_max = batch_inp.n_seq_max; - this->vocab = &vocab; + const int32_t n_tok = (int32_t) batch_inp.tokens.size(); - GGML_ASSERT(batch.n_tokens > 0); + GGML_ASSERT(n_tok > 0); - // - // validate input batch - // - - if (n_seq_max > LLAMA_MAX_SEQ) { + if ((uint32_t) n_seq_max > LLAMA_MAX_SEQ) { LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ); return false; } - if (batch.token) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.token[i] < 0 || (uint32_t) batch.token[i] >= vocab.n_tokens()) { - LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, batch.token[i]); + const bool has_embd = !batch_inp.embd.empty(); + const llama_memory_i * memory = batch_inp.memory; + + // + // build flat token/embd array + // + + if (!has_embd) { + token_vec.resize(n_tok); + for (int32_t i = 0; i < n_tok; ++i) { + const llama_token id = batch_inp.tokens[i].id; + if (id < 0 || id >= batch_inp.n_vocab) { + LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, id); return false; } + token_vec[i] = id; } + } else { + embd_vec = batch_inp.embd; } - if (batch.seq_id) { - for (int32_t i = 0; i < batch.n_tokens; ++i) { - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) { - LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, batch.seq_id[i][s], (llama_seq_id) n_seq_max); - return false; + // + // build flat pos array + // token batch: pos[i] = tokens[i].pos[0] + // embedding batch: pos[j*n_tok + i] = tokens[i].pos[j] (section-major) + // + + { + const int32_t n_pos_total = has_embd ? n_tok * (int32_t) n_pos_per_embd : n_tok; + pos.resize(n_pos_total); + if (!has_embd) { + for (int32_t i = 0; i < n_tok; ++i) { + pos[i] = batch_inp.tokens[i].pos[0]; + } + } else { + for (int32_t i = 0; i < n_tok; ++i) { + for (uint32_t j = 0; j < n_pos_per_embd; ++j) { + pos[(int32_t) j * n_tok + i] = batch_inp.tokens[i].pos[j]; } } } } // - // auto-generate missing fields + // build n_seq_id / seq_id arrays // - if (!batch.n_seq_id) { - n_seq_id.resize(batch.n_tokens); - for (int32_t i = 0; i < batch.n_tokens; i++) { - n_seq_id[i] = seq_id_0.size(); - } - batch.n_seq_id = n_seq_id.data(); - } + n_seq_id.resize(n_tok); + seq_id.resize(n_tok + 1); + seq_id[n_tok] = nullptr; - if (!batch.seq_id) { - seq_id.resize(batch.n_tokens + 1); - seq_id[batch.n_tokens] = NULL; - for (int32_t i = 0; i < batch.n_tokens; i++) { - seq_id[i] = seq_id_0.data(); + { + size_t total = 0; + for (int32_t i = 0; i < n_tok; ++i) { + total += batch_inp.tokens[i].seq_ids.size(); } - batch.seq_id = seq_id.data(); - } + seq_id_data.reserve(total); - if (!batch.pos) { - pos.resize(batch.n_tokens); - - // initialize the starting position for each sequence based on the positions in the memory - llama_pos p0[LLAMA_MAX_SEQ]; - for (uint32_t s = 0; s < n_seq_max; ++s) { - if (!memory) { - // if no memory -> start from 0 - p0[s] = 0; - } else { - p0[s] = memory->seq_pos_max(s) + 1; + for (int32_t i = 0; i < n_tok; ++i) { + for (auto sid : batch_inp.tokens[i].seq_ids) { + seq_id_data.push_back(sid); } } - for (int32_t i = 0; i < batch.n_tokens; i++) { - const llama_seq_id seq_id = batch.seq_id[i][0]; - - pos[i] = p0[seq_id]; + size_t off = 0; + for (int32_t i = 0; i < n_tok; ++i) { + n_seq_id[i] = (int32_t) batch_inp.tokens[i].seq_ids.size(); + seq_id[i] = seq_id_data.data() + off; + off += n_seq_id[i]; - // update the starting position for all sequences that are assigned to the this token - for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { - const llama_seq_id seq_id = batch.seq_id[i][s]; - - p0[seq_id] = pos[i] + 1; + for (int32_t s = 0; s < n_seq_id[i]; ++s) { + if (seq_id[i][s] < 0 || seq_id[i][s] >= (llama_seq_id) n_seq_max) { + LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, seq_id[i][s], (llama_seq_id) n_seq_max); + return false; + } } } - - batch.pos = pos.data(); } - if (!batch.logits) { - if (output_all) { - // return the output for all tokens - output.resize(batch.n_tokens, true); - } else { - // return the output only for the last token - output.resize(batch.n_tokens, false); - output[output.size() - 1] = true; - } + // + // build output/logits array + // - batch.logits = output.data(); - } else if (output_all) { - bool warn = false; + { + output.resize(n_tok, 0); + for (int32_t i = 0; i < n_tok; ++i) { + output[i] = batch_inp.tokens[i].output ? 1 : 0; + } - for (int32_t i = 0; i < batch.n_tokens; ++i) { - if (batch.logits[i] == 0) { - warn = true; + if (output_all) { + bool warn = false; + for (int32_t i = 0; i < n_tok; ++i) { + if (!output[i]) { warn = true; break; } + } + if (warn) { + LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); + std::fill(output.begin(), output.end(), 1); } } + } - if (warn) { - LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__); + // + // set up the internal llama_batch to point to our owned arrays + // - output.resize(batch.n_tokens, true); - batch.logits = output.data(); - } - } + batch.n_tokens = n_tok; + batch.token = has_embd ? nullptr : token_vec.data(); + batch.embd = has_embd ? embd_vec.data() : nullptr; + batch.pos = pos.data(); + batch.n_seq_id = n_seq_id.data(); + batch.seq_id = seq_id.data(); + batch.logits = output.data(); // // compute stats // - this->n_embd = n_embd; - this->n_seq_max = n_seq_max; - // count the outputs in this batch for (int32_t i = 0; i < batch.n_tokens; ++i) { n_outputs += batch.logits[i] != 0; @@ -659,11 +665,14 @@ void llama_batch_allocr::clear() { batch = {}; - pos .clear(); - n_seq_id .clear(); - seq_id .clear(); - seq_id_unq.clear(); - output .clear(); + token_vec .clear(); + embd_vec .clear(); + seq_id_data .clear(); + pos .clear(); + n_seq_id .clear(); + seq_id .clear(); + seq_id_unq .clear(); + output .clear(); for (auto & cur : seq_pos) { cur.clear(); @@ -936,8 +945,9 @@ llama_batch_ext::llama_batch_ext(llama_context * ctx) : void llama_batch_ext::clear() { tokens.clear(); embd .clear(); + pos_max.resize(n_seq_max); for (llama_seq_id i = 0; i < n_seq_max; ++i) { - pos_max[i] = llama_memory_seq_pos_max(memory, i); + pos_max[i] = memory ? llama_memory_seq_pos_max(memory, i) : 0; } } @@ -1082,15 +1092,63 @@ bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id return batch->add_seq(idx, seq_id); } -bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, bool value) { - return batch->set_output(idx, value); -} - bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * pos) { return batch->set_token_pos(idx, pos); } -int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { - // for now, we simply translate the llama_batch_ext into a llama_batch_allocr - return -1; +// llama_batch_compat + +llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp) { + batch_ext = new llama_batch_ext(ctx); + + const bool is_embd = batch_inp.embd != nullptr; + + static const llama_seq_id default_seq_id = 0; + static const int32_t default_n_seq_id = 1; + + for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { + const int32_t n_sid = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : default_n_seq_id; + const llama_seq_id * sids = batch_inp.seq_id ? batch_inp.seq_id[i] : &default_seq_id; + + llama_batch_ext::token t; + + // seq_ids + for (int32_t s = 0; s < n_sid; ++s) { + t.seq_ids.insert(sids[s]); + } + + // position(s) + if (batch_inp.pos) { + if (!is_embd) { + // token batch: one position per token + t.pos[0] = batch_inp.pos[i]; + } else { + // embedding batch (M-RoPE): section-major layout pos[j*n_tokens + i] + for (uint32_t j = 0; j < batch_ext->n_pos_per_embd; ++j) { + t.pos[j] = batch_inp.pos[(int32_t) j * batch_inp.n_tokens + i]; + } + } + } else { + // auto-generate position from the first seq_id + t.pos[0] = batch_ext->advance_pos(sids[0]); + } + + // token id or embeddings + if (!is_embd) { + t.id = batch_inp.token[i]; + } else { + t.embd_off = batch_ext->embd.size(); + const float * src = batch_inp.embd + (size_t) i * batch_ext->n_embd_inp; + batch_ext->embd.insert(batch_ext->embd.end(), src, src + batch_ext->n_embd_inp); + } + + // output flag + t.output = batch_inp.logits ? (batch_inp.logits[i] != 0) : false; + + batch_ext->tokens.push_back(t); + } +} + +llama_batch_compat::~llama_batch_compat() { + delete batch_ext; } diff --git a/src/llama-batch.h b/src/llama-batch.h index 93809767d209..fc5c1a9ded54 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -110,14 +110,10 @@ class llama_batch_allocr { public: llama_batch_allocr(uint32_t n_pos_per_embd); - // sanitize and auto-gen missing data in the input batch - // memory is optional. if provided will be used to check for sequence continuity and to determine the positions + // convert a llama_batch_ext to internal llama_batch and sanitize it bool init( - const llama_batch & batch_inp, + const llama_batch_ext & batch_inp, const llama_vocab & vocab, - const llama_memory_i * memory, - uint32_t n_embd, - uint32_t n_seq_max, bool output_all); const llama_batch & get_batch() const; @@ -173,7 +169,9 @@ class llama_batch_allocr { uint32_t n_seq_max; uint32_t n_outputs; - std::array seq_id_0 = {{ 0 }}; // default sequence id + std::vector token_vec; // owned token IDs built from llama_batch_ext + std::vector embd_vec; // owned embeddings built from llama_batch_ext + std::vector seq_id_data; // flat storage for seq_id pointers below std::vector pos; std::vector n_seq_id; @@ -208,3 +206,10 @@ class llama_batch_allocr { int debug; }; + +// RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext +struct llama_batch_compat { + llama_batch_ext * batch_ext; + llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp); + ~llama_batch_compat(); +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 168dbabd7667..bc00af6aaa5c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1369,24 +1369,18 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return res; } -int llama_context::encode(const llama_batch & batch_inp) { - // MTP hook batches carry both token (next-token id) and embd (h_nextn row), - // so accept either present rather than requiring exactly one. - GGML_ASSERT(batch_inp.token || batch_inp.embd); - - if (batch_inp.n_tokens == 0) { +int llama_context::encode(const llama_batch_ext & batch_inp) { + if (batch_inp.tokens.empty()) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); return -1; } const auto & hparams = model.hparams; - // eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim - const int64_t n_embd = hparams.n_embd_inp(); const int64_t n_vocab = model.vocab.n_tokens(); - // note: during encode, we always pass the full sequence starting from pos = 0 - if (!balloc->init(batch_inp, model.vocab, nullptr, n_embd, cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { + // note: during encode, we always output all tokens and skip position continuity checks (output_all=true) + if (!balloc->init(batch_inp, model.vocab, true)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -1673,17 +1667,13 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map seq_output_count(n_seq_max, 0); - for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { - if (batch_inp.logits[i] == 0) { + for (const auto & tok : batch_inp.tokens) { + if (!tok.output) { continue; } - const int ns = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : 1; - - for (int32_t s = 0; s < ns; ++s) { - const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; - + for (auto seq_id : tok.seq_ids) { seq_output_count[seq_id]++; if (seq_output_count[seq_id] > 1) { LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n", @@ -1724,7 +1709,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } } - if (!balloc->init(batch_inp, vocab, memory.get(), n_embd, n_seq_max, output_all)) { + if (!balloc->init(batch_inp, vocab, output_all)) { LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); return -1; } @@ -3307,9 +3292,13 @@ void llama_context::opt_epoch_iter( batch.logits [pos_batch] = true; } - if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, true)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); - return; + // TODO: use llama_batch_ext here + { + llama_batch_compat compat(this, batch); + if (!balloc->init(*compat.batch_ext, model.vocab, true)) { + LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + return; + } } const uint32_t n_tokens_all = balloc->get_n_tokens(); @@ -4030,6 +4019,18 @@ size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, lla } } +// compat: llama_batch -> llama_batch_ext -> encode/decode + +int llama_context::encode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp); + return encode(*compat.batch_ext); +} + +int llama_context::decode(const llama_batch & batch_inp) { + llama_batch_compat compat(this, batch_inp); + return decode(*compat.batch_ext); +} + /// int32_t llama_encode( @@ -4119,6 +4120,14 @@ void llama_opt_epoch( callback_eval); } +int32_t llama_process(llama_context * ctx, llama_process_type type, llama_batch_ext * batch) { + switch (type) { + case LLAMA_PROCESS_TYPE_ENCODE: return ctx->encode(*batch); + case LLAMA_PROCESS_TYPE_DECODE: return ctx->decode(*batch); + } + return -1; +} + // // ext // diff --git a/src/llama-context.h b/src/llama-context.h index 853052be2cad..00e6e20a1607 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -139,6 +139,10 @@ struct llama_context { llama_memory_context_i * mctx, ggml_status & ret); + int encode(const llama_batch_ext & batch_inp); + int decode(const llama_batch_ext & batch_inp); + + // compat version int encode(const llama_batch & batch_inp); int decode(const llama_batch & batch_inp); From bf372b380cc84496462e0c6cf238d9c4cd705b05 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 17:41:57 +0200 Subject: [PATCH 05/23] change signature --- include/llama.h | 4 ++-- src/llama-batch.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/llama.h b/include/llama.h index f8589bbd23a0..e2631d5b8710 100644 --- a/include/llama.h +++ b/include/llama.h @@ -994,8 +994,8 @@ extern "C" { LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); // Add an input token to the batch, with a specified token ID or token embedding - LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_token id, llama_seq_id seq_id); - LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, float * embd, llama_seq_id seq_id); + LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); + LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, llama_seq_id seq_id, float * embd); // Add the token at index idx in the batch to another sequence id. The position will stays the same. // Note: this should be called before other _set() functions diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 6ac9ee5ec01f..0ce1f8737547 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1134,7 +1134,7 @@ int32_t llama_batch_ext_add(llama_batch_ext * batch, llama_seq_id seq_id) { return batch->add_token(seq_id); } -int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_token id, llama_seq_id seq_id) { +int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, llama_token id) { int32_t idx = batch->add_token(seq_id); if (idx < 0) { return idx; @@ -1145,7 +1145,7 @@ int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_token id, llama return idx; } -int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, float * embd, llama_seq_id seq_id) { +int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, float * embd) { int32_t idx = batch->add_token(seq_id); if (idx < 0) { return idx; From 901ed6859f993e954519b4680fb3ff92ffbea8bf Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 17:56:56 +0200 Subject: [PATCH 06/23] unused var --- src/llama-context.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e364932abdf2..c53501d0504b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1364,7 +1364,6 @@ int llama_context::encode(const llama_batch_ext & batch_inp) { const auto & hparams = model.hparams; // eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim - const int64_t n_embd = hparams.n_embd_inp_enc(); const int64_t n_vocab = model.vocab.n_tokens(); // note: during encode, we always output all tokens and skip position continuity checks (output_all=true) From 231af77e9dcd60c7f5356d33311376203eff1120 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 18:44:06 +0200 Subject: [PATCH 07/23] demo common_prompt_batch_decode --- common/common.cpp | 32 +++++++++++++++++++++++++------- common/common.h | 6 +++++- include/llama-cpp.h | 5 +++++ include/llama.h | 14 +++++++------- src/llama-batch.cpp | 8 ++++++++ 5 files changed, 50 insertions(+), 15 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 8f13217ab442..f33795e96772 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2013,9 +2013,23 @@ bool common_replay_last_token(struct llama_context * ctx, llama_token last_token return true; } +llama_batch_ext_ptr common_batch_ext_get_one(llama_context * ctx, const llama_tokens & tokens) { + llama_batch_ext_ptr batch(llama_batch_ext_init(ctx)); + + for (size_t i = 0; i < tokens.size(); ++i) { + llama_batch_ext_add_token(batch.get(), 0, tokens[i]); + } + + if (!tokens.empty()) { + llama_batch_ext_set_output_logits(batch.get(), (int32_t) tokens.size() - 1, true); + } + + return batch; +} + bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, @@ -2036,7 +2050,9 @@ bool common_prompt_batch_decode( // Memory implementations in recurrent/hybrid models don't support removing tokens from their // memory, so we can't just remove the last token from the memory and replay the last token which // is the reason for this logic. - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_tokens_before_last))) { + llama_tokens prefix_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_tokens_before_last); + llama_batch_ext_ptr batch_prefix = common_batch_ext_get_one(ctx, prefix_tokens); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_prefix.get())) { COM_ERR("%s", "failed to eval\n"); return false; } @@ -2046,17 +2062,19 @@ bool common_prompt_batch_decode( COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size()); llama_token last_token = all_tokens.back(); - llama_batch batch = llama_batch_get_one(&last_token, 1); - int32_t pos = n_past; - batch.pos = &pos; + llama_batch_ext_ptr batch_last = common_batch_ext_get_one(ctx, { last_token }); + llama_pos pos = n_past; + llama_batch_ext_set_pos(batch_last.get(), 0, &pos); - if (llama_decode(ctx, batch)) { + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_last.get())) { COM_ERR("%s", "failed to eval last token\n"); return false; } n_past++; } else { - if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_new))) { + llama_tokens new_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_new); + llama_batch_ext_ptr batch = common_batch_ext_get_one(ctx, new_tokens); + if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) { COM_ERR("%s", "failed to eval\n"); return false; } diff --git a/common/common.h b/common/common.h index 7ed1a982756b..797f5d21e16a 100644 --- a/common/common.h +++ b/common/common.h @@ -948,6 +948,10 @@ void common_batch_add( const std::vector & seq_ids, bool logits); +// create a single-sequence batch from a list of tokens +// last token always have output_logits set to true +llama_batch_ext_ptr common_batch_ext_get_one(struct llama_context * ctx, const llama_tokens & tokens); + // decodes a single batch of tokens for a prompt and manages session tokens // // Note: We save state before the last token so that we can replay it to ensure @@ -955,7 +959,7 @@ void common_batch_add( // tokens from memory, so this approach works across all model architectures. bool common_prompt_batch_decode( struct llama_context * ctx, - const std::vector & all_tokens, + const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch, diff --git a/include/llama-cpp.h b/include/llama-cpp.h index 8f6368177de0..880a6a5fae8f 100644 --- a/include/llama-cpp.h +++ b/include/llama-cpp.h @@ -24,7 +24,12 @@ struct llama_adapter_lora_deleter { void operator()(llama_adapter_lora * adapter) { llama_adapter_lora_free(adapter); } }; +struct llama_batch_ext_deleter { + void operator()(llama_batch_ext * batch) { llama_batch_ext_free(batch); } +}; + typedef std::unique_ptr llama_model_ptr; typedef std::unique_ptr llama_context_ptr; typedef std::unique_ptr llama_sampler_ptr; typedef std::unique_ptr llama_adapter_lora_ptr; +typedef std::unique_ptr llama_batch_ext_ptr; diff --git a/include/llama.h b/include/llama.h index e2631d5b8710..1eb652237977 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1004,25 +1004,25 @@ extern "C" { int32_t idx, llama_seq_id seq_id); - /* - // Set nextn embedding for the token at index idx in the batch - LLAMA_API bool llama_batch_ext_set_embd_nextn( - struct llama_batch_ext * batch, - int32_t idx, - float * embd_nextn); + // // Set nextn embedding for the token at index idx in the batch + // LLAMA_API bool llama_batch_ext_set_embd_nextn( + // struct llama_batch_ext * batch, + // int32_t idx, + // float * embd_nextn); // Set output embedding for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output logits LLAMA_API bool llama_batch_ext_set_output_embd( struct llama_batch_ext * batch, int32_t idx, bool value); // Set output logits for the token at index idx in the batch + // Note: for now, this is equivalent to setting the output embd LLAMA_API bool llama_batch_ext_set_output_logits( struct llama_batch_ext * batch, int32_t idx, bool value); - */ // Set custom position for the token at index idx in the batch // For M-RoPE models: diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 0ce1f8737547..5246057420ba 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1164,6 +1164,14 @@ bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * p return batch->set_token_pos(idx, pos); } +bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + +bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, bool value) { + return batch->set_output(idx, value); +} + // llama_batch_compat llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp) { From 4cd8c268eb30683d6073f5e888a0799f51f406ad Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 18:50:05 +0200 Subject: [PATCH 08/23] fix pos --- src/llama-batch.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 5246057420ba..5187063dc151 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1015,7 +1015,7 @@ void llama_batch_ext::clear() { embd .clear(); pos_max.resize(n_seq_max); for (llama_seq_id i = 0; i < n_seq_max; ++i) { - pos_max[i] = memory ? llama_memory_seq_pos_max(memory, i) : 0; + pos_max[i] = memory ? llama_memory_seq_pos_max(memory, i) + 1 : 0; } } @@ -1100,9 +1100,9 @@ bool llama_batch_ext::set_token_pos(int32_t idx, llama_pos * pos_in) { } // also update seq pos_max - auto new_temporal_pos = pos_in[0]; + auto next_pos = pos_in[0] + 1; for (llama_seq_id seq : t.seq_ids) { - pos_max[seq] = std::max(pos_max[seq], new_temporal_pos); + pos_max[seq] = std::max(pos_max[seq], next_pos); } return true; From e4c474f3f6a182efd70d43656d00f2e8ffc07a15 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 19:26:05 +0200 Subject: [PATCH 09/23] tmp disable test-batch-alloc --- tests/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 855295c152fa..08a12a5ea1f8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -155,7 +155,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-parser.cpp) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) - llama_build_and_test(test-batch-alloc.cpp) + # llama_build_and_test(test-batch-alloc.cpp) // TODO: fix this before merging this PR llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) From 4ba39e5c13714711b7f7f5ce27b9df76a6835d4d Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 13 Jul 2026 21:42:01 +0200 Subject: [PATCH 10/23] fix compat --- src/llama-batch.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 5187063dc151..ebf500557a5c 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1219,7 +1219,10 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & } // output flag - t.output = batch_inp.logits ? (batch_inp.logits[i] != 0) : false; + // if no logits array is given, default to only the last token being an output + t.output = batch_inp.logits + ? (batch_inp.logits[i] != 0) + : (i == batch_inp.n_tokens - 1); batch_ext->tokens.push_back(t); } From 7d626f598deb0408ddb65e76262b9360cf3bc4a0 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 13 Aug 2026 01:34:11 +0200 Subject: [PATCH 11/23] nits: add const --- src/llama-batch.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/llama-batch.h b/src/llama-batch.h index 3d80e3a7e273..2409ab95a7a1 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -70,12 +70,12 @@ struct llama_ubatch { }; struct llama_batch_ext { - size_t n_tokens_max; // max number of tokens that can be stored in the batch - size_t n_embd_inp; // number of embedding dimensions per token - llama_seq_id n_seq_max; // max number of sequences + const size_t n_tokens_max; // max number of tokens that can be stored in the batch + const size_t n_embd_inp; // number of embedding dimensions per token + const llama_seq_id n_seq_max; // max number of sequences llama_memory_i * memory; // memory for position inference - llama_token n_vocab; // max token ID that we accept - size_t n_pos_per_embd; + const llama_token n_vocab; // max token ID that we accept + const size_t n_pos_per_embd; std::vector pos_max; // keep track of the current position From 3b861104b327e11778b0f83c610c30a751c72840 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Thu, 13 Aug 2026 17:31:54 +0200 Subject: [PATCH 12/23] no more pos_max --- common/common.cpp | 7 ++++++- src/llama-batch.cpp | 44 ++++++++++++++++++-------------------------- src/llama-batch.h | 15 ++++++--------- 3 files changed, 30 insertions(+), 36 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index f56753731334..4b883531718b 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2115,8 +2115,13 @@ bool common_replay_last_token(struct llama_context * ctx, llama_token last_token llama_batch_ext_ptr common_batch_ext_get_one(llama_context * ctx, const llama_tokens & tokens) { llama_batch_ext_ptr batch(llama_batch_ext_init(ctx)); + auto mem = llama_get_memory(ctx); + llama_pos pos = mem ? llama_memory_seq_pos_max(mem, 0) + 1 : 0; + for (size_t i = 0; i < tokens.size(); ++i) { - llama_batch_ext_add_token(batch.get(), 0, tokens[i]); + const int32_t idx = llama_batch_ext_add_token(batch.get(), 0, tokens[i]); + llama_batch_ext_set_pos(batch.get(), idx, &pos); + pos++; } if (!tokens.empty()) { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index ebf500557a5c..ae566ba5b2b4 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -44,7 +44,7 @@ bool llama_batch_allocr::init( } const bool has_embd = !batch_inp.embd.empty(); - const llama_memory_i * memory = batch_inp.memory; + const llama_memory_i * mem = batch_inp.mem; // // build flat token/embd array @@ -267,7 +267,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (batch.token) { if (p0 >= 0 && p0 >= seq_pos_min(s)) { @@ -300,7 +300,7 @@ bool llama_batch_allocr::init( continue; } - const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1; + const llama_pos p0 = mem ? mem->seq_pos_max(s) : -1; if (p0 >= 0) { bool ok = true; @@ -328,12 +328,12 @@ bool llama_batch_allocr::init( } } - if (memory) { + if (mem) { for (uint32_t s0 = 0; s0 < n_seq_max; ++s0) { for (uint32_t s1 = 0; s1 < n_seq_max; ++s1) { if (seq_cpl[s0][s1]) { - if (memory->seq_pos_min(s0) != memory->seq_pos_min(s1) || - memory->seq_pos_max(s0) != memory->seq_pos_max(s1)) { + if (mem->seq_pos_min(s0) != mem->seq_pos_min(s1) || + mem->seq_pos_max(s0) != mem->seq_pos_max(s1)) { LLAMA_LOG_ERROR("%s: sequence %d is coupled to %d in the input batch, but have divereged\n", __func__, s0, s1); return false; } @@ -1004,25 +1004,15 @@ llama_batch_ext::llama_batch_ext(llama_context * ctx) : n_tokens_max(llama_n_batch(ctx)), n_embd_inp(llama_model_n_embd_inp(llama_get_model(ctx))), n_seq_max(llama_n_seq_max(ctx)), - memory(llama_get_memory(ctx)), + mem(llama_get_memory(ctx)), n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), n_pos_per_embd(llama_get_model(ctx)->hparams.n_pos_per_embd()) { - clear(); // initialize pos_max + clear(); } void llama_batch_ext::clear() { tokens.clear(); embd .clear(); - pos_max.resize(n_seq_max); - for (llama_seq_id i = 0; i < n_seq_max; ++i) { - pos_max[i] = memory ? llama_memory_seq_pos_max(memory, i) + 1 : 0; - } -} - -// advance the position and return the post-incremented value -llama_pos llama_batch_ext::advance_pos(llama_seq_id seq_id) { - GGML_ASSERT(seq_id >= 0 && seq_id < n_seq_max); - return pos_max[seq_id]++; } int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { @@ -1033,9 +1023,10 @@ int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { return -3; // invalid sequence id } + // position is left undefined (default-initialized) - the caller must set it + // explicitly via set_token_pos() before the batch is processed token t; t.seq_ids.insert(seq_id); - t.pos = { advance_pos(seq_id), 0, 0, 0 }; tokens.push_back(t); @@ -1099,12 +1090,6 @@ bool llama_batch_ext::set_token_pos(int32_t idx, llama_pos * pos_in) { t.pos[i] = pos_in[i]; } - // also update seq pos_max - auto next_pos = pos_in[0] + 1; - for (llama_seq_id seq : t.seq_ids) { - pos_max[seq] = std::max(pos_max[seq], next_pos); - } - return true; } @@ -1182,6 +1167,13 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & static const llama_seq_id default_seq_id = 0; static const int32_t default_n_seq_id = 1; + // local position tracker for auto-generating positions when batch_inp.pos is null, + // seeded from the memory's current max position per sequence (mirrors the logic in `master`) + std::vector pos_next(batch_ext->n_seq_max); + for (llama_seq_id s = 0; s < (llama_seq_id) batch_ext->n_seq_max; ++s) { + pos_next[s] = llama_memory_seq_pos_max(batch_ext->mem, s) + 1; // assume next pos + } + for (int32_t i = 0; i < batch_inp.n_tokens; ++i) { const int32_t n_sid = batch_inp.n_seq_id ? batch_inp.n_seq_id[i] : default_n_seq_id; const llama_seq_id * sids = batch_inp.seq_id ? batch_inp.seq_id[i] : &default_seq_id; @@ -1206,7 +1198,7 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & } } else { // auto-generate position from the first seq_id - t.pos[0] = batch_ext->advance_pos(sids[0]); + t.pos[0] = pos_next[sids[0]]++; } // token id or embeddings diff --git a/src/llama-batch.h b/src/llama-batch.h index 2409ab95a7a1..6bb83a91d448 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -70,15 +70,13 @@ struct llama_ubatch { }; struct llama_batch_ext { - const size_t n_tokens_max; // max number of tokens that can be stored in the batch - const size_t n_embd_inp; // number of embedding dimensions per token + const size_t n_tokens_max; // max number of tokens that can be stored in the batch + const size_t n_embd_inp; // number of embedding dimensions per token const llama_seq_id n_seq_max; // max number of sequences - llama_memory_i * memory; // memory for position inference - const llama_token n_vocab; // max token ID that we accept + llama_memory_i * mem; // memory for position inference + const llama_token n_vocab; // max token ID that we accept const size_t n_pos_per_embd; - std::vector pos_max; // keep track of the current position - struct token { llama_token id = LLAMA_TOKEN_NULL; size_t embd_off = 0; // index offset in the embd array @@ -93,9 +91,8 @@ struct llama_batch_ext { void clear(); - // advance the position and return the post-incremented value - llama_pos advance_pos(llama_seq_id seq_id); - + // add an entry with an undefined position + // the caller must set it explicitly via set_token_pos() int32_t add_token(llama_seq_id seq_id); bool add_seq(int32_t idx, llama_seq_id seq_id); From 51b471f09e30308e3c8191379815144f94986773 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 13:04:50 +0200 Subject: [PATCH 13/23] add comment about llama_batch_ext_set_embd_state --- include/llama.h | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/include/llama.h b/include/llama.h index 99976dbba92a..4458196a3199 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1019,11 +1019,14 @@ extern "C" { int32_t idx, llama_seq_id seq_id); - // // Set nextn embedding for the token at index idx in the batch - // LLAMA_API bool llama_batch_ext_set_embd_nextn( + // Set the "state" embedding for the token at index idx in the batch + // "state" here means extra hidden state carried over from a previous stage, e.g.: + // - MTP: state from N layers of the target model + // - Qwen3 VL (deepstack): state from N layers of vision encoder + // LLAMA_API bool llama_batch_ext_set_embd_state( // struct llama_batch_ext * batch, // int32_t idx, - // float * embd_nextn); + // float * embd_state); // Set output embedding for the token at index idx in the batch // Note: for now, this is equivalent to setting the output logits From 95956eee0c09bb75b0212878f7ba06ff0aba3855 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 13:20:57 +0200 Subject: [PATCH 14/23] handle n_embd_out properly --- src/llama-batch.cpp | 13 ++++++++----- src/llama-batch.h | 6 ++++++ src/llama-context.cpp | 3 --- tests/test-batch-alloc.cpp | 33 +++++++++++++++++++++++++++++---- 4 files changed, 43 insertions(+), 12 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index ae566ba5b2b4..cee600e47595 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -5,6 +5,7 @@ #include "llama-memory.h" #include "llama-hparams.h" #include "llama-model.h" +#include "llama-context.h" #include #include @@ -1000,9 +1001,13 @@ void llama_batch_free(struct llama_batch batch) { // llama_batch_ext +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, const llama_hparams & hparams) { + return ctx_type == LLAMA_CONTEXT_TYPE_MTP ? hparams.n_embd_out() : hparams.n_embd_inp(); +} + llama_batch_ext::llama_batch_ext(llama_context * ctx) : n_tokens_max(llama_n_batch(ctx)), - n_embd_inp(llama_model_n_embd_inp(llama_get_model(ctx))), + n_embd_inp(llama_batch_ext_select_n_embd_inp(ctx->get_cparams().ctx_type, llama_get_model(ctx)->hparams)), n_seq_max(llama_n_seq_max(ctx)), mem(llama_get_memory(ctx)), n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), @@ -1023,8 +1028,7 @@ int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { return -3; // invalid sequence id } - // position is left undefined (default-initialized) - the caller must set it - // explicitly via set_token_pos() before the batch is processed + // position is left undefined; call set_token_pos() before decoding token t; t.seq_ids.insert(seq_id); @@ -1167,8 +1171,7 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & static const llama_seq_id default_seq_id = 0; static const int32_t default_n_seq_id = 1; - // local position tracker for auto-generating positions when batch_inp.pos is null, - // seeded from the memory's current max position per sequence (mirrors the logic in `master`) + // auto-generates positions locally when batch_inp.pos is null, continuing from memory std::vector pos_next(batch_ext->n_seq_max); for (llama_seq_id s = 0; s < (llama_seq_id) batch_ext->n_seq_max; ++s) { pos_next[s] = llama_memory_seq_pos_max(batch_ext->mem, s) + 1; // assume next pos diff --git a/src/llama-batch.h b/src/llama-batch.h index 6bb83a91d448..98648148cda5 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -69,6 +69,12 @@ struct llama_ubatch { std::shared_ptr data; }; +struct llama_hparams; + +// MTP hook batches carry the target model's hidden state (n_embd_out size). +// Normal batches carry token embeddings (n_embd_inp size). +size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, const llama_hparams & hparams); + struct llama_batch_ext { const size_t n_tokens_max; // max number of tokens that can be stored in the batch const size_t n_embd_inp; // number of embedding dimensions per token diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 1123df609982..e8089848cd45 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1643,9 +1643,6 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { const int64_t n_vocab = vocab.n_tokens(); - const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && !batch_inp.embd.empty(); - const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp(); - // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; const bool has_samplers = !sampling.samplers.empty(); diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index 66d29d6f5164..aae27933f0c0 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -3,6 +3,7 @@ #include "llama.h" #include "../src/llama-batch.h" +#include "../src/llama-hparams.h" #include "../src/llama-memory.h" #include "../src/llama-vocab.h" @@ -650,6 +651,29 @@ static void test_mrope(testing & t) { }); } +static void test_mtp_embd_width(testing & t) { + t.test("mtp_uses_n_embd_out", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; + hparams.n_deepstack_layers = 2; // makes n_embd_inp() = 64 + 64*2 = 192 + hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96 + + t.assert_equal("default context uses n_embd_inp (deepstack-aware)", + (size_t) 192, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, hparams)); + + t.assert_equal("MTP context uses n_embd_out instead (target-model hidden state width)", + (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, hparams)); + }); + + t.test("mtp_falls_back_to_n_embd_when_no_override", [&](testing & t) { + llama_hparams hparams = {}; + hparams.n_embd = 64; // no deepstack, no n_embd_out_impl override + + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, hparams)); + t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, hparams)); + }); +} + int main(int argc, char ** argv) { testing t; @@ -665,10 +689,11 @@ int main(int argc, char ** argv) { t.set_filter(argv[1]); } - t.test("init", test_init); - t.test("split", test_split); - t.test("keep_tail", test_keep_tail); - t.test("mrope", test_mrope); + t.test("init", test_init); + t.test("split", test_split); + t.test("keep_tail", test_keep_tail); + t.test("mrope", test_mrope); + t.test("mtp_embd_width", test_mtp_embd_width); return t.summary(); } From e6404103fa8909f6e9c933cf42843a6dcff2c15c Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 14:03:35 +0200 Subject: [PATCH 15/23] rename api --> embd_token --- include/llama.h | 5 +++-- src/llama-batch.cpp | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/include/llama.h b/include/llama.h index 4458196a3199..2c79013fa497 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1009,8 +1009,9 @@ extern "C" { LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); // Add an input token to the batch, with a specified token ID or token embedding - LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); - LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, llama_seq_id seq_id, float * embd); + // "token" embedding here also covers mtmd embeddings, i.e. the input of build_inp_embd() + LLAMA_API int32_t llama_batch_ext_add_token (struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); + LLAMA_API int32_t llama_batch_ext_add_embd_token(struct llama_batch_ext * batch, llama_seq_id seq_id, float * embd_token); // Add the token at index idx in the batch to another sequence id. The position will stays the same. // Note: this should be called before other _set() functions diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index cee600e47595..3ea309735c25 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1134,12 +1134,12 @@ int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, return idx; } -int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, float * embd) { +int32_t llama_batch_ext_add_embd_token(llama_batch_ext * batch, llama_seq_id seq_id, float * embd_token) { int32_t idx = batch->add_token(seq_id); if (idx < 0) { return idx; } - if (!batch->set_token_embd(idx, embd)) { + if (!batch->set_token_embd(idx, embd_token)) { return -2; } return idx; From 6d00861cd390ad291745f8c3a85f6562d827387a Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 14:53:28 +0200 Subject: [PATCH 16/23] llama_embd --- include/llama.h | 15 ++++++++++----- src/llama-batch.cpp | 17 ++++++++++++----- src/llama-batch.h | 2 +- 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/include/llama.h b/include/llama.h index 2c79013fa497..ab3024d5ce8c 100644 --- a/include/llama.h +++ b/include/llama.h @@ -993,6 +993,12 @@ extern "C" { struct llama_batch_ext; + struct llama_embd { + const float * data; + size_t n_rows; // number of embedding rows in data + size_t n_embd; // size of one row + }; + LLAMA_API struct llama_batch_ext * llama_batch_ext_init (struct llama_context * ctx); LLAMA_API void llama_batch_ext_free (struct llama_batch_ext * batch); LLAMA_API void llama_batch_ext_clear(struct llama_batch_ext * batch); @@ -1006,12 +1012,11 @@ extern "C" { // -1: batch is full // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) // -3: invalid sequence id - LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); + LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); // Add an input token to the batch, with a specified token ID or token embedding - // "token" embedding here also covers mtmd embeddings, i.e. the input of build_inp_embd() LLAMA_API int32_t llama_batch_ext_add_token (struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); - LLAMA_API int32_t llama_batch_ext_add_embd_token(struct llama_batch_ext * batch, llama_seq_id seq_id, float * embd_token); + LLAMA_API int32_t llama_batch_ext_add_embd_token(struct llama_batch_ext * batch, llama_seq_id seq_id, struct llama_embd embd); // Add the token at index idx in the batch to another sequence id. The position will stays the same. // Note: this should be called before other _set() functions @@ -1023,11 +1028,11 @@ extern "C" { // Set the "state" embedding for the token at index idx in the batch // "state" here means extra hidden state carried over from a previous stage, e.g.: // - MTP: state from N layers of the target model - // - Qwen3 VL (deepstack): state from N layers of vision encoder + // - Qwen3 VL (deepstack): state from N layers of the vision encoder // LLAMA_API bool llama_batch_ext_set_embd_state( // struct llama_batch_ext * batch, // int32_t idx, - // float * embd_state); + // struct llama_embd embd); // Set output embedding for the token at index idx in the batch // Note: for now, this is equivalent to setting the output logits diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 3ea309735c25..3d407448aa84 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1063,18 +1063,25 @@ bool llama_batch_ext::set_token_id(int32_t idx, llama_token id) { return true; } -bool llama_batch_ext::set_token_embd(int32_t idx, float * embd_in) { +bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { if (idx < 0 || idx >= (int32_t) tokens.size()) { return false; } - if (!embd_in) { + if (!embd_in.data) { + return false; + } + + const size_t n_total = embd_in.n_rows * embd_in.n_embd; + if (n_total != n_embd_inp) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd_inp); return false; } token & t = tokens[idx]; t.embd_off = embd.size(); - embd.insert(embd.end(), embd_in, embd_in + n_embd_inp); + embd.insert(embd.end(), embd_in.data, embd_in.data + n_total); return true; } @@ -1134,12 +1141,12 @@ int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, return idx; } -int32_t llama_batch_ext_add_embd_token(llama_batch_ext * batch, llama_seq_id seq_id, float * embd_token) { +int32_t llama_batch_ext_add_embd_token(llama_batch_ext * batch, llama_seq_id seq_id, llama_embd embd) { int32_t idx = batch->add_token(seq_id); if (idx < 0) { return idx; } - if (!batch->set_token_embd(idx, embd_token)) { + if (!batch->set_token_embd(idx, embd)) { return -2; } return idx; diff --git a/src/llama-batch.h b/src/llama-batch.h index 98648148cda5..59f7b9209d7a 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -103,7 +103,7 @@ struct llama_batch_ext { bool add_seq(int32_t idx, llama_seq_id seq_id); bool set_token_id(int32_t idx, llama_token id); - bool set_token_embd(int32_t idx, float * embd_in); + bool set_token_embd(int32_t idx, llama_embd embd_in); bool set_token_pos(int32_t idx, llama_pos * pos_in); bool set_output(int32_t idx, bool output_last); }; From 69258fb7633991a105d25e9bbb539a73ad692cc2 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 18:25:37 +0200 Subject: [PATCH 17/23] stub llama_batch_ext_set_embd_state --- include/llama.h | 8 ++++---- src/llama-batch.cpp | 7 +++++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/include/llama.h b/include/llama.h index ab3024d5ce8c..08f6bf30ab95 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1029,10 +1029,10 @@ extern "C" { // "state" here means extra hidden state carried over from a previous stage, e.g.: // - MTP: state from N layers of the target model // - Qwen3 VL (deepstack): state from N layers of the vision encoder - // LLAMA_API bool llama_batch_ext_set_embd_state( - // struct llama_batch_ext * batch, - // int32_t idx, - // struct llama_embd embd); + LLAMA_API bool llama_batch_ext_set_embd_state( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); // Set output embedding for the token at index idx in the batch // Note: for now, this is equivalent to setting the output logits diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 3d407448aa84..57501751ecb7 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1160,6 +1160,13 @@ bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * p return batch->set_token_pos(idx, pos); } +bool llama_batch_ext_set_embd_state(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + // TODO + GGML_UNUSED(batch); + GGML_UNUSED(idx); + GGML_UNUSED(embd); +} + bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) { return batch->set_output(idx, value); } From 2b50dabbb4d070744206f4504a73fb5007978f69 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Mon, 17 Aug 2026 18:39:23 +0200 Subject: [PATCH 18/23] support both token + embd + state in batch --- include/llama.h | 7 +++++ src/llama-batch.cpp | 63 ++++++++++++++++++++++++++++++++++++--------- src/llama-batch.h | 1 + 3 files changed, 59 insertions(+), 12 deletions(-) diff --git a/include/llama.h b/include/llama.h index 08f6bf30ab95..c2853fcb4245 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1025,6 +1025,13 @@ extern "C" { int32_t idx, llama_seq_id seq_id); + // Set the token embedding for the token at index idx in the batch + // use it after llama_batch_ext_add_token() to have an entry with both a token id and an embedding + LLAMA_API bool llama_batch_ext_set_embd_token( + struct llama_batch_ext * batch, + int32_t idx, + struct llama_embd embd); + // Set the "state" embedding for the token at index idx in the batch // "state" here means extra hidden state carried over from a previous stage, e.g.: // - MTP: state from N layers of the target model diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 57501751ecb7..b611784cdc9b 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -44,14 +44,35 @@ bool llama_batch_allocr::init( return false; } - const bool has_embd = !batch_inp.embd.empty(); const llama_memory_i * mem = batch_inp.mem; + // + // determine the content types of the batch + // an entry can carry a token id, a token embedding, or both (e.g. MTP hook batches) + // all entries must carry the same combination + // + + const bool has_token = batch_inp.tokens[0].id != LLAMA_TOKEN_NULL; + const bool has_embd = batch_inp.tokens[0].has_embd; + + for (int32_t i = 1; i < n_tok; ++i) { + if ((batch_inp.tokens[i].id != LLAMA_TOKEN_NULL) != has_token || + batch_inp.tokens[i].has_embd != has_embd) { + LLAMA_LOG_ERROR("%s: all entries in the batch must have the same content types\n", __func__); + return false; + } + } + + if (!has_token && !has_embd) { + LLAMA_LOG_ERROR("%s: batch has neither token ids nor embeddings\n", __func__); + return false; + } + // // build flat token/embd array // - if (!has_embd) { + if (has_token) { token_vec.resize(n_tok); for (int32_t i = 0; i < n_tok; ++i) { const llama_token id = batch_inp.tokens[i].id; @@ -61,7 +82,9 @@ bool llama_batch_allocr::init( } token_vec[i] = id; } - } else { + } + + if (has_embd) { embd_vec = batch_inp.embd; } @@ -72,9 +95,9 @@ bool llama_batch_allocr::init( // { - const int32_t n_pos_total = has_embd ? n_tok * (int32_t) n_pos_per_embd : n_tok; + const int32_t n_pos_total = has_token ? n_tok : n_tok * (int32_t) n_pos_per_embd; pos.resize(n_pos_total); - if (!has_embd) { + if (has_token) { for (int32_t i = 0; i < n_tok; ++i) { pos[i] = batch_inp.tokens[i].pos[0]; } @@ -150,8 +173,8 @@ bool llama_batch_allocr::init( // batch.n_tokens = n_tok; - batch.token = has_embd ? nullptr : token_vec.data(); - batch.embd = has_embd ? embd_vec.data() : nullptr; + batch.token = has_token ? token_vec.data() : nullptr; + batch.embd = has_embd ? embd_vec.data() : nullptr; batch.pos = pos.data(); batch.n_seq_id = n_seq_id.data(); batch.seq_id = seq_id.data(); @@ -1080,6 +1103,12 @@ bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { token & t = tokens[idx]; + if (t.has_embd) { + LLAMA_LOG_ERROR("%s: embedding for token %d is already set\n", __func__, idx); + return false; + } + + t.has_embd = true; t.embd_off = embd.size(); embd.insert(embd.end(), embd_in.data, embd_in.data + n_total); @@ -1160,11 +1189,16 @@ bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * p return batch->set_token_pos(idx, pos); } +bool llama_batch_ext_set_embd_token(llama_batch_ext * batch, int32_t idx, llama_embd embd) { + return batch->set_token_embd(idx, embd); +} + bool llama_batch_ext_set_embd_state(llama_batch_ext * batch, int32_t idx, llama_embd embd) { // TODO GGML_UNUSED(batch); GGML_UNUSED(idx); GGML_UNUSED(embd); + return false; } bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) { @@ -1180,7 +1214,9 @@ bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, boo llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp) { batch_ext = new llama_batch_ext(ctx); - const bool is_embd = batch_inp.embd != nullptr; + // a batch can carry both, for example the MTP hook batches + const bool has_token = batch_inp.token != nullptr; + const bool has_embd = batch_inp.embd != nullptr; static const llama_seq_id default_seq_id = 0; static const int32_t default_n_seq_id = 1; @@ -1204,7 +1240,7 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & // position(s) if (batch_inp.pos) { - if (!is_embd) { + if (has_token) { // token batch: one position per token t.pos[0] = batch_inp.pos[i]; } else { @@ -1218,10 +1254,13 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & t.pos[0] = pos_next[sids[0]]++; } - // token id or embeddings - if (!is_embd) { + // token id and/or embeddings + if (has_token) { t.id = batch_inp.token[i]; - } else { + } + + if (has_embd) { + t.has_embd = true; t.embd_off = batch_ext->embd.size(); const float * src = batch_inp.embd + (size_t) i * batch_ext->n_embd_inp; batch_ext->embd.insert(batch_ext->embd.end(), src, src + batch_ext->n_embd_inp); diff --git a/src/llama-batch.h b/src/llama-batch.h index 59f7b9209d7a..8eba06f38356 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -85,6 +85,7 @@ struct llama_batch_ext { struct token { llama_token id = LLAMA_TOKEN_NULL; + bool has_embd = false; // whether embd_off is set size_t embd_off = 0; // index offset in the embd array bool output = false; // TODO: have dedicated output flags std::unordered_set seq_ids; From de88a4a3ce1552a6e8cbcb32e41abc7c81c6f616 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 29 Aug 2026 00:13:49 +0200 Subject: [PATCH 19/23] llama_batch_ext_add_embd --- include/llama.h | 6 +++--- src/llama-batch.cpp | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/include/llama.h b/include/llama.h index 4e1e2d3dea3a..e5348c339ed2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1021,11 +1021,11 @@ extern "C" { // -1: batch is full // -2: token is invalid (id == LLAMA_TOKEN_NULL or invalid embd) // -3: invalid sequence id - LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); + LLAMA_API int32_t llama_batch_ext_add (struct llama_batch_ext * batch, llama_seq_id seq_id); // Add an input token to the batch, with a specified token ID or token embedding - LLAMA_API int32_t llama_batch_ext_add_token (struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); - LLAMA_API int32_t llama_batch_ext_add_embd_token(struct llama_batch_ext * batch, llama_seq_id seq_id, struct llama_embd embd); + LLAMA_API int32_t llama_batch_ext_add_token(struct llama_batch_ext * batch, llama_seq_id seq_id, llama_token id); + LLAMA_API int32_t llama_batch_ext_add_embd (struct llama_batch_ext * batch, llama_seq_id seq_id, struct llama_embd embd); // Add the token at index idx in the batch to another sequence id. The position will stays the same. // Note: this should be called before other _set() functions diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index b611784cdc9b..daafb6977456 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1170,7 +1170,7 @@ int32_t llama_batch_ext_add_token(llama_batch_ext * batch, llama_seq_id seq_id, return idx; } -int32_t llama_batch_ext_add_embd_token(llama_batch_ext * batch, llama_seq_id seq_id, llama_embd embd) { +int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, llama_embd embd) { int32_t idx = batch->add_token(seq_id); if (idx < 0) { return idx; From e67da541ad21cf5fa1aa63186b64814e45ebeb10 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 12 Sep 2026 10:57:52 +0200 Subject: [PATCH 20/23] upstream some changes --- include/llama.h | 4 ++-- src/llama-batch.cpp | 30 ++++++++++++++++++++++-------- src/llama-batch.h | 12 +++++++++--- src/llama-context.cpp | 14 +++++++++++++- 4 files changed, 46 insertions(+), 14 deletions(-) diff --git a/include/llama.h b/include/llama.h index 26290a979140..3ab653092fb1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1015,7 +1015,7 @@ extern "C" { // Add an input token to the batch, with default values: // id = LLAMA_TOKEN_NULL // embd = nullptr - // pos = next available position for the sequence id + // pos = not set, the caller must set it with llama_batch_ext_set_pos() // Returns the batch index (>= 0) // On error: // -1: batch is full @@ -1071,7 +1071,7 @@ extern "C" { LLAMA_API bool llama_batch_ext_set_pos( struct llama_batch_ext * batch, int32_t idx, - llama_pos * pos); + const llama_pos * pos); // TODO: implement get_embeddings() and get_logits() for llama_batch_ext diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index fa49370046c2..72e902e289e9 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -32,7 +32,7 @@ bool llama_batch_allocr::init( clear(); this->vocab = &vocab; - this->n_embd = batch_inp.n_embd_inp; + this->n_embd = batch_inp.n_embd > 0 ? batch_inp.n_embd : batch_inp.n_embd_inp; this->n_seq_max = batch_inp.n_seq_max; const int32_t n_tok = (int32_t) batch_inp.tokens.size(); @@ -1037,6 +1037,7 @@ size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch a llama_batch_ext::llama_batch_ext(llama_context * ctx) : n_tokens_max(llama_n_batch(ctx)), n_embd_inp(llama_batch_ext_select_n_embd_inp(ctx->get_cparams().ctx_type, llama_get_model(ctx)->arch, llama_get_model(ctx)->hparams)), + n_embd_inp_enc(llama_get_model(ctx)->hparams.n_embd_inp_enc()), n_seq_max(llama_n_seq_max(ctx)), mem(llama_get_memory(ctx)), n_vocab(llama_vocab_n_tokens(llama_model_get_vocab(llama_get_model(ctx)))), @@ -1047,6 +1048,7 @@ llama_batch_ext::llama_batch_ext(llama_context * ctx) : void llama_batch_ext::clear() { tokens.clear(); embd .clear(); + n_embd = 0; } int32_t llama_batch_ext::add_token(llama_seq_id seq_id) { @@ -1101,9 +1103,16 @@ bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { } const size_t n_total = embd_in.n_rows * embd_in.n_embd; - if (n_total != n_embd_inp) { + if (n_embd == 0) { + if (n_total != n_embd_inp && n_total != n_embd_inp_enc) { + LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu or %zu\n", + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd_inp, n_embd_inp_enc); + return false; + } + n_embd = n_total; + } else if (n_total != n_embd) { LLAMA_LOG_ERROR("%s: embedding size mismatch, got %zu rows x %zu = %zu, expected %zu\n", - __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd_inp); + __func__, embd_in.n_rows, embd_in.n_embd, n_total, n_embd); return false; } @@ -1121,7 +1130,7 @@ bool llama_batch_ext::set_token_embd(int32_t idx, llama_embd embd_in) { return true; } -bool llama_batch_ext::set_token_pos(int32_t idx, llama_pos * pos_in) { +bool llama_batch_ext::set_token_pos(int32_t idx, const llama_pos * pos_in) { if (idx < 0 || idx >= (int32_t) tokens.size()) { return false; } @@ -1191,7 +1200,7 @@ bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id return batch->add_seq(idx, seq_id); } -bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, llama_pos * pos) { +bool llama_batch_ext_set_pos(llama_batch_ext * batch, int32_t idx, const llama_pos * pos) { return batch->set_token_pos(idx, pos); } @@ -1217,9 +1226,13 @@ bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, boo // llama_batch_compat -llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp) { +llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row) { batch_ext = new llama_batch_ext(ctx); + if (n_embd_row == 0) { + n_embd_row = batch_ext->n_embd_inp; + } + // a batch can carry both, for example the MTP hook batches const bool has_token = batch_inp.token != nullptr; const bool has_embd = batch_inp.embd != nullptr; @@ -1268,8 +1281,9 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & if (has_embd) { t.has_embd = true; t.embd_off = batch_ext->embd.size(); - const float * src = batch_inp.embd + (size_t) i * batch_ext->n_embd_inp; - batch_ext->embd.insert(batch_ext->embd.end(), src, src + batch_ext->n_embd_inp); + const float * src = batch_inp.embd + (size_t) i * n_embd_row; + batch_ext->embd.insert(batch_ext->embd.end(), src, src + n_embd_row); + batch_ext->n_embd = n_embd_row; } // output flag diff --git a/src/llama-batch.h b/src/llama-batch.h index ac18d8af31a6..da16398b69d0 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -79,12 +79,17 @@ size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch a struct llama_batch_ext { const size_t n_tokens_max; // max number of tokens that can be stored in the batch - const size_t n_embd_inp; // number of embedding dimensions per token + const size_t n_embd_inp; // decoder embd row width + const size_t n_embd_inp_enc; // encoder embd row width (e.g. eagle3/dflash extracted features) const llama_seq_id n_seq_max; // max number of sequences llama_memory_i * mem; // memory for position inference const llama_token n_vocab; // max token ID that we accept const size_t n_pos_per_embd; + // actual embd row width of this batch, set by the first set_token_embd() + // must be either n_embd_inp or n_embd_inp_enc; encode/decode verify it against the graph input + size_t n_embd = 0; + struct token { llama_token id = LLAMA_TOKEN_NULL; bool has_embd = false; // whether embd_off is set @@ -107,7 +112,7 @@ struct llama_batch_ext { bool add_seq(int32_t idx, llama_seq_id seq_id); bool set_token_id(int32_t idx, llama_token id); bool set_token_embd(int32_t idx, llama_embd embd_in); - bool set_token_pos(int32_t idx, llama_pos * pos_in); + bool set_token_pos(int32_t idx, const llama_pos * pos_in); bool set_output(int32_t idx, bool output_last); }; @@ -217,6 +222,7 @@ class llama_batch_allocr { // RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext struct llama_batch_compat { llama_batch_ext * batch_ext; - llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp); + // n_embd_row is the embd row width of batch_inp, 0 = use the decoder width + llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0); ~llama_batch_compat(); }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 69bfdb84e0d4..3e4e54c71fbe 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1411,6 +1411,12 @@ int llama_context::encode(const llama_batch_ext & batch_inp) { const auto & hparams = model.hparams; + if (batch_inp.n_embd > 0 && batch_inp.n_embd != hparams.n_embd_inp_enc()) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the encoder input %u\n", + __func__, batch_inp.n_embd, hparams.n_embd_inp_enc()); + return -1; + } + // eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim const int64_t n_vocab = model.vocab.n_tokens(); @@ -1647,6 +1653,12 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { return -1; } + if (batch_inp.n_embd > 0 && batch_inp.n_embd != batch_inp.n_embd_inp) { + LLAMA_LOG_ERROR("%s: embd row width %zu does not match the decoder input %zu\n", + __func__, batch_inp.n_embd, batch_inp.n_embd_inp); + return -1; + } + const auto & vocab = model.vocab; const auto & hparams = model.hparams; @@ -4222,7 +4234,7 @@ size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, lla // compat: llama_batch -> llama_batch_ext -> encode/decode int llama_context::encode(const llama_batch & batch_inp) { - llama_batch_compat compat(this, batch_inp); + llama_batch_compat compat(this, batch_inp, model.hparams.n_embd_inp_enc()); return encode(*compat.batch_ext); } From 12691fb333b6fa41c81a39911dc2ad9579b1a958 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 23 Sep 2026 13:43:44 +0200 Subject: [PATCH 21/23] nits --- include/llama.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/llama.h b/include/llama.h index 91807fe355ad..1805ed0559f9 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1058,7 +1058,7 @@ extern "C" { int32_t idx, struct llama_embd embd); - // Set output embedding for the token at index idx in the batch + // Set if output embeddings should be available for the token at index idx in the batch // Note: for now, this is equivalent to setting the output logits LLAMA_API bool llama_batch_ext_set_output_embd( struct llama_batch_ext * batch, From 82a63e3080173d9eb726aa54c79f9ad008e85f73 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 23 Sep 2026 14:19:49 +0200 Subject: [PATCH 22/23] fix test-batch-alloc --- src/llama-batch.cpp | 18 ++ src/llama-batch.h | 10 + tests/CMakeLists.txt | 2 +- tests/test-batch-alloc.cpp | 363 ++++++++++++++++++++++++++----------- 4 files changed, 282 insertions(+), 111 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 72e902e289e9..589b0e70e594 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1045,6 +1045,24 @@ llama_batch_ext::llama_batch_ext(llama_context * ctx) : clear(); } +llama_batch_ext::llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd) : + n_tokens_max(n_tokens_max), + n_embd_inp(n_embd_inp), + n_embd_inp_enc(n_embd_inp_enc), + n_seq_max(n_seq_max), + mem(mem), + n_vocab(n_vocab), + n_pos_per_embd(n_pos_per_embd) { + clear(); +} + void llama_batch_ext::clear() { tokens.clear(); embd .clear(); diff --git a/src/llama-batch.h b/src/llama-batch.h index da16398b69d0..dcb6d0e61049 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -103,6 +103,16 @@ struct llama_batch_ext { llama_batch_ext(llama_context * ctx); + // build without a llama_context, used by tests + llama_batch_ext( + size_t n_tokens_max, + size_t n_embd_inp, + size_t n_embd_inp_enc, + llama_seq_id n_seq_max, + llama_memory_i * mem, + llama_token n_vocab, + size_t n_pos_per_embd); + void clear(); // add an entry with an undefined position diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 89f39a06e0b3..9b3a4fcc4bbf 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -159,7 +159,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-parser.cpp) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) - # llama_build_and_test(test-batch-alloc.cpp) // TODO: fix this before merging this PR + llama_build_and_test(test-batch-alloc.cpp) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index c4a863559d14..ba837e97f4f1 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -49,49 +49,55 @@ struct mock_memory : public llama_memory_i { void state_read (llama_io_read_i &, llama_seq_id, llama_state_seq_flags) override { GGML_ASSERT(false && "not implemented"); } }; -// builds embedding batches - an empty llama_vocab rejects all token ids, so -// the tests use embeddings everywhere except the token validation tests +// builds a llama_batch_ext without a llama_context +// n_vocab = 0 by default, so every token id is invalid and the tests use embeddings unless stated otherwise struct batch_builder { - uint32_t n_embd; - - std::vector embd; - std::vector pos; - std::vector n_seq_id; - std::vector logits; - - std::vector> seq; - std::vector seq_ptr; - - batch_builder(uint32_t n_embd = 2) : n_embd(n_embd) {} - - // embd values are 100*i + k so that ubatch contents can be traced back to batch indices - void add(llama_pos p, std::initializer_list seq_ids, bool output) { - const int32_t i = (int32_t) seq.size(); - for (uint32_t k = 0; k < n_embd; ++k) { - embd.push_back(100.0f*i + k); + const uint32_t n_embd; + + llama_batch_ext b; + + batch_builder( + uint32_t n_embd = 2, + llama_memory_i * mem = nullptr, + llama_seq_id n_seq_max = 4, + uint32_t n_pos_per_embd = 1, + llama_token n_vocab = 0, + uint32_t n_embd_inp_enc = 0) + : n_embd(n_embd), + b(/*n_tokens_max*/ 64, n_embd, n_embd_inp_enc > 0 ? n_embd_inp_enc : n_embd, n_seq_max, mem, n_vocab, n_pos_per_embd) {} + + // one embedding row for batch index i, values 100*i + k so ubatch contents can be traced back + std::vector row(int32_t i, uint32_t width) const { + std::vector r(width); + for (uint32_t k = 0; k < width; ++k) { + r[k] = 100.0f*i + k; } - pos.push_back(p); - n_seq_id.push_back((int32_t) seq_ids.size()); - seq.emplace_back(seq_ids); - logits.push_back(output ? 1 : 0); + return r; } - llama_batch make(bool with_pos = true, bool with_seq = true, bool with_logits = true) { - seq_ptr.clear(); - for (auto & s : seq) { - seq_ptr.push_back(s.data()); + // embedding entry with full M-RoPE positions + int32_t add_embd(const llama_pos * pos, std::initializer_list seq_ids, bool output, uint32_t width = 0) { + width = width > 0 ? width : n_embd; + + auto it = seq_ids.begin(); + const int32_t idx = b.add_token(*it); + GGML_ASSERT(idx >= 0); + for (++it; it != seq_ids.end(); ++it) { + GGML_ASSERT(b.add_seq(idx, *it)); } - seq_ptr.push_back(nullptr); - llama_batch res = {}; - res.n_tokens = (int32_t) seq.size(); - res.embd = embd.data(); - res.pos = with_pos ? pos.data() : nullptr; - res.n_seq_id = with_seq ? n_seq_id.data() : nullptr; - res.seq_id = with_seq ? seq_ptr.data() : nullptr; - res.logits = with_logits ? logits.data() : nullptr; + const auto r = row(idx, width); + GGML_ASSERT(b.set_token_embd(idx, { r.data(), 1, width })); + GGML_ASSERT(b.set_token_pos(idx, pos)); + GGML_ASSERT(b.set_output(idx, output)); - return res; + return idx; + } + + // embedding entry with a single sequential position + int32_t add(llama_pos p, std::initializer_list seq_ids, bool output) { + const llama_pos pos[GGML_MROPE_SECTIONS] = { p, 0, 0, 0 }; + return add_embd(pos, seq_ids, output); } }; @@ -99,22 +105,31 @@ static void test_init(testing & t) { llama_vocab vocab; t.test("rejects_n_seq_max_too_large", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, LLAMA_MAX_SEQ + 1); bb.add(0, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, LLAMA_MAX_SEQ + 1, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_token", [&](testing & t) { - llama_token tok = 0; // empty vocab -> every token id is out of range - llama_batch batch = llama_batch_get_one(&tok, 1); + // n_vocab = 0 -> every token id is out of range + // set_token_id() refuses such ids, so the token is poked directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); llama_batch_allocr ba(1); - t.assert_true("token id >= n_tokens", !ba.init(batch, vocab, nullptr, 0, 1, false)); - tok = -1; - t.assert_true("negative token id", !ba.init(batch, vocab, nullptr, 0, 1, false)); + t.assert_true("set_token_id refuses out of range id", !bb.b.set_token_id(idx, 0)); + + bb.b.tokens[idx].id = 0; + t.assert_true("token id >= n_vocab", !ba.init(bb.b, vocab, false)); + + bb.b.tokens[idx].id = -1; + t.assert_true("negative token id", !ba.init(bb.b, vocab, false)); }); t.test("rejects_invalid_seq_id", [&](testing & t) { @@ -122,33 +137,44 @@ static void test_init(testing & t) { { batch_builder bb; - bb.add(0, {4}, true); - t.assert_true("seq_id >= n_seq_max", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true("add_token refuses seq_id >= n_seq_max", bb.b.add_token(4) == -3); + t.assert_true("add_token refuses negative seq_id", bb.b.add_token(-1) == -3); + } + { + // poke the seq_ids directly to reach the init() check + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { 4 }; + t.assert_true("seq_id >= n_seq_max", !ba.init(bb.b, vocab, false)); } { batch_builder bb; - bb.add(0, {-1}, true); - t.assert_true("negative seq_id", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + const int32_t idx = bb.add(0, {0}, true); + bb.b.tokens[idx].seq_ids = { -1 }; + t.assert_true("negative seq_id", !ba.init(bb.b, vocab, false)); } }); - t.test("autofill_defaults", [&](testing & t) { + t.test("copies_pos_seq_output", [&](testing & t) { batch_builder bb; for (int i = 0; i < 4; ++i) { - bb.add(0, {0}, false); + bb.add(i, {0}, i == 3); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, false, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); const llama_batch & batch = ba.get_batch(); t.assert_equal(4u, ba.get_n_tokens()); + t.assert_true("embedding batch", batch.embd != nullptr); + t.assert_true("no token ids", batch.token == nullptr); for (int i = 0; i < 4; ++i) { - t.assert_equal("pos defaults to 0..n-1", i, batch.pos[i]); - t.assert_equal("n_seq_id defaults to 1", 1, batch.n_seq_id[i]); - t.assert_equal("seq_id defaults to 0", 0, batch.seq_id[i][0]); + t.assert_equal(i, batch.pos[i]); + t.assert_equal(1, batch.n_seq_id[i]); + t.assert_equal(0, batch.seq_id[i][0]); + t.assert_equal(100.0f*i, batch.embd[i*bb.n_embd]); } t.assert_equal("only the last token is an output", 1u, ba.get_n_outputs()); @@ -167,7 +193,7 @@ static void test_init(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, true)); + t.assert_true(ba.init(bb.b, vocab, true)); t.assert_equal(4u, ba.get_n_outputs()); }); @@ -178,7 +204,7 @@ static void test_init(testing & t) { bb.add(2, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal(2u, ba.get_n_outputs()); llama_ubatch ub = ba.split_simple(10); @@ -193,17 +219,17 @@ static void test_init(testing & t) { t.assert_equal(2, out_ids[1]); }); - t.test("pos_from_memory", [&](testing & t) { + t.test("pos_after_memory", [&](testing & t) { mock_memory mem; mem.ranges[0] = {0, 9}; - batch_builder bb; + batch_builder bb(2, &mem); for (int i = 0; i < 3; ++i) { - bb.add(0, {0}, false); + bb.add(10 + i, {0}, false); } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(false, true, false), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); t.assert_equal("pos continues after memory", 10, ba.seq_pos_min(0)); t.assert_equal(12, ba.seq_pos_max(0)); @@ -216,22 +242,22 @@ static void test_init(testing & t) { llama_batch_allocr ba(1); { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(10, {0}, false); bb.add(11, {0}, true); - t.assert_true("pos_max + 1 is accepted", ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("pos_max + 1 is accepted", ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(11, {0}, false); bb.add(12, {0}, true); - t.assert_true("gap after memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("gap after memory is rejected", !ba.init(bb.b, vocab, false)); } { - batch_builder bb; + batch_builder bb(2, &mem); bb.add(9, {0}, false); bb.add(10, {0}, true); - t.assert_true("overlap with memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + t.assert_true("overlap with memory is rejected", !ba.init(bb.b, vocab, false)); } }); @@ -242,7 +268,7 @@ static void test_init(testing & t) { bb.add(3, {0}, true); llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("rejects_decreasing_positions", [&](testing & t) { @@ -255,7 +281,7 @@ static void test_init(testing & t) { // seq 0 sees positions 4,5,6,3 in batch order -> the trailing 3 decreases llama_batch_allocr ba(1); - t.assert_true(!ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(!ba.init(bb.b, vocab, false)); }); t.test("allows_equal_positions_in_seq", [&](testing & t) { @@ -265,23 +291,143 @@ static void test_init(testing & t) { bb.add(1, {0}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); }); - t.test("rejects_coupled_diverged_seqs", [&](testing & t) { - batch_builder bb; - bb.add(6, {0, 1}, true); - llama_batch_allocr ba(1); mock_memory mem; mem.ranges[0] = {0, 5}; mem.ranges[1] = {2, 5}; // same pos_max, different pos_min -> diverged - t.assert_true(!ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(!ba.init(bb.b, vocab, false)); + } mem.ranges[1] = {0, 5}; - t.assert_true(ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false)); + { + batch_builder bb(2, &mem); + bb.add(6, {0, 1}, true); + t.assert_true(ba.init(bb.b, vocab, false)); + } + }); +} + +static void test_content_types(testing & t) { + llama_vocab vocab; + + t.test("token_and_embd_together", [&](testing & t) { + // e.g. MTP hook batches: a token id and its embedding on the same entry + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + const int32_t idx = bb.b.add_token(0); + t.assert_true(bb.b.set_token_id(idx, 3)); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + const llama_batch & batch = ba.get_batch(); + t.assert_true("token ids are kept", batch.token != nullptr); + t.assert_true("embeddings are kept", batch.embd != nullptr); + t.assert_equal(3, batch.token[0]); + t.assert_equal(0.0f, batch.embd[0]); + t.assert_equal(1.0f, batch.embd[1]); + + llama_ubatch ub = ba.split_simple(1); + t.assert_true(ub.token != nullptr && ub.embd != nullptr); + t.assert_equal(3, ub.token[0]); + }); + + t.test("rejects_mixed_content_types", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 10); + + // entry 0: token only, entry 1: token + embd + const llama_pos p0 = 0; + const llama_pos p1 = 1; + + int32_t i0 = bb.b.add_token(0); + bb.b.set_token_id(i0, 1); + bb.b.set_token_pos(i0, &p0); + + int32_t i1 = bb.b.add_token(0); + bb.b.set_token_id(i1, 2); + const auto r = bb.row(i1, bb.n_embd); + bb.b.set_token_embd(i1, { r.data(), 1, bb.n_embd }); + bb.b.set_token_pos(i1, &p1); + bb.b.set_output(i1, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_neither_token_nor_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.b.add_token(0); + const llama_pos pos = 0; + bb.b.set_token_pos(idx, &pos); + bb.b.set_output(idx, true); + + llama_batch_allocr ba(1); + t.assert_true(!ba.init(bb.b, vocab, false)); + }); + + t.test("rejects_embd_size_mismatch", [&](testing & t) { + batch_builder bb; // n_embd = 2, n_embd_inp_enc = 2 + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 8); + + t.assert_true("too small", !bb.b.set_token_embd(idx, { r.data(), 1, 1 })); + t.assert_true("too large", !bb.b.set_token_embd(idx, { r.data(), 1, 3 })); + t.assert_true("zero rows", !bb.b.set_token_embd(idx, { r.data(), 0, 2 })); + t.assert_true("null data", !bb.b.set_token_embd(idx, { nullptr, 1, 2 })); + t.assert_true("same total via a different split is accepted", bb.b.set_token_embd(idx, { r.data(), 2, 1 })); + }); + + t.test("rejects_double_embd", [&](testing & t) { + batch_builder bb; + const int32_t idx = bb.add(0, {0}, true); + const auto r = bb.row(idx, bb.n_embd); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, bb.n_embd })); + }); + + t.test("encoder_width", [&](testing & t) { + // e.g. eagle3/dflash: extracted features are wider than the decoder input + const uint32_t n_embd_enc = 6; + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + + const llama_pos p0 = 0; + const llama_pos p1 = 1; + bb.add_embd(&p0, {0}, false, n_embd_enc); + bb.add_embd(&p1, {0}, true, n_embd_enc); + + t.assert_equal("batch width follows the first embedding", (size_t) n_embd_enc, bb.b.n_embd); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + + // the ubatch uses the encoder stride: token 1 starts at offset n_embd_enc + llama_ubatch ub = ba.split_simple(2); + t.assert_equal(2u, ub.n_tokens); + t.assert_equal(100.0f, ub.embd[n_embd_enc]); + t.assert_equal(105.0f, ub.embd[n_embd_enc + 5]); + }); + + t.test("rejects_mixing_widths", [&](testing & t) { + batch_builder bb(2, nullptr, 4, 1, 0, /*n_embd_inp_enc*/ 6); + + const llama_pos p0 = 0; + bb.add_embd(&p0, {0}, false, 2); // first entry fixes the batch width to 2 + + const int32_t idx = bb.b.add_token(0); + const auto r = bb.row(idx, 6); + t.assert_true(!bb.b.set_token_embd(idx, { r.data(), 1, 6 })); }); } @@ -295,7 +441,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); @@ -338,7 +484,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); while (ba.split_simple(1).n_tokens > 0) { } @@ -361,7 +507,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(8, false, 0); t.assert_true(ub.equal_seqs()); @@ -397,7 +543,7 @@ static void test_split(testing & t) { bb.add(1, {0, 1}, true); llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, true, 0); t.assert_equal("sequential split rejects coupled seqs", 0u, ub.n_tokens); @@ -419,7 +565,7 @@ static void test_split(testing & t) { } llama_batch_allocr ba(1); - t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); for (llama_seq_id s = 0; s < 3; ++s) { llama_ubatch ub = ba.split_seq(8); @@ -461,14 +607,14 @@ static void test_keep_tail(testing & t) { } ++s; } - return bb.make(); }; t.test("noop_when_seqs_complete", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 2}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 2}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("both seqs fit whole", 4u, ub.n_tokens); @@ -480,9 +626,10 @@ static void test_keep_tail(testing & t) { t.test("defers_seq_with_short_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {2, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {2, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq: seq 0 completes, seq 1 would be left // with 1 < n_keep_tail remaining, so it is deferred entirely @@ -506,9 +653,10 @@ static void test_keep_tail(testing & t) { t.test("completes_first_seq_when_all_violate", [&](testing & t) { batch_builder bb; + make_batch(bb, {3, 3}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 3}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 2 tokens per seq, leaving both with 1 < n_keep_tail remaining; // seq 0 still fits in n_ubatch, so it is extended to completion and emitted alone @@ -530,9 +678,10 @@ static void test_keep_tail(testing & t) { t.test("truncates_to_preserve_tail", [&](testing & t) { batch_builder bb; + make_batch(bb, {5}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {5}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // 4 tokens would leave a remainder of 1, and the seq does not fit in n_ubatch, // so the ubatch is truncated until n_keep_tail tokens remain @@ -553,9 +702,10 @@ static void test_keep_tail(testing & t) { t.test("keeps_full_ubatch_with_sufficient_remainder", [&](testing & t) { batch_builder bb; + make_batch(bb, {6}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {6}), vocab, nullptr, bb.n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_equal(4, false, 2); t.assert_equal("remainder >= n_keep_tail, no truncation", 4u, ub.n_tokens); @@ -569,10 +719,11 @@ static void test_keep_tail(testing & t) { }); t.test("multi_seq_prefix_kept", [&](testing & t) { - batch_builder bb; + batch_builder bb(2, nullptr, 6); + make_batch(bb, {3, 4}); llama_batch_allocr ba(1); - t.assert_true(ba.init(make_batch(bb, {3, 4}), vocab, nullptr, bb.n_embd, 6, false)); + t.assert_true(ba.init(bb.b, vocab, false)); // expansion stops at 3 tokens per seq: seq 0 completes, seq 1 has 1 < n_keep_tail // remaining and is deferred even though its tokens were already gathered @@ -593,32 +744,26 @@ static void test_mrope(testing & t) { llama_vocab vocab; t.test("pos_layout_and_split", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; - batch_builder bb(n_embd); - bb.add(10, {0}, false); - bb.add(11, {0}, true); - - // M-RoPE positions for embeddings are laid out [n_pos][n_tokens] - std::vector pos = { - 10, 11, // temporal - 5, 6, // y - 7, 8, // x - 0, 0, - }; + batch_builder bb(n_embd, nullptr, 4, n_pos); - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); + // M-RoPE positions per embedding: [temporal, y, x, other] + const llama_pos pos0[n_pos] = { 10, 5, 7, 0 }; + const llama_pos pos1[n_pos] = { 11, 6, 8, 0 }; + bb.add_embd(pos0, {0}, false); + bb.add_embd(pos1, {0}, true); llama_batch_allocr ba(n_pos); - t.assert_true(ba.init(batch, vocab, nullptr, n_embd, 4, false)); + t.assert_true(ba.init(bb.b, vocab, false)); llama_ubatch ub = ba.split_simple(2); t.assert_equal(2u, ub.n_tokens); t.assert_equal(n_pos, ub.n_pos); t.assert_true(ub.is_pos_2d()); + // the ubatch stores positions section-major: [n_pos][n_tokens] const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; for (int i = 0; i < 8; ++i) { t.assert_equal(expected[i], ub.pos[i]); @@ -626,7 +771,7 @@ static void test_mrope(testing & t) { }); t.test("pos_jump_allowed", [&](testing & t) { - const uint32_t n_pos = 4; + const uint32_t n_pos = 4; const uint32_t n_embd = 2; mock_memory mem; @@ -635,15 +780,12 @@ static void test_mrope(testing & t) { llama_batch_allocr ba(n_pos); auto try_pos = [&](llama_pos p0) { - batch_builder bb(n_embd); - bb.add(p0, {0}, true); - - std::vector pos = {p0, 1, 1, 0}; + batch_builder bb(n_embd, &mem, 4, n_pos); - llama_batch batch = bb.make(false, true, true); - batch.pos = pos.data(); + const llama_pos pos[n_pos] = { p0, 1, 1, 0 }; + bb.add_embd(pos, {0}, true); - return ba.init(batch, vocab, &mem, n_embd, 4, false); + return ba.init(bb.b, vocab, false); }; t.assert_true("gap after memory is allowed", try_pos(15)); @@ -707,6 +849,7 @@ int main(int argc, char ** argv) { } t.test("init", test_init); + t.test("content_types", test_content_types); t.test("split", test_split); t.test("keep_tail", test_keep_tail); t.test("mrope", test_mrope); From 40704338f054f90e4568cc59335e4ee6dc75afba Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Wed, 23 Sep 2026 18:42:26 +0200 Subject: [PATCH 23/23] add test for compat --- src/llama-batch.cpp | 9 +- src/llama-batch.h | 5 + tests/test-batch-alloc.cpp | 210 +++++++++++++++++++++++++++++++++++++ 3 files changed, 222 insertions(+), 2 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 589b0e70e594..89a1f3f37c7b 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1244,8 +1244,8 @@ bool llama_batch_ext_set_output_logits(llama_batch_ext * batch, int32_t idx, boo // llama_batch_compat -llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row) { - batch_ext = new llama_batch_ext(ctx); +void llama_batch_compat::init(llama_batch_ext & dst, const llama_batch & batch_inp, size_t n_embd_row) { + llama_batch_ext * batch_ext = &dst; if (n_embd_row == 0) { n_embd_row = batch_ext->n_embd_inp; @@ -1314,6 +1314,11 @@ llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & } } +llama_batch_compat::llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row) { + batch_ext = new llama_batch_ext(ctx); + init(*batch_ext, batch_inp, n_embd_row); +} + llama_batch_compat::~llama_batch_compat() { delete batch_ext; } diff --git a/src/llama-batch.h b/src/llama-batch.h index dcb6d0e61049..201d48cce18d 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -232,7 +232,12 @@ class llama_batch_allocr { // RAII translation layer: converts a llama_batch (old API) into a llama_batch_ext struct llama_batch_compat { llama_batch_ext * batch_ext; + // n_embd_row is the embd row width of batch_inp, 0 = use the decoder width llama_batch_compat(llama_context * ctx, const llama_batch & batch_inp, size_t n_embd_row = 0); ~llama_batch_compat(); + + // fill an existing llama_batch_ext from a llama_batch (old API) + // note: this is called directly by the tests, skipping llama_context creation + static void init(llama_batch_ext & batch_ext, const llama_batch & batch_inp, size_t n_embd_row = 0); }; diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index ba837e97f4f1..ad186c6936d7 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -794,6 +794,215 @@ static void test_mrope(testing & t) { }); } +// conversion from the old llama_batch API (llama_batch_compat::init) +static void test_compat(testing & t) { + llama_vocab vocab; + + t.test("token_batch_explicit_fields", [&](testing & t) { + llama_token token[3] = { 5, 6, 7 }; + llama_pos pos[3] = { 3, 4, 5 }; + int32_t n_seq_id[3] = { 1, 1, 2 }; + llama_seq_id s0[1] = { 1 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id s2[2] = { 1, 2 }; + llama_seq_id * seq_id[4] = { s0, s1, s2, nullptr }; + int8_t logits[3] = { 0, 1, 0 }; + + llama_batch lb = {}; + lb.n_tokens = 3; + lb.token = token; + lb.pos = pos; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + lb.logits = logits; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + t.assert_true("no embeddings", bb.b.embd.empty() && bb.b.n_embd == 0); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_equal(pos[i], bb.b.tokens[i].pos[0]); + t.assert_true(!bb.b.tokens[i].has_embd); + t.assert_equal(logits[i] != 0, bb.b.tokens[i].output); + } + t.assert_equal((size_t) 1, bb.b.tokens[0].seq_ids.size()); + t.assert_true(bb.b.tokens[0].seq_ids.count(1) == 1); + t.assert_equal((size_t) 2, bb.b.tokens[2].seq_ids.size()); + t.assert_true(bb.b.tokens[2].seq_ids.count(1) == 1 && bb.b.tokens[2].seq_ids.count(2) == 1); + + // round trip through the allocator + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true(batch.token != nullptr && batch.embd == nullptr); + for (int i = 0; i < 3; ++i) { + t.assert_equal(token[i], batch.token[i]); + t.assert_equal(pos[i], batch.pos[i]); + } + t.assert_equal(1u, ba.get_n_outputs()); + }); + + t.test("defaults_for_null_fields", [&](testing & t) { + // llama_batch_get_one: only token and n_tokens are set + mock_memory mem; + mem.ranges[0] = {0, 9}; + + llama_token token[3] = { 5, 6, 7 }; + llama_batch lb = llama_batch_get_one(token, 3); + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 3, bb.b.tokens.size()); + for (int i = 0; i < 3; ++i) { + t.assert_equal("pos continues after memory", 10 + i, bb.b.tokens[i].pos[0]); + t.assert_equal("seq_id defaults to 0", (size_t) 1, bb.b.tokens[i].seq_ids.size()); + t.assert_true(bb.b.tokens[i].seq_ids.count(0) == 1); + } + t.assert_true("only the last token is an output", !bb.b.tokens[0].output && !bb.b.tokens[1].output && bb.b.tokens[2].output); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + t.assert_equal(10, ba.seq_pos_min(0)); + t.assert_equal(12, ba.seq_pos_max(0)); + }); + + t.test("auto_pos_starts_at_zero_without_memory", [&](testing & t) { + llama_token token[2] = { 5, 6 }; + llama_batch lb = llama_batch_get_one(token, 2); + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal(0, bb.b.tokens[0].pos[0]); + t.assert_equal(1, bb.b.tokens[1].pos[0]); + }); + + t.test("auto_pos_is_tracked_per_seq", [&](testing & t) { + mock_memory mem; + mem.ranges[0] = {0, 9}; // seq 1 is empty + + llama_token token[4] = { 5, 6, 7, 8 }; + int32_t n_seq_id[4] = { 1, 1, 1, 1 }; + llama_seq_id s0[1] = { 0 }; + llama_seq_id s1[1] = { 1 }; + llama_seq_id * seq_id[5] = { s0, s1, s0, s1, nullptr }; + + llama_batch lb = {}; + lb.n_tokens = 4; + lb.token = token; + lb.n_seq_id = n_seq_id; + lb.seq_id = seq_id; + + batch_builder bb(2, &mem, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal("seq 0 continues after memory", 10, bb.b.tokens[0].pos[0]); + t.assert_equal("seq 1 starts from 0", 0, bb.b.tokens[1].pos[0]); + t.assert_equal(11, bb.b.tokens[2].pos[0]); + t.assert_equal( 1, bb.b.tokens[3].pos[0]); + }); + + t.test("embd_batch_with_mrope_positions", [&](testing & t) { + const uint32_t n_pos = 4; + const uint32_t n_embd = 2; + + float embd[2*n_embd] = { 0, 1, 100, 101 }; + // section-major layout: pos[j*n_tokens + i] + llama_pos pos[n_pos*2] = { + 10, 11, // temporal + 5, 6, // y + 7, 8, // x + 0, 0, + }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(n_embd, nullptr, 4, n_pos); + llama_batch_compat::init(bb.b, lb); + + t.assert_equal((size_t) 2, bb.b.tokens.size()); + t.assert_equal("batch width", (size_t) n_embd, bb.b.n_embd); + for (int i = 0; i < 2; ++i) { + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal(LLAMA_TOKEN_NULL, bb.b.tokens[i].id); + t.assert_equal((size_t) i*n_embd, bb.b.tokens[i].embd_off); + for (uint32_t j = 0; j < n_pos; ++j) { + t.assert_equal(pos[j*2 + i], bb.b.tokens[i].pos[j]); + } + } + t.assert_equal(100.0f, bb.b.embd[2]); + t.assert_equal(101.0f, bb.b.embd[3]); + + llama_batch_allocr ba(n_pos); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0}; + for (int i = 0; i < 8; ++i) { + t.assert_equal(expected[i], ub.pos[i]); + } + }); + + t.test("token_and_embd_both_set", [&](testing & t) { + // e.g. MTP hook batches + llama_token token[2] = { 5, 6 }; + float embd[4] = { 0, 1, 100, 101 }; + llama_pos pos[2] = { 3, 4 }; + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.token = token; + lb.embd = embd; + lb.pos = pos; + + batch_builder bb(2, nullptr, 4, 1, /*n_vocab*/ 100); + llama_batch_compat::init(bb.b, lb); + + for (int i = 0; i < 2; ++i) { + t.assert_equal(token[i], bb.b.tokens[i].id); + t.assert_true(bb.b.tokens[i].has_embd); + t.assert_equal("one position per token", pos[i], bb.b.tokens[i].pos[0]); + } + t.assert_equal(100.0f, bb.b.embd[2]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + const llama_batch & batch = ba.get_batch(); + t.assert_true("both kept", batch.token != nullptr && batch.embd != nullptr); + }); + + t.test("embd_row_width_override", [&](testing & t) { + // encoder input (e.g. eagle3/dflash) is wider than the decoder input + const uint32_t n_embd_enc = 6; + float embd[2*n_embd_enc]; + for (int i = 0; i < 2*6; ++i) { + embd[i] = (float) i; + } + + llama_batch lb = {}; + lb.n_tokens = 2; + lb.embd = embd; + + batch_builder bb(2, nullptr, 4, 1, 0, n_embd_enc); + llama_batch_compat::init(bb.b, lb, n_embd_enc); + + t.assert_equal((size_t) n_embd_enc, bb.b.n_embd); + t.assert_equal((size_t) 2*n_embd_enc, bb.b.embd.size()); + t.assert_equal((size_t) n_embd_enc, bb.b.tokens[1].embd_off); + t.assert_equal(6.0f, bb.b.embd[n_embd_enc]); + + llama_batch_allocr ba(1); + t.assert_true(ba.init(bb.b, vocab, false)); + llama_ubatch ub = ba.split_simple(2); + t.assert_equal("ubatch uses the encoder stride", 6.0f, ub.embd[n_embd_enc]); + }); +} + static void test_mtp_embd_width(testing & t) { t.test("mtp_uses_n_embd_out", [&](testing & t) { llama_hparams hparams = {}; @@ -850,6 +1059,7 @@ int main(int argc, char ** argv) { t.test("init", test_init); t.test("content_types", test_content_types); + t.test("compat", test_compat); t.test("split", test_split); t.test("keep_tail", test_keep_tail); t.test("mrope", test_mrope);