[GDN] Support ReplaySSM Ring Spec-Verify (#28695)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: vincentzed <207368749+vincentzed@users.noreply.github.com>

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -94,10 +94,16 @@ class MambaAttnBackendBase(AttentionBackend):
             )
             # The ring cursor is a per-slot decode counter shared by all GDN layers;
             # manage it once here (snapshot, hand to layers, advance mod L), not per-layer.
+            # Gate on the linear_replayssm FLAG, not on cursor-tensor presence: the
+            # spec-verify ring (--enable-gdn-replayssm-spec) shares the write_pos
+            # allocation but owns it exclusively via commit_gdn_replayssm_spec
+            # (advance-by-accept-count once per verify step). Advancing it here as
+            # well inserts one phantom/stale ring entry per step and cumulatively
+            # poisons the reconstruction (degenerate repetition at 10k+ tokens).
             mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
             write_pos_buf = (
-                getattr(mamba_pool, "replayssm_write_pos", None)
-                if mamba_pool is not None
+                mamba_pool.replayssm_write_pos
+                if mamba_pool is not None and mamba_pool.enable_linear_replayssm
                 else None
             )
             if write_pos_buf is not None:
@@ -323,12 +329,20 @@ class MambaAttnBackendBase(AttentionBackend):
         )
 
     def _replayssm_enabled(self) -> bool:
-        """True iff --enable-linear-replayssm allocated the ring cursor
-        (MambaPool.replayssm_write_pos doubles as the on/off gate)."""
+        """True iff --enable-linear-replayssm is on for this pool.
+
+        Gate on the FLAG, not on ``replayssm_write_pos is not None``: the
+        spec-verify ring (--enable-gdn-replayssm-spec) also allocates the
+        cursor tensor but owns it exclusively via commit_gdn_replayssm_spec.
+        The decode-ring metadata machinery gated here (per-bs static cursor
+        buffers, the per-replay snapshot + advance-by-one in _replay_metadata,
+        and the decode-kernel ring rerouting downstream) must stay fully
+        dormant for the spec ring.
+        """
         mamba_pool = getattr(self.req_to_token_pool, "mamba_pool", None)
         if mamba_pool is None:
             return False
