upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-27 00:17:27 -05:00
Code Issues Packages Projects Releases Wiki Activity

llama: add llama_batch_ext (#24669)

Browse Source 
* (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
This commit is contained in:
Xuan-Son Nguyen
2026-09-24 16:25:07 +02:00
committed by GitHub
parent 308883b335
commit fc343a84bb
10 changed files with 1210 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+30 -7
View File
@@ -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.h
+5 -1
View File
@@ -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,
include/llama-cpp.h
+5
View File
@@ -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.h
+90
View File
@@ -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.cpp
+440 -103
View File
@@ -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,43 +26,120 @@ 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);
//
// validate input batch
//
if (n_seq_max > LLAMA_MAX_SEQ) {
if ((uint32_t) n_seq_max > LLAMA_MAX_SEQ) {
LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ);
return false;
}
if (batch.token) {
for (int32_t i = 0; i < batch.n_tokens; ++i) {
if (batch.token[i] < 0 || (uint32_t) batch.token[i] >= vocab.n_tokens()) {
LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, batch.token[i]);
const llama_memory_i * mem = batch_inp.mem;
//
// determine the content types of the batch
// an entry can carry a token id, a token embedding, or both (e.g. MTP hook batches)
// all entries must carry the same combination
//
const bool has_token = batch_inp.tokens[0].id != LLAMA_TOKEN_NULL;
const bool has_embd = batch_inp.tokens[0].has_embd;
for (int32_t i = 1; i < n_tok; ++i) {
if ((batch_inp.tokens[i].id != LLAMA_TOKEN_NULL) != has_token ||
batch_inp.tokens[i].has_embd != has_embd) {
LLAMA_LOG_ERROR("%s: all entries in the batch must have the same content types\n", __func__);
return false;
}
}
if (!has_token && !has_embd) {
LLAMA_LOG_ERROR("%s: batch has neither token ids nor embeddings\n", __func__);
return false;
}
//
// build flat token/embd array
//
if (has_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 (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];
}
}
}
}
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);
//
// build n_seq_id / seq_id arrays
//
n_seq_id.resize(n_tok);
seq_id.resize(n_tok + 1);
seq_id[n_tok] = nullptr;
{
size_t total = 0;
for (int32_t i = 0; i < n_tok; ++i) {
total += batch_inp.tokens[i].seq_ids.size();
}
seq_id_data.reserve(total);
for (int32_t i = 0; i < n_tok; ++i) {
for (auto sid : batch_inp.tokens[i].seq_ids) {
seq_id_data.push_back(sid);
}
}
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];
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;
}
}
@@ -67,91 +147,43 @@ bool llama_batch_allocr::init(
}
//
// auto-generate missing fields
// build output/logits array
//
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();
}
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();
}
batch.seq_id = seq_id.data();
}
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;
}
{
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++) {
const llama_seq_id seq_id = batch.seq_id[i][0];
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;
}
}
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;
}
batch.logits = output.data();
} else if (output_all) {
bool warn = false;
for (int32_t i = 0; i < batch.n_tokens; ++i) {
if (batch.logits[i] == 0) {
warn = true;
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__);
output.resize(batch.n_tokens, true);
batch.logits = output.data();
}
}
//
// set up the internal llama_batch to point to our owned arrays
//
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.h
+76 -7
View File
@@ -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);
};
src/llama-context.cpp
+52 -32
View File
@@ -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
//
src/llama-context.h
+4
View File
@@ -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.cpp
+2 -1
View File
@@ -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.cpp
+506 -111
View File
@@ -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;
const 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;
llama_batch_ext b;
std::vector<std::vector<llama_seq_id>> seq;
std::vector<llama_seq_id *> seq_ptr;
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) {}
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);
// 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;
const int32_t idx = bb.add(0, {0}, true);
bb.b.tokens[idx].seq_ids = { 4 };
t.assert_true("seq_id >= n_seq_max", !ba.init(bb.b, vocab, false));
}
{
batch_builder bb;
bb.add(0, {-1}, true);
t.assert_true("negative seq_id", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
const int32_t idx = bb.add(0, {0}, true);
bb.b.tokens[idx].seq_ids = { -1 };
t.assert_true("negative seq_id", !ba.init(bb.b, vocab, false));
}
});
t.test("autofill_defaults", [&](testing & t) {
t.test("copies_pos_seq_output", [&](testing & t) {
batch_builder bb;
for (int i = 0; i < 4; ++i) {
bb.add(0, {0}, false);
bb.add(i, {0}, i == 3);
}
llama_batch_allocr ba(1);
t.assert_true(ba.init(bb.make(false, false, false), vocab, nullptr, bb.n_embd, 4, false));
t.assert_true(ba.init(bb.b, vocab, false));
const llama_batch & batch = ba.get_batch();
t.assert_equal(4u, ba.get_n_tokens());
t.assert_true("embedding batch", batch.embd != nullptr);
t.assert_true("no token ids", batch.token == nullptr);
for (int i = 0; i < 4; ++i) {
t.assert_equal("pos defaults to 0..n-1", i, batch.pos[i]);
t.assert_equal("n_seq_id defaults to 1", 1, batch.n_seq_id[i]);
t.assert_equal("seq_id defaults to 0", 0, batch.seq_id[i][0]);
t.assert_equal(i, batch.pos[i]);
t.assert_equal(1, batch.n_seq_id[i]);
t.assert_equal(0, batch.seq_id[i][0]);
t.assert_equal(100.0f*i, batch.embd[i*bb.n_embd]);
}
t.assert_equal("only the last token is an output", 1u, ba.get_n_outputs());
@@ -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);
batch_builder bb(n_embd, nullptr, 4, n_pos);
// 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,
};
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);
batch_builder bb(n_embd, &mem, 4, n_pos);
std::vector<llama_pos> pos = {p0, 1, 1, 0};
const llama_pos pos[n_pos] = { p0, 1, 1, 0 };
bb.add_embd(pos, {0}, true);
llama_batch batch = bb.make(false, true, true);
batch.pos = pos.data();
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();
}