Pull requests / #743
#743 qsa_select: coalesced streaming top-k for long Turing prompts (+23% 248K prefill)
closed · @konijiwa110 · 0 Kommentare · Auf GitHub
BenchmarksNVIDIA / CUDAModels & quants
Beschreibung
## What The register top-k holds at most `TK_T * TK_PER` blocks (~135K cells). Past that the prompt path falls back to `block_topk_kernel`, whose per-thread runs read the scores uncoalesced: 12.8 ms per 256 queries at 250K on an RTX 2080 Ti, the largest prefill phase there at that length. `block_topk_stream_kernel` reads the keys from memory on every radix pass with block `t + i * TK_T` (coalesced), so it has no register limit, then emits the cells in ascending chunks of `TK_T` blocks with two block scans (ties to the lowest index), giving `block_topk_kernel`'s ids. It is taken only where the active bound is already used (Turing prefill, or `STRATA_TOPK_ACTIVE_ANY=1` in tests) above 22,528 blocks (~90K cells). Decode and other cards are unchanged. `STRATA_TOPK_STREAM=0` keeps the previous dispatch. `qsa_topk_active_parity` gains cases either side of the 22,528-block bound, heavily tied scores at 200K and all-zero scores at 262K (20 cases), comparing against the capacity-dispatch reference. ## Measured RTX 2080 Ti, 256 queries, capacity 262144, top-k with the active bound (`STRATA_TOPK_STREAM=0` / unset): | context | before | stream | |---|---|---| | 90K | 0.49 ms | 0.48 ms | | 131K | 1.03 ms | 0.75 ms | | 250K | 12.8 ms | 1.13 ms | Selections identical to the reference kernel 256/256 at each length. `qsa_topk_active_parity`: PASS (20 cases, ids bitwise identical), with the stream on and off. End to end, Swift 1.5 IQ3_XXS, same build A/B, prefill tok/s: 248K 951 -> 1167 (+23%), 155K 1166 -> 1210 (+4%), 128K and below unchanged. `qsa select` at 248K: 70.9 s -> 22.2 s. (These runs had #742's scorer on as well.) Only tested on sm_75. Independent of #741 and #742 (merges cleanly with both). Developed with an AI coding assistant; all numbers measured on an RTX 2080 Ti 22 GB / i5-12490F / 64 GB. 🤖 Generated with [Claude Code](https://claude.com/claude-code)
Mehr auf der Site
Links zu Install, Modellen, Releases.