[hsthe29]

Motivation

How it works today. On hybrid GDN models (Qwen3.5 / 3.6 / 3.8, Qwen3-Next), speculative target-verify runs the recurrent kernel with disable_state_update=True and writes a full [HV, V, K] SSM snapshot for every draft position into intermediate_ssm. After acceptance, scatter_mamba_states_after_mtp_verify copies two snapshots back into the persistent pool: the last accepted position, and, with the extra_buffer radix strategy, the position that crossed a mamba-track boundary. Every linear-chain algorithm goes through this path: EAGLE/MTP, DFLASH and DSPARK.

The problem. That buffer is large, and it grows with the draft width. On Qwen3.6-35B-A3B one snapshot is 60 MiB (30 GDN layers × 32 heads × 128 × 128 × fp32), so each request needs 60 MiB × D. _handle_max_mamba_cache reserves this for every request slot before the KV pool is sized, so all of it comes out of KV capacity:

setup (1× H100)	intermediate_ssm
Qwen3.6-35B-A3B-FP8, MTP D=4, 25 requests	6.09 GB
Qwen3.6-35B-A3B-FP8, DFLASH block 16, 8 requests	8.44 GB
Qwen3.8-27B-FP8, DFlash2 / DSpark D=8, 8 requests	10.12 GB

In the first setup the KV pool ends up at 498,880 tokens. That is not enough for 25 concurrent 26k-token requests, so 7 of them sit in the queue.

--enable-linear-replayssm-spec already gets rid of these snapshots, but it changes how verify is computed: a frozen checkpoint plus a compact ring, a causal-masked window, and compensated materialization. Its results are therefore not bit-identical to the default verify. On Qwen3.6-35B-A3B with MTP (measured on v0.5.21), greedy outputs diverged from the default path after 23–143 tokens on 8 of 8 prompts. Many deployments can live with that. Some cannot: evals, A/B rollouts and regression testing all need the outputs to match the default path exactly.

The idea. Each recurrent step is

h_t = exp(g_t) * h_{t-1} + k_t ⊗ u_t,    u_t = beta_t * (v_t - (exp(g_t) * h_{t-1})^T k_t)

The verify kernel already computes u_t, the l2-normalized k_t and g_t. And h_0 is still in the pool, because verify does not update state. So instead of keeping every h_t, we can keep (u_t, k_t, g_t) and rebuild the accepted h_t at commit time. That costs HV * (V + K + 1) floats per step instead of HV * V * K, which is 63.8× less at K = V = 128.

The rebuild has to round exactly the way the verify kernel does. Inside the kernel the decayed state h * exp(g) is used twice, by the delta reduction and by the update. So the compiler keeps it as a separately rounded multiply and folds the update into a single FMA: fma(k, u, round(h * exp(g))). The replay writes out the same two instructions. A naive replay written as one expression lets the compiler contract a different multiply and drifts by 1–2 ULP. The tests include that case, to show they would catch it.

Modifications

Everything sits behind a new flag, --enable-linear-lossless-verify, off by default. With the flag off, no code path and no memory reservation changes.

Verify kernel (fused_sigmoid_gating_recurrent.py, gdn_triton.py)

• New optional u_states_buffer / k_states_buffer / g_states_buffer arguments and a CACHE_U_TUPLES constexpr. For each draft step, the kernel stores u (after the beta scaling), the l2-normalized k, and g, at the row given by intermediate_state_indices. k and g do not depend on the V tile, so only the i_v == 0 program writes them.
• The per-request pitch comes from the buffer shape, so it stays correct under --speculative-adaptive.
• The tuple path picks the same launch config as the dense verify (target_verify=True, which includes the SM90 BV=4 tile). Verify outputs are therefore identical by construction, not just in testing.

Verify call site (gdn_backend.py)

• When the pool has a tuple ring, the ring is passed to the kernel instead of intermediate_ssm.
• A one-time check that the dispatcher picked the Triton verify kernel. The FlashInfer verify kernel accepts **kwargs: it would silently ignore the buffers, and the commit would replay zeros.
• An assert that no tree metadata (retrieve_parent_token) reaches the tuple path.

Commit (hybrid_linear_attn_backend.py, mamba_state_scatter_triton.py)

• _recon_kernel replays all GDN layers of a batch in one launch, starting from ssm_states[slot] and stopping at the accepted step.
  • On CUDA, the rounding is pinned with mul.rn.f32 / fma.rn.f32 inline PTX.
  • PTX does not exist on ROCm: the AMDGPU backend rejects the f register constraint. HIP builds instead use a plain multiply feeding tl.fma, which the compiler cannot contract.
  • Both forms are bit-exact against verify.
