[GDN] Amortize ReplaySSM checkpoint materialization (#35544)

python/sglang/srt/arg_groups/attention_hook.py

@@ -371,9 +371,9 @@ def handle_linear_attn_backend(server_args: Any):
             )
 
     # ReplaySSM spec-verify (Part B of #28511): linear-chain target verify via
-    # fold-every-commit -- the verify stores each draft step's raw inputs into
+    # compact cached replay. Verify stores normalized keys, update vectors,
-    # the per-slot (rawv, rawk, g, beta) window and the commit replays the
+    # and fp32 log-decays; accepted BF16 windows are materialized with
-    # accepted prefix into the fp32 checkpoint. The intra-window interaction
+    # compensated hi/lo accumulation. 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.
@@ -439,8 +439,8 @@ def handle_linear_attn_backend(server_args: Any):
         if cfg.mamba_ssm_dtype is None:
             logger.info(
                 "--enable-linear-replayssm-spec: setting --mamba-ssm-dtype "
-                "float32 (the closed-loop exact fold keeps the SSM checkpoint "
+                "float32 (cached replay uses compensated checkpoint "
-                "bit-identical to the recurrent baseline)."
+                "projection and materialization)."
             )
             declare_resolution(
                 server_args,
@@ -450,10 +450,8 @@ def handle_linear_attn_backend(server_args: Any):
         elif cfg.mamba_ssm_dtype != "float32":
             logger.warning(
                 "--enable-linear-replayssm-spec with --mamba-ssm-dtype=%s: the "
-                "closed-loop fold re-quantizes the committed state each "
+                "compact checkpoint is materialized after each accepted "
-                "commit/flush (fp32 keeps it bit-exact to the fp32 recurrent "
+                "verify window; validate long-sequence accuracy and throughput.",
-                "baseline), so it may drift over long sequences. Validate "
-                "accuracy for your model.",
                 cfg.mamba_ssm_dtype,
             )
 

python/sglang/srt/configs/mamba_utils.py

@@ -125,29 +125,35 @@ class BaseLinearStateParams(ABC):
         ) * len(self.layers)
 
     def replayssm_ring_bytes_per_req(self, record_len: int) -> int:
-        """Per-slot bytes of the ReplaySSM spec-verify fold window (all
+        """ReplaySSM spec-verify scratch bytes across all layers.
-        layers). Not part of ``mamba_cache_per_req``, so the memory solver
+
-        must charge it separately. MUST mirror the ``MambaPool`` allocation:
+        GDN keeps compact d/k/g plus low parts for the activation-dtype d/k
-        raw v/k in the conv dtype + fp32 beta, plus the fp32 gate ring
+        rings. KDA keeps its raw-input fold window and d/k rings.
-        (per-head scalar for GDN, per-K vector for KDA). KDA additionally
+        """
-        keeps the chunked d/k rings under spec (its forward_decode routes on
-        their presence), also in the conv dtype."""
         hv, v_dim, k_dim = self.shape.temporal
         h_k = self.shape.num_k_heads_per_tp
         conv_b = self.dtype.conv.itemsize
         fp32_b = 4
-        per_layer = (
-            hv * record_len * v_dim * conv_b  # rawv
-            + h_k * record_len * k_dim * conv_b  # rawk
-            + hv * record_len * fp32_b  # beta (fp32)
-            # g (fp32): GDN per-head scalar, KDA per-K vector
-            + hv * record_len * (k_dim if self.is_kda else 1) * fp32_b
-        )
         if self.is_kda:
-            per_layer += (
+            per_layer = (
-                hv * record_len * v_dim * conv_b  # d
+                hv * record_len * v_dim * conv_b  # rawv
+                + h_k * record_len * k_dim * conv_b  # rawk
+                + hv * record_len * fp32_b  # beta
+                + hv * record_len * k_dim * fp32_b  # vector g
+                + hv * record_len * v_dim * conv_b  # d
                 + h_k * record_len * k_dim * conv_b  # k
             )
+        else:
+            per_layer = (
+                hv * record_len * v_dim * conv_b  # d
+                + h_k * record_len * k_dim * conv_b  # normalized k
+                + hv * record_len * fp32_b  # scalar g
+            )
+            if self.dtype.conv != torch.float32:
+                per_layer += (
+                    hv * record_len * v_dim * conv_b  # d low part
+                    + h_k * record_len * k_dim * conv_b  # normalized-k low part
+                )
         return per_layer * len(self.layers)
 
     @property

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

@@ -853,6 +853,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
                     mamba_pool=mamba_pool,
                     layer_cache=mamba_cache_params,
                     cache_indices=cache_indices,
+                    replay_indices=forward_batch.req_pool_indices,
                     query_start_loc=query_start_loc,
                     draft_token_num=forward_batch.spec_info.draft_token_num,
                 )
