Skip to content

kv-cells: hoist the sequence list out of the for_each_token_in cell loop #154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
danielhanchen wants to merge 1 commit into from
Closed

kv-cells: hoist the sequence list out of the for_each_token_in cell loop #154

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions src/llama-kv-cells.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,20 +314,28 @@ class llama_kv_cells {
// note: used by n-gram input embeddings to recover the tokens preceding a ubatch
template<typename F>
void for_each_token_in(const std::bitset<LLAMA_MAX_SEQ> & 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);
Comment on lines +336 to +338

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve the sparse-cell fast path for large sequence batches

When a PLE ubatch contains many unique sequences, get_prev_tokens() puts every ubatch.seq_id_unq into seqs, so this loop now tests all selected sequences for every in-range cell. For example, with 256 selected sequences and a cell belonging only to sequence 0, the old intersection/count loop stopped after one sequence test, while this loop performs 256 tests; across a long KV cache this can substantially regress the multi-sequence decode workload that the optimization is meant to accelerate. Retain the previous sparse-cell path, or iterate whichever of the selected and cell sequence sets is smaller.

Useful? React with 👍 / 👎.

}
}
}
Expand Down
Loading