-        return getattr(mamba_pool, "replayssm_write_pos", None) is not None
+        return bool(mamba_pool.enable_linear_replayssm)
 
     def _replayssm_track_flush_mask(
         self, seq_lens_cpu: torch.Tensor, bs: int
@@ -544,7 +558,10 @@ class MambaAttnBackendBase(AttentionBackend):
                     static_ff.copy_(force_flush_dev)
                 else:
                     static_ff.zero_()
-                if not in_capture:
+                # Defense in depth: the decode-ring advance is only meaningful for
+                # decode/idle forwards (mirrors the eager path's gating). A
+                # TARGET_VERIFY replay must never advance the cursor.
+                if not in_capture and forward_mode.is_decode_or_idle():
                     L = mamba_pool.linear_replayssm_cache_len
                     # Advance only valid (non-padded) slots; a forced flush empties
                     # the ring -> next write_pos 0, like the natural L-1 wrap.

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -565,22 +565,62 @@ class GDNAttnBackend(MambaAttnBackendBase):
             value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
 
         if is_target_verify:
-            core_attn_out = self.kernel_dispatcher.target_verify(
-                A_log=layer.A_log,
-                dt_bias=layer.dt_bias,
-                q=query,
-                k=key,
-                v=value,
-                a=a,
-                b=b,
-                ssm_states=ssm_states,
-                cache_indices=cache_indices,
-                query_start_loc=query_start_loc,
-                intermediate_states_buffer=intermediate_state_cache,
-                intermediate_state_indices=intermediate_state_indices,
-                cache_steps=forward_batch.spec_info.draft_token_num,
-                retrieve_parent_token=retrieve_parent_token,
+            # ReplaySSM spec-verify (Part B of #28511): when the per-slot ring is
+            # allocated (--enable-gdn-replayssm-spec, GDN + linear-chain topk<=1),
+            # reconstruct the verify output for the whole draft window from the
+            # frozen checkpoint (`temporal`) + the per-slot circular (d, k, g) ring
+            # instead of the recurrent verify that snapshots a full state per draft
+            # token. The cursors are advanced once per decode step by the worker
+            # (commit_gdn_replayssm_spec in spec_utils). GDN-only: KDA (per-K gate)
+            # routes through kda_backend and never reaches here; we additionally
+            # guard on `not replayssm_is_kda` for safety. Falls back to the
+            # recurrent verify when the ring is absent.
+            mamba_pool = self.req_to_token_pool.mamba_pool
+            use_replayssm_spec = (
+                mamba_cache_params.replayssm_d is not None
+                and getattr(mamba_pool, "replayssm_cache_base", None) is not None
+                and not getattr(mamba_pool, "replayssm_is_kda", False)
             )
+            if use_replayssm_spec:
+                core_attn_out = self._replayssm_target_verify(
+                    layer=layer,
+                    query=query,
+                    key=key,
+                    value=value,
+                    a=a,
+                    b=b,
+                    mamba_pool=mamba_pool,
+                    layer_cache=mamba_cache_params,
+                    cache_indices=cache_indices,
+                    query_start_loc=query_start_loc,
+                    draft_token_num=forward_batch.spec_info.draft_token_num,
+                )
+            else:
+                # The recurrent fallback needs the per-draft snapshots, which
+                # the pool gates OFF under --enable-gdn-replayssm-spec (the
+                # same flag that makes `use_replayssm_spec` true above), so
+                # this branch is unreachable with a None buffer by
+                # construction -- keep it loud rather than silently frozen.
+                assert intermediate_state_cache is not None, (
+                    "recurrent target_verify fallback requires intermediate_ssm, "
+                    "which is not allocated under --enable-gdn-replayssm-spec"
+                )
+                core_attn_out = self.kernel_dispatcher.target_verify(
+                    A_log=layer.A_log,
+                    dt_bias=layer.dt_bias,
+                    q=query,
+                    k=key,
+                    v=value,
+                    a=a,
+                    b=b,
+                    ssm_states=ssm_states,
+                    cache_indices=cache_indices,
+                    query_start_loc=query_start_loc,
+                    intermediate_states_buffer=intermediate_state_cache,
+                    intermediate_state_indices=intermediate_state_indices,
+                    cache_steps=forward_batch.spec_info.draft_token_num,
+                    retrieve_parent_token=retrieve_parent_token,
+                )
         else:
             g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
             core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
@@ -612,3 +652,92 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 )
 
         return core_attn_out
+
+    def _replayssm_target_verify(
+        self,
+        *,
+        layer: RadixLinearAttention,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        a: torch.Tensor,
+        b: torch.Tensor,
+        mamba_pool: MambaPool,
+        layer_cache: "MambaPool.SpeculativeState",
+        cache_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        draft_token_num: int,
+    ) -> torch.Tensor:
+        """ReplaySSM GDN spec-verify (Part B of #28511).
+
+        Reconstructs the verify output for the whole draft window from the frozen
+        checkpoint (``temporal``) + the per-slot circular ``(d, k, g)`` ring, and
+        appends this window's drafts to the rings (chunked ``d`` for output
+        reconstruction; raw ``v`` / pre-norm ``k`` / fp32 ``beta`` for the
+        closed-loop exact fold that replays the recurrent update into the fp32
+        checkpoint at flush). The rings are PER-LAYER
+        (sliced via ``mamba2_layer_cache``), while the cursors (write_pos,
+        cache_base, is_flush) are PER-SLOT pool attributes shared by all GDN layers
+        of the step; the cursors persist across steps and are advanced once per step
+        by the worker (commit_gdn_replayssm_spec) -- here we only read them and
+        write this step's ring entries. GDN has K == V, so ``temporal``
+        ([slots, HV, K, V]) is consumed directly as the kernel's [slots, HV, V, K]
+        checkpoint.
+        """
+        from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_decode import (
+            gdn_replayssm_spec_decode,
+        )
+
+        H, K = layer.num_k_heads, layer.head_k_dim
+        HV, V = layer.num_v_heads, layer.head_v_dim
+        # q/k/v may be [1, seq, *] (fallback split) or [seq, *] (fused split);
+        # derive the packed token count from numel so both layouts flatten.
+        seq_len = query.numel() // (H * K)
+        q = query.reshape(seq_len, H, K)
+        k = key.reshape(seq_len, H, K)
+        v = value.reshape(seq_len, HV, V)
+        a = a.reshape(seq_len, HV)
+        b = b.reshape(seq_len, HV)
+        d_cache = layer_cache.replayssm_d  # [slots, HV, L, V]
+        max_cache_len = d_cache.shape[-2]  # ring length L
+        out = q.new_empty(seq_len, HV, V)
+        gdn_replayssm_spec_decode(
+            q=q,
+            k=k,
+            v=v,
+            a=a,
+            b=b,
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            checkpoint_state=layer_cache.temporal,
+            d_cache=d_cache,
+            k_cache=layer_cache.replayssm_k,
+            g_cache=layer_cache.replayssm_g,
+            # Closed-loop exact-fold rings: raw v / raw pre-norm k / fp32 beta.
+            # The flush replays these through the recurrent update (bit-identical
+            # to the recurrent baseline) instead of folding `d` open-loop.
+            rawv_cache=layer_cache.replayssm_rawv,
+            rawk_cache=layer_cache.replayssm_rawk,
+            beta_cache=layer_cache.replayssm_beta,
+            out=out,
+            query_start_loc=query_start_loc,
+            ssm_state_indices=cache_indices,
+            # Per-slot cursors live on the pool (shared across all GDN layers),
+            # NOT in forward_metadata: the verify kernel reads/writes them
+            # block-keyed via ssm_state_indices and must NOT advance write_pos
+            # (the worker does that after acceptance), so the decode-path
+            # forward_metadata.replayssm_write_pos snapshot is not used here.
+            write_pos=mamba_pool.replayssm_write_pos,
+            cache_base=mamba_pool.replayssm_cache_base,
+            is_flush=mamba_pool.replayssm_is_flush,
+            max_cache_len=max_cache_len,
+            max_spec_len=draft_token_num,
+            scale=K**-0.5,
+            use_qk_l2norm_in_kernel=True,
+            # SGLang marks invalid/padding requests with a negative mamba slot
+            # index (valid slots start at 0), so the kernel's "null block"
+            # sentinel is -1, not the vLLM default of 0.
+            null_block_id=-1,
+        )
+        # Match the recurrent target_verify output shape (== value.shape).
+        return out.reshape(value.shape)

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -723,6 +723,14 @@ class KVCacheConfigurator:
             enable_linear_replayssm=self.server_args.enable_linear_replayssm,
             linear_replayssm_cache_len=self.server_args.linear_replayssm_cache_len,
             mamba_envelope_layout=self.server_args.enable_page_major_kv_layout,
+            # ReplaySSM spec-verify is GDN-only: activate the pool machinery
+            # (rings + cursors + the intermediate_ssm gate) only for GDN-hybrid
+            # models, so any other mamba-ish model (Mamba2/Nemotron, lightning,
+            # ...) run with the flag set stays byte-identical to flag-off.
+            enable_gdn_replayssm_spec=(
+                self.server_args.enable_gdn_replayssm_spec
+                and self.hybrid_gdn_config is not None
+            ),
         )
         return req_to_token_pool
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -329,9 +329,20 @@ class MambaPool:
         #   replayssm_d: [num_layers, num_slots, HV, L, V]
         #   replayssm_k: [num_layers, num_slots, H,  L, K]
         #   replayssm_g: [num_layers, num_slots, HV, L]  (fp32)