• Rows with a negative step, or an out-of-range slot or row index, are skipped. This is the same validity rule fused_mamba_state_scatter_with_mask uses.
• Order matters. The mamba-track replay runs before the in-place replay of the accepted state, because both must start from the pre-verify state. If the in-place replay runs first, the track replay starts from the already-advanced state and stores a wrong prefix-cache state. This bug does not show up in throughput benchmarks. It shows up as the second turn of a conversation diverging from the default path; the server-level test below covers it.
• Conv windows still go through the existing scatter. scatter_mamba_states_after_mtp_verify gets a skip_ssm argument, so the SSM half is left to the replay.
• PP spec stable rows (SGLANG_ENABLE_PP_SPEC) work: the replay reads the same rows verify wrote to (src_indices_raw).

Pool and memory accounting (memory_pool.py, kv_cache_configurator.py, disaggregation/decode.py, unified_memory_pool.py)

• MambaPool.SpeculativeState gets intermediate_ssm_u / _k / _g, shaped [layers, spec_slots + 1, D, HV, V | K | -]. In this mode intermediate_ssm is None.
• The ring is always fp32, whatever --mamba-ssm-dtype is, so bf16 state pools stay exact too.
• The ring is excluded from state transfer, like the other spec-only buffers.
• spec_intermediate_bytes_per_req computes the scratch size from the same shapes MambaPool allocates: the tuple ring plus the deduplicated conv window.
• With the flag on, _handle_max_mamba_cache reserves exactly that in all three sizing branches, including the joint solve for max_mamba_cache_size. The freed memory becomes KV cache instead of an idle reservation. With the flag off, the existing formula is used unchanged.
• The flag is plumbed through HybridReqToTokenPool and the PD decode pool. The unified memory pool rejects it.

Gating (attention_hook.py, pipeline.py)

The replay walks one linear chain, so the flag requires --speculative-eagle-topk 1. This is enforced in three places:

• At startup: validate_linear_lossless_verify runs right after handle_speculative_decoding. It cannot live in handle_linear_attn_backend, because that runs earlier, before DFLASH and DSPARK pin topk to 1.
  • Accepted: EAGLE/MTP with topk 1, DFLASH, DSPARK.
  • Rejected with a clear message: tree EAGLE (topk > 1), NGRAM, no speculative decoding, either ReplaySSM flag, a non-Triton verify backend, --enable-unified-memory, and any SGLANG_RAGGED_VERIFY_MODE other than static.
• In the pool: MambaPool rejects topk ≠ 1 and KDA layers, and the KV cache configurator rejects non-GDN models.
• At runtime: forward_extend asserts that no tree metadata reaches the tuple path.

Tests: test/registered/attention/unittests/gdn/test_gdn_lossless_verify.py, registered for CUDA and AMD CI. Details below.

Not changed: the default verify path, KDA, tree verify, and ReplaySSM.

Accuracy Tests

Unit tests (test_gdn_lossless_verify.py, 8 tests). They drive the real MambaPool, the real verify kernel and the real commit entry point (HybridLinearAttnBackend.update_mamba_state_after_mtp_verify), and require torch.equal against the snapshot verify for:

• verify outputs, every SSM slot, every conv slot and every mamba-track slot, for D ∈ {2, 3, 4, 8, 16} and heads (16, 32) and (4, 16);
• a bf16 state pool;
• PP stable request rows;
• out-of-range rows, which must be left untouched;
• the PTX and portable replay forms (CUDA);
• the scratch-size helper against the real allocation, byte for byte;
• rejection of tree verify;
• a control: replaying the accepted state before the track state gives a wrong track state, so the ordering above is actually exercised.

On H100 these 8 pass, and so do the existing GDN tests: test_triton.py, test_gdn_replayssm_spec_fold.py, test_fused_verify_triton_gdn.py, test_gdn_noncontiguous_stride.py.

Server-level, greedy. Stock main (d75d5b3) against this branch with --enable-linear-lossless-verify, same flags otherwise (--max-running-requests 8 --max-mamba-cache-size 40 --kv-cache-dtype bf16 --page-size 64), 1× H100. Each workload runs twice per server, with a cache flush in between, so the server's own run-to-run determinism is measured too:

• single: 8 prompts, one at a time, 300 tokens each.
• multi-turn: turn 1 decodes 700 tokens, crossing the 256 and 512 mamba-track boundaries. Turn 2 is turn 1 plus a follow-up, so its prefill resumes from the state tracked during turn-1 decode (cached_tokens = 512). This is the workload that catches a wrong track state.
• batched: the 8 prompts in one request, so the batch composition is fixed.

