Pull requests / #575
#575 qsa select: split the top-k past the register fit - 9% faster long-prompt prefill
closed · @orangeswim · 0 comentarios · En GitHub
AMD / HIPNVIDIA / CUDAModels & quantsWindows
Descripción
**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
En el sitio
Enlaces a install, modelos, releases.