@@ -1017,6 +1018,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
         mamba_pool: MambaPool,
         layer_cache: "MambaPool.SpeculativeState",
         cache_indices: torch.Tensor,
+        replay_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
         draft_token_num: int,
     ) -> torch.Tensor:
@@ -1024,15 +1026,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
 
         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
+        appends this window's drafts to the compact ``(d, k, g)`` rings. BF16
-        reconstruction; raw ``v`` / pre-norm ``k`` / fp32 ``beta`` for the
+        checkpoints also keep low parts for compensated materialization. The rings
-        closed-loop exact fold that replays the recurrent update into the fp32
+        are per-layer (sliced via ``mamba2_layer_cache``), while the cursors
-        checkpoint at flush). The rings are PER-LAYER
+        (write_pos, cache_base, is_flush) are request-slot pool attributes shared
-        (sliced via ``mamba2_layer_cache``), while the cursors (write_pos,
+        by all GDN layers of the step. The cursors advance once per accepted step;
-        cache_base, is_flush) are PER-SLOT pool attributes shared by all GDN layers
+        here we only read them and write this step's ring entries. GDN has K == V,
-        of the step; the cursors persist across steps and are advanced once per step
+        so ``temporal``
-        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.
         """
@@ -1065,21 +1065,18 @@ class GDNAttnBackend(MambaAttnBackendBase):
             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.
+            # BF16 compact D/K low parts; None for fp32 checkpoints.
-            # 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),
+            replay_indices=replay_indices,
-            # NOT in forward_metadata: the verify kernel reads/writes them
+            # Request-slot cursors live on the pool (shared across all GDN layers)
-            # block-keyed via ssm_state_indices and must NOT advance write_pos
+            # and advance only after acceptance, so the decode-path
-            # (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,
+            write_pos=mamba_pool.replayssm_spec_write_pos,
             cache_base=mamba_pool.replayssm_cache_base,
             is_flush=mamba_pool.replayssm_is_flush,
             max_cache_len=max_cache_len,
@@ -1090,6 +1087,9 @@ class GDNAttnBackend(MambaAttnBackendBase):
             # 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,
+            # Capacity folds are committed once across every GDN layer after
+            # acceptance; the active path is therefore a single launch/layer.
+            launch_mode="verify",
         )
         # Match the recurrent target_verify output shape (== value.shape).
         return out.reshape(value.shape)

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -2276,22 +2276,14 @@ class KVCacheConfigurator:
         # no longer reserves the (1 + D/ratio) intermediate factor -- the whole
         # budget goes to persistent slots (K sized like non-spec), which is how the
         # freed ~9GB turns into higher max_running.
-        # The ring is allocated per slot but is not part of mamba_cache_per_req;
+        # The ring is not part of mamba_cache_per_req. GDN replay is fixed-size
-        # the solve must charge it too or num_slots is over-provisioned.
+        # request scratch; KDA replay remains attached to each mamba slot.
         replayssm_active = get_exec().mamba.enable_linear_replayssm_spec and (
             self.hybrid_gdn_config is not None
             or kimi_linear_config(self.model_config) is not None
         )
         if replayssm_active:
-            # GDN sizes the fold window to the draft maximum; the KDA ring
+            record_len = get_exec().mamba.linear_replayssm_cache_len
-            # stays --linear-replayssm-cache-len long (mirrors MambaPool).
-            max_draft_tokens = max_speculative_num_draft_tokens()
-            if kimi_linear_config(self.model_config) is not None:
-                record_len = get_exec().mamba.linear_replayssm_cache_len
-            elif max_draft_tokens is not None:
-                record_len = max_draft_tokens
-            else:
-                record_len = get_exec().mamba.linear_replayssm_cache_len
             replayssm_ring_per_req = (
                 config.mamba2_cache_params.replayssm_ring_bytes_per_req(
                     record_len=record_len
@@ -2300,6 +2292,15 @@ class KVCacheConfigurator:
         else:
             replayssm_ring_per_req = 0
         replayssm_ring_per_req = int(replayssm_ring_per_req * pp_layer_scale)
+        if replayssm_active and kimi_linear_config(self.model_config) is None:
+            replay_req_slots = (
+                get_schedule().max_running_requests // self.ps.attn_dp_size + 1
+            )
+            replayssm_fixed_bytes = replayssm_ring_per_req * replay_req_slots
+            replayssm_ring_per_slot = 0
+        else:
+            replayssm_fixed_bytes = 0
+            replayssm_ring_per_slot = replayssm_ring_per_req
         if has_spec_dec:
             assert get_spec().speculative_num_draft_tokens is not None
             assert get_schedule().max_running_requests is not None
@@ -2380,11 +2381,12 @@ class KVCacheConfigurator:
                 intermediate_size = per_req * (capped_reqs + 1) * D
                 total_rest_memory = total_rest_memory - (intermediate_size / (1 << 30))
             else:
-                per_slot = per_req + replayssm_ring_per_req
+                per_slot = per_req + replayssm_ring_per_slot
                 get_context().override(
                     "mamba_pool.memory_budget",
                     max_mamba_cache_size=int(
-                        (mamba_budget_bytes - per_slot) // per_slot
+                        (mamba_budget_bytes - replayssm_fixed_bytes - per_slot)
+                        // per_slot
                     ),
                 )
 
@@ -2404,14 +2406,12 @@ class KVCacheConfigurator:
                 f"(4) use GPUs with more memory."
             )
 
-        # +1: the pool's padding slot is allocated alongside the request slots.
+        # +1 accounts for each pool's padding slot.
-        # ReplaySSM ring rides on every slot too (replayssm_ring_per_req is 0 when
-        # the ring is not allocated).
         mamba_state_memory = (
             (get_schedule().max_mamba_cache_size + 1)
-            * (stage_per_req + replayssm_ring_per_req)
+            * (stage_per_req + replayssm_ring_per_slot)
-            / (1 << 30)
+            + replayssm_fixed_bytes
-        )
+        ) / (1 << 30)
         return total_rest_memory - mamba_state_memory
 
 

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -568,10 +568,23 @@ class MambaRadixCache(BasePrefixCache):
                 # donate to the last flush boundary (where temporal is current)
                 # and reset the cursor, keeping the donated checkpoint consistent
                 # with its key length. page_size is asserted == 1, so no realign.
-                write_pos_buf = self.req_to_token_pool.mamba_pool.replayssm_write_pos
+                mamba_pool = self.req_to_token_pool.mamba_pool
+                write_pos_buf = mamba_pool.replayssm_write_pos
+                cursor_idx = req.kv.mamba_pool_idx
+                if write_pos_buf is None:
+                    write_pos_buf = getattr(
+                        mamba_pool, "replayssm_spec_write_pos", None
+                    )
+                    cursor_idx = req.kv.req_pool_idx
                 if write_pos_buf is not None:
-                    cache_len -= int(write_pos_buf[req.kv.mamba_pool_idx].item())
+                    cache_len -= int(write_pos_buf[cursor_idx].item())
-                    write_pos_buf[req.kv.mamba_pool_idx] = 0
+                    write_pos_buf[cursor_idx] = 0
+                    if (
+                        getattr(mamba_pool, "replayssm_spec_write_pos", None)
+                        is not None
+                    ):
+                        mamba_pool.replayssm_cache_base[cursor_idx] = 0
+                        mamba_pool.replayssm_is_flush[cursor_idx] = 0
             if cache_len is None:
                 cache_len = 0
             if cache_len != len(token_ids):
@@ -688,6 +701,11 @@ class MambaRadixCache(BasePrefixCache):
             if self.enable_mamba_extra_buffer
             else len(token_ids)
         )
+        spec_write_pos = getattr(
+            self.req_to_token_pool.mamba_pool, "replayssm_spec_write_pos", None
+        )
+        if not self.enable_mamba_extra_buffer and spec_write_pos is not None:
+            cache_len -= int(spec_write_pos[req.kv.req_pool_idx].item())
         if self.disable or cache_len is None:
             return _skip_cache_unfinished_req(req)
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -383,14 +383,8 @@ 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)
+        # Under GDN spec verify, rawv/rawk hold compact D/K low parts and beta is
-        #   replayssm_rawk: [num_layers, num_slots, H,  L, K]  (conv/activation dtype)
+        # None. KDA uses the same fields for its raw-input fold window.
-        #   replayssm_beta: [num_layers, num_slots, HV, L]     (fp32)
-        # The raw rings + beta exist only under --enable-linear-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
@@ -529,12 +523,13 @@ class MambaPool:
         self.linear_replayssm_cache_len = linear_replayssm_cache_len
         # ReplaySSM: the decode ring (--enable-linear-replayssm) allocates the
         # chunked (d, k) records + write_pos; the spec-verify flag
-        # (--enable-linear-replayssm-spec) always uses fold-every-commit and
+        # (--enable-linear-replayssm-spec) uses compact replay for GDN and raw
-        # allocates only the raw (v, k, g, beta) window -- no chunked records,
+        # fold-every-commit for KDA. The shared g allocation gates on
-        # no cursors (KDA additionally keeps d/k, see the allocation below).
+        # `_replayssm_on`.
-        # The shared g allocation gates on `_replayssm_on`.
         self.enable_linear_replayssm_spec = enable_linear_replayssm_spec
-        self.replayssm_spec_fold = bool(enable_linear_replayssm_spec)
+        self.replayssm_spec_fold = bool(
+            enable_linear_replayssm_spec and cache_params.is_kda
+        )
         _replayssm_on = enable_linear_replayssm or enable_linear_replayssm_spec
 
         # for disagg with nvlink
@@ -623,16 +618,15 @@ class MambaPool:
                 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
+                # GDN speculative replay is request-lifetime scratch. Size it by
-                # Ring dtype. DECODE ring (--enable-linear-replayssm): records
+                # active requests instead of every persistent radix-cache slot.
-                # follow the SSM dtype -- its flush folds `d` directly into the
+                num_slots = (
-                # state. SPEC-verify ring (--enable-linear-replayssm-spec): d/k feed
+                    spec_state_size + 1
-                # ONLY the one-shot output reconstruction (the closed-loop exact
+                    if enable_linear_replayssm_spec and not cache_params.is_kda
-                # fold replays the raw rings for state instead), so their
+                    else size + 1
-                # quantization noise stays below the bf16 output cast; keep them
+                )
-                # in the conv/activation dtype instead of the (fp32-enforced)
+                # Decode records follow the SSM dtype. Spec-verify compact d/k
-                # SSM dtype to halve the ring traffic. g stays fp32 everywhere
+                # records follow the activation dtype; g stays fp32.
-                # (exact-fold input). The two flags are mutually exclusive.
                 ring_dtype = conv_dtype if enable_linear_replayssm_spec else ssm_dtype
                 # Fold-every-commit: one verify window, no chunked (d, k)
                 # records. KDA is the exception on both counts: its window
@@ -674,14 +668,10 @@ class MambaPool:
                     dtype=torch.float32,
                     device=device,
                 )
-                # Closed-loop exact-fold rings (spec-verify only). Raw v / raw
+                # KDA still uses raw-input fold-every-commit. GDN materializes
-                # pre-norm k live in the conv (activation) dtype -- they are born
+                # its compact d/k/g history directly and needs no duplicate ring.
-                # there, so storage round-trips losslessly -- beta in fp32. The
+                if enable_linear_replayssm_spec and cache_params.is_kda:
-                # flush replays these through the recurrent update sequentially
+                    if cache_params.is_kda or not self.replayssm_spec_fold:
-                # (bit-identical to the recurrent baseline) instead of folding
-                # the chunked `d` records open-loop.
-                if enable_linear_replayssm_spec:
-                    if cache_params.is_kda:
                         # Backstop for the KDA ring invariants; this pool is
                         # sized with the final adaptive-aware draft maximum.
                         if L & (L - 1) != 0:
@@ -711,6 +701,20 @@ class MambaPool:
                         dtype=torch.float32,
                         device=device,
                     )
+                elif enable_linear_replayssm_spec and ring_dtype != torch.float32:
+                    # Low parts of compact D and normalized K. The rings follow
+                    # the activation dtype regardless of checkpoint dtype, so
+                    # materialization always needs both parts.
+                    replayssm_rawv = torch.zeros(
+                        size=(num_mamba_layers, num_slots, hv, record_len, v_dim),
+                        dtype=conv_dtype,
+                        device=device,
+                    )
+                    replayssm_rawk = torch.zeros(
+                        size=(num_mamba_layers, num_slots, h_k, record_len, k_dim),
+                        dtype=conv_dtype,
+                        device=device,
+                    )
 
             if speculative_num_draft_tokens is not None:
                 if _is_npu:
@@ -875,8 +879,12 @@ class MambaPool:
                     + (
                         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_linear_replayssm_spec
+                            f"beta={get_tensor_size_bytes(replayssm_beta) / GB:.3f}GB "
+                            if replayssm_beta is not None
+                            else ""
+                        )
+                        if replayssm_rawv is not None
                         else ""
                     )
                 )
@@ -884,26 +892,24 @@ class MambaPool:
             # 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(_replayssm_on and cache_params.is_kda)
-            # Persistent per-slot decode-position cursor for ReplaySSM. Shared
+            # Decode ReplaySSM remains keyed by persistent mamba slots.
-            # across all linear-attn layers; advanced once per decode forward by
-            # 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 _replayssm_on and not self.replayssm_spec_fold
+                if enable_linear_replayssm
+                else None
+            )
+            self.replayssm_spec_write_pos = (
+                torch.zeros((spec_state_size + 1,), dtype=torch.int32, device=device)
+                if enable_linear_replayssm_spec and not self.replayssm_spec_fold
                 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.
             self.replayssm_cache_base = (
-                torch.zeros((size + 1,), dtype=torch.int32, device=device)
+                torch.zeros((spec_state_size + 1,), dtype=torch.int32, device=device)
                 if enable_linear_replayssm_spec and not self.replayssm_spec_fold
                 else None
             )
             self.replayssm_is_flush = (
-                torch.zeros((size + 1,), dtype=torch.int8, device=device)
+                torch.zeros((spec_state_size + 1,), dtype=torch.int8, device=device)
                 if enable_linear_replayssm_spec and not self.replayssm_spec_fold
                 else None
             )
@@ -1031,12 +1037,6 @@ 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()
@@ -1047,46 +1047,22 @@ 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):
-        # Accept both the legacy 2-tuple (conv, temporal) and the 3-tuple that also
+        # Accept historical 3-tuples, but request-keyed replay scratch is not
-        # carries the ReplaySSM spec-verify cursors.
+        # restored with a physical checkpoint slot.
         if len(mamba_cache_cpu) == 3:
-            conv_cpu, temporal_cpu, cursors_cpu = mamba_cache_cpu
+            conv_cpu, temporal_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()
 
     _NON_TRANSFER_STATE_FIELDS = frozenset(
@@ -1349,10 +1325,21 @@ class HybridReqToTokenPool(ReqToTokenPool):
     # For chunk prefill req, we do not need to allocate mamba cache,
     # We could use allocated mamba cache instead.
     def alloc(self, reqs: List[Req]) -> Optional[List[int]]:
+        fresh_req_rows = [req.kv.req_pool_idx is None for req in reqs]
         select_index = super().alloc(reqs)
         if select_index is None:
             return None
 
+        spec_write_pos = getattr(self.mamba_pool, "replayssm_spec_write_pos", None)
+        if spec_write_pos is not None:
+            fresh_indices = [
+                idx for idx, fresh in zip(select_index, fresh_req_rows) if fresh
+            ]
+            if fresh_indices:
+                spec_write_pos[fresh_indices] = 0
+                self.mamba_pool.replayssm_cache_base[fresh_indices] = 0
+                self.mamba_pool.replayssm_is_flush[fresh_indices] = 0
+
         mamba_indices: list[torch.Tensor] = []
         mamba_ping_pong_track_buffers: list[torch.Tensor] = []
         for req in reqs:
@@ -1371,12 +1358,6 @@ 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.kv.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.kv.mamba_pool_idx] = 0
-                    self.mamba_pool.replayssm_is_flush[req.kv.mamba_pool_idx] = 0
             mamba_indices.append(req.kv.mamba_pool_idx)
             if self.enable_mamba_extra_buffer:
                 if req.kv.mamba_ping_pong_track_buffer is None:

python/sglang/srt/speculative/spec_utils.py

@@ -913,6 +913,7 @@ def commit_mamba_states_after_verify(
         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_circular,
             commit_gdn_replayssm_spec,
         )
         from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
@@ -922,37 +923,59 @@ def commit_mamba_states_after_verify(
         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
+        replay_indices = batch.req_pool_indices
-        # bonus token). max_cache_len = ring length L = replayssm_d.shape[-2].
+        last_correct_step_indices, mamba_steps_to_track = _verify_commit_step_indices(
-        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).
-        last_correct_step_indices, _ = _verify_commit_step_indices(
             batch=batch,
             accept_index=accept_index,
             accept_lens=accept_lens,
             draft_token_num=draft_token_num,
         )
+        # 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(
+            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],
+            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,
         )
-        # NOTE: radix mamba prefix-caching (mamba_track / extra_buffer) would need
+        if batch.mamba_track_indices is not None:
-        # a device-side force-flush so `temporal` reflects the ring before a
+            fused_conv_window_scatter_with_mask(
-        # snapshot; not wired for Part B (server_args forbids extra_buffer with
+                spec_state.conv[0],
-        # --enable-linear-replayssm-spec), so the per-track scatters are intentionally
+                spec_state.intermediate_conv_window[0],
-        # skipped here.
+                batch.mamba_track_indices,
+                mamba_steps_to_track,
+            )
         return
 
     # KDA ReplaySSM (fold-every-commit): KDA keeps its own recurrent verify kernel

test/registered/unit/mem_cache/test_replayssm_ring_accounting.py

@@ -28,10 +28,8 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 # temporal = (hv=4, v_dim=8, k_dim=8), num_k_heads_per_tp = 4, record_len = 8,
 # 2 layers. conv bf16 (2B), fp32 gate/beta (4B). Ring tensors (per slot, per
 # layer):
-#   rawv hv*RL*v_dim, rawk h_k*RL*k_dim  -> conv dtype
+# GDN stores compact d/k plus scalar g. KDA stores raw v/k, vector g, beta,
-#   g    hv*RL (GDN) / hv*RL*k_dim (KDA) -> fp32
+# and its existing d/k rings.
-#   beta hv*RL                           -> fp32
-#   d/k  like rawv/rawk                  -> conv dtype (KDA only)
 DTYPE = Mamba2StateDType(conv=torch.bfloat16, temporal=torch.float32)
 RL = 8
 LAYERS = [0, 1]
@@ -44,7 +42,7 @@ def _kda_params():
     return KimiLinearCacheParams(shape=shape, dtype=DTYPE, layers=LAYERS)
 
 
-def _gdn_params():
+def _gdn_params(temporal_dtype=torch.float32):
     # Only shape.temporal and shape.num_k_heads_per_tp are read here; the rest
     # are dummy (the accounting does not depend on them).
     shape = Mamba2StateShape(
@@ -59,15 +57,20 @@ def _gdn_params():
         conv_kernel=0,
         num_k_heads_per_tp=4,
     )
-    return Mamba2CacheParams(shape=shape, dtype=DTYPE, layers=LAYERS)
+    dtype = Mamba2StateDType(conv=torch.bfloat16, temporal=temporal_dtype)
+    return Mamba2CacheParams(shape=shape, dtype=dtype, layers=LAYERS)
 
 
 class TestReplaySSMRingAccounting(CustomTestCase):
     def test_gdn_fold(self):
-        # fold window: rawv 512 + rawk 512 + g(scalar, 4*8*4) 128 + beta 128 = 1280
+        # d 512 + normalized k 512 + scalar g 128 + d/k low parts 1024 = 2176
         self.assertEqual(
             _gdn_params().replayssm_ring_bytes_per_req(record_len=RL),
-            1280 * len(LAYERS),
+            2176 * len(LAYERS),
+        )
+        self.assertEqual(
+            _gdn_params(torch.bfloat16).replayssm_ring_bytes_per_req(record_len=RL),
+            2176 * len(LAYERS),
         )
 
     def test_kda_fold(self):