+        #   replayssm_rawv: [num_layers, num_slots, HV, L, V]  (conv/activation dtype)
+        #   replayssm_rawk: [num_layers, num_slots, H,  L, K]  (conv/activation dtype)
+        #   replayssm_beta: [num_layers, num_slots, HV, L]     (fp32)
+        # The raw rings + beta exist only under --enable-gdn-replayssm-spec: the
+        # closed-loop exact fold sequentially replays them through the recurrent
+        # update at flush -- bit-identical to the recurrent baseline -- instead
+        # of folding the chunked `d` records open-loop (which accumulates error
+        # across flushes). See fla/gdn_replayssm_spec_decode.py.
         replayssm_d: Optional[torch.Tensor] = None
         replayssm_k: Optional[torch.Tensor] = None
         replayssm_g: Optional[torch.Tensor] = None
+        replayssm_rawv: Optional[torch.Tensor] = None
+        replayssm_rawk: Optional[torch.Tensor] = None
+        replayssm_beta: Optional[torch.Tensor] = None
 
         def at_layer_idx(self, layer: int):
             kwargs = {}
@@ -357,7 +368,10 @@ class MambaPool:
 
     @dataclass(frozen=True, kw_only=True)
     class SpeculativeState(State):
-        intermediate_ssm: torch.Tensor
+        # None under --enable-gdn-replayssm-spec: the spec ring owns rollback
+        # (verify writes ring records, commit moves cursors), so the per-draft
+        # full-state snapshots are never produced or consumed.
+        intermediate_ssm: Optional[torch.Tensor]
         intermediate_conv_window: List[torch.Tensor]
 
     def _allocate_deduplicated_conv_window(
@@ -414,6 +428,7 @@ class MambaPool:
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         envelope_layout: bool = False,
+        enable_gdn_replayssm_spec: bool = False,
     ):
         conv_state_shape = cache_params.shape.conv
         temporal_state_shape = cache_params.shape.temporal
@@ -429,6 +444,13 @@ class MambaPool:
         self.debug_memory_pool = envs.SGLANG_DEBUG_MEMORY_POOL.get()
         self.enable_linear_replayssm = enable_linear_replayssm
         self.linear_replayssm_cache_len = linear_replayssm_cache_len
