[YAMY1234]

Motivation

#30775 enabled PP x speculative decoding (SGLANG_ENABLE_PP_SPEC=1). On hybrid linear-attention targets (KDA / GDN / Mamba-style recurrent state) it is currently incorrect in two independent ways, both found while validating PP + MTP across model families (tracking: #39634):

1. Non-last PP stages never commit the accepted recurrent state. The last stage commits KDA/Mamba state inside run_eagle_verify; earlier stages only run target verify and relay the result, so their recurrent state stays at the pre-verify value while the last stage advances. GLM-5.3-Flash PP2 + MTP: GSM8K 0.745 vs 0.990 without spec, accept length looks normal (3.5), so the bug is silent.
2. With more than two stages the relayed accept result is paired with the wrong micro-batch. The commit hook took its destination slots from the backend's active forward_metadata, which on a non-last stage still belongs to the previous forward (slot (i-1) % N), and the KDA verify scratch was written at positional rows 0..bs-1 that in-flight micro-batches overwrite before the relay returns. With PP2 (i-1) % 2 == (i+1) % 2, so both errors are masked; with PP4 the first concurrent batch crashes with indices length mismatch: dst_indices_raw.shape[0]=7 vs step_indices_raw.shape[0]=12, and after fixing only the destination the accuracy is 0.485.

Pipeline-parallel speculative decoding can receive verify acceptance after a later micro-batch has overwritten positional hybrid-state scratch. When SGLANG_ENABLE_PP_SPEC=1, hybrid verify scratch and the delayed recurrent-state commit therefore use request-pool stable rows; with the gate off, the existing positional / active-metadata path is unchanged.

Modifications

• scheduler_pp_mixin: relay accept_index unconditionally (chains too). On non-last stages _pp_spec_commit_relayed_accept commits the acceptance relayed for that same micro-batch (same commit_mamba_states_after_verify the last stage uses) before seq_lens advances, then compacts the accepted KV. The scheduler snapshots each PP-spec micro-batch and clones req_pool_indices (the Mamba-track tensor is rebound, never mutated, so the snapshot shares it).
• Gating: pp_spec_stable_rows_enabled() (layers/attention/linear/utils.py) is the single predicate for the stable-row policy; kernels take stable_rows= from the backend instead of reading the env.
• Hybrid Mamba commit (hybrid_linear_attn_backend, mamba_state_scatter_triton): under the PP-spec gate, when the caller carries req_pool_indices, rebuild and translate destination state indices from the matching forward snapshot instead of the active forward_metadata; with the gate off, or without rows, the existing path runs.
• KDA / GDN verify (kda_backend, gdn_backend, gdn_flashinfer): the request-pool slot is the stable scratch row (padding still goes to the discard row); row selection is shared in select_verify_intermediate_state_indices. The FlashInfer GDN MTP kernel writes positional 0..B-1 and takes no row indices, so the PP-spec path uses a positional scratch and remaps to the stable rows after the kernel. KDA PP verify falls back from the FlashInfer kernel to Triton; ReplaySSM verify reads its source rows from the same stable mapping (src_indices_raw / dst_indices_raw across backends).
• Qwen PLE short-conv / n-gram side state (ple_state_pool, qwen4_exp) gets the same request-row pairing.

No change to the attention math or the acceptance algorithm; only where state is stored and how relayed results are paired.

Accuracy / validation (GB300, GSM8K 200 questions, SGLANG_ENABLE_PP_SPEC=1, EAGLE/MTP steps 3 topk 1)

Model	Config	Before	After
GLM-5.3-Flash (KDA)	PP2xTP2 + MTP	0.745 / AL 3.51	0.990 / 3.39 (TP4+MTP baseline 0.985 / 3.40)
GLM-5.3-Flash	PP2xTP2 + MTP + ReplaySSM	-	0.995 / 3.40
GLM-5.3-Flash	PP4xTP1 + MTP	crash (7 vs 12)	0.990 / 3.39
Qwen3.8-Flash-Next (GDN)	PP4xTP1 + MTP	crash (1 vs 12)	0.965 / 3.02 (TP4+MTP baseline 0.980 / 3.16)
Qwen3.8-Flash-Next	PP2xTP2 + MTP, with #39643 and #40501 on main	-	0.975 / 3.13 (0.955 / 3.02 without this PR)
DeepSeek-V4-Flash (non-hybrid control)	PP2xTP2 + MTP	0.985 / 2.87	0.985 / 2.89
DeepSeek-V4-Flash	PP4xTP1 + MTP	0.980 / 2.87	0.980 / 2.87
Llama-2-7b EAGLE	test/registered/e2e/pp/test_pp_spec.py::TestPPSpecConsistency	2 passed	2 passed

Re-verified after merging main d229952 (with #39643 applied, which the aggregate path needs on main for the GLM / Qwen3.5 draft-embedding key): GLM-5.3-Flash PP2xTP2 + MTP GSM8K-200 0.990 / AL 3.40; Qwen3.5-4B hybrid e2e (TestPPSpecHybridConsistency, GSM8K-256): no-PP 0.8906 / accept 3.33, PP2 0.8906 / accept 3.33.

Tests: test/registered/e2e/pp/test_pp_spec.py::TestPPSpecHybridConsistency (Qwen3.5-4B with its own MTP head, --language-only, PP1 vs PP2 on GSM8K-256 with the shared 2% score and accept-length floors) replaces the earlier pipeline mocks; unit cases stay in test_ngram_mamba_verify_update.py (scatter source rows, FlashInfer row remap, row selection, ReplaySSM fold and circular helpers, PLE rows) and test_kda_helion_dispatcher.py (KDA Triton fallback).

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed.

---

CI States

Latest PR Test (Base): ✅ Run #36697732993
Latest PR Test (Extra): ✅ Run #36697732150
Latest PR Test (AMD ROCm 10): ⏳ Run #36697733207

[YAMY1234]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[YAMY1234]

All NV CI and CI extra Passed
https://github.com/sgl-project/sglang/actions/runs/36697732993/job/109964844295?pr=40001
https://github.com/sgl-project/sglang/actions/runs/36697732150/job/109908510387?pr=40001