[jianc99]

Stacked on #2 (ReplaySSM for DFLASH and DSPARK); this PR's own changes are the commits after it.

Motivation

With a bf16 checkpoint, the GDN ReplaySSM speculative path folds the accepted history into the checkpoint after every verify step (commit_gdn_replayssm_circular).

The commit kernel ran one CTA per (64-column V tile, request, value head x layer). On Qwen3.8-27B that is about 590k CTAs per step at batch 128 (48 GDN layers, 48 value heads per rank at TP1). Every CTA reloaded the request scalars, the ring positions and the key tile. It then folded the history with four chained tensor-core products: key and key residual, each times the high and low update parts.

On GB300 the commit took 5.40 ms per step at batch 128.

Modifications

• New launch grid. The grid is (request, layer, key head). A program handles the HV / H value heads that read its key head, in 32-column V tiles. The group size comes from the per-rank head counts in the tensor shapes (3 for Qwen3.8 at any TP), and the wrapper asserts HV % H == 0.

• One product per 16-step chunk. The history is folded 16 steps at a time. The key and key-residual rows and the high and low update parts of a chunk are stacked along the reduction axis, so each chunk is one tensor-core product with fp32 accumulation.

• Default ring. A ring of one chunk (the default 16 steps) loads its key tile once per program and folds each tile in one product, with at most 37 KB of shared memory and no spills.

• Ring-length fallback. The host picks the kernel from the ring length (--linear-replayssm-cache-len):

  • ring of 16 steps (the default, which fits DFLASH and DSPARK drafts of up to 16 tokens): the head-group kernel;
  • rings of 32 steps and more: the previous per-value-head kernel, kept unchanged.

  On a longer ring the head-group kernel folds chunk by chunk, so its time grows with the history; the per-head kernel folds the whole ring in one product and barely depends on the history. The ring length is a static shape fixed at startup, so the choice needs no sync, is the same on every TP rank and does not depend on the batch. The tests run every case through both kernels.

• fp32 rings carry no low part.

• Both changes are needed. Grouping heads alone gives no speedup, and stacking alone on the old grid is 44% slower.

• Fold and track products are computed for every row and masked at the store. Carrying the track accumulator through a branch on the track flag produced wrong values in the last loop iteration with software pipelining, so the kernel avoids it.

• Unchanged: slot and layer strides, null rows, the early exit for rows that neither fold nor track, the history-load masking and the cursor finish kernel.

• Numerics: the result differs from the previous kernel by fp32 summation order only: at most one bf16 ulp, or 2.8e-6 for fp32 checkpoints with bf16 rings (up to 1.6e-4 with fp32 rings, TF32 operands).

Grouping value heads per CTA follows SubSir's WY ring commit kernel (one CTA per request, layer and head group). Thanks to SubSir for the idea.

Accuracy Tests

• New test_gdn_replayssm_circular_commit.py. It checks folded and tracked checkpoints against a float64 reference of the closed form, through both commit kernels. Coverage:
  • value/key head ratios 1-4
  • fp32 and bf16 checkpoints
  • a partial V tile
  • contiguous, slot-strided and layer-inner state pools
  • null rows
  • rings of 16, 32, 64 and 128 steps
  • track rows on and off
  • bf16 rings with residuals, and fp32 rings without them
  • NaN in ring entries past the committed history
  • It also checks that untouched slots stay bit-identical and that the cursors are reset.
  • A dispatch test checks that ring 16 launches the head-group kernel and rings 32, 64 and 128 launch the per-head kernel.
• Through the spec commit path: in test_gdn_replayssm_direct_commit.py the chain tests use ring 16 (head-group kernel) and the lazy-fold test uses ring 32 (per-head kernel).
• Against the previous kernel: 466 random cases across the same axes. bf16 checkpoints differ by at most one ulp, in under 0.03% of elements; fp32 checkpoints by at most 2.8e-6 with bf16 rings; with fp32 rings both kernels fold TF32 operands and the checkpoints differ by up to 1.6e-4.
• Sanitizers: compute-sanitizer memcheck and initcheck report no device-side errors.

Speed Tests and Profiling

Commit kernel on GB300 (SM clock 2070 MHz): 48 layers, K = V = 128, ring 16, bf16 checkpoint, 4 history steps, every row folding.

heads per rank	batch	before	after
HV 48 / H 16 (TP1)	1	39.8 us	42.0 us
	8	341 us	253 us
	128	5.40 ms	3.85 ms