+        # ReplaySSM spec-verify (Part B of #28511) REUSES the linear_replayssm ring
+        # (replayssm_d/k/g + write_pos) and ADDS two per-slot cursors
+        # (replayssm_cache_base + replayssm_is_flush). Enabling the spec-verify path
+        # therefore implies the ring, so the d/k/g + write_pos allocation gates on
+        # `_replayssm_on` (either flag). GDN-only is enforced upstream + below.
+        self.enable_gdn_replayssm_spec = enable_gdn_replayssm_spec
+        _replayssm_on = enable_linear_replayssm or enable_gdn_replayssm_spec
 
         # for disagg with nvlink
         self.enable_custom_mem_pool, self.custom_mem_pool, _ = (
@@ -504,22 +526,34 @@ class MambaPool:
 
             # GDN ReplaySSM ring buffers (slice 1a). Allocated only when the
             # flag is on; otherwise left as None so the legacy State is
-            # byte-identical. temporal_state_shape == (HV, V, K).
+            # byte-identical. temporal_state_shape == (HV, V, K). Either the decode
+            # ring (--enable-linear-replayssm) or the spec-verify ring
+            # (--enable-gdn-replayssm-spec) shares this allocation.
             replayssm_d = replayssm_k = replayssm_g = None
-            if enable_linear_replayssm:
+            replayssm_rawv = replayssm_rawk = replayssm_beta = None
+            if _replayssm_on:
                 hv, v_dim, k_dim = temporal_state_shape
                 h_k = getattr(cache_params.shape, "num_k_heads_per_tp", hv)
                 L = linear_replayssm_cache_len
                 num_slots = size + 1
-                # Ring records live in the SSM dtype (bf16/fp32) except g (fp32).
+                # Ring dtype. DECODE ring (--enable-linear-replayssm): records
+                # follow the SSM dtype -- its flush folds `d` directly into the
+                # state. SPEC-verify ring (--enable-gdn-replayssm-spec): d/k feed
+                # ONLY the one-shot output reconstruction (the closed-loop exact
+                # fold replays the raw rings for state instead), so their
+                # quantization noise stays below the bf16 output cast; keep them
+                # in the conv/activation dtype instead of the (fp32-enforced)
+                # SSM dtype to halve the ring traffic. g stays fp32 everywhere
+                # (exact-fold input). The two flags are mutually exclusive.
+                ring_dtype = conv_dtype if enable_gdn_replayssm_spec else ssm_dtype
                 replayssm_d = torch.zeros(
                     size=(num_mamba_layers, num_slots, hv, L, v_dim),
-                    dtype=ssm_dtype,
+                    dtype=ring_dtype,
                     device=device,
                 )
                 replayssm_k = torch.zeros(
                     size=(num_mamba_layers, num_slots, h_k, L, k_dim),
-                    dtype=ssm_dtype,
+                    dtype=ring_dtype,
                     device=device,
                 )
                 # The log-decay gate ring (fp32): per-head SCALAR for the GDN
@@ -535,6 +569,42 @@ class MambaPool:
                     dtype=torch.float32,
                     device=device,
                 )
+                # Closed-loop exact-fold rings (spec-verify only). Raw v / raw
+                # pre-norm k live in the conv (activation) dtype -- they are born
+                # there, so storage round-trips losslessly -- beta in fp32. The
+                # flush replays these through the recurrent update sequentially
+                # (bit-identical to the recurrent baseline) instead of folding
+                # the chunked `d` records open-loop.
+                if enable_gdn_replayssm_spec:
+                    # Backstop for the spec-verify ring invariants; this pool
+                    # is sized with the final adaptive-aware draft maximum.
+                    if L & (L - 1) != 0:
+                        raise ValueError(
+                            f"spec-verify ring length must be a power of two, got {L}"
+                        )
+                    if (
+                        speculative_num_draft_tokens is not None
+                        and L < 2 * speculative_num_draft_tokens
+                    ):
+                        raise ValueError(
+                            f"spec-verify ring too small: {L} < "
+                            f"2 * {speculative_num_draft_tokens} (early-flush margin)"
+                        )
+                    replayssm_rawv = torch.zeros(
+                        size=(num_mamba_layers, num_slots, hv, L, v_dim),
+                        dtype=conv_dtype,
+                        device=device,
+                    )
+                    replayssm_rawk = torch.zeros(
+                        size=(num_mamba_layers, num_slots, h_k, L, k_dim),
+                        dtype=conv_dtype,
+                        device=device,
+                    )
+                    replayssm_beta = torch.zeros(
+                        size=(num_mamba_layers, num_slots, hv, L),
+                        dtype=torch.float32,
+                        device=device,
+                    )
 
             if speculative_num_draft_tokens is not None:
                 if _is_npu:
@@ -546,18 +616,30 @@ class MambaPool:
                     )
                 # Cache intermediate SSM states per draft token during target verify
                 # Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V]
-                intermediate_ssm_state_cache = torch.zeros(
-                    size=(
-                        num_mamba_layers,
-                        spec_state_size + 1,
-                        speculative_num_draft_tokens,
-                        temporal_state_shape[0],
-                        temporal_state_shape[1],
-                        temporal_state_shape[2],
-                    ),
-                    dtype=ssm_dtype,
-                    device="cuda",
-                )
+                #
+                # ReplaySSM spec-verify owns rollback via the ring + cursors (the
+                # verify kernel never writes per-draft snapshots; the commit never
+                # reads them), so this buffer -- the dominant spec scratch, ~46x
+                # the conv state -- is dead weight there and is skipped. The conv
+                # intermediate windows below STAY (conv rollback consumes them).
+                # The recurrent-verify fallback cannot be reached under the flag
+                # (GDN + linear chain + triton enforced in server_args; the
+                # backend asserts loudly if it ever is).
+                if enable_gdn_replayssm_spec:
+                    intermediate_ssm_state_cache = None
+                else:
+                    intermediate_ssm_state_cache = torch.zeros(
+                        size=(
+                            num_mamba_layers,
+                            spec_state_size + 1,
+                            speculative_num_draft_tokens,
+                            temporal_state_shape[0],
+                            temporal_state_shape[1],
+                            temporal_state_shape[2],
+                        ),
+                        dtype=ssm_dtype,
+                        device="cuda",
+                    )
                 # Cache intermediate conv windows (last K-1 inputs) per draft token
                 # during target verify.
                 #
@@ -625,13 +707,21 @@ class MambaPool:
                     replayssm_d=replayssm_d,
                     replayssm_k=replayssm_k,
                     replayssm_g=replayssm_g,
