upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-26 07:57:24 -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
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();
}