| 504 | } |
| 505 | |
| 506 | llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) { |
| 507 | if (sequential && has_cpl) { |
| 508 | 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__); |
| 509 | |
| 510 | return {}; |
| 511 | } |
| 512 | |
| 513 | std::vector<seq_set_t> cur_seq_set; |
| 514 | |
| 515 | llama_seq_id last_seq_id = -1; |
| 516 | |
| 517 | // determine the non-overlapping sequence sets participating in this ubatch |
| 518 | for (int32_t i = 0; i < batch.n_tokens; ++i) { |
| 519 | if (used[i]) { |
| 520 | continue; |
| 521 | } |
| 522 | |
| 523 | bool add = true; |
| 524 | |
| 525 | for (uint32_t s = 0; s < cur_seq_set.size(); ++s) { |
| 526 | // no overlap with existing sequence sets: |
| 527 | if (!(cur_seq_set[s] & seq_set[i]).none()) { |
| 528 | add = false; |
| 529 | break; |
| 530 | } |
| 531 | } |
| 532 | |
| 533 | // accept only increasing sequence ids |
| 534 | if (sequential) { |
| 535 | add = add && (cur_seq_set.empty() || batch.seq_id[i][0] == last_seq_id + 1); |
| 536 | } |
| 537 | |
| 538 | if (add) { |
| 539 | cur_seq_set.push_back(seq_set[i]); |
| 540 | |
| 541 | last_seq_id = batch.seq_id[i][0]; |
| 542 | |
| 543 | if (cur_seq_set.size() > n_ubatch) { |
| 544 | break; |
| 545 | } |
| 546 | } |
| 547 | } |
| 548 | |
| 549 | const uint32_t n_seqs = cur_seq_set.size(); |
| 550 | |
| 551 | // we are done |
| 552 | if (n_seqs == 0) { |
| 553 | return {}; |
| 554 | } |
| 555 | |
| 556 | // the current batch index of each sequence set |
| 557 | std::vector<int32_t> cur_idx(n_seqs, 0); |
| 558 | |
| 559 | for (uint32_t s = 0; s < n_seqs; ++s) { |
| 560 | while (used[seq_set_map[cur_seq_set[s]][cur_idx[s]]]) { |
| 561 | ++cur_idx[s]; |
| 562 | } |
| 563 | } |
no test coverage detected