+                    replayssm_rawv=replayssm_rawv,
+                    replayssm_rawk=replayssm_rawk,
+                    replayssm_beta=replayssm_beta,
+                )
+                intermediate_ssm_gb = (
+                    get_tensor_size_bytes(intermediate_ssm_state_cache) / GB
+                    if intermediate_ssm_state_cache is not None
+                    else 0.0
                 )
                 logger.info(
                     f"Mamba Cache is allocated. "
                     f"max_mamba_cache_size: {size}, "
                     f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
                     f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
-                    f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB "
+                    f"intermediate_ssm_state_cache size: {intermediate_ssm_gb:.2f}GB "
                     # Report the deduplicated PHYSICAL conv-window buffers (the view
                     # over-reports its logical, un-deduplicated size).
                     f"intermediate_conv_window_cache size: {get_tensor_size_bytes(self._intermediate_conv_window_phys) / GB:.2f}GB "
@@ -643,6 +733,9 @@ class MambaPool:
                     replayssm_d=replayssm_d,
                     replayssm_k=replayssm_k,
                     replayssm_g=replayssm_g,
+                    replayssm_rawv=replayssm_rawv,
+                    replayssm_rawk=replayssm_rawk,
+                    replayssm_beta=replayssm_beta,
                 )
                 logger.info(
                     f"Mamba Cache is allocated. "
@@ -650,26 +743,48 @@ class MambaPool:
                     f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
                     f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
                 )
-            if enable_linear_replayssm:
+            if _replayssm_on:
                 logger.info(
                     f"GDN ReplaySSM ring buffers allocated (L="
                     f"{linear_replayssm_cache_len}): "
                     f"d={get_tensor_size_bytes(replayssm_d) / GB:.3f}GB, "
                     f"k={get_tensor_size_bytes(replayssm_k) / GB:.3f}GB, "
                     f"g={get_tensor_size_bytes(replayssm_g) / GB:.3f}GB "
+                    + (
+                        f"rawv={get_tensor_size_bytes(replayssm_rawv) / GB:.3f}GB, "
+                        f"rawk={get_tensor_size_bytes(replayssm_rawk) / GB:.3f}GB, "
+                        f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB "
+                        if enable_gdn_replayssm_spec
+                        else ""
+                    )
                 )
             # Gate granularity of the linear-attn layers (drives the kernel's
             # IS_KDA path + the g_cache layout). Read by the backend metadata to
             # decide the per-K (KDA) vs scalar (GDN) flush/advance handling.
-            self.replayssm_is_kda = bool(
-                enable_linear_replayssm and cache_params.is_kda
-            )
+            self.replayssm_is_kda = bool(_replayssm_on and cache_params.is_kda)
             # Persistent per-slot decode-position cursor for ReplaySSM. Shared
             # across all linear-attn layers; advanced once per decode forward by
-            # the backend metadata build. Index 0..size; reset on slot (re)alloc.
+            # the backend metadata build (decode ring) or once per verify step by
+            # the worker (spec-verify ring). Index 0..size; reset on slot (re)alloc.
             self.replayssm_write_pos = (
                 torch.zeros((size + 1,), dtype=torch.int32, device=device)
-                if enable_linear_replayssm
+                if _replayssm_on
+                else None
+            )
+            # ReplaySSM spec-verify (Part B of #28511) extra per-slot cursors. The
+            # circular ring's rolling origin (cache_base) + the per-slot flush flag
+            # (is_flush). Block-keyed (indexed by the physical mamba slot), shared by
+            # all GDN layers of one verify step; advanced by commit_gdn_replayssm_spec.
+            # Only allocated for the spec-verify ring (the decode ring does not use
+            # a circular buffer); None otherwise.
+            self.replayssm_cache_base = (
+                torch.zeros((size + 1,), dtype=torch.int32, device=device)
+                if enable_gdn_replayssm_spec
+                else None
+            )
+            self.replayssm_is_flush = (
+                torch.zeros((size + 1,), dtype=torch.int8, device=device)
+                if enable_gdn_replayssm_spec
                 else None
             )
             mem_usage_bytes = self.mamba_cache.mem_usage_bytes()
@@ -804,6 +919,12 @@ class MambaPool:
             ]
         if self.replayssm_write_pos is not None:
             self.replayssm_write_pos[dst_indices] = 0
+        # ReplaySSM spec-verify ring: a copied checkpoint has no pending ring
+        # entries, so its rolling origin + flush flag reset alongside write_pos.
+        if self.replayssm_cache_base is not None:
+            self.replayssm_cache_base[dst_indices] = 0
+        if self.replayssm_is_flush is not None:
+            self.replayssm_is_flush[dst_indices] = 0
 
     def get_cpu_copy(self, indices):
         current_platform.synchronize()
@@ -814,17 +935,46 @@ class MambaPool:
         temporal_cpu = self.mamba_cache.temporal[:, indices].to(
             "cpu", non_blocking=True
         )
