Pull requests / #575

#575 qsa select: split the top-k past the register fit - 9% faster long-prompt prefill

closed · @orangeswim · 0 Kommentare · Auf GitHub

AMD / HIPNVIDIA / CUDAModels & quantsWindows

Beschreibung

**The problem.** `block_topk_reg_kernel` holds one query's keys in registers, which caps
it at `TK_T × TK_PER` = 33,792 blocks — contexts past ~135K cells, and the 204,800
context's 51,202 blocks. Past that, `qsa_block_topk` falls back to `block_topk_kernel`,
which re-reads the whole score row six times (four radix passes + count + emit) against
the register kernel's one read. `qsa select` was 19% of a cold 198K prefill on the 3090,
and every prompt pays it whatever its length: the dispatch keys off the context's
capacity, not the prompt. The sm_75 counted dispatch (0.1.38) bounds by the prompt's
active blocks; nothing covers prompts that genuinely exceed the registers.

**The change.** Past the fit but within two of them, the range splits at the fit: the
register kernel runs on each half (its block range is now a parameter; whole-range calls
are unchanged), and `block_topk_merge_kernel` rebuilds the whole-range selection from the
halves' candidates — same radix rule, one count per cell instead of per weighted block,
ties bounded by a single id, each selected cell placed at its rank by id. The rule is
unchanged, so the ids must be. `block_excl_scan` gains a per-call zero of the warp slots
a 512-thread CTA leaves unfilled (1024-thread callers unaffected); the scratch is one
stream-ordered allocation per call, shared by nothing and held past the merge by nothing.

**Scope.** Prompt path only — its caller passes an active block count, decode's captured
windows pass none. HIP keeps its own dispatch. `STRATA_TOPK_OLD=1` still forces the
reference as the A/B arm. Geometries past a 2,051-wide selection or a 256-query batch
fall back.

**Correctness.** `qsa_topk_active_parity` (CMake target, `add_test`, exit 77 without a
GPU) drives both paths and diverges exactly at this branch: active-bound → counted →
split, no-bound → capacity → reference. 19 cases, bitwise identical, on both sides of
the fit — 135,164 sits at the fit, 135,168 and 135,171 split (the last puts the final
query's half B at one block, 3 cells, and leaves the earlier queries' half B empty) —
with NaN keys on every fifth block, where order_key maps NaN to 0 and a NaN tail loses
its +1e9 boost. End to end, the 26-needle suite's finds and misses are unchanged against
stock.

**Measured** (RTX 3090, Ryzen 7 5800X, 32 GB RAM, Coder IQ1_M pack, ctx 204,800, cold,
seeds fixed; this tree ± the commit):

| | stock | + split | |
|---|---|---|---|
| 100K prompt | 2,394 t/s | 2,437 t/s | +1.8% |
| 198K prompt | 2,190 t/s | 2,368–2,388 (4 runs) | +8–9% |

Retrievals and decode (~101 t/s) identical on both builds. The gain scales with the
reference kernel's re-reads, and because the old fallback keyed on capacity, it covers
every prompt size on an affected context — including short prompts, which on sm_86 have
no counted dispatch and still take the split.

## AI Tools
Claude + GLM-5.3, Strata qwen3.8-flash-next-coder-iq1_m

Mehr auf der Site

Links zu Install, Modellen, Releases.