[huangwei021230]

Motivation

ReplaySSM speculative-state folding already supports GDN when speculative decoding commits through commit_mamba_states_after_verify(). DFlash and DSpark use a different path and call HybridLinearAttnBackend.update_mamba_state_after_mtp_verify() directly after target verification.

Before this change, the direct commit path recognized ReplaySSM only for KDA. A GDN target therefore fell back to the regular intermediate-state scatter path, even though --enable-linear-replayssm-spec intentionally removes intermediate_ssm. The KV-cache configurator prevented this invalid path by rejecting every non-KDA DFlash/DSpark target.

This PR routes the built-in GDN backend through the existing ReplaySSM fold implementation from #28695. It does not change the ReplaySSM kernels, model outputs, public flags, or the generic speculative-decoding commit path.

Modifications

• Resolve the direct ReplaySSM verify-commit handler when HybridLinearAttnBackend is initialized.
• Select the KDA handler only when speculative folding is enabled, and select the GDN handler only for a concrete GDNAttnBackend.
• Keep the existing intermediate-state scatter fallback for every other backend and for ReplaySSM-disabled execution.
• Replace the pool-local KDA-only rejection with validation scoped to DFlash/DSpark target workers before memory-pool sizing.
• Fail fast for unsupported direct-verify configurations: non-GDN/KDA targets, non-CUDA/ROCm devices, unified memory, and non-static ragged verification for GDN. KDA compact-ragged verification remains supported.
• Add a focused CPU regression test for GDN/KDA handler selection, constructor wiring, direct commit dispatch, and configuration guards.

Accuracy Tests

The following focused and neighboring regression tests pass:

Test	Result
New GDN direct-commit regression test	3/3 passed
GDN ReplaySSM speculative fold parity	4/4 passed
GDN/KDA ReplaySSM packed-decode parity	6/6 passed
GDN attention backend regression suite	9/9 passed
Mamba2 attention backend regression suite	9/9 passed
Hybrid FlashInfer metadata regression suite	2/2 passed
NGRAM/Mamba verify-update regression suite	14/14 passed
ReplaySSM ring accounting	3/3 passed

The GPU tests were run on a single NVIDIA H20. The clean end-to-end A/B test used Qwen3.8-27B as the target, Qwen3.8-27B-DFlash2 as the draft model, DFlash with one speculative step and eight draft tokens, TP=1, Triton attention backends, CUDA graphs disabled, five-shot GSM8K with 200 questions, parallelism 4, temperature 0, and top-p 1. The only server-side difference was --enable-linear-replayssm-spec.

Configuration	Accuracy	Invalid outputs	Exact output match
Baseline	194/200 (97.0%)	0	Reference
ReplaySSM	194/200 (97.0%)	0	200/200

Speed Tests and Profiling

N/A. This PR fixes routing to the existing ReplaySSM implementation and does not change its kernels. The GSM8K run above is reported only as an accuracy test; its single-run wall-clock timing is intentionally not used as a performance claim.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. N/A: this PR adds no public flag or API and fixes an already documented GDN ReplaySSM path.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, and /rerun-failed-ci.
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #32688323568
Latest PR Test (Extra): ❌ Run #32688323390
Latest PR Test (AMD ROCm 7.2): ❌ Run #32688323552

[yuan-luo]

This fix is solid. The only concern is that we once observed ReplaySSM accuracy drop for extra long output. Could you paste the AIME 2025 test result to confirm it also works correctly for DSpark?

[yuan-luo]

KDA supports compact-ragged, but GDN requires static ragged, why are the implementions inconsistent?

[huangwei021230]

Thanks for pointing this out. The base GDN backend currently does not have a compact-ragged target-verification implementation, while KDA already supports it. I prototyped compact GDN verification in this PR by following KDA’s packed-to-dense causal-conv path, and the prototype passes the correctness tests. I have also included the ragged mRoPE position fix from #35586 that was used for the Qwen3.8 non-uniform compact tests. However, the current implementation does not appear performance-ready.

On a single H20 with GSM8K-200, B=12/D=8, ReplaySSM enabled, FA3 full attention, Triton GDN verify, CUDA graphs disabled, and the same forced-50% verify-budget policy, cap-accept achieved 483.117 output tok/s in 63.138 s, while compact achieved 426.462 output tok/s in 71.575 s (-11.73% throughput and +13.36% wall time). Both scored 196/200 with zero invalid outputs, and all 200 extracted answers and correctness decisions matched. A separate four-repeat fixed-token workload showed the same direction: compact was 18.20%, 11.37%, and 10.33% slower at B=8, 12, and 16, respectively. Would you prefer that I retain this correctness prototype in the current PR for future optimization, or remove it and leave optimized compact GDN support to a follow-up?

My preference is to keep this PR within its original scope: fix the direct DFlash/DSpark ReplaySSM commit routing, allow GDN static and cap-accept (cap-accept still performs dense target verification and only caps acceptance/commit), reject GDN compact, and leave supports_ragged_verify_graph disabled. A performant compact GDN path likely requires dedicated native-varlen or fused kernel work. However, I am happy to follow your preference if you would rather keep the current prototype in this PR.

[yuan-luo]

replayssm_is_kda is pool level flag, while later on isinstance in Line 994 is backend type. If GDN backend with replayssm_is_kda=True will fallback to KDA handler. Here we need to add an assertion for this case.

[huangwei021230]

Thanks for pointing this out. I’ve added this assertion.

[huangwei021230]

This fix is solid. The only concern is that we once observed ReplaySSM accuracy drop for extra long output. Could you paste the AIME 2025 test result to confirm it also works correctly for DSpark?

Sorry for the delayed follow-up. I ran an AIME 2025 ReplaySSM off/on comparison using Qwen3.8-27B with the DSpark draft model, TP=1 on H20, Triton backends, CUDA graphs disabled, greedy decoding, single-request concurrency, a 65,536-token context, and a 32,768-token output limit.

Both arms scored 26/30 (86.7%) with the same terminal-answer scorer, and all 30 paired outputs were byte-identical. The four longest completed responses contained approximately 15.6K–28.5K tokens.

One caveat: the runs used compact mode without a profiled SPS table, so DSpark used a uniform full D=8 verify window (verify-all). This validates long-output ReplaySSM state/commit correctness, but not non-uniform compact scheduling.