speculative setup	single	multi-turn (turn 1 · turn 2)	batched
EAGLE/MTP 3/1/4, Qwen3.6-35B-A3B-FP8	8/8	4/4 · 4/4	8/8
DFLASH, Qwen3.6-35B-A3B-FP8 + z-lab/Qwen3.6-35B-A3B-DFlash (D=16)	8/8	4/4 · 4/4	8/8
DFLASH with a DFlash2 draft, Qwen3.8-27B-FP8 + z-lab/Qwen3.8-27B-DFlash2 (D=8)	8/8	4/4 · 4/4	8/8
DSPARK, Qwen3.8-27B-FP8 + RadixArk/Qwen3.8-27B-DSpark (D=8)	8/8	4/4 · 4/4	8/8

Every cell matched stock on both runs. When the same 8 prompts are fired from 8 threads instead, outputs occasionally differ in 2 of 8. Stock differs from itself in exactly the same way between runs, because arrival order changes how requests are batched; it is not specific to this change.

ROCm. The same kernels and commit logic, backported to v0.5.21 and run on an MI325X (gfx942, ROCm 7.15, Triton 3.8):

• A standalone harness passed 144 of 144 exact checks with the portable replay.
• Server-level, with --enable-deterministic-inference --attention-backend triton and SGLANG_USE_AITER=0, all four workloads matched stock on both runs.
• With the default aiter backends, stock itself is not reproducible from run to run (even bs=1 greedy differs), so that configuration cannot be compared. Accept length was the same (2.852 vs 2.850).
• gfx950: test_gdn_lossless_verify.py passes at this PR's head on an MI350X (gfx950, ROCm 10, Triton 3.8). Every check is bit-exact for one request and for a batch, which run gfx950's own verify tiles (BV=8 / 4 warps and BV=16 / 2 warps). Server-level runs on gfx950 are not done yet.

Speed Tests and Profiling

End to end. Qwen3.6-35B-A3B-FP8, MTP 3/1/4, 1× H100, with --mem-fraction-static 0.8 --kv-cache-dtype bf16 --page-size 64 --max-running-requests 25 --max-mamba-cache-size 150 --enable-hierarchical-cache. bench_serving on the random dataset, 25 requests × (26k in / 2k out), concurrency 25:

	verify scratch	max_total_num_tokens	max running	duration	output tok/s	accept length
main	6.09 GB	498,880	18	36.75 s	1,360.6	3.51
this PR, flag on	0.10 GB	788,224 (+58%)	24	33.03 s	1,513.7 (+11%)	3.45

Accept length is within run-to-run noise: it ranged from 3.45 to 3.51 across repeated runs of both configurations. That is expected, since the outputs are bit-identical.

Memory with other algorithms (8 requests, 40 mamba slots, so the absolute numbers are not comparable with the table above):

setup	verify scratch	max_total_num_tokens
MTP, Qwen3.6-35B-A3B-FP8 (D=4)	2.11 → 0.03 GB	1,007,680 → 1,107,776 (+10%)
DFLASH, Qwen3.6-35B-A3B-FP8 (D=16)	8.44 → 0.13 GB	390,208 → 591,040 (+51%)
DFlash2, Qwen3.8-27B-FP8 (D=8)	10.12 → 0.16 GB	176,640 → 302,464 (+71%)
DSpark, Qwen3.8-27B-FP8 (D=8)	10.12 → 0.16 GB	178,944 → 304,832 (+70%)

Per-step kernel cost. Microbenchmark on the Qwen3.6-35B-A3B GDN geometry (30 layers, 16/32 heads, 128 × 128), H100, CUDA events, about 25% of rows crossing a track boundary. The verify kernel gets cheaper, because it writes HV * (V + K + 1) instead of HV * V * K floats per step. The replay is cheaper than the scatter at D=4 and costs a little more at D=8 and D=16, where it walks a longer chain. Net, every configuration is faster per verify step:

requests × draft tokens	verify, per GDN layer	commit, 30 layers	net per verify step
25 × 4	99.7 → 53.9 µs	1868 → 1489 µs	−1.75 ms
8 × 4	41.0 → 41.8 µs	605 → 413 µs	−0.17 ms
8 × 8	64.9 → 42.0 µs	527 → 548 µs	−0.67 ms
8 × 16	120.9 → 60.3 µs	470 → 690 µs	−1.60 ms

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. The flag is described in --help; there is no docs page that lists the linear-attention flags.
• 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, /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 #42155370412
Latest PR Test (Extra): ❌ Run #42155370523
Latest PR Test (AMD ROCm 10): ⏳ Run #42155370457

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[hsthe29]

Hi @hnyls2002 @yizhang2077 @Qiaolin-Yu and others, could one of you trigger CI on this PR when you have a moment?

