Felipe97/llama-cpp-compiled
01.1k
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 ranges17struct 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, so51// the tests use embeddings everywhere except the token validation tests52struct 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 indices66 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 range109 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 decreases254 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 -> diverged279 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 seq454 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 left486 // with 1 < n_keep_tail remaining, so it is deferred entirely487 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 alone513 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 remain537 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_tail576 // remaining and is deferred even though its tokens were already gathered577 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, // temporal604 5, 6, // y605 7, 8, // x606 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}675