[imanu86]

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.

[imanu86]

@Niko1221 Companion to #2058, independent of it. STRATA_QSA_ROWS=1 computes the decode window's QSA block scores (up to 4 queries) with the queries' slices in registers and six shuffles per block instead of twenty, keeping warp_sum's pairings so scores and top-k ids are bit-exact (qsa_rows_bench --rows, memcmp, 16 cases on sm_75 and sm_86). On the 2080 Ti the call goes from 101 to 55 us at T=4; end to end on our forced 131k bench +3.2% / +0.1% over two B arms, text unchanged. The prompt path and HIP are untouched (#1825's territory). Ported and re-checked on v0.1.43. Off by default; glad to refine whatever you would like changed.

[Niko1221]

Thanks. Ported as an opt-in (STRATA_QSA_ROWS=1), text identical in our gate. On an RTX 3060 and an RTX A4000 it is level with the default (0.997 and 1.003 at 30K context), so the gain you measured is specific to the 2080 Ti. We keep it off by default.

[Niko1221]

Thanks for the pull request.

Your change is already in main for 0.2.0-beta, with your credit, so there is nothing left for you to do here. We're closing it because we're resetting the tracker for 0.2.0-beta: there are more than 1,300 open issues and pull requests, and we're closing everything that isn't being worked on right now.

If part of your work did not make it in and you want it considered, open a new pull request on top of main that follows CONTRIBUTING.md, or comment /reopen after fixing the title area prefix and the template.