It adds an opt-in --enable-linear-lossless-verify for GDN speculative verify: per-draft (u, k, g) tuples replace the full SSM snapshots, outputs stay bit-identical to the default verify, and the freed memory goes to the KV cache (+58% tokens on Qwen3.6-35B-A3B with MTP). I tested it on H100 with MTP, DFLASH (incl. a DFlash2 draft) and DSPARK, and the ROCm path on an MI325X; details are in the description. Thanks!

[hsthe29]

Hi @hnyls2002 @yizhang2077 @Qiaolin-Yu, @ShangmingCai, @HaiShaw, @BBuf and others, could one of you trigger CI on this PR when you have a moment?

It adds an opt-in --enable-linear-lossless-verify for GDN speculative verify: per-draft (u, k, g) tuples replace the full SSM snapshots, outputs stay bit-identical to the default verify, and the freed memory goes to the KV cache (+58% tokens on Qwen3.6-35B-A3B with MTP). I tested it on H100 with MTP, DFLASH (incl. a DFlash2 draft) and DSPARK, and the ROCm path on an MI325X; details are in the description. Thanks!

[BBuf]

Please address the [P1] items, then re-ping.

_{Agent-assisted pre-review, findings verified against 24b47a2ce7ac.}

[BBuf]

[P1] These keywords are passed even when lossless verify is disabled. On XPU this function is rebound to the platform wrapper, which accepts neither these arguments nor **kwargs, so ordinary GDN speculative verification now raises TypeError. Preserve the flag-off signature and gate tuple forwarding on backend support.

[BBuf]

[P1] Checking the backend name does not exclude Ascend: this configuration passes validation, but the registry selects the Ascend verify and commit implementations. The pool removes intermediate_ssm while the Ascend commit still consumes it and accesses its shape, so the first commit fails. Reject unsupported platforms before allocating the tuple-only pool.

[BBuf]

@hsthe29 I pushed d1097b76 to your branch: Preserve the flag-off XPU verify signature and reject lossless verification on unsupported platforms.
Please pull before your next push. The review runs again on the new head.

_{Automated pre-review for 24b47a2ce7ac; it re-checks when the PR is updated.}

[BBuf]

Please address the [P1] items, then re-ping.

_{Agent-assisted pre-review, findings verified against d1097b76955e.}

[BBuf]

[P1] This AMD lane defaults to MI300, so the (4, 16) case never reaches the gfx950-specific verify tiling. The flag admits MI35x, but the description says gfx950 is untested and the ROCm evidence comes from a v0.5.21 backport. Please register/run these parity checks on MI35x at this head, covering both N=1 and N>1, or gate that architecture out until validated.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-test test/registered/attention/unittests/gdn/test_gdn_lossless_verify.py

Targeted validation for d1097b76955e.

[github-actions]

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/attention/unittests/gdn/test_gdn_lossless_verify.py

[hsthe29]

@BBuf Thanks for the review, and for pushing d1097b7 to preserve the flag-off XPU verify signature and reject unsupported platforms.

I've renamed the PR to follow the title convention.

For the remaining P1 (running the parity checks on MI35x at this head, covering both N=1 and N>1), I don't have access to the MI35x machine over the weekend, so I'll pick it up on Monday. Sorry for the wait, and thanks for your patience.

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 413bdb6.

Approve preview

[hsthe29]

@BBuf Pushed 413bdb6:

• Merged main and fixed the three base-a-test-cpu failures: documented --enable-linear-lossless-verify in server_arguments.mdx, appended the field to POSITIONAL_FIELD_ORDER instead of inserting it mid-list, and dropped the unified-memory guard in _init_unified_pools (validate_linear_lossless_verify already rejects --enable-unified-memory, and the guard broke test_unified_tri_pool's mocked args).
• MI35x: test_gdn_lossless_verify.py is now registered on stage-b-test-1-gpu-small-amd-mi35x, with a new single-request sweep. On an MI350X (gfx950, ROCm 10, Triton 3.8) at this head, the (4, 16) heads run gfx950's own tiles, (BV=8, 4 warps) for N=1 and (BV=16, 2 warps) for N=5. All 9 tests pass (the PTX check is skipped as CUDA-only), and every check is bit-exact (torch.equal) against the snapshot verify.

Could you take another look?

[BBuf]

[BBuf Review Agent] @hsthe29 The docs, positional-order and unified-pool fixes are in place, and the MI35x tests cover both gfx950 tile configurations. I’m fine with the reported parity results at this head.

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_fused_verify_triton_gdn.py

[github-actions]

🚀 1-gpu-h100 (1 test): ⏳ View workflow run

cd test/ && python3 registered/kernels/ops/attention/test_fused_verify_triton_gdn.py  # base-b-kernel-unit