[Feat] gdn: run DFLASH and DSPARK verify on the ReplaySSM spec ring

python/sglang/srt/arg_groups/attention_hook.py

@@ -443,9 +443,10 @@ def handle_linear_attn_backend(server_args: Any):
     # 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 sizes the window to the draft maximum; KDA (kda_backend) keeps a
-    # --linear-replayssm-cache-len window and folds via its own fused
-    # verify ring-write + commit_kda_replayssm_after_verify.
+    # GDN keeps a --linear-replayssm-cache-len ring that must hold the draft
+    # window (MambaPool checks it); KDA (kda_backend) keeps a window of the
+    # same length and folds via its own fused verify ring-write +
+    # commit_kda_replayssm_after_verify.
     if cfg.enable_linear_replayssm_spec:
         if cfg.speculative_eagle_topk not in (None, 1):
             raise ValueError(

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

@@ -547,14 +547,15 @@ class ExecMamba(msgspec.Struct):
     ] = False
     linear_replayssm_cache_len: A[
         int,
-        "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps.",
+        "Ring-buffer length L for ReplaySSM linear-attn decode. The full recurrent state is flushed to HBM every L decode steps. With --enable-linear-replayssm-spec on GDN models it is also the spec-verify ring length, which must be a power of two, at least 16 and at least the number of draft tokens.",
     ] = 16
     # ReplaySSM spec-verify (Part B of RFC #28511): linear-attn target-verify via
     # fold-every-commit instead of per-draft full-state snapshots -- the verify
     # stores each draft step's raw inputs into a per-slot window and the commit
