kv-cache: scan only the predecessor window in get_prev_tokens - #153
kv-cache: scan only the predecessor window in get_prev_tokens#153danielhanchen 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: c49b634dce
ℹ️ 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".
| // a lookup landed in an M-RoPE gap, so now scan below the window | ||
| if (p_scan0 > 0) { | ||
| for (uint32_t s = 0; s < n_stream; ++s) { | ||
| v_cells[s].for_each_token_in(seqs, 0, p_scan0, collect); |
There was a problem hiding this comment.
When an M-RoPE lookup misses, this fallback makes a second complete traversal of the cache: llama_kv_cells::for_each_token_in() iterates every entry in used and only then filters by position (src/llama-kv-cells.h:317-319). The previous implementation traversed used once, whereas the new miss path traverses it twice while invoking collect for the same total set of cells, so multimodal batches with position gaps regress at large contexts; use a position-indexed range traversal or preserve the below-window result from the first traversal.
Useful? React with 👍 / 👎.
| // only [w0, p_max] is read back from hist. cells below w0 only fill below[], which is used | ||
| // when a lookup finds nothing in the window. so scan the window first, the rest only if needed |
There was a problem hiding this comment.
This new comment hard-wraps one sentence across two source lines, and the same pattern recurs in the comments below. AGENTS.md explicitly forbids hard-wrapping comments and splitting a line mid-sentence, so keep each sentence on one line.
AGENTS.md reference: AGENTS.md:L77-L81
Useful? React with 👍 / 👎.
get_prev_tokens rebuilds the n-gram predecessor history by walking every cell of the cache, but only positions in [w0, p_max] are ever read back out of it. The cells below w0 exist solely to fill below[], which is the fallback for a lookup that lands in an M-RoPE position gap. Scan the window first and fall back to the full scan only when a lookup actually misses, which for text never happens: positions are unique, so every lookup resolves inside the window. Assisted-by: Claude
c49b634 to
dbac7c4
Compare
1 participant
get_prev_tokensrebuilds the n-gram predecessor history for the PLE input by walking every cell of the cache, at every decode step. Only positions in[w0, p_max]are ever read back out of the map it builds. The cells beloww0exist for one reason: to fillbelow[], which is the fallback for a lookup that lands in an M-RoPE position gap.For text that fallback is never used, because text positions are unique and every lookup resolves inside the window. So at 64k context the scan visits 65536 cells to return two tokens.
This scans the window first and falls back to the full scan only when a lookup actually misses.
The M-RoPE path is why this is not a one line change
An earlier version simply narrowed the range and dropped
below[]. It was byte identical on text, which made the fallback look like dead weight. It is not: a multimodal ubatch repeats one position across a whole image, so a lookup can legitimately land in a gap, and that version produced a different, still plausible caption.lookuptherefore reports a miss instead of silently returningbelow[], the missed(index, seq_id)pairs are collected, and the[0, w0)scan runs only if that list is non empty, through the samecollectlambda.Measurement
Note that
62acc89c2(ggml-org#28011) landed first and already stops the inner sequence loop once every sequence in a cell has been seen. That removed most of this cost. Any number measured before it, including the ones I had from earlier work, no longer describes this patch.Re-measured against current master,
llama-bench -d 65536 -n 32, arms alternated inside one job, 8 reps each:7 of 8 reps beat every base rep. This is a couple of percent at 64k context, not the large win it was before ggml-org#28011, and it is worth stating plainly.
The two changes attack different loops, this one the number of cells visited and ggml-org#28011 the work per cell, but they land on the same total because what remains after ggml-org#28011 is small.
Testing
Built at
9723942ad, CUDA, B200.llama-mtmd-cli, which is the gate the earlier version failed.test-llama-archsfor qwen4exp, llama, qwen3next, gemma3n and qwen3moe.