Commit fc343a84b for llama.cpp
commit fc343a84bbd925b37dde3219de35ea0bed50d630
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Thu Sep 24 16:25:07 2026 +0200
llama: add llama_batch_ext (#24669)
* (wip) add llama_batch_ext
* wip
* updated design
* updated impl
* change signature
* unused var
* demo common_prompt_batch_decode
* fix pos
* tmp disable test-batch-alloc
* fix compat
* nits: add const
* no more pos_max
* add comment about llama_batch_ext_set_embd_state
* handle n_embd_out properly
* rename api --> embd_token
* llama_embd
* stub llama_batch_ext_set_embd_state
* support both token + embd + state in batch
* llama_batch_ext_add_embd
* upstream some changes
* nits
* fix test-batch-alloc
* add test for compat
diff --git a/common/common.cpp b/common/common.cpp
index d8319cd9a..a316b2ae1 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -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;
}
diff --git a/common/common.h b/common/common.h
index 7afc266ac..39f855c4f 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1021,6 +1021,10 @@ 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
@@ -1028,7 +1032,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<llama_token> & 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 8f6368177..880a6a5fa 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, 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;
diff --git a/include/llama.h b/include/llama.h
index 31bbf8b0d..1805ed055 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -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)
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 2b98a552f..89a1f3f37 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -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;
+}
diff --git a/src/llama-batch.h b/src/llama-batch.h
index a3d1889d4..201d48cce 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -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 @@ private:
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 @@ private:
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);
+};
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index fcd4dfb13..da9c558d2 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -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,29 +1702,27 @@ 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;
@@ -1731,20 +1730,17 @@ int llama_context::decode(const llama_batch & batch_inp) {
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
//
diff --git a/src/llama-context.h b/src/llama-context.h
index 77ef92fc6..b403b099b 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -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);
diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp
index 34b3c6880..b83f45ba9 100644
--- a/src/llama-hparams.cpp
+++ b/src/llama-hparams.cpp
@@ -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 {
diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp
index 66d29d6f5..ad186c693 100644
--- a/tests/test-batch-alloc.cpp
+++ b/tests/test-batch-alloc.cpp
@@ -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,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<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);
}
};
@@ -97,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) {
@@ -120,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;
- 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,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<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]);
@@ -624,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;
@@ -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();
}