[Feat] gdn: Add --enable-linear-lossless-verify for bit-exact tuple replay in spec verify

docs/docs/advanced_features/server_arguments.mdx

@@ -2124,6 +2124,12 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>triton</code>, <code>cutedsl</code>, <code>flashinfer</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-linear-lossless-verify`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Keep per-draft (u, k, g) tuples instead of full SSM-state snapshots during GDN speculative target-verify and replay the accepted state at commit. Bit-identical to the default verify with ~64x less verify scratch; the freed memory goes to the KV cache. GDN hybrid models with a linear draft chain (<code>--speculative-eagle-topk 1</code>: EAGLE/MTP, DFLASH, DSPARK) and the Triton verify kernel, on CUDA or ROCm.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
+    </tr>
 </tbody>
 </table>
 

python/sglang/kernels/ops/attention/linear/fused_sigmoid_gating_recurrent.py

@@ -238,6 +238,14 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     USE_GDC: tl.constexpr = False,
     MATCH_CUTEDSL_DECODE: tl.constexpr = False,
     ROUND_STATE_TO_BF16: tl.constexpr = False,
+    # Lossless verify (--enable-linear-lossless-verify): cache per-position corrected
+    # values u (=beta*(v - k^T h)), the l2-normed k and the decay g instead of
+    # full [K,V] states. Per-layer shapes [slots, cache_steps, HV, {V,K,-}] (no
+    # K x V product -> ~K x smaller). S_accept is replayed post-accept.
+    u_states_buffer=None,
+    k_states_buffer=None,
+    g_states_buffer=None,
+    CACHE_U_TUPLES: tl.constexpr = False,
 ):
     """
     Fused kernel that combines sigmoid gating computation with recurrent delta rule update.
@@ -338,7 +346,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     # contiguous but `cache_idx * cache_steps * HV * K * V` can exceed int32 for
     # large slot counts.
     cache_idx = -1
-    if CACHE_INTERMEDIATE_STATES:
+    if CACHE_INTERMEDIATE_STATES or CACHE_U_TUPLES:
         cache_idx = tl.load(intermediate_state_indices + i_n).to(tl.int64)
 
     step_idx = 0
@@ -548,6 +556,37 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
                 )
                 tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
 
+        # Lossless-verify tuple cache: store the corrected value u (=b_v here) per
+        # position. k (l2-normed) and g are also kept, so any accepted S_p can be
+        # replayed post-accept as h = exp(g) * h + k outer u.
+        if CACHE_U_TUPLES:
+            if cache_idx >= 0:
+                u_ptr = (
+                    u_states_buffer
+                    + cache_idx * cache_steps * HV * V
+                    + step_idx * HV * V
+                    + i_hv * V
+                    + o_v
+                )
+                tl.store(u_ptr, b_v.to(u_ptr.dtype.element_ty), mask=mask_v)
+                # k and g are independent of the V-block, so store them once per
+                # (step, v-head) from the i_v==0 program.
+                if i_v == 0:
+                    k_ptr = (
+                        k_states_buffer
+                        + cache_idx * cache_steps * HV * K
+                        + step_idx * HV * K
+                        + i_hv * K
+                        + o_k
+                    )
+                    tl.store(k_ptr, b_k.to(k_ptr.dtype.element_ty), mask=mask_k)
+                    g_ptr = (
+                        g_states_buffer
+                        + cache_idx * cache_steps * HV
+                        + step_idx * HV
+                        + i_hv
+                    )
+                    tl.store(g_ptr, b_g.to(g_ptr.dtype.element_ty))
         # A BF16 recurrent-state decode persists its state after every token.
         # Multi-token verification must make that precision boundary visible to
         # the next recurrence step instead of retaining the FP32 accumulator for
@@ -619,6 +658,10 @@ def fused_sigmoid_gating_delta_rule_update(
     # decays and updates through FMA, and a BF16 state pool is rounded after
     # every token.
     match_cutedsl_decode: bool = False,
+    # Lossless verify tuple ring (per-layer views [slots, steps, HV, ...]).
+    u_states_buffer: Optional[torch.Tensor] = None,
+    k_states_buffer: Optional[torch.Tensor] = None,
+    g_states_buffer: Optional[torch.Tensor] = None,
 ):
     """
     Fused triton implementation of sigmoid gating delta rule update.
@@ -669,7 +712,9 @@ def fused_sigmoid_gating_delta_rule_update(
         K,
         V,
         is_kda,
-        target_verify=target_verify,
+        # The lossless tuple verify must tile exactly like the dense one it
+        # replaces, so its outputs stay bit-identical on every arch.
+        target_verify=(target_verify or u_states_buffer is not None),
         cache_steps=cache_steps,
     )
     if is_kda and _is_sm90 and envs.SGLANG_ENABLE_FP8_BLOCK_GEMV.get():
@@ -703,6 +748,10 @@ def fused_sigmoid_gating_delta_rule_update(
     # allocated per-request pitch, which is preserved in stride(0).
     if intermediate_states_buffer is not None:
         cache_stride_steps = intermediate_states_buffer.stride(0) // (HV * K * V)
+    elif u_states_buffer is not None:
+        # The tuple ring keeps the allocated per-request pitch in dim 1, which
+        # can exceed the runtime draft count under --speculative-adaptive.
+        cache_stride_steps = u_states_buffer.shape[1]
     elif cache_steps is not None and cache_steps > 0:
         cache_stride_steps = cache_steps
     else:
@@ -803,6 +852,10 @@ def fused_sigmoid_gating_delta_rule_update(
         stride_beta_slot=stride_beta_slot,
         MAX_CACHE_LEN=max_cache_len,
         CACHE_RING=cache_ring,
+        u_states_buffer=u_states_buffer,
+        k_states_buffer=k_states_buffer,
+        g_states_buffer=g_states_buffer,
+        CACHE_U_TUPLES=u_states_buffer is not None,
         ring_indices=ring_indices
         if ring_indices is not None
         else initial_state_indices,

python/sglang/kernels/ops/mamba/mamba_state_scatter_triton.py

@@ -785,13 +785,17 @@ def scatter_mamba_states_after_mtp_verify(
     mamba_track_indices: torch.Tensor | None,
     mamba_steps_to_track: torch.Tensor | None,
     src_indices_raw: torch.Tensor | None = None,
+    skip_ssm: bool = False,
 ) -> None:
     """Scatter per-step verify states (ssm + all conv types) into the
-    persistent caches, plus the interval-crossing track slots."""
+    persistent caches, plus the interval-crossing track slots.
+
+    ``skip_ssm`` leaves the SSM state to the caller: the lossless GDN tuple
+    verify keeps no per-draft SSM snapshot and replays the accepted state."""
     ssm_states = mamba_caches.temporal
     intermediate_state_cache = mamba_caches.intermediate_ssm
 
-    if ssm_states.numel() > 0:
+    if ssm_states.numel() > 0 and not skip_ssm:
         fused_mamba_state_scatter_with_mask(
             ssm_states,
             intermediate_state_cache,

python/sglang/srt/arg_groups/attention_hook.py

@@ -537,6 +537,69 @@ def handle_linear_attn_backend(server_args: Any):
             )
 
 
+def validate_linear_lossless_verify(server_args: Any) -> None:
+    """Gate --enable-linear-lossless-verify once the speculative args are final.
+
+    Registered after ``handle_speculative_decoding``: DFLASH and DSPARK pin
+    ``speculative_eagle_topk`` to 1 there and EAGLE/MTP resolve their defaults,
+    so the top-k seen here is the one the verify will run with.
+    """
+    cfg = resolving_view(server_args)
+    if not cfg.enable_linear_lossless_verify:
+        return
+    if not (get_platform().is_cuda or get_platform().is_hip):
+        raise ValueError(
+            "--enable-linear-lossless-verify is only supported on CUDA/HIP: "
+            "other platforms do not implement both tuple-writing verification "
+            "and tuple reconstruction."
+        )
+    if cfg.speculative_algorithm is None:
+        raise ValueError(
+            "--enable-linear-lossless-verify only changes speculative "
+            "target-verify; it requires --speculative-algorithm."
+        )
+    if cfg.speculative_eagle_topk != 1:
+        raise ValueError(
+            "--enable-linear-lossless-verify requires a linear draft chain "
+            "(--speculative-eagle-topk 1): the tuple replay applies the "
+            "accepted steps in order and cannot follow an EAGLE tree. Got "
+            f"--speculative-eagle-topk={cfg.speculative_eagle_topk!r}."
+        )
+    if cfg.enable_linear_replayssm_spec or cfg.enable_linear_replayssm:
+        raise ValueError(
+            "--enable-linear-lossless-verify is mutually exclusive with "
+            "--enable-linear-replayssm and --enable-linear-replayssm-spec: "
+            "they replace the same per-draft state snapshots."
+        )
+    decode = cfg.linear_attn_decode_backend or cfg.linear_attn_backend
+    verify = cfg.linear_attn_verify_backend
+    if verify not in (None, "triton") or (verify is None and decode == "flashinfer"):
+        raise ValueError(
+            "--enable-linear-lossless-verify needs the Triton GDN verify "
+            "kernel, which is the one that writes the (u, k, g) tuples. Pass "
+            "--linear-attn-verify-backend triton. Got "
+            f"--linear-attn-verify-backend={verify!r} with decode backend "
+            f"{decode!r}."
+        )
+    if cfg.enable_unified_memory:
+        raise ValueError(
+            "--enable-linear-lossless-verify is not supported with "
+            "--enable-unified-memory: the unified pool builds its own "
+            "per-draft state snapshots."
+        )
+    from sglang.srt.speculative.ragged_verify import (
+        RaggedVerifyMode,
+        read_ragged_verify_mode,
+    )
+
+    if read_ragged_verify_mode() is not RaggedVerifyMode.STATIC:
+        raise ValueError(
+            "--enable-linear-lossless-verify requires "
+            "SGLANG_RAGGED_VERIFY_MODE=static: the GDN verify reshapes the batch "
+            "to a fixed draft width."
+        )
+
+
 def handle_multi_item_scoring(server_args: Any):
     """Setup and validate multi-item scoring constraints.
 

python/sglang/srt/arg_groups/field_order.py

@@ -560,5 +560,6 @@
     "uocr_doc_page_window",
     "uocr_max_page_tokens",
     "enable_fp4_indexer",
+    "enable_linear_lossless_verify",
 )
 # fmt: on

python/sglang/srt/arg_groups/fields/exec_.py

@@ -623,6 +623,14 @@ class ExecMamba(msgspec.Struct):
         bool,
         "Enable the ReplaySSM spec-verify: fold-every-commit -- a per-slot raw-input window replaces the recurrent verify's per-draft full-state snapshots. GDN or KDA hybrid linear-attn models, linear-chain (--speculative-eagle-topk in {None, 1}) only.",
     ] = False
+    # Lossless GDN tuple verify: target-verify keeps each draft step's
+    # (u, k, g) tuple instead of a full [HV, V, K] state snapshot and replays the
+    # accepted prefix at commit with the verify kernel's exact rounding, so the
+    # committed state and every output stay bit-identical to the snapshot verify.
+    enable_linear_lossless_verify: A[
+        bool,
+        "Keep per-draft (u, k, g) tuples instead of full SSM-state snapshots during GDN speculative target-verify and replay the accepted state at commit. Bit-identical to the default verify with ~64x less verify scratch; the freed memory goes to the KV cache. GDN hybrid models with a linear draft chain (--speculative-eagle-topk 1: EAGLE/MTP, DFLASH, DSPARK) and the Triton verify kernel.",
+    ] = False
 
 
 class ExecGraph(msgspec.Struct):

python/sglang/srt/arg_groups/pipeline.py

@@ -322,6 +322,11 @@ def run_resolution_pipeline(server_args: Any) -> None:
 
     run_hook(handle_layernorm_sp, server_args)
 
+    # After the speculative hook so speculative_eagle_topk is final.
+    from sglang.srt.arg_groups.attention_hook import validate_linear_lossless_verify
+
+    run_hook(validate_linear_lossless_verify, server_args)
+
     # Validate the FlashInfer A2A token budget now that draft-token counts are final.
     run_hook(validate_flashinfer_a2a_token_budget, server_args)
 

python/sglang/srt/arg_groups/resolution_hooks.py

@@ -83,6 +83,7 @@
         "validate_experimental_sgl_marlin",
         "handle_speculative_decoding",
         "handle_layernorm_sp",
+        "validate_linear_lossless_verify",
         "validate_flashinfer_a2a_token_budget",
         "validate_mega_moe_token_budget_for_model",
         "handle_load_format",

python/sglang/srt/disaggregation/decode.py

@@ -335,6 +335,7 @@ def __init__(
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
         enable_linear_replayssm_spec: bool = False,
+        enable_linear_lossless_verify: bool = False,
         short_conv_layer_ids: Optional[List[int]] = None,
         short_conv_state_shape: Optional[Tuple[int, int]] = None,
         ngram_context_len: int = 0,
@@ -391,6 +392,7 @@ def __init__(
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             mamba_envelope_layout=mamba_envelope_layout,
             enable_linear_replayssm_spec=enable_linear_replayssm_spec,
+            enable_linear_lossless_verify=enable_linear_lossless_verify,
             short_conv_layer_ids=short_conv_layer_ids,
             short_conv_state_shape=short_conv_state_shape,
             replayssm_request_ring=replayssm_request_ring,