From 31e432e758f8cc4b2c5f27902721500173bf39db Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 31 Aug 2026 01:31:53 +0000 Subject: [PATCH] kv-cells: hoist the sequence list out of the for_each_token_in cell loop seqs is constant for the whole scan, so its set bits can be listed once instead of intersecting a LLAMA_MAX_SEQ-wide bitset against every cell and counting the result. Builds on 62acc89c2, which stops the inner loop once every sequence in the cell has been seen. That removed the fixed LLAMA_MAX_SEQ inner loop; this removes the per-cell bitset AND and popcount that remained. Assisted-by: Claude --- src/llama-kv-cells.h | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index a4292c79e44..6b4e6f5de40 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -314,20 +314,28 @@ class llama_kv_cells { // note: used by n-gram input embeddings to recover the tokens preceding a ubatch template void for_each_token_in(const std::bitset & seqs, llama_pos p0, llama_pos p1, F && f) const { + // hoisted: intersecting a LLAMA_MAX_SEQ-wide bitset per cell is the cost being removed + llama_seq_id sel[LLAMA_MAX_SEQ]; + int n_sel = 0; + + for (llama_seq_id s = 0; s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) { + if (seqs.test(s)) { + sel[n_sel++] = s; + } + } + + if (n_sel == 0) { + return; + } + for (const auto & i : used) { if (pos[i] < p0 || pos[i] >= p1) { continue; } - const auto m = seq[i] & seqs; - - // a cell carries a handful of sequences at most, out of LLAMA_MAX_SEQ - size_t left = m.count(); - - for (llama_seq_id s = 0; left > 0 && s < (llama_seq_id) LLAMA_MAX_SEQ; ++s) { - if (m.test(s)) { - f(s, pos[i], ext[i].tok); - --left; + for (int k = 0; k < n_sel; ++k) { + if (seq[i].test(sel[k])) { + f(sel[k], pos[i], ext[i].tok); } } }