Repository navigation
qsa : keep the incremental QSA prefix across mrope images (rebase of #163) - #235
Merged
Merged
Conversation
With one image anywhere in the cache, qsa_scalar_visibility_cells rejected every ubatch, so all later decode and verify steps took the masked top-k path. The mrope mask compares 2-D extents only between a cell and a query at the same linear position, and with one sequence in the cache the selection already ranks the cells in the mask's own order (position, then y, then x). Such a cell is therefore harmless unless a query of the ubatch sits at its position. Several sequences are not ranked, so any 2-D cell still keeps the masked path there, as do 2-D queries. Text-only caches are unchanged. Assisted-by: Claude
…mages An mrope image puts all its cells at one position, so qsa_apply rejected its ubatch and qsa_recover refused its cells: after the first image in a conversation every later decode and verify step re-pooled all indexer blocks on the GPU and rebuilt the block metadata with the per-cell scan on the host, both in proportion to the context. The tracked prefix now keeps the cells in the mask's order (position, then y, then x), the order the selection ranks them in once an image is present, and the ranks of a 1-D ubatch after the image are the rank blocks the selection uses. Tails and block starts are ranks; a block's pooled key takes the rope position the selection gives it (its first cell's key for rank blocks, the block start otherwise). Image ubatches are tracked but do not match, so they keep the scan. A prefix with a position hole and no image still matches no selection layout and stays invalid, the prefill scoring bound stays position-only, and a ranked prefix only matches while its sequence is all the cache holds. Recovery (qsa_rebuild_prefix, which skips the cells of other sequences since #116) orders the cells of the sequence the same way: a counting sort by position, then y and x among an image's cells. It builds the prefix through apply, so a recovered prefix has the same ranks, keys and blocks as one kept ubatch by ubatch. A cell shared with another sequence is still refused. A ranked prefix is used only while its sequence is alone in the cache; with other sequences resident the selection does not rank, so those ubatches keep the scan. Text-only caches keep rank == position and are unchanged. Assisted-by: Claude
…image) An MTP draft context never receives an image's cells, so its positions skip each image's extent. Its prefix was invalid from the first image on, and every later draft step re-pooled all indexer blocks and rebuilt the metadata with the per-cell scan, in proportion to the context. Without an image in the cache the selection pools position blocks and numbers the complete ones in order; a block with a hole is never pooled, and a query's tail is its own position block with -1 for missing cells. The tracked prefix now keeps that list of complete blocks (first rank, and the start the selection compares and ropes with), so the closed-form metadata and the incremental pooling work on such a prefix too. Rank blocks after an image are the same list. The pooled keys are rebuilt if the numbering changes under existing blocks, and a prefix with holes only matches while its sequence is all the cache holds. The prefill scoring bound keeps its old rule. Recovery accepts such a prefix too: test-qsa-recover's hole case now expects position blocks instead of a refusal. Text-only caches (no holes, no image) are unchanged. Assisted-by: Claude
test-qsa-recover now rebuilds seq 0 after every random step of a cache it shares with two other conversations, cells in any order: text, mrope images, position holes (an MTP draft), rollbacks, and a cell shared for a while (refused). The rebuilt tracker must equal the incremental one in every field the selection metadata reads (cells, ranks, mrope keys, first duplicate and gap, complete blocks and their starts). Assisted-by: Claude
pwilkin
approved these changes
Oct 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Rebase of #163 (
qsa : keep the incremental QSA prefix across mrope images, by David Hauf) onto current master54e9dccd4. #163 is based on6aa81b88eand no longer applies: commits 2 and 3 conflict insrc/llama-memory-hybrid-idx.cppandtests/test-qsa-prefix.cpp. This PR carries the same four commits, authorship kept. The design, the selection policy for image histories and the earlier measurements are described in #163 and are not repeated here.This rebase is opened because the bug that #163 fixes is now also a crash on master, not only a slowdown:
qsa_scalar_visibility_cellsrejects every later ubatch, so prefill after an image takes the masked route. Its compute buffer grows with ubatch x cache depth (about 40 bytes per query x key).-ub 8192, a multimodal request on Signal-3.8-Flash-Next fails at about 78K tokens of depth:allocating 28058.43 MiB on device 0: cudaMalloc failed: out of memory->decode: Compute error. A user reported the same failure at 148K (allocating 46195.12 MiB).What changed in the rebase
git range-diff 6aa81b88e..0a3e66dd2 54e9dccd4..HEAD:=).qsa_text_pos()(an M-RoPE text token is[p, p, p, 0], axis 3 stays zero). The qsa : keep the incremental QSA prefix across mrope images (target and MTP draft) #163 code compared all axes withp, which on current master rejects every text token. Both places inqsa_prefix_matchesnow callqsa_text_pos(). Theqsa_applyloop takes the qsa : keep the incremental QSA prefix across mrope images (target and MTP draft) #163 version (it readsyandxfrom axes 1 and 2 only, so axis 3 does not matter there; it also tracks the image rows of a mixed ubatch, which master added after qsa : keep the incremental QSA prefix across mrope images (target and MTP draft) #163).tests/test-qsa-prefix.cpp: master'stest_scan_foreign_imageis kept next to the newtest_mrope_ranksandtest_position_holes.No other change.
Measurements
Server, the reporter's command with the context reduced to 261888 (see Additional information):
Requests, one server per build, sent in this order, non-thinking,
max_tokens32. Each prompt is numbered text lines and asks for the number on the last line. The image is a 1024x1024 ImageMagickplasma:fractalPNG (sha256 ae985c46...).54e9dccd42749(correct), 1149 t/s2749, 1151 t/s2749, 1153 t/s12749(correct), 855 t/s12749, 1130 t/s12749, 1131 t/sallocating 28058.43 MiB,Compute error26999(correct), 1065 t/s26999, 1066 t/sPrompt t/s is the server's
timings.prompt_per_second, one run per cell. On master the time per 8192-token ubatch grows with depth on B (8 s at 25K -> 10 s at 58K) and stays flat on A (7 s from 16K to 61K).These are single runs of an end-to-end reproduction, not a benchmark: no performance claim is made from them beyond "master fails, the change does not". The decode gains reported in #163 were not re-measured on this rebase.
Correctness:
test-qsa-prefix,test-qsa-fast,test-qsa-recover: PASS on the rebased head (ROCm build).Additional information
git range-diffabove shows the only differences). The maintainer decision requested in qsa : keep the incremental QSA prefix across mrope images (target and MTP draft) #163 on the selection for image histories (complete blocks plus own tail, vs master's 2051 cells with a top-k tie-break) still applies here unchanged.--ctx-size 262144disables QSA block selection for every ubatch above 512 tokens, because the prefill depth bound inqwen4exp_use_block_selectionisn_kv <= 262140and the context is padded to a multiple of 256 (262140 -> 262144). With-ub 8192the masked route then needs an 83,796.50 MiB compute buffer at startup and the model does not load.--ctx-size 261888avoids it (2,806 MiB compute buffer at-ub 8192). qwen4exp : 1M context (uint32 QSA block ids, depth-bounded indexer scratch, RoPE-scaled slots) #215 and cuda : widen QSA selected-key block ids to uint32 and lift the depth bound to 1M #124 change that bound.Requirements
tests/corpus/correctness-*.txtcorpora,llama-perplexity, deterministic multi-step decode,test-backend-ops(no ggml change in this PR) and the full ctest suite;llama-benchPP2048 at depths 0 / 12000 / 32000 / 64000 andllama-benchywith MTP;-ubother than 8192,--kv-unified/ several slots, quantized KV;