ggml-org/llama.cpp · #24669
llama: add llama_batch_ext
common/common.cpp30 + / 7 −
@@ -2198,9 +2198,28 @@ 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));++ 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) {+ 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()) {+ 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<llama_token> & all_tokens,+ const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch,@@ -2221,7 +2240,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<llama_token*>(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; }@@ -2231,17 +2252,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<llama_token*>(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; }common/common.h5 + / 1 −
@@ -1021,14 +1021,18 @@ void common_batch_add( const std::vector<llama_seq_id> & 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 // compatibility with all memory types. Recurrent/hybrid models cannot remove // tokens from memory, so this approach works across all model architectures. bool common_prompt_batch_decode( struct llama_context * ctx,- const std::vector<llama_token> & all_tokens,+ const llama_tokens & all_tokens, int n_new, int & n_past, int n_batch,include/llama-cpp.h5 + / 0 −
@@ -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, llama_model_deleter> llama_model_ptr; typedef std::unique_ptr<llama_context, llama_context_deleter> llama_context_ptr; typedef std::unique_ptr<llama_sampler, llama_sampler_deleter> llama_sampler_ptr; typedef std::unique_ptr<llama_adapter_lora, llama_adapter_lora_deleter> llama_adapter_lora_ptr;+typedef std::unique_ptr<llama_batch_ext, llama_batch_ext_deleter> llama_batch_ext_ptr;include/llama.h90 + / 0 −
@@ -293,6 +293,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; @@ -999,6 +1004,91 @@ extern "C" { struct llama_context * ctx, struct llama_batch batch); + //+ // Extended batch API+ //++ 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);++ // Add an input token to the batch, with default values:+ // id = LLAMA_TOKEN_NULL+ // embd = nullptr+ // 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+ // -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_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+ LLAMA_API bool llama_batch_ext_add_seq(+ struct llama_batch_ext * batch,+ 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+ // - 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);++ // 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,+ 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:+ // - 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,+ const 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);+ // 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)src/llama-batch.cpp428 + / 91 −
@@ -3,6 +3,9 @@ #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h"+#include "llama-hparams.h"+#include "llama-model.h"+#include "llama-context.h" #include <cassert> #include <cstring>@@ -23,135 +26,164 @@ 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 > 0 ? batch_inp.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);++ 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;+ }++ const llama_memory_i * mem = batch_inp.mem; //- // validate input batch+ // 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 // - if (n_seq_max > LLAMA_MAX_SEQ) {- LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ);+ 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; } - 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]);+ //+ // build flat token/embd array+ //++ 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;+ 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; } } - 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;+ if (has_embd) {+ embd_vec = batch_inp.embd;+ }++ //+ // 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_token ? n_tok : n_tok * (int32_t) n_pos_per_embd;+ pos.resize(n_pos_total);+ if (has_token) {+ 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];+ 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]; - pos[i] = p0[seq_id];-- // 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_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();+ 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;@@ -259,7 +291,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)) {@@ -292,7 +324,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;@@ -320,12 +352,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; }@@ -725,11 +757,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();@@ -985,3 +1020,305 @@ void llama_batch_free(struct llama_batch batch) { } if (batch.logits) free(batch.logits); }+++// llama_batch_ext++size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, const llama_hparams & hparams) {+ if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) {+ return hparams.n_embd_out();+ }+ if (arch == LLM_ARCH_DFLASH) {+ return hparams.n_embd_inp_enc();+ }+ return hparams.n_embd_inp();+}++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)))),+ n_pos_per_embd(llama_get_model(ctx)->hparams.n_pos_per_embd()) {+ 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();+ n_embd = 0;+}++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+ }++ // position is left undefined; call set_token_pos() before decoding+ token t;+ t.seq_ids.insert(seq_id);++ 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, llama_embd embd_in) {+ if (idx < 0 || idx >= (int32_t) tokens.size()) {+ return false;+ }+ if (!embd_in.data) {+ return false;+ }++ const size_t n_total = embd_in.n_rows * embd_in.n_embd;+ 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);+ return false;+ }++ 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);++ return true;+}++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;+ }+ 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];+ }++ 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);+}++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(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_seq_id seq_id, llama_token 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, 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)) {+ 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_pos(llama_batch_ext * batch, int32_t idx, const llama_pos * pos) {+ 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) {+ 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++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;+ }++ // 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;++ // auto-generates positions locally when batch_inp.pos is null, continuing from memory+ std::vector<llama_pos> 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;++ 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 (has_token) {+ // 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] = pos_next[sids[0]]++;+ }++ // token id and/or embeddings+ if (has_token) {+ t.id = batch_inp.token[i];+ }++ if (has_embd) {+ t.has_embd = true;+ t.embd_off = batch_ext->embd.size();+ 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+ // 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);+ }+}++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;+}src/llama-batch.h76 + / 7 −
@@ -2,6 +2,7 @@ #include "llama.h" +#include "llama-arch.h" #include "llama-cparams.h" #include <array>@@ -10,6 +11,7 @@ #include <bitset> #include <memory> #include <unordered_map>+#include <unordered_set> // keep this struct lightweight struct llama_ubatch {@@ -68,19 +70,71 @@ struct llama_ubatch { std::shared_ptr<data_t> data; }; +struct llama_hparams;++// MTP hook batches carry the target model's hidden state (n_embd_out size).+// DFlash batches carry the fused target features at the encoder input width (n_embd_inp_enc size).+// Normal batches carry token embeddings (n_embd_inp size).+size_t llama_batch_ext_select_n_embd_inp(llama_context_type ctx_type, llm_arch arch, 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; // 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+ size_t embd_off = 0; // index offset in the embd array+ bool output = false; // TODO: have dedicated output flags+ std::unordered_set<llama_seq_id> seq_ids;+ std::array<llama_pos, GGML_MROPE_SECTIONS> pos = {0, 0, 0, 0};+ };+ std::vector<token> tokens;+ std::vector<float> embd;++ 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+ // 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);+ 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, const llama_pos * pos_in);+ bool set_output(int32_t idx, bool output_last);+};+ // a helper for sanitizing, fulfilling and splitting a batch 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;@@ -137,7 +191,9 @@ class llama_batch_allocr { uint32_t n_seq_max; uint32_t n_outputs; - std::array<llama_seq_id, 1> seq_id_0 = {{ 0 }}; // default sequence id+ std::vector<llama_token> token_vec; // owned token IDs built from llama_batch_ext+ std::vector<float> embd_vec; // owned embeddings built from llama_batch_ext+ std::vector<llama_seq_id> seq_id_data; // flat storage for seq_id pointers below std::vector<llama_pos> pos; std::vector<int32_t> n_seq_id;@@ -172,3 +228,16 @@ 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;++ // 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);+};src/llama-context.cpp52 + / 32 −
@@ -1463,24 +1463,25 @@ 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; + 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_embd = hparams.n_embd_inp_enc(); 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; }@@ -1701,50 +1702,45 @@ static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map<llama_s return false; // all sequences use backend sampling } -int llama_context::decode(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);-+int llama_context::decode(const llama_batch_ext & batch_inp) { if (!memory) { LLAMA_LOG_DEBUG("%s: cannot decode batches with this context (calling encode() instead)\n", __func__); return encode(batch_inp); } - if (batch_inp.n_tokens == 0) {+ if (batch_inp.tokens.empty()) { LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); 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; const int64_t n_vocab = vocab.n_tokens();- const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;- // DFlash embd batches carry the fused target features at the encoder input width- const bool dflash_embd = model.arch == LLM_ARCH_DFLASH && batch_inp.embd;- const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : dflash_embd ? hparams.n_embd_inp_enc() : hparams.n_embd_inp(); // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; const bool has_samplers = !sampling.samplers.empty(); const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; - // embedding contexts output every token even when batch.logits is not set- if (has_samplers && (output_all || batch_inp.logits)) {+ // TODO: avoid this workaround in the future+ // embedding contexts output every token even when no token is explicitly marked as output+ if (has_samplers) { std::vector<int32_t> seq_output_count(n_seq_max, 0); - for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {- if (!output_all && batch_inp.logits[i] == 0) {+ for (const auto & tok : batch_inp.tokens) {+ if (!output_all && !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) { if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { continue; }@@ -1762,7 +1758,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; }@@ -3557,9 +3553,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();@@ -4310,6 +4310,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, model.hparams.n_embd_inp_enc());+ 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(@@ -4399,6 +4411,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 //src/llama-context.h4 + / 0 −
@@ -141,6 +141,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); src/llama-hparams.cpp2 + / 1 −
@@ -283,7 +283,8 @@ bool llama_hparams::is_ple(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 {tests/test-batch-alloc.cpp509 + / 114 −
@@ -3,6 +3,8 @@ #include "llama.h" #include "../src/llama-batch.h"+#include "../src/llama-arch.h"+#include "../src/llama-hparams.h" #include "../src/llama-memory.h" #include "../src/llama-vocab.h" @@ -47,106 +49,132 @@ 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<float> embd;- std::vector<llama_pos> pos;- std::vector<int32_t> n_seq_id;- std::vector<int8_t> logits;-- std::vector<std::vector<llama_seq_id>> seq;- std::vector<llama_seq_id *> 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<llama_seq_id> 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<float> row(int32_t i, uint32_t width) const {+ std::vector<float> 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<llama_seq_id> 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<llama_seq_id> seq_ids, bool output) {+ const llama_pos pos[GGML_MROPE_SECTIONS] = { p, 0, 0, 0 };+ return add_embd(pos, seq_ids, output); } }; 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) { llama_batch_allocr ba(1); { 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;- 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 = { 4 };+ t.assert_true("seq_id >= n_seq_max", !ba.init(bb.b, vocab, false));+ }+ {+ batch_builder bb;+ 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());@@ -165,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()); }); @@ -176,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);@@ -191,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));@@ -214,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)); } }); @@ -240,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) {@@ -253,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) {@@ -263,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 })); }); } @@ -293,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);@@ -336,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) { }@@ -359,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());@@ -395,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);@@ -417,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);@@ -459,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);@@ -478,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@@ -504,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@@ -528,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@@ -551,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);@@ -567,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@@ -591,40 +744,34 @@ 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<llama_pos> 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]); } }); 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;@@ -633,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<llama_pos> 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));@@ -650,6 +794,254 @@ 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 = {};+ 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, LLM_ARCH_LLAMA, 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, LLM_ARCH_LLAMA, 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, LLM_ARCH_LLAMA, hparams));+ t.assert_equal((size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_LLAMA, hparams));+ });++ t.test("dflash_uses_n_embd_inp_enc", [&](testing & t) {+ llama_hparams hparams = {};+ hparams.n_embd = 64;+ hparams.n_embd_inp_enc_impl = 128; // makes n_embd_inp_enc() = 128+ hparams.n_embd_out_impl = 96; // makes n_embd_out() = 96++ t.assert_equal("DFlash uses the encoder input width",+ (size_t) 128, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_DFLASH, hparams));++ t.assert_equal("other archs ignore n_embd_inp_enc",+ (size_t) 64, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_DEFAULT, LLM_ARCH_LLAMA, hparams));++ t.assert_equal("MTP takes precedence over DFlash",+ (size_t) 96, llama_batch_ext_select_n_embd_inp(LLAMA_CONTEXT_TYPE_MTP, LLM_ARCH_DFLASH, hparams));+ });+}+ int main(int argc, char ** argv) { testing t; @@ -665,10 +1057,13 @@ 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("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);+ t.test("mtp_embd_width", test_mtp_embd_width); return t.summary(); }