-    # replays the accepted prefix into the fp32 checkpoint. GDN sizes the window
-    # to the draft maximum; KDA folds a (raw v, pre-norm k, gate, beta) ring of
-    # length --linear-replayssm-cache-len. Linear-chain (topk <= 1) only.
+    # replays the accepted prefix into the fp32 checkpoint. GDN keeps a ring of
+    # length --linear-replayssm-cache-len (a power of two, at least 16, that
+    # holds the draft window); KDA folds a (raw v, pre-norm k, gate, beta) ring
+    # of the same length. Linear-chain (topk <= 1) only.
     enable_linear_replayssm_spec: A[
         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.",

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

@@ -1572,6 +1572,45 @@ def update_mamba_state_after_mtp_verify(
             )
             return
 
+        # ReplaySSM-GDN (the request cursors exist only there): the verify
+        # appended this window's drafts to the compact request ring and wrote
+        # no snapshots. Fold every accepted window, so `temporal` is the
+        # committed state after each step as with the KDA fold above; radix
+        # donation and state copies read the slot directly. Chain layout only.
+        if getattr(mamba_pool, "replayssm_cache_base", None) is not None:
+            assert req_pool_indices is not None, (
+                "the GDN ReplaySSM commit needs the request rows of the ring"
+            )
+            if request_number == 0:
+                return
+            from sglang.kernels.ops.attention.linear.gdn_replayssm_spec_decode import (
+                commit_gdn_replayssm_after_verify,
+            )
+
+            commit_gdn_replayssm_after_verify(
+                spec_state=mamba_caches,
+                write_pos=mamba_pool.replayssm_spec_write_pos,
+                cache_base=mamba_pool.replayssm_cache_base,
+                is_flush=mamba_pool.replayssm_is_flush,
+                state_batch_indices=state_indices_tensor,
+                replay_indices=req_pool_indices[:request_number],
+                accept_lens=last_correct_step_indices + 1,
+                last_correct_step_indices=last_correct_step_indices,
+                fold_every_commit=True,
+                mamba_track_indices=mamba_track_indices,
+                mamba_steps_to_track=mamba_steps_to_track,
+                src_indices_raw=src_indices_raw,
+                null_block_id=-1,
+            )
+            self._update_ple_state_after_mtp_verify(
+                state_indices_tensor,
+                last_correct_step_indices,
+                mamba_track_indices,
+                mamba_steps_to_track,
+                src_indices_raw,
+            )
+            return
+
         # KDA fused-accept: the next verify seeds itself in-kernel from the
         # accepted checkpoint slot (recurrent_kda's num_accepted_tokens), so the
         # SSM state never round-trips through `temporal` and only the conv
@@ -1598,6 +1637,13 @@ def update_mamba_state_after_mtp_verify(
             )
             return
 
+        assert (
+            mamba_caches.temporal.numel() == 0
+            or mamba_caches.intermediate_ssm is not None
+        ), (
+            "per-draft SSM snapshots are not allocated (ReplaySSM spec verify), "
+            "but no ReplaySSM commit handled this pool"
+        )
         scatter_mamba_states_after_mtp_verify(
             mamba_caches,
             state_indices_tensor,

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

@@ -25,6 +25,7 @@
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.runtime_context import get_exec, get_memory, get_schedule
+from sglang.srt.speculative.ragged_verify import ragged_verify_compact_enabled
 from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu, is_xpu
 from sglang.srt.utils.common import is_pin_memory_available, rank0_log
 
@@ -248,6 +249,22 @@ def _validate_gdn_linear_attn_backends(backends: LinearAttnBackends) -> None:
         )
 
 
+def _validate_gdn_target_verify_layout(model_runner: ModelRunner) -> None:
+    """Compact verify packs per-request verify lengths into one token list,
+    but GDN target verify views a dense [batch, draft tokens] window from the
+    conv on. Cap-accept verifies the dense window and only trims acceptance."""
+    if model_runner.is_draft_worker:
+        return
+    if not model_runner.spec_algorithm.supports_ragged_verify():
+        return
+    if ragged_verify_compact_enabled():
+        raise ValueError(
+            "SGLANG_RAGGED_VERIFY_MODE=compact is not supported on GDN models: "
+            "GDN target verify views a dense [batch, draft tokens] window. Use "
+            "SGLANG_RAGGED_VERIFY_MODE=static or cap-accept."
+        )
+
+
 class GDNKernelDispatcher:
     """Dispatches GDN kernel calls to the appropriate backend per mode."""
 
@@ -521,6 +538,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
 
     def __init__(self, model_runner: ModelRunner):
         _validate_gdn_linear_attn_backends(model_runner.linear_attn_backends)
+        _validate_gdn_target_verify_layout(model_runner)
         super().__init__(model_runner)
         self.enable_mis = get_exec().features.enable_mis
         self.mis_metadata: Optional[GDNMISMetadata] = None

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -1661,18 +1661,6 @@ def _build_hybrid_req_pool(
         max_num_reqs: int,
         extra_max_context_len: int,
     ) -> ReqToTokenPool:
-        # DSPARK/DFLASH commit routes through the backend fold (KDA-only); a
-        # non-KDA model there would scatter a None intermediate_ssm and crash.
-        _algo = (get_spec().speculative_algorithm or "").upper()
-        if (
-            get_exec().mamba.enable_linear_replayssm_spec
-            and _algo in ("DSPARK", "DFLASH")
-            and self.hybrid_kda_config is None
-        ):
-            raise ValueError(
-                "--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA "
-                "model; got a non-KDA model."
-            )
         req_to_token_pool = HybridReqToTokenPool(
             size=max_num_reqs,
             mamba_size=get_schedule().max_mamba_cache_size,

python/sglang/srt/mem_cache/memory_pool.py

@@ -762,6 +762,25 @@ def __init__(
                 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
+                if enable_linear_replayssm_spec and not cache_params.is_kda:
+                    # The GDN verify writes its window at ring positions
+                    # [write_pos, write_pos + T) and drops stores past L; the
+                    # commit folds early enough to keep the window inside the
+                    # ring only if a whole window fits. The compact commit
+                    # reduces over the whole ring with tl.dot, which needs at
+                    # least 16 positions.
+                    if L & (L - 1) != 0:
+                        raise ValueError(
+                            "spec-verify ring length --linear-replayssm-cache-len "
+                            f"must be a power of two, got {L}"
+                        )
+                    min_len = max(16, speculative_num_draft_tokens or 0)
+                    if L < min_len:
+                        raise ValueError(
+                            "spec-verify ring too small: --linear-replayssm-cache-len "
+                            f"{L} < {min_len} (at least 16 and the draft tokens, "
+                            f"{speculative_num_draft_tokens})"
+                        )
                 # GDN speculative replay is request-lifetime scratch. Size it by
                 # active requests instead of every persistent radix-cache slot.
                 num_slots = (

python/sglang/srt/speculative/spec_utils.py

@@ -982,16 +982,10 @@ def commit_mamba_states_after_verify(
         if batch.forward_mode.is_idle() or accept_index.numel() == 0:
             return
         from sglang.kernels.ops.attention.linear.gdn_replayssm_spec_decode import (
-            commit_gdn_replayssm_circular,
-            commit_gdn_replayssm_spec,
-        )
-        from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
-            fused_conv_window_scatter_with_mask,
+            commit_gdn_replayssm_after_verify,
         )
 
         spec_state = req_pool.get_speculative_mamba2_params_all_layers()
-        state_batch_indices = req_pool.get_mamba_indices(batch.req_pool_indices)
-        replay_indices = batch.req_pool_indices
         if prepared_step_indices is None:
             last_correct_step_indices, mamba_steps_to_track = (
                 _verify_commit_step_indices(
@@ -1003,54 +997,24 @@ def commit_mamba_states_after_verify(
             )
         else:
             last_correct_step_indices, mamba_steps_to_track = prepared_step_indices
-        # Advance the per-request circular cursors by the accepted count (incl. the
-        # bonus token). max_cache_len = ring length L = replayssm_d.shape[-2].
-        commit_gdn_replayssm_spec(
+        # Advance the request cursors by the accepted count (incl. the bonus
+        # token), fold capacity rows into `temporal`, and roll back conv.
+        commit_gdn_replayssm_after_verify(
+            spec_state=spec_state,
             write_pos=mamba_pool.replayssm_spec_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
-            replay_indices=replay_indices,
-            max_cache_len=spec_state.replayssm_d.shape[-2],
+            state_batch_indices=req_pool.get_mamba_indices(batch.req_pool_indices),
+            replay_indices=batch.req_pool_indices,
+            accept_lens=accept_lens,
+            last_correct_step_indices=last_correct_step_indices,
             max_spec_len=draft_token_num,
             fold_every_commit=spec_state.temporal.dtype != torch.float32,
-            null_block_id=-1,  # SGLang: valid slots >= 0, padding == -1
-        )
-        # Capacity rows fold all layers in one launch; track rows snapshot the
-        # exact crossing state without disturbing the active circular history.
-        commit_gdn_replayssm_circular(
-            checkpoint_state=spec_state.temporal,
-            d_cache=spec_state.replayssm_d,
-            k_cache=spec_state.replayssm_k,
-            g_cache=spec_state.replayssm_g,
-            d_residual_cache=spec_state.replayssm_rawv,
-            k_residual_cache=spec_state.replayssm_rawk,
-            state_batch_indices=state_batch_indices,
-            replay_indices=replay_indices,
-            write_pos=mamba_pool.replayssm_spec_write_pos,
-            cache_base=mamba_pool.replayssm_cache_base,
-            is_flush=mamba_pool.replayssm_is_flush,
-            accept_lens=accept_lens,
             mamba_track_indices=batch.mamba_track_indices,
             mamba_steps_to_track=mamba_steps_to_track,
-            null_block_id=-1,
-        )
-        # Roll back active conv state and snapshot its interval-crossing window.
-        fused_conv_window_scatter_with_mask(
-            spec_state.conv[0],
-            spec_state.intermediate_conv_window[0],
-            state_batch_indices,
-            last_correct_step_indices,
-            src_indices_raw,
+            src_indices_raw=src_indices_raw,
+            null_block_id=-1,  # SGLang: valid slots >= 0, padding == -1
         )
-        if batch.mamba_track_indices is not None:
-            fused_conv_window_scatter_with_mask(
-                spec_state.conv[0],
-                spec_state.intermediate_conv_window[0],
-                batch.mamba_track_indices,
-                mamba_steps_to_track,
-                src_indices_raw,
-            )
         return
 
     # KDA ReplaySSM (fold-every-commit): KDA keeps its own recurrent verify kernel