Felipe97/llama-cpp-compiled
01.1k
1#include "llama-batch.h"2 3#include "llama-impl.h"4#include "llama-vocab.h"5#include "llama-memory.h"6 7#include <cassert>8#include <cstring>9#include <algorithm>10#include <sstream>11 12llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {13 const char * LLAMA_BATCH_DEBUG = getenv("LLAMA_BATCH_DEBUG");14 debug = LLAMA_BATCH_DEBUG ? atoi(LLAMA_BATCH_DEBUG) : 0;15 16 seq_pos.resize(LLAMA_MAX_SEQ);17 seq_cpl.resize(LLAMA_MAX_SEQ);18 for (auto & cur : seq_cpl) {19 cur.resize(LLAMA_MAX_SEQ);20 }21 22 seq_idx.resize(LLAMA_MAX_SEQ, -1);23}24 25bool llama_batch_allocr::init(26 const llama_batch & batch_inp,27 const llama_vocab & vocab,28 const llama_memory_i * memory,29 uint32_t n_embd,30 uint32_t n_seq_max,31 bool output_all) {32 clear();33 34 batch = batch_inp;35 36 this->vocab = &vocab;37 38 GGML_ASSERT(batch.n_tokens > 0);39 40 //41 // validate input batch42 //43 44 if (n_seq_max > LLAMA_MAX_SEQ) {45 LLAMA_LOG_ERROR("%s: n_seq_max = %d > %d\n", __func__, n_seq_max, LLAMA_MAX_SEQ);46 return false;47 }48 49 if (batch.token) {50 for (int32_t i = 0; i < batch.n_tokens; ++i) {51 if (batch.token[i] < 0 || (uint32_t) batch.token[i] >= vocab.n_tokens()) {52 LLAMA_LOG_ERROR("%s: invalid token[%d] = %d\n", __func__, i, batch.token[i]);53 return false;54 }55 }56 }57 58 if (batch.seq_id) {59 for (int32_t i = 0; i < batch.n_tokens; ++i) {60 for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {61 if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) {62 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);63 return false;64 }65 }66 }67 }68 69 //70 // auto-generate missing fields71 //72 73 if (!batch.n_seq_id) {74 n_seq_id.resize(batch.n_tokens);75 for (int32_t i = 0; i < batch.n_tokens; i++) {76 n_seq_id[i] = seq_id_0.size();77 }78 batch.n_seq_id = n_seq_id.data();79 }80 81 if (!batch.seq_id) {82 seq_id.resize(batch.n_tokens + 1);83 seq_id[batch.n_tokens] = NULL;84 for (int32_t i = 0; i < batch.n_tokens; i++) {85 seq_id[i] = seq_id_0.data();86 }87 batch.seq_id = seq_id.data();88 }89 90 if (!batch.pos) {91 pos.resize(batch.n_tokens);92 93 // initialize the starting position for each sequence based on the positions in the memory94 llama_pos p0[LLAMA_MAX_SEQ];95 for (uint32_t s = 0; s < n_seq_max; ++s) {96 if (!memory) {97 // if no memory -> start from 098 p0[s] = 0;99 } else {100 p0[s] = memory->seq_pos_max(s) + 1;101 }102 }103 104 for (int32_t i = 0; i < batch.n_tokens; i++) {105 const llama_seq_id seq_id = batch.seq_id[i][0];106 107 pos[i] = p0[seq_id];108 109 // update the starting position for all sequences that are assigned to the this token110 for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {111 const llama_seq_id seq_id = batch.seq_id[i][s];112 113 p0[seq_id] = pos[i] + 1;114 }115 }116 117 batch.pos = pos.data();118 }119 120 if (!batch.logits) {121 if (output_all) {122 // return the output for all tokens123 output.resize(batch.n_tokens, true);124 } else {125 // return the output only for the last token126 output.resize(batch.n_tokens, false);127 output[output.size() - 1] = true;128 }129 130 batch.logits = output.data();131 } else if (output_all) {132 bool warn = false;133 134 for (int32_t i = 0; i < batch.n_tokens; ++i) {135 if (batch.logits[i] == 0) {136 warn = true;137 }138 }139 140 if (warn) {141 LLAMA_LOG_WARN("%s: embeddings required but some input tokens were not marked as outputs -> overriding\n", __func__);142 143 output.resize(batch.n_tokens, true);144 batch.logits = output.data();145 }146 }147 148 //149 // compute stats150 //151 152 this->n_embd = n_embd;153 this->n_seq_max = n_seq_max;154 155 // count the outputs in this batch156 for (int32_t i = 0; i < batch.n_tokens; ++i) {157 n_outputs += batch.logits[i] != 0;158 }159 160 has_cpl = false;161 162 // determine coupled sequences163 // these are pairs of sequences that have at least one token in the input batch that is assigned to both of them164 for (int32_t i = 0; i < batch.n_tokens; ++i) {165 const llama_seq_id s0 = batch.seq_id[i][0];166 167 for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {168 const llama_seq_id s1 = batch.seq_id[i][s];169 170 seq_pos[s1].insert(batch.pos[i]);171 172 if (s > 0) {173 // mark that sequence s1 is coupled to s0174 seq_cpl[s1][s0] = true;175 176 // note: tracking the other way around is not necessary for now177 //seq_cpl[s0][s1] = true;178 179 has_cpl = true;180 }181 }182 }183 184 // precompute the sequence sets for each token and determine the unique sequence ids that participate in the batch185 {186 seq_set_t seq_set_unq;187 188 for (int32_t i = 0; i < batch.n_tokens; ++i) {189 seq_set_t cur;190 for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {191 const llama_seq_id seq_id = batch.seq_id[i][s];192 193 cur .set(seq_id);194 seq_set_unq.set(seq_id);195 }196 197 seq_set.push_back(cur);198 seq_set_map[cur].push_back(i);199 }200 201 for (uint32_t s = 0; s < n_seq_max; ++s) {202 if (seq_set_unq.test(s)) {203 seq_idx[s] = seq_id_unq.size();204 seq_id_unq.push_back(s);205 }206 }207 }208 209 if (debug > 0) {210 LLAMA_LOG_DEBUG("%s: input batch info:\n", __func__);211 212 llama_ubatch ubatch {213 /*.b_equal_seqs =*/ false,214 /*.n_tokens =*/ (uint32_t) batch.n_tokens,215 /*.n_seq_tokens =*/ (uint32_t) 1,216 /*.n_seqs =*/ (uint32_t) batch.n_tokens,217 /*.n_seqs_unq =*/ (uint32_t) this->seq_id_unq.size(),218 /*.n_pos =*/ n_pos_per_embd,219 /*.token =*/ batch.token,220 /*.embd =*/ batch.embd,221 /*.pos =*/ batch.pos,222 /*.n_seq_id =*/ batch.n_seq_id,223 /*.seq_id =*/ batch.seq_id,224 /*.seq_id_unq =*/ this->seq_id_unq.data(),225 /*.seq_idx =*/ this->seq_idx.data(),226 /*.output =*/ batch.logits,227 /*.data =*/ {},228 };229 230 ubatch_print(ubatch, debug);231 232 LLAMA_LOG_DEBUG("%s: seq = [\n", __func__);233 for (int s0 = 0; s0 < (int) seq_pos.size(); ++s0) {234 if (seq_pos[s0].empty()) {235 continue;236 }237 238 std::stringstream ss;239 for (int s1 = 0; s1 < (int) seq_cpl[s0].size(); ++s1) {240 if (seq_cpl[s0][s1]) {241 ss << s1 << " ";242 }243 }244 245 LLAMA_LOG_DEBUG("%s: %4d: pos = [%4d, %4d], cpl = %s\n",246 __func__, s0, seq_pos_min(s0), seq_pos_max(s0), ss.str().empty() ? "-" : ss.str().c_str());247 }248 LLAMA_LOG_DEBUG("%s: ]\n", __func__);249 }250 251 //252 // consistency checks253 //254 255 if (n_pos_per_embd > 1) {256 // M-RoPE case: allow position to "jump" forward only (non-continuous positions are allowed)257 for (uint32_t s = 0; s < n_seq_max; ++s) {258 if (seq_pos[s].empty()) {259 continue;260 }261 262 const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1;263 264 if (batch.token) {265 if (p0 >= 0 && p0 >= seq_pos_min(s)) {266 LLAMA_LOG_ERROR(267 "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"268 " - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n"269 " - the tokens for sequence %d in the input batch have a starting position of Y = %d\n"270 " for M-RoPE, it is required that the position satisfies: X < Y\n",271 __func__, s, s, p0, s, seq_pos_min(s));272 273 return false;274 }275 } else {276 // embedding inputs can have overlapping positions277 if (p0 >= 0 && p0 > seq_pos_min(s)) {278 LLAMA_LOG_ERROR(279 "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"280 " - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n"281 " - the tokens for sequence %d in the input batch have a starting position of Y = %d\n"282 " for M-RoPE, it is required that the position satisfies: X <= Y\n",283 __func__, s, s, p0, s, seq_pos_min(s));284 285 return false;286 }287 }288 }289 } else {290 for (uint32_t s = 0; s < n_seq_max; ++s) {291 if (seq_pos[s].empty()) {292 continue;293 }294 295 const llama_pos p0 = memory ? memory->seq_pos_max(s) : -1;296 297 if (p0 >= 0) {298 bool ok = true;299 300 if (seq_pos_min(s) != p0 + 1) {301 ok = false;302 }303 304 if (!ok) {305 LLAMA_LOG_ERROR(306 "%s: the tokens of sequence %d in the input batch have inconsistent sequence positions:\n"307 " - the last position stored in the memory module of the context (i.e. the KV cache) for sequence %d is X = %d\n"308 " - the tokens for sequence %d in the input batch have a starting position of Y = %d\n"309 " it is required that the sequence positions remain consecutive: Y = X + 1\n",310 __func__, s, s, p0, s, seq_pos_min(s));311 312 return false;313 }314 }315 316 if (seq_pos_max(s) - seq_pos_min(s) + 1 > (int) seq_pos[s].size()) {317 LLAMA_LOG_ERROR("%s: sequence %d positions are not continuous\n", __func__, s);318 return false;319 }320 }321 }322 323 if (memory) {324 for (uint32_t s0 = 0; s0 < n_seq_max; ++s0) {325 for (uint32_t s1 = 0; s1 < n_seq_max; ++s1) {326 if (seq_cpl[s0][s1]) {327 if (memory->seq_pos_min(s0) != memory->seq_pos_min(s1) ||328 memory->seq_pos_max(s0) != memory->seq_pos_max(s1)) {329 LLAMA_LOG_ERROR("%s: sequence %d is coupled to %d in the input batch, but have divereged\n", __func__, s0, s1);330 return false;331 }332 }333 }334 }335 }336 337 // disallow partial sequence sub-sets:338 //339 // invalid: x340 // i: 0 1 2 ...341 // ---------------------------------------342 // seq_id[i][0]: 0 0 1343 // seq_id[i][1]: 1 1 2344 // seq_id[i][2]: 2345 //346 // disallow decreasing sequence positions:347 //348 // invalid: x349 // i: 0 1 2 3 4 5 6 ...350 // ---------------------------------------351 // pos[i]: 4 5 0 1 6 2 3352 // seq_id[i][0]: 0 0 1 1 0 1 0353 //354 {355 seq_set_t cur_seq_set[LLAMA_MAX_SEQ];356 for (uint32_t s = 0; s < n_seq_max; ++s) {357 cur_seq_set[s].set();358 }359 360 llama_pos cur_seq_pos[LLAMA_MAX_SEQ];361 for (uint32_t s = 0; s < n_seq_max; ++s) {362 cur_seq_pos[s] = -1;363 }364 365 for (int32_t i = 0; i < batch.n_tokens; ++i) {366 const llama_pos pos = batch.pos[i];367 368 for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {369 const llama_seq_id seq_id = batch.seq_id[i][s];370 371 cur_seq_set[seq_id] &= seq_set[i];372 373 if (cur_seq_set[seq_id].none()) {374 LLAMA_LOG_ERROR("%s: sequence %d belongs to incompatible sequence sets (not allowed)\n", __func__, seq_id);375 return false;376 }377 378 if (pos < cur_seq_pos[seq_id]) {379 LLAMA_LOG_ERROR("%s: sequence %d positions are decreasing (not allowed)\n", __func__, seq_id);380 return false;381 }382 383 cur_seq_pos[seq_id] = pos;384 }385 }386 }387 388 split_reset();389 390 return true;391}392 393llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t n_seqs) {394 const uint32_t n_tokens = n_seq_tokens*n_seqs;395 396 clear();397 split_reset();398 399 const int64_t n_pos_all = (int64_t) n_tokens*n_pos_per_embd;400 401 auto udata = std::make_shared<llama_ubatch::data_t>();402 403 udata->token .resize(n_tokens);404 udata->embd .clear();405 udata->pos .resize(n_pos_all);406 udata->n_seq_id .resize(n_tokens);407 udata->seq_id .resize(n_tokens);408 udata->seq_id_unq.resize(0);409 udata->seq_idx .resize(LLAMA_MAX_SEQ, -1);410 udata->output .resize(n_tokens);411 412 for (uint32_t s = 0; s < n_seqs; ++s) {413 udata->seq_idx[s] = s;414 udata->seq_id_unq.push_back(s);415 }416 417 llama_ubatch res {418 /*.b_equal_seqs =*/ true,419 /*.n_tokens =*/ n_tokens,420 /*.n_seq_tokens =*/ n_seq_tokens,421 /*.n_seqs =*/ n_seqs,422 /*.n_seqs_unq =*/ n_seqs,423 /*.n_pos =*/ n_pos_per_embd,424 425 /*.token =*/ udata->token.data(),426 /*.embd =*/ nullptr,427 /*.pos =*/ udata->pos.data(),428 /*.n_seq_id =*/ udata->n_seq_id.data(),429 /*.seq_id =*/ udata->seq_id.data(),430 /*.seq_id_unq =*/ udata->seq_id_unq.data(),431 /*.seq_idx =*/ udata->seq_idx.data(),432 /*.output =*/ udata->output.data(),433 /*.data =*/ std::move(udata),434 };435 436 return res;437}438 439const llama_batch & llama_batch_allocr::get_batch() const {440 return batch;441}442 443uint32_t llama_batch_allocr::get_n_tokens() const {444 return batch.n_tokens;445}446 447uint32_t llama_batch_allocr::get_n_outputs() const {448 return n_outputs;449}450 451uint32_t llama_batch_allocr::get_n_used() const {452 return n_used;453}454 455std::vector<int32_t> & llama_batch_allocr::get_out_ids() {456 return out_ids;457}458 459llama_pos llama_batch_allocr::seq_pos_min(llama_seq_id seq_id) const {460 return seq_pos[seq_id].empty() ? -1 : *seq_pos[seq_id].begin();461}462 463llama_pos llama_batch_allocr::seq_pos_max(llama_seq_id seq_id) const {464 return seq_pos[seq_id].empty() ? -1 : *seq_pos[seq_id].rbegin();465}466 467void llama_batch_allocr::split_reset() {468 out_ids.clear();469 470 n_used = 0;471 472 used.clear();473 used.resize(get_n_tokens(), false);474}475 476llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {477 // find the first unused token478 uint32_t cur_idx = 0;479 while (cur_idx < used.size() && used[cur_idx]) {480 ++cur_idx;481 }482 483 // we are done484 if (cur_idx >= used.size()) {485 return {};486 }487 488 std::vector<int32_t> idxs;489 490 while (true) {491 idxs.push_back(cur_idx);492 493 used[cur_idx] = true;494 ++n_used;495 496 ++cur_idx;497 498 if (cur_idx >= used.size()) {499 break;500 }501 502 if (idxs.size() >= n_ubatch) {503 break;504 }505 }506 507 return ubatch_add(idxs, idxs.size(), false);508}509 510llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) {511 if (sequential && has_cpl) {512 LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);513 514 return {};515 }516 517 std::vector<seq_set_t> cur_seq_set;518 519 llama_seq_id last_seq_id = -1;520 521 // determine the non-overlapping sequence sets participating in this ubatch522 for (int32_t i = 0; i < batch.n_tokens; ++i) {523 if (used[i]) {524 continue;525 }526 527 bool add = true;528 529 for (uint32_t s = 0; s < cur_seq_set.size(); ++s) {530 // no overlap with existing sequence sets:531 if (!(cur_seq_set[s] & seq_set[i]).none()) {532 add = false;533 break;534 }535 }536 537 // accept only increasing sequence ids538 if (sequential) {539 add = add && (cur_seq_set.empty() || batch.seq_id[i][0] == last_seq_id + 1);540 }541 542 if (add) {543 cur_seq_set.push_back(seq_set[i]);544 545 last_seq_id = batch.seq_id[i][0];546 547 if (cur_seq_set.size() > n_ubatch) {548 break;549 }550 }551 }552 553 uint32_t n_seqs = cur_seq_set.size();554 555 // we are done556 if (n_seqs == 0) {557 return {};558 }559 560 // the current batch index of each sequence set561 std::vector<int32_t> cur_idx(n_seqs, 0);562 563 for (uint32_t s = 0; s < n_seqs; ++s) {564 while (used[seq_set_map[cur_seq_set[s]][cur_idx[s]]]) {565 ++cur_idx[s];566 }567 }568 569 // the list of batch indices for each sequence set570 // at the end we will concat these to get the final ubatch571 std::vector<idx_vec_t> idxs_per_seq(n_seqs);572 573 while (true) {574 // we can only add new n_seq_tokens tokens if all the sequence sets have at least 1 more unused tokens and575 // if we haven't reached n_ubatch576 bool can_expand = true;577 578 for (uint32_t s = 0; s < n_seqs; ++s) {579 if (cur_idx[s] >= (int32_t) seq_set_map[cur_seq_set[s]].size()) {580 can_expand = false;581 break;582 }583 }584 585 if (!can_expand) {586 break;587 }588 589 for (uint32_t s = 0; s < n_seqs; ++s) {590 const int32_t idx = seq_set_map[cur_seq_set[s]][cur_idx[s]];591 592 idxs_per_seq[s].push_back(idx);593 594 used[idx] = true;595 ++n_used;596 597 ++cur_idx[s];598 }599 600 if ((idxs_per_seq[0].size() + 1)*n_seqs > n_ubatch) {601 break;602 }603 }604 605 // if n_keep_tail > 0, keep only the seqs that either finish in this ubatch or have at least606 // n_keep_tail tokens remaining for a future ubatch, so that the trailing n_keep_tail tokens607 // of each seq are never split across ubatches608 if (n_keep_tail > 0) {609 GGML_ASSERT(n_ubatch > n_keep_tail);610 611 auto n_remaining = [&](uint32_t s) {612 return (uint32_t) (seq_set_map[cur_seq_set[s]].size() - cur_idx[s]);613 };614 615 // keep the longest prefix of seqs that satisfy the constraint, to preserve sequential seq ids616 uint32_t n_keep = 0;617 while (n_keep < n_seqs) {618 const uint32_t remaining = n_remaining(n_keep);619 620 if (remaining != 0 && remaining < n_keep_tail) {621 break;622 }623 624 n_keep++;625 }626 627 // all seqs violate the constraint - resolve the first one directly and emit it alone628 if (n_keep == 0) {629 auto & idxs = idxs_per_seq[0];630 631 const auto & seq_idxs = seq_set_map[cur_seq_set[0]];632 633 if (idxs.size() + n_remaining(0) <= n_ubatch) {634 // extend the seq to completion635 while (n_remaining(0) > 0) {636 const int32_t idx = seq_idxs[cur_idx[0]];637 638 idxs.push_back(idx);639 640 used[idx] = true;641 ++n_used;642 643 ++cur_idx[0];644 }645 } else {646 // truncate the seq so that at least n_keep_tail tokens remain647 while (n_remaining(0) < n_keep_tail) {648 used[idxs.back()] = false;649 --n_used;650 651 idxs.pop_back();652 653 --cur_idx[0];654 }655 }656 657 n_keep = 1;658 }659 660 // return the tokens of the deferred seqs back to the pool661 for (uint32_t s = n_keep; s < n_seqs; ++s) {662 for (const int32_t idx : idxs_per_seq[s]) {663 used[idx] = false;664 --n_used;665 }666 }667 668 n_seqs = n_keep;669 }670 671 // concat the per-sequence-set lists672 std::vector<int32_t> idxs;673 674 for (uint32_t s = 0; s < n_seqs; ++s) {675 idxs.insert(idxs.end(), idxs_per_seq[s].begin(), idxs_per_seq[s].end());676 }677 678 return ubatch_add(idxs, n_seqs, true);679}680 681llama_ubatch llama_batch_allocr::split_seq(uint32_t n_ubatch) {682 // find the first unused token683 uint32_t cur_idx = 0;684 while (cur_idx < used.size() && used[cur_idx]) {685 ++cur_idx;686 }687 688 // we are done689 if (cur_idx >= used.size()) {690 return {};691 }692 693 // this is the starting sequence set694 // we allow adding tokens only if their sequence set is a subset of the current sequence set695 auto cur_seq_set = seq_set[cur_idx];696 697 std::vector<int32_t> idxs;698 699 while (true) {700 idxs.push_back(cur_idx);701 702 used[cur_idx] = true;703 ++n_used;704 705 if (idxs.size() >= n_ubatch) {706 break;707 }708 709 do {710 ++cur_idx;711 } while (cur_idx < get_n_tokens() && (used[cur_idx] || ((cur_seq_set & seq_set[cur_idx]) != seq_set[cur_idx])));712 713 if (cur_idx == get_n_tokens()) {714 break;715 }716 717 cur_seq_set = seq_set[cur_idx];718 }719 720 return ubatch_add(idxs, 1, true);721}722 723void llama_batch_allocr::clear() {724 n_outputs = 0;725 726 batch = {};727 728 pos .clear();729 n_seq_id .clear();730 seq_id .clear();731 seq_id_unq.clear();732 output .clear();733 734 for (auto & cur : seq_pos) {735 cur.clear();736 }737 738 for (auto & cur : seq_cpl) {739 std::fill(cur.begin(), cur.end(), false);740 }741 742 seq_set.clear();743 744 seq_set_map.clear();745 746 std::fill(seq_idx.begin(), seq_idx.end(), -1);747}748 749llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, uint32_t n_seqs, bool equal_seqs) {750 const uint32_t n_tokens = idxs.size();751 752 assert(n_tokens%n_seqs == 0);753 754 auto udata = std::make_shared<llama_ubatch::data_t>();755 756 const int64_t n_embd_all = batch.embd ? (int64_t) n_tokens*n_embd : 0;757 const int64_t n_pos_all = (int64_t) n_tokens*n_pos_per_embd;758 759 udata->token .resize(n_tokens);760 udata->embd .resize(n_embd_all);761 udata->pos .resize(n_pos_all);762 udata->n_seq_id .resize(n_tokens);763 udata->seq_id .resize(n_tokens);764 udata->seq_id_unq.resize(0);765 udata->seq_idx .resize(LLAMA_MAX_SEQ, -1);766 udata->output .resize(n_tokens);767 768 udata->seq_id_data.reserve(n_tokens);769 770 seq_set_t seq_set_unq;771 772 for (size_t i = 0; i < idxs.size(); ++i) {773 if (batch.token) {774 udata->token[i] = batch.token[idxs[i]];775 }776 777 if (batch.embd) {778 memcpy(udata->embd.data() + i*n_embd, batch.embd + (int64_t) idxs[i]*n_embd, n_embd*sizeof(float));779 }780 781 for (size_t j = 0; j < (size_t)n_pos_per_embd; ++j) {782 // if we are using M-RoPE783 // if the current batch is text, we need to broadcast the same position across all RoPE sections784 // otherwise, the input batch is image embeddings, we copy the positions as-is785 // if we are not using M-RoPE, there is only one position per token (this loop runs only once)786 size_t src_off = batch.token ? 0 : j*batch.n_tokens;787 udata->pos[j*n_tokens + i] = batch.pos[src_off + idxs[i]];788 }789 790 udata->n_seq_id[i] = batch.n_seq_id[idxs[i]];791 udata->output[i] = batch.logits[idxs[i]];792 793 for (int s = 0; s < udata->n_seq_id[i]; ++s) {794 const llama_seq_id seq_id = batch.seq_id[idxs[i]][s];795 796 udata->seq_id_data.push_back(seq_id);797 seq_set_unq.set(seq_id);798 }799 800 if (udata->output[i]) {801 out_ids.push_back(idxs[i]);802 }803 }804 805 llama_seq_id * seq_id_ptr = udata->seq_id_data.data();806 for (size_t i = 0; i < idxs.size(); ++i) {807 udata->seq_id[i] = seq_id_ptr;808 seq_id_ptr += udata->n_seq_id[i];809 }810 811 for (uint32_t s = 0; s < n_seq_max; ++s) {812 if (seq_set_unq.test(s)) {813 udata->seq_idx[s] = udata->seq_id_unq.size();814 udata->seq_id_unq.push_back(s);815 }816 }817 818 llama_ubatch res {819 /*.b_equal_seqs =*/ equal_seqs,820 /*.n_tokens =*/ n_tokens,821 /*.n_seq_tokens =*/ n_tokens/n_seqs,822 /*.n_seqs =*/ n_seqs,823 /*.n_seqs_unq =*/ (uint32_t) udata->seq_id_unq.size(),824 /*.n_pos =*/ n_pos_per_embd,825 826 /*.token =*/ batch.token ? udata->token.data() : nullptr,827 /*.embd =*/ batch.embd ? udata->embd.data() : nullptr,828 /*.pos =*/ udata->pos.data(),829 /*.n_seq_id =*/ udata->n_seq_id.data(),830 /*.seq_id =*/ udata->seq_id.data(),831 /*.seq_id_unq =*/ udata->seq_id_unq.data(),832 /*.seq_idx =*/ udata->seq_idx.data(),833 /*.output =*/ udata->output.data(),834 /*.data =*/ std::move(udata),835 };836 837 if (debug > 0) {838 LLAMA_LOG_DEBUG("%s: added ubatch to split:\n", __func__);839 840 ubatch_print(res, debug);841 }842 843 return res;844}845 846void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {847 if (debug > 0) {848 LLAMA_LOG_DEBUG("%s: equal_seqs = %d\n", __func__, ubatch.equal_seqs());849 LLAMA_LOG_DEBUG("%s: n_tokens = %d\n", __func__, ubatch.n_tokens);850 LLAMA_LOG_DEBUG("%s: n_seq_tokens = %d\n", __func__, ubatch.n_seq_tokens);851 LLAMA_LOG_DEBUG("%s: n_seqs = %d\n", __func__, ubatch.n_seqs);852 LLAMA_LOG_DEBUG("%s: n_seqs_unq = %d\n", __func__, ubatch.n_seqs_unq);853 854 std::stringstream ss_seq_id_unq;855 std::stringstream ss_seq_idx;856 857 ss_seq_id_unq << "[ ";858 ss_seq_idx << "[";859 860 for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {861 ss_seq_id_unq << ubatch.seq_id_unq[s] << " ";862 }863 864 for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {865 if (ubatch.seq_idx[s] >= 0) {866 ss_seq_idx << ubatch.seq_idx[s]%10;867 } else {868 ss_seq_idx << ".";869 }870 }871 872 ss_seq_id_unq << "]";873 ss_seq_idx << "]";874 875 LLAMA_LOG_DEBUG("%s: token = %p\n", __func__, (void *) ubatch.token);876 LLAMA_LOG_DEBUG("%s: embd = %p\n", __func__, (void *) ubatch.embd);877 LLAMA_LOG_DEBUG("%s: pos = %p\n", __func__, (void *) ubatch.pos);878 LLAMA_LOG_DEBUG("%s: n_seq_id = %p\n", __func__, (void *) ubatch.n_seq_id);879 LLAMA_LOG_DEBUG("%s: seq_id = %p\n", __func__, (void *) ubatch.seq_id);880 LLAMA_LOG_DEBUG("%s: seq_id_unq = %s\n", __func__, ss_seq_id_unq.str().c_str());881 LLAMA_LOG_DEBUG("%s: seq_idx = %s\n", __func__, ss_seq_idx.str().c_str());882 LLAMA_LOG_DEBUG("%s: output = %p\n", __func__, (void *) ubatch.output);883 LLAMA_LOG_DEBUG("%s: n_outputs = %d\n", __func__, n_outputs);884 885 if (debug > 0) {886 int seq_id_max = 0;887 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {888 for (int s = 0; s < ubatch.n_seq_id[i]; ++s) {889 for (int s = 0; s < ubatch.n_seq_id[i]; ++s) {890 seq_id_max = std::max(seq_id_max, ubatch.seq_id[i][s]);891 }892 }893 }894 ++seq_id_max;895 896 LLAMA_LOG_DEBUG("%s: token = [\n", __func__);897 for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {898 std::vector<int8_t> seq_id(seq_id_max);899 900 for (int s = 0; s < ubatch.n_seq_id[i]; ++s) {901 seq_id[ubatch.seq_id[i][s]] = 1;902 }903 904 std::stringstream ss;905 for (int s = 0; s < seq_id_max; ++s) {906 if (seq_id[s]) {907 ss << s%10;908 } else {909 ss << ".";910 }911 }912 913 if (ubatch.token) {914 LLAMA_LOG_DEBUG("%s: %4d: id = %6d (%16s), pos = %4d, n_seq_id = %2d, seq_id = [%s], output = %d\n",915 __func__, i, ubatch.token[i], vocab->token_to_piece(ubatch.token[i]).c_str(),916 ubatch.pos[i], ubatch.n_seq_id[i], ss.str().c_str(), ubatch.output[i]);917 } else {918 LLAMA_LOG_DEBUG("%s: %4d: [embd], pos = %4d, n_seq_id = %2d, seq_id = [%s], output = %d\n",919 __func__, i, ubatch.pos[i], ubatch.n_seq_id[i], ss.str().c_str(), ubatch.output[i]);920 }921 }922 LLAMA_LOG_DEBUG("%s: ]\n", __func__);923 }924 }925}926 927//928// interface implementation929//930 931struct llama_batch llama_batch_get_one(932 llama_token * tokens,933 int32_t n_tokens) {934 return {935 /*n_tokens =*/ n_tokens,936 /*tokens =*/ tokens,937 /*embd =*/ nullptr,938 /*pos =*/ nullptr,939 /*n_seq_id =*/ nullptr,940 /*seq_id =*/ nullptr,941 /*logits =*/ nullptr,942 };943}944 945struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_t n_seq_max) {946 llama_batch batch = {947 /*n_tokens =*/ 0,948 /*tokens =*/ nullptr,949 /*embd =*/ nullptr,950 /*pos =*/ nullptr,951 /*n_seq_id =*/ nullptr,952 /*seq_id =*/ nullptr,953 /*logits =*/ nullptr,954 };955 956 if (embd) {957 batch.embd = (float *) malloc(sizeof(float) * n_tokens_alloc * embd);958 } else {959 batch.token = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc);960 }961 962 batch.pos = (llama_pos *) malloc(sizeof(llama_pos) * n_tokens_alloc);963 batch.n_seq_id = (int32_t *) malloc(sizeof(int32_t) * n_tokens_alloc);964 batch.seq_id = (llama_seq_id **) malloc(sizeof(llama_seq_id *) * (n_tokens_alloc + 1));965 for (int i = 0; i < n_tokens_alloc; ++i) {966 batch.seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id) * n_seq_max);967 }968 batch.seq_id[n_tokens_alloc] = nullptr;969 970 batch.logits = (int8_t *) malloc(sizeof(int8_t) * n_tokens_alloc);971 972 return batch;973}974 975void llama_batch_free(struct llama_batch batch) {976 if (batch.token) free(batch.token);977 if (batch.embd) free(batch.embd);978 if (batch.pos) free(batch.pos);979 if (batch.n_seq_id) free(batch.n_seq_id);980 if (batch.seq_id) {981 for (int i = 0; batch.seq_id[i] != nullptr; ++i) {982 free(batch.seq_id[i]);983 }984 free(batch.seq_id);985 }986 if (batch.logits) free(batch.logits);987}988