HV 12 / H 4 (TP4)	1	11.5 us	17.9 us
	8	90 us	72 us
	128	1.35 ms	0.97 ms

Batch 1 is slower: +2.2 us at TP1 and +6.4 us at TP4, because there are fewer, longer programs (with an fp32 checkpoint at TP4, 16.8 -> 26.5 us). The kernel choice deliberately ignores the batch: the two kernels differ in fp32 summation order, so a batch rule would make a request's checkpoint depend on the batch it runs in, to save at most about 12 us per step (TP4 heads, fp32 checkpoint, track rows: 45.1 vs 33.5 us for the per-head kernel).

Track snapshots: with half of the rows taking a snapshot, batch 128 at TP1 goes from 10.07 to 7.37 ms.

Ring-length fallback. Both kernels on every ring, batch 128, TP1 heads, bf16 checkpoint, bf16 ring with residuals, every row folding (CUDA graphs, best of 3). The kernel that runs is in bold. Without the fallback the head-group kernel ran on every ring; with it, rings above 16 run the per-head kernel, so their commit time is the same as before this PR.

ring	history steps	head-group kernel	per-head kernel
16	4	3.87 ms	5.45 ms
16	16	4.09 ms	5.42 ms
32	4	7.58 ms	5.79 ms
32	32 (full)	15.84 ms	5.84 ms
64	4	8.17 ms	8.47 ms
64	16	9.25 ms	8.55 ms
64	64 (full)	34.89 ms	9.62 ms
128	4	8.32 ms	13.81 ms
128	16	9.15 ms	14.00 ms
128	128 (full)	65.43 ms	17.30 ms

fp32 checkpoints, track rows, batches 1 and 8 and TP4 heads follow the same pattern (264 cases): on ring 32 the per-head kernel is 1.2-3.9x faster for every history, and on full rings 2.4-7x faster. A nearly full ring is what a fold at capacity commits; only fp32 checkpoints with the lazy fold (generic EAGLE/MTP) fold that much at once.

Known cost: on rings 64 and 128 with short histories (4-16 steps), the per-head kernel is up to 1.7x slower than the head-group kernel alone would be (in the table above, ring 128 with 4 steps: 13.81 ms instead of 8.32 ms; with an fp32 checkpoint up to 16.87 instead of 10.77 ms). A raised ring folds only short histories when every commit folds, where the ring only has to fit one draft window, or in a track snapshot soon after a lazy fold. The history length lives on the device, so the host cannot pick the kernel by it.

End to end, on top of the DFLASH ReplaySSM commit branch: Qwen3.8-27B, TP1, DFLASH with an 8-token draft block, bf16 checkpoint, default 16-step ring (head-group kernel), radix cache off, 128 max running requests, T=1, three runs each. These were measured with an earlier revision of this kernel that has the same ring-16 arithmetic and the same kernel time.

Decode tok/s = per-request decode rate (output tokens after the first / time after the first token), times concurrency.

concurrency	before, decode tok/s	after, decode tok/s	change
64	6701 / 6685 / 6675, mean 6687	6780 / 6757 / 6755, mean 6764	+1.2% (run ranges do not overlap)
128	7741 / 7802 / 7767, mean 7770	7849 / 7950 / 7862, mean 7887	+1.5% (run ranges do not overlap)

Accept length is unchanged. At c128 the running cap equals the concurrency, so about one request waits in the queue; #8 shows that this setup reads 7-15% higher at that level than a larger cap, on both sides, so compare c128 only within this table.

Related upstream work

• [Fix] Mask unused ReplaySSM history before checkpoint arithmetic sgl-project/sglang#42688 masks the history loads of the same commit kernel ([Feat] gdn: run DFLASH and DSPARK verify on the ReplaySSM spec ring #2 carries the same masks), and also masks the keys and key residuals past the tracked prefix before the track product. The head-group kernel zeroes the update parts past the tracked prefix but multiplies them with keys masked only to the committed history, so a non-finite key between the tracked prefix and the end of the history can still reach a track snapshot. What to do: if [Fix] Mask unused ReplaySSM history before checkpoint arithmetic sgl-project/sglang#42688 lands first, this PR rebases onto it and carries its track-prefix key masking into the head-group kernel; if this PR lands first, [Fix] Mask unused ReplaySSM history before checkpoint arithmetic sgl-project/sglang#42688 needs the same change in both kernels.