Repository navigation
qsa : enable the incremental QSA cache on Vulkan - #117
Conversation
|
Tested #116 + #117 on top of #133 (selected-key WMMA, with #115/#119 under it) on a Strix Halo (gfx1151, RADV, stock Mesa), Flash-Next REAP-320 Q3_K_XL,
So the cache stays exact on the maskless path too. One gap: our second-turn prompts were re-processed in full, so rebuilding a partly cached prefix after another conversation's turn wasn't exercised. With this, #115/#119 (tested earlier) and #116/#117 all check out on gfx1151. Would you be OK taking the stack out of draft so we can start merging, #115 first? Assisted-by: Claude Opus 5.5 |
|
Thanks for testing the stack on gfx1151, and for #133. The gap you mention, rebuilding a partly cached prefix after another conversation's turn, is what #116 adds. I've taken #115–#119 out of draft. #115 is the one to start with; it also fixes #130. Assisted-by: Claude Opus 5.5 |
|
The conflict with master is only this branch's copy of #115 ( Please also re-check after rebasing: your bit-identical runs were on #115, where Vulkan took the masked top-k. On master, Vulkan now takes the maskless selected-key path, which also reads Comment generated via automated review of clanker. |
…ed cache The incremental QSA state (qsa_prefix: one sequence's cells in position order, the pooled indexer-key cache, the compact metadata) could only be rebuilt when the tracked sequence was the only occupant of the cache: qsa_recover() gave up at the first cell of another sequence. A multi-slot server with a unified KV pool therefore never regained the incremental path once a second conversation had cells in the pool. - qsa_rebuild_prefix() (src/prefix.h, unit-tested in test-qsa-recover) skips other sequences' cells; holes, shared cells and duplicate positions are still refused. - qsa_apply() recovers the prefix of whichever single-sequence ubatch comes next, including right on a switch between sequences, and tries a sequence that cannot form a prefix only once until the state is invalidated again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Assisted-by: Claude Opus 5.5
The incremental QSA cache keeps the pooled indexer keys of the tracked sequence and scores only new cells. It was enabled only when the QSA layers run on ROCm. Validated on Vulkan (Radeon 890M, gfx1150, Mesa 26.1.6) with Qwen3.8-Flash-Next UD-IQ4_XS, on the maskless selected-key path: - exact: the four repository corpora, as is and repeated to 12.7k-53.4k tokens, give bit-identical prompt tokens, complete next-token distributions and 64 greedy steps with top-20 log-probabilities with the cache on and off - decode with one conversation (llama-bench tg128): +7.4% / +20.8% / +44.6% at 12k / 32k / 64k cells, unchanged on an empty cache - with three other 16k-token conversations in the unified cache: the same 64 tokens as with an empty cache; decode at 97% of the empty-cache speed (72% without the cache), and at 102% (75%) once the other conversation has stopped LLAMA_QSA_INCREMENTAL=0 disables it, =1 enables it on any backend. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Rebased onto master 4e0d917. c20ba6f is dropped, so the branch is now #116 plus the gate,
Decode at depth vs master ( Assisted-by: Claude Opus 5.5 |
…e kernel (halo-box #117), dropping the six visibility ops per strip On the maskless path qwen4exp_apply_compact_visibility turns the raw head sum into score + log(step(tail[q] - start[b])) with cast(starts), cast(tails), reshape, repeat, sub, step, log and add (plus a get_rows with several sequence rows): seven or eight dispatches and a [n_blocks, n_query] f32 temporary per QSA layer per strip, in front of the block select. The block-select kernel now takes the limits buffer (binding 4) and applies the same test per element while it loads the row (positions stay below 2^24, so the f32 subtraction the chain does is exact and the integer compare is the same test); the fused range starts at the first visibility cast and ends at the chain's CONCAT, so none of those ops runs. The scores it reads are the raw QSA_SCORE output; the matcher checks the whole visibility subgraph (one or several rows) and that every visibility op sits inside the range, and still accepts the TOP_K entry (scores carrying the visibility, v1). Chain validation no longer uses ggml_can_fuse_subgraph_ext: it rejects a view of a tensor outside the subgraph unless that is a weight, and does not count a cast's self-reference in src[1], so the chain (views of the score and of the block cells, three casts) could never pass it. The matcher counts the readers of every node itself: all inside the chain, only the CONCAT read outside. QSA_SCORE on the same path: its keep pattern required a trailing bias ADD, which the compact path does not have adjacent to the head sum, so the scorer chain was not kept together there; qsa_score_pattern_nobias (the chain without the ADD) is kept and protected too, and the scorer matcher leaves a log(step)-fed ADD to the block select when that fits the row. Knobs: GGML_VK_QSA_BLOCK_SELECT (default on), GGML_VK_QSA_BLOCK_SELECT_LIMITS=0 keeps the visibility ops and enters at the TOP_K (the v1 arm). Tests: five limits-mode cases (one and two sequence rows, budget past the visible blocks, a decode row), exact compare against the CPU chain, each asserting exactly one QSA_BLOCK_SELECT dispatch. Assisted-by: Claude Opus 5.5
The comments the QSA select / score / radix / flat GET_ROWS changes added named ledger items (A2, E11, B-SEL, Hi 4, halo-box #117, fork hashes, dated debug runs) and ran to 3-11 hard-wrapped lines restating the commit messages. They now give the what, the why and the measurement in 1-2 lines; the attribution for the QSA_SCORE idea stays (review C1, C2). Assisted-by: Claude Opus 5.5
3 participants
Overview
Depends on #116. This branch contains it; the change here is the last commit.
The incremental QSA cache keeps the tracked sequence's pooled, normalised and rotated indexer keys between ubatches. It then scores only the new cells, instead of pooling the whole cache again every time. Until now it was enabled only when the QSA layers run on ROCm. On Vulkan it gives bit-identical results on master's maskless selected-key path (below), so this PR enables it there too.
LLAMA_QSA_INCREMENTAL=0disables the cache, and=1enables it on any backend.Measurements
Decode at depth is what this PR changes. Master and this PR were run in the order master, this PR, this PR, master, each time with
-r 2:llama-bench -m MODEL -p 0 -n 128 -d 0,12000,32000,64000 -b 2048 -ub 512 -r 2 -ngl 99 -fa on -ctk f16 -ctv f16 -lm none -lzm on-direct -o jsonlWithout the cache, decode slows from 10.7 to 6.6 t/s between an empty cache and 64k cells. With it, decode only goes
from 10.7 to 9.5.
Prompt speed of a single conversation does not change beyond about 1%. For example, a 22.5k-token prompt runs at 151.8 t/s with the cache off and 152.0 t/s with it, and a 53.4k-token prompt at 147.3 t/s on master and 146.4 t/s with this PR.
Several conversations in one unified cache:
llama-server -c 65536 -np 4 -kvu --cache-ram 0 --no-cache-idle-slots, no MTP. Slot 0 reads a 2048-token prompt and generates 64 tokens with the cache empty, then again with three 16k-token conversations resident:LLAMA_QSA_INCREMENTAL=0)Without the cache, every decode step pools the indexer keys of the whole unified cache, so the other conversations slow it down. With the cache, a step scores only its own new cells. The middle column is this PR's gate without #116: master cannot rebuild the tracked prefix while another conversation has cells in the cache, so next to the residents the cache is not used.
Correctness, with the cache on vs
LLAMA_QSA_INCREMENTAL=0on the same build, all on the maskless selected-key path:llama-server(1 slot, no MTP), each as is (570-876 tokens) and repeated to 12.7k-53.4k tokens. We compare the prompt token IDs, the complete next-token distribution after prefill (all 248,320 vocabulary entries) and 64 greedy steps with top-20 log-probabilities. The cache on and off are bit-identical on all eight prompts. With the cache off, the build is also bit-identical to master.leak-diag.py). With this PR, X and Y give identical tokens, with no log-probability difference over the top 20 (0.0000), the same as on master.test-qsa-recover(qsa : track a sequence's prefix while other sequences share the unified cache #116) passes.test-backend-opsdoes not apply.Requirements
qwen4exp incremental QSA cache) on the Vulkan backend. Our measurements come from Strix Point (gfx1150), not
Strix Halo.
measurement. Every commit carries
Assisted-by: Claude Opus 5.5.