+        # ReplaySSM spec-verify ring: round-trip the per-slot cursors with the
+        # checkpoint so a restored slot reconstructs exactly. Only the spec ring
+        # adds the 3rd tuple element; every other config keeps the legacy 2-tuple
+        # so those paths stay byte-identical.
+        if self.replayssm_cache_base is not None:
+            cursors_cpu = (
+                self.replayssm_write_pos[indices].to("cpu", non_blocking=True),
+                self.replayssm_cache_base[indices].to("cpu", non_blocking=True),
+                self.replayssm_is_flush[indices].to("cpu", non_blocking=True),
+            )
+            current_platform.synchronize()
+            return conv_cpu, temporal_cpu, cursors_cpu
         current_platform.synchronize()
         return conv_cpu, temporal_cpu
 
     def load_cpu_copy(self, mamba_cache_cpu, indices):
-        conv_cpu, temporal_cpu = mamba_cache_cpu
+        # Accept both the legacy 2-tuple (conv, temporal) and the 3-tuple that also
+        # carries the ReplaySSM spec-verify cursors.
+        if len(mamba_cache_cpu) == 3:
+            conv_cpu, temporal_cpu, cursors_cpu = mamba_cache_cpu
+        else:
+            conv_cpu, temporal_cpu = mamba_cache_cpu
+            cursors_cpu = None
         current_platform.synchronize()
         for i, conv in enumerate(self.mamba_cache.conv):
             conv[:, indices] = conv_cpu[i].to(conv.device, non_blocking=True)
         self.mamba_cache.temporal[:, indices] = temporal_cpu.to(
             self.mamba_cache.temporal.device, non_blocking=True
         )
+        if cursors_cpu is not None and self.replayssm_cache_base is not None:
+            wp_cpu, cb_cpu, fl_cpu = cursors_cpu
+            self.replayssm_write_pos[indices] = wp_cpu.to(
+                self.replayssm_write_pos.device, non_blocking=True
+            )
+            self.replayssm_cache_base[indices] = cb_cpu.to(
+                self.replayssm_cache_base.device, non_blocking=True
+            )
+            self.replayssm_is_flush[indices] = fl_cpu.to(
+                self.replayssm_is_flush.device, non_blocking=True
+            )
         current_platform.synchronize()
 
     def get_contiguous_buf_infos(self):
@@ -841,7 +991,14 @@ class MambaPool:
                 continue
             # Skip GDN ReplaySSM ring buffers: they are derived/transient decode
             # scratch, not part of the persistent transferable state.
-            if field in ("replayssm_d", "replayssm_k", "replayssm_g"):
+            if field in (
+                "replayssm_d",
+                "replayssm_k",
+                "replayssm_g",
+                "replayssm_rawv",
+                "replayssm_rawk",
+                "replayssm_beta",
+            ):
                 continue
             value = getattr(self.mamba_cache, field)
             if value is None:
@@ -882,6 +1039,9 @@ class MambaPool:
                 "replayssm_d",
                 "replayssm_k",
                 "replayssm_g",
+                "replayssm_rawv",
+                "replayssm_rawk",
+                "replayssm_beta",
             ):
                 continue
             value = getattr(self.mamba_cache, field)
@@ -964,6 +1124,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
+        enable_gdn_replayssm_spec: bool = False,
     ):
         super().__init__(
             size=size,
@@ -990,6 +1151,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_linear_replayssm=enable_linear_replayssm,
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             mamba_envelope_layout=mamba_envelope_layout,
+            enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
         )
 
     def _init_mamba_pool(
@@ -1005,6 +1167,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
         enable_linear_replayssm: bool = False,
         linear_replayssm_cache_len: int = 16,
         mamba_envelope_layout: bool = False,
+        enable_gdn_replayssm_spec: bool = False,
     ):
         self.mamba_pool = self.mamba_pool_cls(
             size=mamba_size,
@@ -1018,6 +1181,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
             enable_linear_replayssm=enable_linear_replayssm,
             linear_replayssm_cache_len=linear_replayssm_cache_len,
             envelope_layout=mamba_envelope_layout,
+            enable_gdn_replayssm_spec=enable_gdn_replayssm_spec,
         )
         self.mamba_allocator = MambaSlotAllocator(
             size=mamba_size,
@@ -1116,6 +1280,12 @@ class HybridReqToTokenPool(ReqToTokenPool):
                 # (the post-prefill state that prefill wrote into this slot).
                 if self.mamba_pool.replayssm_write_pos is not None:
                     self.mamba_pool.replayssm_write_pos[req.mamba_pool_idx] = 0
+                # ReplaySSM spec-verify ring: an empty ring also resets the
+                # circular origin + flush flag so the first verify step on this
+                # freshly-prefilled slot reconstructs from the checkpoint alone.
+                if self.mamba_pool.replayssm_cache_base is not None:
+                    self.mamba_pool.replayssm_cache_base[req.mamba_pool_idx] = 0
+                    self.mamba_pool.replayssm_is_flush[req.mamba_pool_idx] = 0
             mamba_indices.append(req.mamba_pool_idx)
             if self.enable_mamba_extra_buffer:
                 if req.mamba_ping_pong_track_buffer is None:

python/sglang/srt/server_args.py

@@ -2164,6 +2164,16 @@ class ServerArgs:
         int,
         "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.",
     ] = 16
