kv-cells: stop the sequence scan once all sequences are seen

[ServeurpersoCom]

Overview

Reduce the generation slowdown as context grows. Split out of #27977 as requested, one PR per change. This is the quick win of the series: a handful of lines, and the largest gain of the whole set.

Looking up the previous tokens of an n-gram was checking all 256 possible sequences for every cell of the KV cache, when a cell almost always belongs to just one. It now stops as soon as it has seen the ones the cell actually holds, which is a lot cheaper as the cache fills up.

Additional information

for_each_token_in tested all LLAMA_MAX_SEQ sequences for every used cell, while a cell almost always belongs to one. The scan now stops once the cell's own sequences have been seen. Same visit order, same callback arguments, so behaviour is unchanged.

get_prev_tokens is the only caller, so this affects the n-gram path.

RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs:

55k context generation 56.3 -> 74.3 t/s
132k context generation 33.6 -> 50.9 t/s

Prompt processing is unchanged, the scan is amortised over the ubatch there. The gain follows the number of used cells, so it grows with context and is invisible on short prompts.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES

[drluoto]

Isolated datapoint for this 5-line PR alone, on plain master cc231cb (no other qwen4exp patches), gfx1151 / Strix Halo, ROCm 7.1, UD-IQ4_XS, llama-bench tg64 r=2, A/B in the same tree:

depth	master	+this PR
4096	19.33	19.56 (+1.2%)
16384	15.21	15.65 (+2.9%)
32768	13.30	13.85 (+4.1%)

Monotone with depth, no regressions anywhere. The gain is smaller here than the CUDA numbers in #27977 for a knowable reason: on stock master this box's TOP_K still runs on the CPU fallback (>1024, no CUB on HIP), so the cores are already saturated by that — the ngram scan isn't the sole binding constraint. With GPU TOP_K in place (#26592/#27466) the same scan removal was worth noticeably more in our stack testing.

Five lines, measurable everywhere, grows with context: easy +1 from this hardware for landing it.

[tvanderka]

This PR easily solves most of get_prev_tokens bottleneck. But I can still measure it being ~8% slower at 128k ctx vs poc in #27992

[ServeurpersoCom]

This PR easily solves most of get_prev_tokens bottleneck. But I can still measure it being ~8% slower at 128k ctx vs poc in #27992

I ran the two side by side on the same stack, only swapping my PR for yours, RTX PRO 6000, Qwen3.8-Flash-Next UD-Q4_K_XL, fa on, warm runs:

                      #28011      #27992
  55k   generation    74.7        77.5    (+3.7%)
  132k  generation    48.1        56.7    (+17.9%)
  55k   prompt      2638        2570
  132k  prompt      1889        1871

So your 8% is if anything understated here, and the gap widens with depth, which is what you would expect from log n versus a cheaper linear scan. Output checked correct in all runs.

The prompt processing difference is within my run to run drift, I would not read anything into it without an alternated A/B.

They are not complementary: once the index is there my scan never runs, so there is nothing left for it to save. That probably explains the regression sammcj saw carrying both.

[ggerganov]

I think we should get this small change in since it is simple.

The get_prev_tokens is something that will be reworked for sure, but I expect it would be a more complicated effort and will take some time. #27992 could be an option, though I have a feeling that the best approach has to be confined into llama_kv_cells to avoid bloating the llama_kv_cache with extra logic.

[ServeurpersoCom]

I'm making a note to POC a log n approach that doesn't touch the llama-kv-cache.

[ServeurpersoCom]

This PR easily solves most of get_prev_tokens bottleneck. But I can still measure it being ~8% slower at 128k ctx vs poc in #27992

Out of curiosity, I took your unit test and ran it against the index in #28040, adapting the reference since that one drops for_each_token_in. 9480 lookups, no failures.