Pull requests / #2059
#2059 QSA decode: the window's block scores with the queries in registers and six shuffles per block (STRATA_QSA_ROWS=1, opt-in, bit-exact)
open · @imanu86 · 0 comments · View on GitHub
BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows
Description
## Title Issue: none (an opt-in speed change). ## Summary `STRATA_QSA_ROWS=1` (opt-in, off by default) computes the verify window's QSA **block scores** (decode, up to 4 queries) with a kernel that keeps the queries' slices in registers and reduces the four heads with six shuffles per block instead of twenty. Scores and top-k ids are **bit for bit** the multi kernel's (`qsa_rows_bench --rows`, `memcmp`, 16 cases per card). CUDA only; the prompt path (#1825's territory) and HIP are not touched. ## Why In a 131k decode the QSA kernel that grows most with the window's rows is `block_scores_multi_kernel`: ~20 us more per call for every row of the window on both an RTX 2080 Ti and an RTX 3060 (nsys, 8 QSA layers on the second card of a layer split). The keys are already read once for all the queries; the cost per row is in the shared-memory and shuffle pipes - per (key block, query) 4 x `LDS.128` of the query slices and 5 shuffles per head (20 per block and query), at 32,769 blocks ~15-20 us of issue per query. ## What changed - `src/kernels/cuda/qsa_select.cu`: `block_scores_rows_kernel<NQ>` (NQ = 1..4): the NQ queries' slices (4 heads x 4 floats per lane) in registers for the whole kernel (44 / 63 / 78 / 94 registers, no stack); the four heads' lane products reduced with a scatter that keeps `warp_sum`'s pairings (xor 16 and 8 leave one head per lane, then 4, 2, 1) and gathered on lane 0 in head order. Tail block and the `+1e9` as before. `qsa_block_scores` takes it for the calls without `active_blocks` with up to 4 queries when `STRATA_QSA_ROWS=1`; otherwise, or beyond 4 queries, the multi kernel runs as before. - `include/strata/kernels/qsa_select.hpp`: the `rows` entry. - `src/kernels/qsa_rows_bench.cpp` (`EXCLUDE_FROM_ALL`): what the T rows of a verification window cost in a QSA layer's decode kernels at a long context (block scores, top-k, attention over each row's selection, indexer append; synthetic data at the real geometry, int8 pools, no model); `--rows` runs the new kernel against the multi kernel, scores and top-k ids bit for bit at T = 1..4 at 131k and at a short context with a full tail block, and prints both kernels' microseconds. - `docs/QSAROWS.md`. ## Bit for bit Every addition in the new reduction adds the same two operands of the same `warp_sum` level (float addition commutes), the product expression is the same text in the same file, the relus are summed in head order. `qsa_rows_bench --rows`: 8 cases per card (T 1-4 x two contexts), 0 differ, on sm_75 and sm_86. In the engine the greedy text and the drafter's acceptance are unchanged. ## Measurements Hardware: RTX 2080 Ti 22 GB (sm_75) + RTX 3060 12 GB (sm_86), Windows 11, CUDA 12.6, layer split, Qwen3.8-Flash-Next IQ3_XXS pack, 131k context, `--max-context 262144 --kv int8`. `qsa_block_scores`, us per call at 131k (both kernels read the same pooled keys): | | T=1 | T=2 | T=3 | T=4 | per row | |---|---|---|---|---|---| | RTX 2080 Ti, multi (reference) | 47.2 | 56.9 | 78.2 | 101.2 | +18.0 | | RTX 2080 Ti, `STRATA_QSA_ROWS=1` | **36.2** | **38.0** | **49.2** | **54.6** | **+6.1** | | RTX 3060, multi | 73.4 | 109.4 | 151.9 | 193.5 | +40.0 | | RTX 3060, `STRATA_QSA_ROWS=1` | **63.3** | **75.5** | **98.0** | **130.2** | **+22.3** | In the engine (nsys): per window the second card saves 0.09 ms at T=1 and 0.37 ms at T=4 (8 QSA layers), the first 0.04 / 0.25 ms (4 layers). End to end, deterministic 131k bench (forced tokens and window sizes, 1024 decode tokens, three interleaved runs lever on / off / on): **+3.2% and +0.1%** decode tokens/s in the two B arms, text identical. Together with the companion experts lever (`STRATA_EXPDEC=1`, separate PR): +2.9%. Ported and validated on v0.1.43 (micro-bench bit-exact on both cards; our full stack on 0.1.43 matches our 0.1.42 daily on the forced 131k bench). ## Risks / limits - Off by default; `STRATA_QSA_ROWS` must be exactly `1`. Only the decode calls with <= 4 queries and no `active_blocks` change; everything else runs the kernels it ran before. - On the RTX 3060 at T=4 the kernel is at 94 registers (2 CTAs per SM); a two-pass version of 2 queries would stay at 63 if that card matters more. - Not measured on sm_80/89/120 or HIP. **Scope.** The gain comes from VRAM-resident routed experts / decode QSA scores on Turing/Ampere setups where these kernels are a large share of the verify window. On rigs where most routed experts are off-GPU (e.g. 4x RTX 5080 16 GB, see the test in #2058) the window is dominated by stage hand-offs and host experts, and the levers measure neutral.
Related on strata.com
Editorial links to help you install, pick models, or read release notes — not part of the upstream thread.