+    # ReplaySSM spec-verify (Part B of RFC #28511): GDN linear-chain target-verify
+    # via a per-slot circular (d, k, g) ring + periodic flush instead of per-draft
+    # full-state snapshots. GDN only; linear-chain (topk <= 1) only. Reuses the
+    # `linear_replayssm` ring (replayssm_d/k/g + write_pos) and adds two per-slot
+    # cursors (cache_base, is_flush); the ring length reuses
+    # `linear_replayssm_cache_len`.
+    enable_gdn_replayssm_spec: A[
+        bool,
+        "Enable the ReplaySSM GDN spec-verify kernel (Part B of RFC #28511): a per-slot circular (d, k, g) ring + periodic flush replacing the recurrent verify's per-draft full-state snapshots. GDN only, linear-chain (--speculative-eagle-topk in {None, 1}) only. Reuses --linear-replayssm-cache-len for the ring length.",
+    ] = False
 
     # -------------------------------------------------------------------------
     # Hierarchical cache
@@ -2998,6 +3008,9 @@ class ServerArgs:
 
         handle_speculative_decoding(self)
 
+        # Needs the draft-token count derived just above.
+        self._validate_gdn_replayssm_spec_ring()
+
         # Validate the CuteDSL A2A token budget now that num_tokens_per_req is final.
         self._validate_cutedsl_a2a_token_budget()
 
@@ -5307,6 +5320,110 @@ class ServerArgs:
                     f"{self.linear_replayssm_cache_len}."
                 )
 
