feat: add Qwen3.5 GDN ReplaySSM

[tuanzhangCS]

Summary

• add ReplaySSM support for Qwen3.5 GDN during MTP decoding
• retain per-draft-token replay inputs instead of materializing speculative SSM states
• replay the accepted prefix after verification with one all-layer commit kernel
• fuse replay payload capture into the QKV split and use a cooperative 4-warp CTA for 128x128 GDN state heads
• keep the existing Kimi K3 KDA lazy-replay path unchanged

Motivation

The previous MTP path saved a full draft SSM state for every speculative token. This makes scratch memory scale with the full recurrent state and adds state-copy work even though verification only commits an accepted prefix.

ReplaySSM stores the per-token K/V/a/b inputs needed to reconstruct the recurrence. Once verification determines the accepted length, the commit path recomputes only the accepted state transitions.

Implementation

• add Qwen3.5 GDN replay payload allocation and paging
• capture replay K/V/a/b as part of the fused QKV split
• store replay payloads and state pools in layer-indexed buffers
• launch a single replay kernel over all Qwen3.5 GDN layers
• split 128x128 state-head work across a cooperative 4-warp CTA
• add runtime and kernel correctness coverage plus model documentation

Performance

Model: nvidia/Qwen3.5-122B-A10B-NVFP4 on NVIDIA B200.

Common settings: input length 128, requested output length 128, greedy decoding, MTP 3 steps / 4 draft tokens, CUDA Graph, FP8 KV cache, three measured runs with the median reported. MTP produced 133 output tokens per sequence in both paths; throughput uses the actual token count.

TP1

Batch	Baseline latency	ReplaySSM latency	Latency reduction	Baseline throughput	ReplaySSM throughput	Throughput gain
1	565.933 ms	550.037 ms	2.81%	235.0 tok/s	241.8 tok/s	2.89%
4	650.307 ms	591.259 ms	9.08%	818.1 tok/s	899.8 tok/s	9.99%
8	732.119 ms	659.910 ms	9.86%	1453.3 tok/s	1612.3 tok/s	10.94%
16	904.146 ms	790.027 ms	12.62%	2353.6 tok/s	2693.6 tok/s	14.45%

TP4

Batch	Baseline latency	ReplaySSM latency	Latency reduction	Baseline throughput	ReplaySSM throughput	Throughput gain
1	561.955 ms	546.627 ms	2.73%	236.7 tok/s	243.3 tok/s	2.80%
4	624.124 ms	607.500 ms	2.66%	852.4 tok/s	875.7 tok/s	2.74%
8	659.460 ms	630.253 ms	4.43%	1613.4 tok/s	1688.2 tok/s	4.63%
16	726.453 ms	697.810 ms	3.94%	2929.3 tok/s	3049.5 tok/s	4.11%

For TP4, the target NVFP4 MoE backend remained FlashInfer TRT-LLM. The BF16 MTP draft MoE used Triton, deterministic communication was enabled, and all-reduce fusion was disabled because the available FlashInfer draft GEMM/custom communication path independently hit an illegal-memory-access error. ReplaySSM was the only A/B difference.

Output token counts and output-ID digests matched between baseline and ReplaySSM for every TP1 and TP4 batch.

Validation

• pre-commit run --all-files
• pytest -q test/runtime/layers/test_gdn_qkv_split_fused.py test/runtime/test_cache_setup.py test/runtime/test_gdn_state_paging.py test/runtime/test_qwen35_gdn_replay.py test/runtime/test_kimi_k3_kda.py — 58 passed, 3 skipped
• pytest -q test/ops/attention/test_gdn_replay_commit.py from tokenspeed-kernel — 11 passed
• full-model Qwen3.5 TP1 and TP4 ReplaySSM on/off A/B benchmarks with matching outputs