perf(dspark): overlap KV staging with compute via cp.async (1.23x dspark-decode@4k)

[James-CUDA]

Summary

The GQA-6 KV staging loop is strictly serial — stage -> __syncthreads -> compute -> __syncthreads
— so DRAM latency is exposed once per tile. Staging a deeper tile put more loads in flight per
barrier but left the two phases serial. Double-buffer the staged K/V and issue the next tile's
global->shared copies with cp.async before consuming the current one, so the copy runs underneath
the compute.

Proof of speedup

• Tested on RTX 5090 (sm_120)

Decode tok/s (end-to-end, ctx=4096):

	decode tok/s
before (main)	30.31
after (this PR)	37.32

Prefill pp tok/s — n/a, this PR targets decode only.

Numbers come from runtime/examples/dspark_tau_check.cpp, not bench/scripts/bench.sh: the latter
benchmarks a GGUF, and the path this PR touches is the NVFP4 DSpark decode. dspark_tau_check is
end-to-end (a full generation, AR and speculative timed in one process), not an isolated-kernel
microbenchmark, and it reports the losslessness gate alongside the throughput.

Same binary on every arm, env-selected, so build and machine variance are out of the comparison.
SPARKINFER_FA_PIPE=0 is main's behaviour exactly; the shallow-tile arm is included only to show
where the remaining headroom was:

$ SPARKINFER_FA_TILE_DEEP=1 SPARKINFER_FA_PIPE=0 build/runtime/dspark_tau_check \
      models_q38_modelopt dspark 4096 <4096 prompt ids>      # main
$ SPARKINFER_FA_TILE_DEEP=1 SPARKINFER_FA_PIPE=1 ...          # this PR

arm              dspark_tps  ar_tps    accept    lossless
tile8            26.5467     26.6901   1.0828    1
tile32 (main)    30.3137     30.6892   1.0828    1
tile32+pipe      37.3181     38.3079   1.0828    1

1.23x over main. LOSSLESS 1 on every arm and acceptance identical to four decimals (1.0828) —
pure speed, no numerical drift.

Why there was still headroom

nsys at ctx=4096 attributes 61.2% of decode GPU time to fa_split_gqa_kernel at 1.11 ms per
call. At n_splits=1 the grid is one block per (seq, kv_head, split) = 4 blocks on 170 SMs, and the
staging loop is plain __ldg: blockDim(192) x 16 B is ~12 KB of loads in flight before each
barrier against ~700 ns of DRAM latency, so the block sustains ~7.5 GB/s. The roofline for the bytes
it reads is 18.7 us (4096 keys x 4 kv-heads x 256 x 2 x 2 B = 33.6 MB at 1792 GB/s). A deeper tile
closed part of that; overlapping the copy with the compute closes more.

Why it is bit-identical

cp.async changes when bytes land in shared memory, not their values, and the token walk is
unchanged: tiles are consecutive, the mainloop still walks start..end in order, and the
online-softmax update is per token in that order. Every row therefore accumulates the same keys in
the same sequence.

Cost and fallbacks

A separate kernel (fa_split_gqa_pipe_kernel), so the synchronous paths keep their exact code.
bf16 KV only — cp.async copies bytes verbatim and cannot dequantize, so an int8 pool still takes
the dequant-into-smem path.

Double buffering needs 4 * TILE * 256 * 2 B = 64 KB at TILE=32, past the 48 KB default, so the
launcher opts in once via cudaFuncSetAttribute. If that is refused the error is cleared and the
synchronous tile runs instead — the fallback is prior behaviour, not a failed launch.
SPARKINFER_FA_PIPE=0 restores synchronous staging.

[skyrocket2026]

sparkinfer DSpark auto-eval — eval-dspark:XL

metric	value
label	eval-dspark:XL
scored at	DSpark speculative decode @ ctx=4k on the ModelOpt NVFP4 checkpoint
PR DSpark tok/s	43.05
main DSpark tok/s	36.28
speedup vs main	1.187× (+18.6%)
PR AR tok/s (floor)	47.37
main AR tok/s (floor)	39.24
AR vs main (floor)	+20.7%
DSpark vs AR	0.909× — 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
qwen3.8 (ModelOpt) guard @16k	✅ no regression (decode+prefill) — decode 89.5 tok/s · prefill 13551 pp
qwen3.6 guard @16k	✅ no regression (decode+prefill) — decode 473.7 tok/s · prefill 28348 pp
PPL PR / main	3.204 / 3.204
Polaris receipt	collected, not signed (no key configured)
commit	3d2bbed5a

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.