perf(qwen3.8): split the flash-decode KV range across warp-groups

[widecloud]

Summary

The GQA-6 flash-decode block runs six warps. At a small split count the launcher issues
dim3(num_kv_heads * n_splits, num_seqs) = ~4 CTAs, so the whole decode attention runs on four
SMs of 170. #874 hid the staging latency; it did not change how few warps do the walking. Split
the block's KV range across three warp-groups, each with its own tile and its own (m, l, acc),
merged once at the end.

Proof of speedup

• Tested on RTX 5090 (sm_120)

Decode tok/s (dspark_tau_check, ctx=4096, 128 new tokens — the scored harness):

	decode tok/s
before (main)	42.39
after (this PR)	47.63

gate
mean accept τ	1.0847 → 1.0847, unchanged to 4 dp
LOSSLESS	1 on every run
AR_TPS	47.89 → 55.41, +15.7% (no regression; this is the kernel speedup)

main                 DSPARK 42.3629  AR 47.7918  MEAN_ACCEPT 1.0847  LOSSLESS 1
KVG off (control)    DSPARK 42.4143  AR 47.8603  MEAN_ACCEPT 1.0847  LOSSLESS 1
KVG=2 KT=16          DSPARK 46.4792  AR 53.8132  MEAN_ACCEPT 1.0847  LOSSLESS 1   +9.7%
KVG=2 KT=20          DSPARK 46.8122  AR 54.1784  MEAN_ACCEPT 1.0847  LOSSLESS 1  +10.5%
KVG=3 KT=12 (env)    DSPARK 47.6699  AR 55.4242  MEAN_ACCEPT 1.0847  LOSSLESS 1  +12.5%

and as the COMPILED DEFAULT with no env set anywhere -- which is what the bot runs --
over two ALTERNATING pairs (alternating, not blocked: the GPU drifts thermally across a
session, so a blocked A/A/B/B would manufacture a delta from the drift):
main     42.4048, 42.3751   median 42.3900   AR 47.8903, 47.8301   tau 1.0847
this PR  47.6138, 47.6457   median 47.6298   AR 55.4147, 55.4473   tau 1.0847

  ratio 1.1236 (+12.4%), ranges DISJOINT (main best 42.4048 < PR worst 47.6138)
  tau identical on all four runs; LOSSLESS=1 on all four

Control and main agree to 0.12%, so the kernel is inert when disabled and the
comparison is not measuring a build difference. Every arm is the SAME binary,
selected by SPARKINFER_FA_KVG / SPARKINFER_FA_KVG_TILE.

The gain is monotone in group count, which is the diagnosis restated: the block was
warp-starved, and each warp-group added recovers more of the walk.

τ is deliberately the first row. This PR cannot buy throughput with acceptance: it never touches
the draft, the proposal depth, the engage thresholds or the idle logic. AR_TPS — which does not
depend on acceptance at all — rises 12.1%, and that is the honest measure of the change.

Mechanism

Each of the KVG warp-groups walks a contiguous stripe of the block's KV range with its own
double-buffered cp.async tile and its own running (m, l, acc). At the end the groups are folded
with the standard log-sum-exp rescale, in ascending group index. The block still emits exactly one
split's partials per q-head, so n_splits, the part_m/part_l/part_acc layout and the combine
kernel are all untouched.

Exactness. Within a stripe the token order is unchanged and the merge formula is fixed, so the
result is deterministic run to run. It is a different summation order than the single-group walk —
the same reassociation the existing n_splits > 1 combine already performs — and LOSSLESS holds on
every run because both legs of the harness share it.

Shared memory

sm_120 has 128 KB of L1/shared per SM, ~100 KB addressable as dynamic smem (NOT the 228 KB of
datacenter Blackwell). KVG=3 at KT=12 needs 72 KB of tiles plus ~19 KB of merge scratch = 91 KB.
That is what bounds the group count: KVG=4 needs KT<=8, and KVG=3 at KT>=16 does not fit. Opt-in is via
cudaFuncSetAttribute and falls back to #874's kernel if refused, so a device with a smaller cap
keeps exactly today's behaviour.

Scope

• GQA-6 hd256 bf16 KV only, num_seqs == 1, and only where perf(dspark): overlap KV staging with compute via cp.async (1.23x dspark-decode@4k) #874's pipeline is already selected.
• int8 KV is untouched (it takes the MMA path).
• SPARKINFER_FA_KVG=1 restores main's kernel for A/B in one binary; the launcher also
  cascades 3 groups -> 2 -> perf(dspark): overlap KV staging with compute via cp.async (1.23x dspark-decode@4k) #874's kernel if an opt-in is refused.

Overlap with recent PRs

Builds on #872 (deep tile) and #874 (cp.async pipeline) rather than overlapping them: both changed
what one block does with its stripe, this changes how many warps a block puts on the range. No
shared code beyond the kernel both are in.

[skyrocket2026]

sparkinfer DSpark auto-eval — eval-dspark:L

metric	value
label	eval-dspark:L
scored at	DSpark speculative decode @ ctx=4k on the ModelOpt NVFP4 checkpoint
PR DSpark tok/s	48.13
main DSpark tok/s	43.06
speedup vs main	1.118× (+11.7%)
PR AR tok/s (floor)	54.33
main AR tok/s (floor)	47.39
AR vs main (floor)	+14.6%
DSpark vs AR	0.886× — above 1.0 means speculation finally pays
mean accept τ	1.085 (main 1.085, ceiling 7)
accuracy gate	✅ top1=1.000 (bar >=0.9) · KL=0.0000 (bar <=0.1)
losslessness gate	✅ DSpark matches the AR reference token-for-token, verified across 3 independent runs
mean accept τ floor	✅ 1.0847 vs main 1.0847 (bar ≥95%)
qwen3.8 (ModelOpt) guard @16k	✅ no regression (decode+prefill) — decode 89.4 tok/s · prefill 13544 pp
qwen3.6 guard @16k	✅ no regression (decode+prefill) — decode 474.3 tok/s · prefill 28358 pp
PPL PR / main	3.204 / 3.204
Polaris receipt	collected, not signed (no key configured)
commit	9a13964df

ok

_{Scored on the pinned eval box vs same-box origin/main: DSpark speculative decode throughput at ctx=4k on the ModelOpt NVFP4 checkpoint, with the AR reference measured in the same process and the same model load. Both a regression in AR decode and any divergence from the AR token sequence are hard REJECTs — a speculative decoder that is fast because it skips verification is not faster, it is wrong. τ is the lever, and the row above reports it against a block_size of 7. This is informational, not a judgment on your PR: a none label just means no measurable DSpark decode@4k speedup was verified, which is expected and fine if that isn't what your change is about. Automated — merge behaviour depends on SPARKINFER_DSPARK_AUTOMERGE.}

[skyrocket2026]

Auto-merged as the round's dspark-merge-first winner — verified same-box 128-token decode speedup over main, accuracy-gated vs llama.cpp.