| 450 | } |
| 451 | |
| 452 | bool llama_memory_recurrent::find_slot(const llama_ubatch & ubatch) { |
| 453 | const uint32_t n_seq_tokens = ubatch.n_seq_tokens; |
| 454 | const uint32_t n_seqs = ubatch.n_seqs; |
| 455 | |
| 456 | // if we have enough unused cells before the current head -> |
| 457 | // better to start searching from the beginning of the cache, hoping to fill it |
| 458 | if (head > used + 2*n_seqs) { |
| 459 | head = 0; |
| 460 | } |
| 461 | |
| 462 | // For recurrent state architectures (like Mamba or RWKV), |
| 463 | // each cache cell can store the state for a whole sequence. |
| 464 | // A slot should be always be contiguous. |
| 465 | |
| 466 | // can only process batches with an equal number of new tokens in each sequence |
| 467 | GGML_ASSERT(ubatch.equal_seqs()); |
| 468 | |
| 469 | int32_t min = size - 1; |
| 470 | int32_t max = 0; |
| 471 | |
| 472 | // everything should fit if all seq_ids are smaller than the max |
| 473 | for (uint32_t s = 0; s < n_seqs; ++s) { |
| 474 | const uint32_t i = s*n_seq_tokens; // first token of sequence set s |
| 475 | const uint32_t n_seq_id = ubatch.n_seq_id[i]; |
| 476 | |
| 477 | for (uint32_t j = 0; j < n_seq_id; ++j) { |
| 478 | const llama_seq_id seq_id = ubatch.seq_id[i][j]; |
| 479 | |
| 480 | if (seq_id < 0 || (uint32_t) seq_id >= size) { |
| 481 | // too big seq_id |
| 482 | // TODO: would it be possible to resize the cache instead? |
| 483 | LLAMA_LOG_ERROR("%s: seq_id=%d >= n_seq_max=%u Try using a bigger --parallel value\n", __func__, seq_id, n_seq_max); |
| 484 | return false; |
| 485 | } |
| 486 | if (j > 0) { |
| 487 | auto & seq = cells[seq_id]; |
| 488 | if (seq.tail >= 0) { |
| 489 | auto & cell = cells[seq.tail]; |
| 490 | // clear cells from seq_ids that become shared |
| 491 | // (should not normally happen, but let's handle it anyway) |
| 492 | cell.seq_id.erase(seq_id); |
| 493 | seq.tail = -1; |
| 494 | if (cell.seq_id.empty()) { |
| 495 | cell.pos = -1; |
| 496 | cell.src = -1; |
| 497 | used -= 1; |
| 498 | } |
| 499 | } |
| 500 | } |
| 501 | } |
| 502 | } |
| 503 | |
| 504 | #ifndef NDEBUG |
| 505 | { |
| 506 | std::vector<int32_t> tails_verif; |
| 507 | tails_verif.assign(size, -1); |
| 508 | for (uint32_t i = 0; i < size; ++i) { |
| 509 | auto & cell = cells[i]; |