char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aman Guptallama-batch: add unit test (#25471)961e4b26a

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
22.7 KiB674 linesraw
1#include "testing.h"
2
3#include "llama.h"
4
5#include "../src/llama-batch.h"
6#include "../src/llama-memory.h"
7#include "../src/llama-vocab.h"
8
9#include <cstdlib>
10#include <initializer_list>
11#include <map>
12#include <string>
13#include <utility>
14#include <vector>
15
16// mock memory that only provides per-sequence position ranges
17struct mock_memory : public llama_memory_i {
18    std::map<llama_seq_id, std::pair<llama_pos, llama_pos>> ranges; // seq_id -> [pos_min, pos_max]
19
20    llama_memory_context_ptr init_batch(llama_batch_allocr &, uint32_t, bool) override {  GGML_ASSERT(false && "not implemented"); }
21    llama_memory_context_ptr init_full() override {  GGML_ASSERT(false && "not implemented"); }
22    llama_memory_context_ptr init_update(llama_context *, bool) override { GGML_ASSERT(false && "not implemented"); }
23
24    bool get_can_shift() const override { GGML_ASSERT(false && "not implemented"); }
25
26    void clear(bool) override { GGML_ASSERT(false && "not implemented"); }
27
28    bool seq_rm  (llama_seq_id, llama_pos, llama_pos) override { GGML_ASSERT(false && "not implemented"); }
29    void seq_cp  (llama_seq_id, llama_seq_id, llama_pos, llama_pos) override { GGML_ASSERT(false && "not implemented"); }
30    void seq_keep(llama_seq_id) override { GGML_ASSERT(false && "not implemented"); }
31    void seq_add (llama_seq_id, llama_pos, llama_pos, llama_pos) override { GGML_ASSERT(false && "not implemented"); }
32    void seq_div (llama_seq_id, llama_pos, llama_pos, int) override { GGML_ASSERT(false && "not implemented");  }
33
34    llama_pos seq_pos_min(llama_seq_id seq_id) const override {
35        auto it = ranges.find(seq_id);
36        return it == ranges.end() ? -1 : it->second.first;
37    }
38
39    llama_pos seq_pos_max(llama_seq_id seq_id) const override {
40        auto it = ranges.find(seq_id);
41        return it == ranges.end() ? -1 : it->second.second;
42    }
43
44    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override { return {}; }
45
46    void state_write(llama_io_write_i &, llama_seq_id, llama_state_seq_flags) const override { GGML_ASSERT(false && "not implemented"); }
47    void state_read (llama_io_read_i &,  llama_seq_id, llama_state_seq_flags) override { GGML_ASSERT(false && "not implemented"); }
48};
49
50// builds embedding batches - an empty llama_vocab rejects all token ids, so
51// the tests use embeddings everywhere except the token validation tests
52struct batch_builder {
53    uint32_t n_embd;
54
55    std::vector<float>     embd;
56    std::vector<llama_pos> pos;
57    std::vector<int32_t>   n_seq_id;
58    std::vector<int8_t>    logits;
59
60    std::vector<std::vector<llama_seq_id>> seq;
61    std::vector<llama_seq_id *>            seq_ptr;
62
63    batch_builder(uint32_t n_embd = 2) : n_embd(n_embd) {}
64
65    // embd values are 100*i + k so that ubatch contents can be traced back to batch indices
66    void add(llama_pos p, std::initializer_list<llama_seq_id> seq_ids, bool output) {
67        const int32_t i = (int32_t) seq.size();
68        for (uint32_t k = 0; k < n_embd; ++k) {
69            embd.push_back(100.0f*i + k);
70        }
71        pos.push_back(p);
72        n_seq_id.push_back((int32_t) seq_ids.size());
73        seq.emplace_back(seq_ids);
74        logits.push_back(output ? 1 : 0);
75    }
76
77    llama_batch make(bool with_pos = true, bool with_seq = true, bool with_logits = true) {
78        seq_ptr.clear();
79        for (auto & s : seq) {
80            seq_ptr.push_back(s.data());
81        }
82        seq_ptr.push_back(nullptr);
83
84        llama_batch res = {};
85        res.n_tokens = (int32_t) seq.size();
86        res.embd     = embd.data();
87        res.pos      = with_pos    ? pos.data()      : nullptr;
88        res.n_seq_id = with_seq    ? n_seq_id.data() : nullptr;
89        res.seq_id   = with_seq    ? seq_ptr.data()  : nullptr;
90        res.logits   = with_logits ? logits.data()   : nullptr;
91
92        return res;
93    }
94};
95
96static void test_init(testing & t) {
97    llama_vocab vocab;
98
99    t.test("rejects_n_seq_max_too_large", [&](testing & t) {
100        batch_builder bb;
101        bb.add(0, {0}, true);
102
103        llama_batch_allocr ba(1);
104        t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, LLAMA_MAX_SEQ + 1, false));
105    });
106
107    t.test("rejects_invalid_token", [&](testing & t) {
108        llama_token tok = 0; // empty vocab -> every token id is out of range
109        llama_batch batch = llama_batch_get_one(&tok, 1);
110
111        llama_batch_allocr ba(1);
112        t.assert_true("token id >= n_tokens", !ba.init(batch, vocab, nullptr, 0, 1, false));
113
114        tok = -1;
115        t.assert_true("negative token id", !ba.init(batch, vocab, nullptr, 0, 1, false));
116    });
117
118    t.test("rejects_invalid_seq_id", [&](testing & t) {
119        llama_batch_allocr ba(1);
120
121        {
122            batch_builder bb;
123            bb.add(0, {4}, true);
124            t.assert_true("seq_id >= n_seq_max", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
125        }
126        {
127            batch_builder bb;
128            bb.add(0, {-1}, true);
129            t.assert_true("negative seq_id", !ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
130        }
131    });
132
133    t.test("autofill_defaults", [&](testing & t) {
134        batch_builder bb;
135        for (int i = 0; i < 4; ++i) {
136            bb.add(0, {0}, false);
137        }
138
139        llama_batch_allocr ba(1);
140        t.assert_true(ba.init(bb.make(false, false, false), vocab, nullptr, bb.n_embd, 4, false));
141
142        const llama_batch & batch = ba.get_batch();
143
144        t.assert_equal(4u, ba.get_n_tokens());
145
146        for (int i = 0; i < 4; ++i) {
147            t.assert_equal("pos defaults to 0..n-1", i, batch.pos[i]);
148            t.assert_equal("n_seq_id defaults to 1", 1, batch.n_seq_id[i]);
149            t.assert_equal("seq_id defaults to 0",   0, batch.seq_id[i][0]);
150        }
151
152        t.assert_equal("only the last token is an output", 1u, ba.get_n_outputs());
153        t.assert_equal(0, (int) batch.logits[0]);
154        t.assert_equal(1, (int) batch.logits[3]);
155
156        t.assert_equal(0, ba.seq_pos_min(0));
157        t.assert_equal(3, ba.seq_pos_max(0));
158        t.assert_equal(-1, ba.seq_pos_min(1));
159    });
160
161    t.test("output_all", [&](testing & t) {
162        batch_builder bb;
163        for (int i = 0; i < 4; ++i) {
164            bb.add(i, {0}, false);
165        }
166
167        llama_batch_allocr ba(1);
168        t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, true));
169        t.assert_equal(4u, ba.get_n_outputs());
170    });
171
172    t.test("explicit_logits", [&](testing & t) {
173        batch_builder bb;
174        bb.add(0, {0}, true);
175        bb.add(1, {0}, false);
176        bb.add(2, {0}, true);
177
178        llama_batch_allocr ba(1);
179        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
180        t.assert_equal(2u, ba.get_n_outputs());
181
182        llama_ubatch ub = ba.split_simple(10);
183        t.assert_equal(3u, ub.n_tokens);
184        t.assert_equal(1, (int) ub.output[0]);
185        t.assert_equal(0, (int) ub.output[1]);
186        t.assert_equal(1, (int) ub.output[2]);
187
188        const auto & out_ids = ba.get_out_ids();
189        t.assert_equal((size_t) 2, out_ids.size());
190        t.assert_equal(0, out_ids[0]);
191        t.assert_equal(2, out_ids[1]);
192    });
193
194    t.test("pos_from_memory", [&](testing & t) {
195        mock_memory mem;
196        mem.ranges[0] = {0, 9};
197
198        batch_builder bb;
199        for (int i = 0; i < 3; ++i) {
200            bb.add(0, {0}, false);
201        }
202
203        llama_batch_allocr ba(1);
204        t.assert_true(ba.init(bb.make(false, true, false), vocab, &mem, bb.n_embd, 4, false));
205
206        t.assert_equal("pos continues after memory", 10, ba.seq_pos_min(0));
207        t.assert_equal(12, ba.seq_pos_max(0));
208    });
209
210    t.test("pos_continuity_with_memory", [&](testing & t) {
211        mock_memory mem;
212        mem.ranges[0] = {0, 9};
213
214        llama_batch_allocr ba(1);
215
216        {
217            batch_builder bb;
218            bb.add(10, {0}, false);
219            bb.add(11, {0}, true);
220            t.assert_true("pos_max + 1 is accepted", ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false));
221        }
222        {
223            batch_builder bb;
224            bb.add(11, {0}, false);
225            bb.add(12, {0}, true);
226            t.assert_true("gap after memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false));
227        }
228        {
229            batch_builder bb;
230            bb.add(9, {0}, false);
231            bb.add(10, {0}, true);
232            t.assert_true("overlap with memory is rejected", !ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false));
233        }
234    });
235
236    t.test("rejects_non_continuous_positions", [&](testing & t) {
237        batch_builder bb;
238        bb.add(0, {0}, false);
239        bb.add(1, {0}, false);
240        bb.add(3, {0}, true);
241
242        llama_batch_allocr ba(1);
243        t.assert_true(!ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
244    });
245
246    t.test("rejects_decreasing_positions", [&](testing & t) {
247        batch_builder bb;
248        const llama_pos    pos[7] = {4, 5, 0, 1, 6, 2, 3};
249        const llama_seq_id seq[7] = {0, 0, 1, 1, 0, 1, 0};
250        for (int i = 0; i < 7; ++i) {
251            bb.add(pos[i], {seq[i]}, false);
252        }
253        // seq 0 sees positions 4,5,6,3 in batch order -> the trailing 3 decreases
254
255        llama_batch_allocr ba(1);
256        t.assert_true(!ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false));
257    });
258
259    t.test("allows_equal_positions_in_seq", [&](testing & t) {
260        batch_builder bb;
261        bb.add(0, {0}, false);
262        bb.add(0, {0}, false);
263        bb.add(1, {0}, true);
264
265        llama_batch_allocr ba(1);
266        t.assert_true(ba.init(bb.make(true, true, false), vocab, nullptr, bb.n_embd, 4, false));
267    });
268
269
270    t.test("rejects_coupled_diverged_seqs", [&](testing & t) {
271        batch_builder bb;
272        bb.add(6, {0, 1}, true);
273
274        llama_batch_allocr ba(1);
275
276        mock_memory mem;
277        mem.ranges[0] = {0, 5};
278        mem.ranges[1] = {2, 5}; // same pos_max, different pos_min -> diverged
279        t.assert_true(!ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false));
280
281        mem.ranges[1] = {0, 5};
282        t.assert_true(ba.init(bb.make(), vocab, &mem, bb.n_embd, 4, false));
283    });
284}
285
286static void test_split(testing & t) {
287    llama_vocab vocab;
288
289    t.test("split_simple_chunks", [&](testing & t) {
290        batch_builder bb;
291        for (int i = 0; i < 5; ++i) {
292            bb.add(i, {0}, i == 4);
293        }
294
295        llama_batch_allocr ba(1);
296        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
297
298        llama_ubatch ub = ba.split_simple(2);
299        t.assert_equal(2u, ub.n_tokens);
300        t.assert_true(!ub.equal_seqs());
301        t.assert_equal(1u, ub.n_seqs_unq);
302        t.assert_equal(0, ub.seq_id_unq[0]);
303        t.assert_equal(0, ub.seq_idx[0]);
304        for (int i = 0; i < 2; ++i) {
305            t.assert_equal(i, ub.pos[i]);
306            t.assert_equal(1, ub.n_seq_id[i]);
307            t.assert_equal(0, ub.seq_id[i][0]);
308            t.assert_equal(100.0f*i, ub.embd[i*bb.n_embd]);
309            t.assert_equal(100.0f*i + 1, ub.embd[i*bb.n_embd + 1]);
310        }
311
312        ub = ba.split_simple(2);
313        t.assert_equal(2u, ub.n_tokens);
314        t.assert_equal(2, ub.pos[0]);
315        t.assert_equal(3, ub.pos[1]);
316
317        ub = ba.split_simple(2);
318        t.assert_equal(1u, ub.n_tokens);
319        t.assert_equal(4, ub.pos[0]);
320        t.assert_equal(1, (int) ub.output[0]);
321
322        t.assert_equal(5u, ba.get_n_used());
323
324        ub = ba.split_simple(2);
325        t.assert_equal("batch is consumed", 0u, ub.n_tokens);
326
327        const auto & out_ids = ba.get_out_ids();
328        t.assert_equal((size_t) 1, out_ids.size());
329        t.assert_equal(4, out_ids[0]);
330    });
331
332    t.test("split_reset_allows_resplit", [&](testing & t) {
333        batch_builder bb;
334        for (int i = 0; i < 3; ++i) {
335            bb.add(i, {0}, i == 2);
336        }
337
338        llama_batch_allocr ba(1);
339        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
340
341        while (ba.split_simple(1).n_tokens > 0) {
342        }
343        t.assert_equal(3u, ba.get_n_used());
344
345        ba.split_reset();
346        t.assert_equal(0u, ba.get_n_used());
347
348        llama_ubatch ub = ba.split_simple(10);
349        t.assert_equal(3u, ub.n_tokens);
350    });
351
352    t.test("split_equal_unequal_lengths", [&](testing & t) {
353        batch_builder bb;
354        for (int i = 0; i < 4; ++i) {
355            bb.add(i, {0}, i == 3);
356        }
357        for (int i = 0; i < 2; ++i) {
358            bb.add(i, {1}, i == 1);
359        }
360
361        llama_batch_allocr ba(1);
362        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
363
364        llama_ubatch ub = ba.split_equal(8, false, 0);
365        t.assert_true(ub.equal_seqs());
366        t.assert_equal("both seqs advance by the shorter length", 4u, ub.n_tokens);
367        t.assert_equal(2u, ub.n_seq_tokens);
368        t.assert_equal(2u, ub.n_seqs);
369        t.assert_equal(2u, ub.n_seqs_unq);
370        // tokens are grouped per sequence set: [s0 s0 s1 s1]
371        t.assert_equal(0, ub.seq_id[0][0]);
372        t.assert_equal(0, ub.seq_id[1][0]);
373        t.assert_equal(1, ub.seq_id[2][0]);
374        t.assert_equal(1, ub.seq_id[3][0]);
375        t.assert_equal(0, ub.pos[0]);
376        t.assert_equal(1, ub.pos[1]);
377        t.assert_equal(0, ub.pos[2]);
378        t.assert_equal(1, ub.pos[3]);
379
380        ub = ba.split_equal(8, false, 0);
381        t.assert_equal("only seq 0 remains", 2u, ub.n_tokens);
382        t.assert_equal(1u, ub.n_seqs);
383        t.assert_equal(2, ub.pos[0]);
384        t.assert_equal(3, ub.pos[1]);
385
386        ub = ba.split_equal(8, false, 0);
387        t.assert_equal(0u, ub.n_tokens);
388
389        t.assert_equal(6u, ba.get_n_used());
390    });
391
392    t.test("split_equal_coupled", [&](testing & t) {
393        batch_builder bb;
394        bb.add(0, {0, 1}, false);
395        bb.add(1, {0, 1}, true);
396
397        llama_batch_allocr ba(1);
398        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
399
400        llama_ubatch ub = ba.split_equal(4, true, 0);
401        t.assert_equal("sequential split rejects coupled seqs", 0u, ub.n_tokens);
402
403        ub = ba.split_equal(4, false, 0);
404        t.assert_equal(2u, ub.n_tokens);
405        t.assert_equal("one sequence set", 1u, ub.n_seqs);
406        t.assert_equal("two unique seq ids", 2u, ub.n_seqs_unq);
407        t.assert_equal(2, ub.n_seq_id[0]);
408        t.assert_equal(0, ub.seq_idx[0]);
409        t.assert_equal(1, ub.seq_idx[1]);
410    });
411
412    t.test("split_seq_per_sequence", [&](testing & t) {
413        batch_builder bb;
414        for (llama_seq_id s = 0; s < 3; ++s) {
415            bb.add(0, {s}, false);
416            bb.add(1, {s}, true);
417        }
418
419        llama_batch_allocr ba(1);
420        t.assert_true(ba.init(bb.make(), vocab, nullptr, bb.n_embd, 4, false));
421
422        for (llama_seq_id s = 0; s < 3; ++s) {
423            llama_ubatch ub = ba.split_seq(8);
424            t.assert_equal(2u, ub.n_tokens);
425            t.assert_equal(1u, ub.n_seqs);
426            t.assert_equal(s, ub.seq_id[0][0]);
427            t.assert_equal(s, ub.seq_id_unq[0]);
428        }
429
430        t.assert_equal(0u, ba.split_seq(8).n_tokens);
431        t.assert_equal(6u, ba.get_n_used());
432    });
433
434    t.test("ubatch_reserve", [&](testing & t) {
435        llama_batch_allocr ba(1);
436
437        llama_ubatch ub = ba.ubatch_reserve(3, 2);
438        t.assert_equal(6u, ub.n_tokens);
439        t.assert_equal(3u, ub.n_seq_tokens);
440        t.assert_equal(2u, ub.n_seqs);
441        t.assert_equal(2u, ub.n_seqs_unq);
442        t.assert_true(ub.equal_seqs());
443        t.assert_equal(0, ub.seq_id_unq[0]);
444        t.assert_equal(1, ub.seq_id_unq[1]);
445        t.assert_true(ub.token != nullptr);
446        t.assert_true(ub.embd == nullptr);
447    });
448}
449
450static void test_keep_tail(testing & t) {
451    llama_vocab vocab;
452
453    // batch with n_tokens[s] tokens for each seq s, output on the last token of each seq
454    auto make_batch = [](batch_builder & bb, std::initializer_list<int> n_tokens) {
455        llama_seq_id s = 0;
456        for (int n : n_tokens) {
457            for (int i = 0; i < n; ++i) {
458                bb.add(i, {s}, i == n - 1);
459            }
460            ++s;
461        }
462        return bb.make();
463    };
464
465    t.test("noop_when_seqs_complete", [&](testing & t) {
466        batch_builder bb;
467
468        llama_batch_allocr ba(1);
469        t.assert_true(ba.init(make_batch(bb, {2, 2}), vocab, nullptr, bb.n_embd, 4, false));
470
471        llama_ubatch ub = ba.split_equal(4, false, 2);
472        t.assert_equal("both seqs fit whole", 4u, ub.n_tokens);
473        t.assert_equal(2u, ub.n_seqs);
474        t.assert_equal(2u, ub.n_seq_tokens);
475
476        t.assert_equal(0u, ba.split_equal(4, false, 2).n_tokens);
477    });
478
479    t.test("defers_seq_with_short_remainder", [&](testing & t) {
480        batch_builder bb;
481
482        llama_batch_allocr ba(1);
483        t.assert_true(ba.init(make_batch(bb, {2, 3}), vocab, nullptr, bb.n_embd, 4, false));
484
485        // expansion stops at 2 tokens per seq: seq 0 completes, seq 1 would be left
486        // with 1 < n_keep_tail remaining, so it is deferred entirely
487        llama_ubatch ub = ba.split_equal(4, true, 2);
488        t.assert_equal(2u, ub.n_tokens);
489        t.assert_equal(1u, ub.n_seqs);
490        t.assert_equal(0, ub.seq_id[0][0]);
491        t.assert_equal(2u, ba.get_n_used());
492
493        ub = ba.split_equal(4, true, 2);
494        t.assert_equal("deferred seq comes back whole", 3u, ub.n_tokens);
495        t.assert_equal(1u, ub.n_seqs);
496        t.assert_equal(1, ub.seq_id[0][0]);
497        for (int i = 0; i < 3; ++i) {
498            t.assert_equal(i, ub.pos[i]);
499        }
500
501        t.assert_equal(5u, ba.get_n_used());
502        t.assert_equal(0u, ba.split_equal(4, true, 2).n_tokens);
503    });
504
505    t.test("completes_first_seq_when_all_violate", [&](testing & t) {
506        batch_builder bb;
507
508        llama_batch_allocr ba(1);
509        t.assert_true(ba.init(make_batch(bb, {3, 3}), vocab, nullptr, bb.n_embd, 4, false));
510
511        // expansion stops at 2 tokens per seq, leaving both with 1 < n_keep_tail remaining;
512        // seq 0 still fits in n_ubatch, so it is extended to completion and emitted alone
513        llama_ubatch ub = ba.split_equal(4, false, 2);
514        t.assert_equal(3u, ub.n_tokens);
515        t.assert_equal(1u, ub.n_seqs);
516        t.assert_equal(3u, ub.n_seq_tokens);
517        t.assert_equal(0, ub.seq_id[0][0]);
518        for (int i = 0; i < 3; ++i) {
519            t.assert_equal(i, ub.pos[i]);
520        }
521        t.assert_equal(3u, ba.get_n_used());
522
523        ub = ba.split_equal(4, false, 2);
524        t.assert_equal(3u, ub.n_tokens);
525        t.assert_equal(1, ub.seq_id[0][0]);
526        t.assert_equal(6u, ba.get_n_used());
527    });
528
529    t.test("truncates_to_preserve_tail", [&](testing & t) {
530        batch_builder bb;
531
532        llama_batch_allocr ba(1);
533        t.assert_true(ba.init(make_batch(bb, {5}), vocab, nullptr, bb.n_embd, 4, false));
534
535        // 4 tokens would leave a remainder of 1, and the seq does not fit in n_ubatch,
536        // so the ubatch is truncated until n_keep_tail tokens remain
537        llama_ubatch ub = ba.split_equal(4, false, 2);
538        t.assert_equal(3u, ub.n_tokens);
539        t.assert_equal(1u, ub.n_seqs);
540        t.assert_equal(2, ub.pos[2]);
541        t.assert_equal(3u, ba.get_n_used());
542
543        ub = ba.split_equal(4, false, 2);
544        t.assert_equal("trailing tokens stay in one ubatch", 2u, ub.n_tokens);
545        t.assert_equal(3, ub.pos[0]);
546        t.assert_equal(4, ub.pos[1]);
547        t.assert_equal(1, (int) ub.output[1]);
548
549        t.assert_equal(5u, ba.get_n_used());
550    });
551
552    t.test("keeps_full_ubatch_with_sufficient_remainder", [&](testing & t) {
553        batch_builder bb;
554
555        llama_batch_allocr ba(1);
556        t.assert_true(ba.init(make_batch(bb, {6}), vocab, nullptr, bb.n_embd, 4, false));
557
558        llama_ubatch ub = ba.split_equal(4, false, 2);
559        t.assert_equal("remainder >= n_keep_tail, no truncation", 4u, ub.n_tokens);
560
561        ub = ba.split_equal(4, false, 2);
562        t.assert_equal(2u, ub.n_tokens);
563        t.assert_equal(4, ub.pos[0]);
564        t.assert_equal(5, ub.pos[1]);
565
566        t.assert_equal(6u, ba.get_n_used());
567    });
568
569    t.test("multi_seq_prefix_kept", [&](testing & t) {
570        batch_builder bb;
571
572        llama_batch_allocr ba(1);
573        t.assert_true(ba.init(make_batch(bb, {3, 4}), vocab, nullptr, bb.n_embd, 6, false));
574
575        // expansion stops at 3 tokens per seq: seq 0 completes, seq 1 has 1 < n_keep_tail
576        // remaining and is deferred even though its tokens were already gathered
577        llama_ubatch ub = ba.split_equal(6, true, 2);
578        t.assert_equal(3u, ub.n_tokens);
579        t.assert_equal(1u, ub.n_seqs);
580        t.assert_equal(0, ub.seq_id[0][0]);
581        t.assert_equal(3u, ba.get_n_used());
582
583        ub = ba.split_equal(6, true, 2);
584        t.assert_equal(4u, ub.n_tokens);
585        t.assert_equal(1, ub.seq_id[0][0]);
586        t.assert_equal(7u, ba.get_n_used());
587    });
588}
589
590static void test_mrope(testing & t) {
591    llama_vocab vocab;
592
593    t.test("pos_layout_and_split", [&](testing & t) {
594        const uint32_t n_pos = 4;
595        const uint32_t n_embd = 2;
596
597        batch_builder bb(n_embd);
598        bb.add(10, {0}, false);
599        bb.add(11, {0}, true);
600
601        // M-RoPE positions for embeddings are laid out [n_pos][n_tokens]
602        std::vector<llama_pos> pos = {
603            10, 11, // temporal
604             5,  6, // y
605             7,  8, // x
606             0,  0,
607        };
608
609        llama_batch batch = bb.make(false, true, true);
610        batch.pos = pos.data();
611
612        llama_batch_allocr ba(n_pos);
613        t.assert_true(ba.init(batch, vocab, nullptr, n_embd, 4, false));
614
615        llama_ubatch ub = ba.split_simple(2);
616        t.assert_equal(2u, ub.n_tokens);
617        t.assert_equal(n_pos, ub.n_pos);
618        t.assert_true(ub.is_pos_2d());
619
620        const llama_pos expected[8] = {10, 11, 5, 6, 7, 8, 0, 0};
621        for (int i = 0; i < 8; ++i) {
622            t.assert_equal(expected[i], ub.pos[i]);
623        }
624    });
625
626    t.test("pos_jump_allowed", [&](testing & t) {
627        const uint32_t n_pos = 4;
628        const uint32_t n_embd = 2;
629
630        mock_memory mem;
631        mem.ranges[0] = {0, 9};
632
633        llama_batch_allocr ba(n_pos);
634
635        auto try_pos = [&](llama_pos p0) {
636            batch_builder bb(n_embd);
637            bb.add(p0, {0}, true);
638
639            std::vector<llama_pos> pos = {p0, 1, 1, 0};
640
641            llama_batch batch = bb.make(false, true, true);
642            batch.pos = pos.data();
643
644            return ba.init(batch, vocab, &mem, n_embd, 4, false);
645        };
646
647        t.assert_true("gap after memory is allowed",     try_pos(15));
648        t.assert_true("overlap is allowed for embd",     try_pos(9));
649        t.assert_true("pos behind memory is rejected",  !try_pos(8));
650    });
651}
652
653int main(int argc, char ** argv) {
654    testing t;
655
656    const char * verbose = getenv("LLAMA_TEST_VERBOSE");
657    if (verbose) {
658        t.verbose = std::string(verbose) == "1";
659    }
660    if (!t.verbose) {
661        llama_log_set([](ggml_log_level, const char *, void *) {}, nullptr);
662    }
663
664    if (argc > 1) {
665        t.set_filter(argv[1]);
666    }
667
668    t.test("init",      test_init);
669    t.test("split",     test_split);
670    t.test("keep_tail", test_keep_tail);
671    t.test("mrope",     test_mrope);
672
673    return t.summary();
674}