[gdn] support replayssm with extra buffer (#32692)

test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py

@@ -0,0 +1,234 @@
+"""GDN ReplaySSM fold-every-commit: fused ring-write + commit fold.
+
+The design contract is BITWISE parity with the recurrent verify + per-draft
+snapshot baseline, so every case asserts ``torch.equal``: ring-write leaves the
+verify output unchanged; the folded checkpoint equals the accepted-step
+snapshot (fp32 and bf16, each vs the same-dtype baseline); HAS_TRACK stores the
+crossing-step state and skips -1 steps / null slots; a 256-iteration
+verify->commit chain stays bitwise equal at every step (no accumulation channel
+on the state path -- the long-decode drift failure mode).
+"""
+
+import sys
+import unittest
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[4]))
+
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
+    fused_sigmoid_gating_delta_rule_update,
+)
+from sglang.kernels.ops.attention.fla.gdn_replayssm_spec_fold import (
+    commit_gdn_replayssm_fold_all_layers,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
+
+B, T = 3, 4
+H, HV = 4, 8
+K = V = 64
+NUM_SLOTS = 8
+DEVICE = "cuda"
+
+
+def _make_window(step_seed: int):
+    gen = torch.Generator(device=DEVICE).manual_seed(step_seed)
+
+    def rand(*shape, dtype=torch.bfloat16):
+        return torch.randn(*shape, device=DEVICE, dtype=dtype, generator=gen)
+
+    return {
+        "q": rand(1, B * T, H, K),
+        "k": rand(1, B * T, H, K),
+        "v": rand(1, B * T, HV, V),
+        "a": rand(B * T, HV),
+        "b": rand(B * T, HV),
+    }
+
+
+def _run_verify(inputs, gating, state, slots, *, snapshots=None, rings=None):
+    kwargs = {}
+    if snapshots is not None:
+        kwargs.update(
+            intermediate_states_buffer=snapshots,
+            intermediate_state_indices=slots,
+            cache_steps=T,
+        )
+    if rings is not None:
+        # Per-layer views, matching the backend's mamba2_layer_cache slices.
+        kwargs.update(
+            cache_ring=True,
+            replayssm_rawv=rings["rawv"][0],
+            replayssm_rawk=rings["rawk"][0],
+            replayssm_g=rings["g"][0],
+            replayssm_beta=rings["beta"][0],
+        )
+    cu_seqlens = torch.arange(0, B * T + 1, step=T, dtype=torch.int32, device=DEVICE)
+    return fused_sigmoid_gating_delta_rule_update(
+        A_log=gating["A_log"],
+        dt_bias=gating["dt_bias"],
+        softplus_beta=1.0,
+        softplus_threshold=20.0,
+        q=inputs["q"],
+        k=inputs["k"],
+        v=inputs["v"],
+        b=inputs["b"],
+        a=inputs["a"],
+        initial_state_source=state,
+        initial_state_indices=slots,
+        cu_seqlens=cu_seqlens,
+        use_qk_l2norm_in_kernel=True,
+        is_kda=False,
+        disable_state_update=True,
+        **kwargs,
+    )
+
+
+def _make_rings(dtype=torch.bfloat16):
+    return {
+        "rawv": torch.zeros(1, NUM_SLOTS, HV, T, V, device=DEVICE, dtype=dtype),
+        "rawk": torch.zeros(1, NUM_SLOTS, H, T, K, device=DEVICE, dtype=dtype),
+        "g": torch.zeros(1, NUM_SLOTS, HV, T, device=DEVICE, dtype=torch.float32),
+        "beta": torch.zeros(1, NUM_SLOTS, HV, T, device=DEVICE, dtype=torch.float32),
+    }
+
+
+def _fold(state, rings, slots, accept_lens, track_slots=None, track_steps=None):
+    commit_gdn_replayssm_fold_all_layers(
+        checkpoint_state=state,
+        rawv_cache=rings["rawv"],
+        rawk_cache=rings["rawk"],
+        g_cache=rings["g"],
+        beta_cache=rings["beta"],
+        ssm_state_indices=slots,
+        accept_lens=accept_lens,
+        max_cache_len=T,
+        num_k_heads=H,
+        mamba_track_indices=track_slots,
+        mamba_steps_to_track=track_steps,
+        null_block_id=-1,
+    )
+
+
+class TestGdnReplayssmSpecFold(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.gating = {
+            "A_log": torch.randn(HV, device=DEVICE) * 0.1,
+            "dt_bias": torch.randn(HV, device=DEVICE) * 0.1,
+        }
+        self.slots = torch.tensor([5, 2, 7], dtype=torch.int32, device=DEVICE)
+        self.accept_lens = torch.tensor([3, 1, 4], dtype=torch.int32, device=DEVICE)
+
+    def _state(self, dtype):
+        gen = torch.Generator(device=DEVICE).manual_seed(1)
+        return torch.randn(
+            NUM_SLOTS, HV, K, V, device=DEVICE, dtype=dtype, generator=gen
+        )
+
+    def test_ring_write_does_not_change_verify_output(self):
+        for dtype in (torch.float32, torch.bfloat16):
+            state = self._state(dtype)
+            inputs = _make_window(11)
+            out_plain = _run_verify(inputs, self.gating, state.clone(), self.slots)
+            out_ring = _run_verify(
+                inputs, self.gating, state.clone(), self.slots, rings=_make_rings()
+            )
+            self.assertTrue(torch.equal(out_plain, out_ring), f"{dtype=}")
+
+    def test_fold_matches_snapshot_baseline_bitwise(self):
+        for dtype in (torch.float32, torch.bfloat16):
+            state = self._state(dtype)
+            inputs = _make_window(22)
+
+            snapshots = torch.zeros(NUM_SLOTS, T, HV, K, V, device=DEVICE, dtype=dtype)
+            _run_verify(
+                inputs, self.gating, state.clone(), self.slots, snapshots=snapshots
+            )
+
+            fold_state = state.clone().unsqueeze(0).contiguous()
+            rings = _make_rings()
+            _run_verify(inputs, self.gating, fold_state[0], self.slots, rings=rings)
+            _fold(fold_state, rings, self.slots, self.accept_lens)
+
+            for s, n in zip(self.slots.tolist(), self.accept_lens.tolist()):
+                self.assertTrue(
+                    torch.equal(snapshots[s, n - 1], fold_state[0, s]),
+                    f"{dtype=} slot={s} accept_len={n}",
+                )
+            untouched = set(range(NUM_SLOTS)) - set(self.slots.tolist())
+            for s in untouched:
+                self.assertTrue(torch.equal(fold_state[0, s], state[s]))
+
+    def test_track_store_and_null_slots(self):
+        dtype = torch.float32
+        state = self._state(dtype)
+        inputs = _make_window(33)
+
+        snapshots = torch.zeros(NUM_SLOTS, T, HV, K, V, device=DEVICE, dtype=dtype)
+        _run_verify(inputs, self.gating, state.clone(), self.slots, snapshots=snapshots)
+
+        fold_state = state.clone().unsqueeze(0).contiguous()
+        rings = _make_rings()
+        _run_verify(inputs, self.gating, fold_state[0], self.slots, rings=rings)
+
+        track_slots = torch.tensor([1, 0, 3], dtype=torch.int64, device=DEVICE)
+        track_steps = torch.tensor([1, -1, 2], dtype=torch.int64, device=DEVICE)
+        slots_with_null = self.slots.clone()
+        slots_with_null[1] = -1
+        _fold(
+            fold_state,
+            rings,
+            slots_with_null,
+            self.accept_lens,
+            track_slots=track_slots,
+            track_steps=track_steps,
+        )
+
+        self.assertTrue(torch.equal(fold_state[0, 1], snapshots[5, 1]))
+        self.assertTrue(torch.equal(fold_state[0, 3], snapshots[7, 2]))
+        # Null slot: neither committed nor tracked (track step 1 is masked
+        # to -1 only for row 1's -1 step; row 1's slot itself was nulled).
+        self.assertTrue(torch.equal(fold_state[0, 2], state[2]))
+        self.assertTrue(torch.equal(fold_state[0, 0], state[0]))
+
+    def test_long_chain_no_accumulation(self):
+        """256 chained verify->commit iterations stay bitwise equal to the
+        baseline chain at every iteration (fp32 + bf16): no error channel
+        can accumulate with sequence length."""
+        num_iters = 256
+        for dtype in (torch.float32, torch.bfloat16):
+            base_state = self._state(dtype)
+            fold_state = base_state.clone().unsqueeze(0).contiguous()
+            snapshots = torch.zeros(NUM_SLOTS, T, HV, K, V, device=DEVICE, dtype=dtype)
+            gen = torch.Generator().manual_seed(7)
+            for it in range(num_iters):
+                inputs = _make_window(1000 + it)
+                accept_lens = torch.randint(1, T + 1, (B,), generator=gen).to(
+                    device=DEVICE, dtype=torch.int32
+                )
+
+                out_base = _run_verify(
+                    inputs, self.gating, base_state, self.slots, snapshots=snapshots
+                )
+                for s, n in zip(self.slots.tolist(), accept_lens.tolist()):
+                    base_state[s] = snapshots[s, n - 1]
+
+                rings = _make_rings()
+                out_fold = _run_verify(
+                    inputs, self.gating, fold_state[0], self.slots, rings=rings
+                )
+                _fold(fold_state, rings, self.slots, accept_lens)
+
+                self.assertTrue(torch.equal(out_base, out_fold), f"{dtype=} {it=}")
+                self.assertTrue(
+                    torch.equal(base_state, fold_state[0]), f"{dtype=} {it=}"
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_replayssm_ring_accounting.py

@@ -0,0 +1,67 @@
+"""ReplaySSM ring per-slot byte accounting (BaseLinearStateParams.replayssm_ring_bytes_per_req).
+
+The memory solver charges this on top of mamba_cache_per_req so num_slots is not
+over-provisioned (the ring is allocated per slot but is NOT part of the state
+cache cost). Pins the arithmetic against hand-computed byte counts for the
+fold window (raw v / pre-norm k / g / beta). If the MambaPool allocation
+changes shape, update both together.
+"""
+
+import pytest
+import torch
+
+from sglang.srt.configs.mamba_utils import (
+    Mamba2CacheParams,
+    Mamba2StateDType,
+    Mamba2StateShape,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+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
+#   g    hv*RL                          -> fp32
+#   beta hv*RL                          -> fp32
+DTYPE = Mamba2StateDType(conv=torch.bfloat16, temporal=torch.float32)
+RL = 8
+LAYERS = [0, 1]
+
+
+def _gdn_params():
+    # 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(
+        conv=[(4, 3)],
+        temporal=(4, 8, 8),
+        intermediate_size=0,
+        conv_dim=0,
+        ssm_state_size=0,
+        num_heads=0,
+        head_dim=0,
+        state_size=0,
+        conv_kernel=0,
+        num_k_heads_per_tp=4,
+    )
+    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
+        self.assertEqual(
+            _gdn_params().replayssm_ring_bytes_per_req(record_len=RL),
+            1280 * len(LAYERS),
+        )
+
+    def test_zero_len_ring(self):
+        self.assertEqual(_gdn_params().replayssm_ring_bytes_per_req(record_len=0), 0)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/spec/test_ngram_mamba_verify_update.py

@@ -99,6 +99,9 @@ class TestNgramMambaVerifyUpdate(CustomTestCase):
         target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify = (
             MagicMock()
         )
+        mamba_pool = target_worker.model_runner.req_to_token_pool.mamba_pool
+        mamba_pool.replayssm_spec_fold = False
+        mamba_pool.replayssm_cache_base = None
         return target_worker
 
     def test_mamba_verify_update_called_with_correct_indices(self):