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 コメント · GitHub で見る

BenchmarksAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

本文

## 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.

関連リンク

インストール・モデル・リリースへの站内リンク。