kv-cells: hoist the sequence list out of the for_each_token_in cell loop - #154
kv-cells: hoist the sequence list out of the for_each_token_in cell loop#154danielhanchen wants to merge 1 commit into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b17ab37a67
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // seqs does not change over the scan, so list its bits once instead of | ||
| // intersecting a LLAMA_MAX_SEQ-wide bitset against every cell |
There was a problem hiding this comment.
Join this newly added explanation into one line; it is hard-wrapped halfway through a sentence, contrary to the repository's explicit comment-formatting rule.
AGENTS.md reference: AGENTS.md:L81-L81
Useful? React with 👍 / 👎.
| for (int k = 0; k < n_sel; ++k) { | ||
| if (seq[i].test(sel[k])) { | ||
| f(sel[k], pos[i], ext[i].tok); |
There was a problem hiding this comment.
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 👍 / 👎.
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 62acc89, 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
b17ab37 to
31e432e
Compare
* unsloth: pin ggml-org#27941, #152 and #154 ggml-org#27941 at 02eb201, no longer a draft and MERGEABLE upstream. It fixes four qwen4exp correctness defects and is the only account of the Flash-Next reporter who is on Vulkan only, where the allocator variable is never set. The previous pin PR #148 carried the stale 8161d11 and is closed. #152 at 258345e and #154 at 31e432e both merge onto b10708 on their own. #137 is already pinned at 4e1865e, which is still its head, so no change. #142 and #144 are NOT pinned here: both predate an upstream refactor that folded lazy_mode and model_shared into a lazy struct, so they conflict on src/llama.cpp and src/llama-model-loader.cpp against b10708. b10708 has 'ml.lazy.mode = params.lazy_mode' where those branches still write 'ml.lazy_mode' and 'ml.model_shared'. They need rebasing onto current upstream before they can be pinned; pinning them now would fail the resolve. * unsloth: repin ggml-org#27941 to 6b2b85c The PR moved on at 10:44Z, after 02eb201 was pinned. Verified to merge onto the current base tag: 6 files changed, 480 insertions, 138 deletions. * unsloth: pin #144, rebased, and leave #142 out as contained in it #144 rebased onto b10709 at 6fc8df1. Two of its nine commits were dropped rather than carried: 'qwen4exp: fix QSA correctness defects and harden metadata loading' and the comment tidy on top of it. That work is what ggml-org#27941 supersedes, and keeping both copies is what made the two pins conflict in llama-kv-cells.h and llama-memory-hybrid-idx.cpp. Listed after ggml-org#27941 so the upstream version of that work lands first. #142 is deliberately NOT pinned: its single commit is byte-identical to #144's c7bd6f2 apart from the lazy API spelling, so #144 already contains it and pinning both would apply the same change twice. Verified: b10709 + ggml-org#27941 + #144 merges clean, 24 files, +612/-94. * unsloth: repin ggml-org#27754 to 949f7ef The PR moved on at 10:39Z; 5796547 was pinned by #159 earlier today. Verified to merge onto b10709: 44 files changed, 2673 insertions, 38 deletions. Every entry in the set is now at its PR's current head. --------- Co-authored-by: danielhanchen <elliegouldingstuff@gmail.com>
1 participant
for_each_token_intakes aseqsbitset that does not change over the scan, then intersects it against every cell and counts the result.62acc89c2(ggml-org#28011) fixed the larger half of this by stopping the inner loop once every sequence in the cell has been seen, which removed the fixedLLAMA_MAX_SEQiteration. What remains is the per cellseq[i] & seqsandm.count(), both of which are computed from an operand that is constant for the whole call.Listing the set bits of
seqsonce, before the cell loop, removes both. A scan typically asks for a single sequence, so the inner work becomes onetest()per cell.Measurement
llama-bench -d 65536 -n 32on Qwen3.8-Flash-Next UD-IQ1_S, arms alternated inside one job, 8 reps each, on top of a master that already contains ggml-org#28011:7 of 8 reps beat every base rep.
This is a small, incremental change on top of ggml-org#28011 rather than an independent win, and it should be read that way.
for_each_token_incurrently has exactly one caller, the n-gram predecessor scan, so the effect is confined to models with per layer token embeddings.Testing
Built at
9723942ad, CUDA, B200.llama-mtmd-cli.test-llama-archsfor qwen4exp, llama, qwen3next, gemma3n and qwen3moe.