+        # ReplaySSM spec-verify (Part B of #28511): GDN-only, linear-chain target
+        # verify. Reuses the `linear_replayssm` ring (replayssm_d/k/g + write_pos)
+        # plus two extra per-slot cursors (cache_base, is_flush) and the chunked
+        # (I+A)^-1 reconstruction verify kernel. The intra-window interaction uses a
+        # strictly-lower causal mask, so it is valid ONLY for a linear draft chain
+        # (speculative_eagle_topk in {None, 1}, i.e. NEXTN / MTP); EAGLE tree verify
+        # (topk > 1) must fall back to the recurrent verify. GDN-only is enforced at
+        # runtime (KDA routes through kda_backend, which never enters this path; the
+        # pool gate also checks `not cache_params.is_kda`). The ring length reuses
+        # --linear-replayssm-cache-len (no separate flag).
+        if self.enable_gdn_replayssm_spec:
+            if self.speculative_eagle_topk not in (None, 1):
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec requires a linear draft chain "
+                    "(--speculative-eagle-topk in {None, 1}); the chunked verify "
+                    "kernel uses a strictly-lower causal mask and is invalid for "
+                    "EAGLE tree verify. Got "
+                    f"--speculative-eagle-topk={self.speculative_eagle_topk!r}."
+                )
+            if decode != "triton":
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec requires the Triton linear-attn "
+                    "decode backend, got "
+                    f"--linear-attn-decode-backend={decode!r}."
+                )
+            if self.enable_mamba_extra_buffer():
+                # The spec-verify path does not yet implement the device-side
+                # force-flush needed to keep `temporal` consistent with the ring at
+                # radix mamba-track boundaries, so it is incompatible with
+                # extra_buffer (radix prefix caching).
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec is not yet compatible with mamba "
+                    "extra_buffer (radix prefix caching); use --disable-radix-cache "
+                    "or --mamba-radix-cache-strategy no_buffer."
+                )
+            if self.disaggregation_mode != "null":
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec is not supported under PD "
+                    "disaggregation yet (follow-up). Got "
+                    f"--disaggregation-mode={self.disaggregation_mode!r}."
+                )
+            if self.linear_replayssm_cache_len < 1:
+                raise ValueError(
+                    "--linear-replayssm-cache-len must be >= 1, got "
+                    f"{self.linear_replayssm_cache_len}."
+                )
+            if self.enable_linear_replayssm:
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec and --enable-linear-replayssm are "
+                    "mutually exclusive: they share the ring storage but drive it "
+                    "with incompatible cursor protocols (per-decode-forward vs "
+                    "per-verify-commit advance)."
+                )
+            ring_len = self.linear_replayssm_cache_len
+            if ring_len & (ring_len - 1) != 0:
+                raise ValueError(
+                    "--linear-replayssm-cache-len must be a power of two for the "
+                    f"circular spec-verify ring, got {ring_len}."
+                )
+            # ring_len >= 2 * max drafts is checked in
+            # _validate_gdn_replayssm_spec_ring() (draft tokens not derived yet).
+            # Closed-loop exact fold: the flush replays raw ring inputs through
+            # the recurrent update into the checkpoint, bit-identical to the
+            # recurrent baseline -- which keeps its state in fp32. A 16-bit
+            # checkpoint would re-quantize the exactly-folded state every flush
+            # and become the dominant residual error source, so require fp32.
+            if self.mamba_ssm_dtype is None:
+                logger.info(
+                    "--enable-gdn-replayssm-spec: setting --mamba-ssm-dtype "
+                    "float32 (the closed-loop exact fold requires the fp32 SSM "
+                    "checkpoint for recurrent-parity)."
+                )
+                self.mamba_ssm_dtype = "float32"
+            elif self.mamba_ssm_dtype != "float32":
+                raise ValueError(
+                    "--enable-gdn-replayssm-spec requires --mamba-ssm-dtype "
+                    f"float32, got {self.mamba_ssm_dtype!r}. The closed-loop "
+                    "exact fold keeps the committed state bit-identical to the "
+                    "recurrent baseline, which is only meaningful against the "
+                    "fp32 checkpoint; a 16-bit checkpoint would re-quantize it "
+                    "every flush."
+                )
+
+    def _validate_gdn_replayssm_spec_ring(self):
+        """Enforce ring_len >= 2 * max draft tokens for the spec-verify ring.
+
+        Early-flush margin: write_pos + spec_len <= ring_len must hold on every
+        verify step (see _advance_gdn_spec_cursors_kernel). Runs after
+        handle_speculative_decoding() so the (adaptive-aware) max is final;
+        MambaPool re-checks at ring allocation as a backstop.
+        """
+        if not self.enable_gdn_replayssm_spec:
+            return
+        max_drafts = self.max_speculative_num_draft_tokens
+        if max_drafts is None:
+            return
+        ring_len = self.linear_replayssm_cache_len
+        if ring_len < 2 * max_drafts:
+            raise ValueError(
+                "--linear-replayssm-cache-len must be >= 2 * the maximum "
+                "speculative draft-token count for the spec-verify ring "
+                f"(early-flush margin), got {ring_len} < {2 * max_drafts}."
+            )
+
     def _handle_legacy_cp_arguments(self):
         legacy_mode_to_strategy = {
             "in-seq-split": "zigzag",

python/sglang/srt/speculative/spec_utils.py

@@ -696,6 +696,73 @@ def commit_mamba_states_after_verify(
     model_runner = target_worker.model_runner
     if mambaish_config(model_runner.model_config) is None:
         return
+
+    # ReplaySSM spec-verify path (Part B of #28511): the accepted drafts already
+    # live in the per-slot circular ring (written during verify). Instead of
+    # scattering an intermediate full SSM state into `temporal`, advance the
+    # block-keyed cursors by the accepted count (the ring owns the SSM state; the
+    # verify/flush kernel folds it into `temporal` periodically). The CONV state
+    # still needs its usual accept-rollback, so we keep the conv-window scatter and
+    # skip only the SSM scatter. GDN-only + linear-chain (topk<=1) -- the runtime
+    # ring is allocated only then; KDA never allocates the cursors.
+    req_pool = model_runner.req_to_token_pool
+    mamba_pool = getattr(req_pool, "mamba_pool", None)
+    if (
+        mamba_pool is not None
+        and getattr(mamba_pool, "replayssm_cache_base", None) is not None
+        and not getattr(mamba_pool, "replayssm_is_kda", False)
+    ):
+        if batch.forward_mode.is_idle() or accept_index.numel() == 0:
+            return
+        from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_decode import (
+            commit_gdn_replayssm_spec,
+        )
+        from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+            fused_conv_window_scatter_with_mask,
+        )
+
+        spec_state = req_pool.get_speculative_mamba2_params_all_layers()
+        bs = accept_lens.shape[0]
+        state_batch_indices = req_pool.get_mamba_indices(batch.req_pool_indices)
+        # Advance the per-slot circular cursors by the accepted count (incl. the
+        # bonus token). max_cache_len = ring length L = replayssm_d.shape[-2].
+        commit_gdn_replayssm_spec(
+            write_pos=mamba_pool.replayssm_write_pos,
+            cache_base=mamba_pool.replayssm_cache_base,
+            is_flush=mamba_pool.replayssm_is_flush,
+            num_accepted=accept_lens,  # [bs], includes the bonus token
+            state_batch_indices=state_batch_indices,
+            max_cache_len=spec_state.replayssm_d.shape[-2],
+            max_spec_len=draft_token_num,
+            null_block_id=-1,  # SGLang: valid slots >= 0, padding == -1
+        )
+        # Roll back / commit the conv state to the last accepted draft step
+        # (same logic as the recurrent commit, but conv-only).
+        accept_indices_offset = torch.arange(
+            0,
+            bs * draft_token_num,
+            step=draft_token_num,
+            dtype=accept_lens.dtype,
+            device=accept_lens.device,
+        )
+        req_idx = torch.arange(bs, dtype=torch.int64, device=accept_lens.device)
+        last_correct_step_indices = (
+            accept_index[req_idx, (accept_lens - 1).to(torch.int64)]
+            - accept_indices_offset
+        )
+        fused_conv_window_scatter_with_mask(
+            spec_state.conv[0],
+            spec_state.intermediate_conv_window[0],
+            state_batch_indices,
+            last_correct_step_indices,
+        )
+        # NOTE: radix mamba prefix-caching (mamba_track / extra_buffer) would need
+        # a device-side force-flush so `temporal` reflects the ring before a
+        # snapshot; not wired for Part B (server_args forbids extra_buffer with
+        # --enable-gdn-replayssm-spec), so the per-track scatters are intentionally
+        # skipped here.
+        return
+
     attn_backend = model_runner.attn_backend
 
     bs = accept_lens.shape[0]