Store mamba prefix-cache checkpoints at the configured SSM state dtype (#34820)

test/registered/kernel/ops/attention/test_kda_track_state.py

@@ -0,0 +1,183 @@
+import unittest
+
+import torch
+
+from sglang.kernels.ops.attention.fla.kda import chunk_kda
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=180, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+CHUNK_SIZE = 64
+
+_BACKENDS = {"triton": chunk_kda}
+HELION_AVAILABLE = True
+try:
+    import helion  # noqa: F401
+except ModuleNotFoundError as error:
+    # A broken install (transitive import failure) must stay loud; only the
+    # absent package downgrades the run to triton-only.
+    if error.name != "helion":
+        raise
+    HELION_AVAILABLE = False
+if HELION_AVAILABLE:
+    from sglang.kernels.ops.attention.helion.kda_prefill import (
+        chunk_kda as helion_chunk_kda,
+    )
+
+    _BACKENDS["helion"] = helion_chunk_kda
+
+
+def _make_varlen_inputs(seed, lens, num_heads=2, head_dim=128):
+    """Packed varlen KDA inputs: [1, sum(lens), H, D] plus a zero fp32 state pool."""
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    total = sum(lens)
+
+    def randn(*shape, dtype=torch.bfloat16):
+        return torch.randn(*shape, generator=generator, device="cuda", dtype=dtype)
+
+    q = randn(1, total, num_heads, head_dim)
+    k = randn(1, total, num_heads, head_dim)
+    v = (0.1 * randn(1, total, num_heads, head_dim, dtype=torch.float32)).to(
+        torch.bfloat16
+    )
+    gate = randn(1, total, num_heads, head_dim)
+    beta = torch.sigmoid(randn(1, total, num_heads, dtype=torch.float32)).to(
+        torch.bfloat16
+    )
+    a_log = randn(num_heads, dtype=torch.float32)
+    dt_bias = randn(num_heads * head_dim, dtype=torch.float32)
+    state = torch.zeros(
+        len(lens), num_heads, head_dim, head_dim, device="cuda", dtype=torch.float32
+    )
+    cu_seqlens = torch.tensor(
+        [0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
+    )
+    return q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens
+
+
+def _run_chunk_kda(
+    chunk_kda_fn, q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens, **kwargs
+):
+    return chunk_kda_fn(
+        # chunk_kda writes in place (the attention output lands in v, the gate
+        # cumsum in g); hand every run fresh copies so runs stay independent.
+        q=q.clone(),
+        k=k.clone(),
+        v=v.clone(),
+        g=gate.clone(),
+        beta=beta.clone(),
+        scale=q.shape[-1] ** -0.5,
+        initial_state=state,
+        initial_state_indices=torch.arange(
+            state.shape[0], device="cuda", dtype=torch.int32
+        ),
+        use_qk_l2norm_in_kernel=True,
+        cu_seqlens=cu_seqlens,
+        A_log=a_log,
+        dt_bias=dt_bias,
+        lower_bound=-5.0,
+        **kwargs,
+    )
+
+
+class TestKdaTrackState(CustomTestCase):
+    def test_helion_backend_ran(self):
+        """Visibility hook: without helion installed the snapshot check above
+        runs triton-only and the Helion track configs go untested — surface
+        that as an explicit skip instead of a silent pass."""
+        if not HELION_AVAILABLE:
+            self.skipTest("helion is not installed; triton backend only")
+
+    @torch.inference_mode()
+    def test_track_state_snapshots_fp32_accumulator(self):
+        """Bug regression: the mamba radix track path snapshots the SSM state at
+        the last chunk boundary of unaligned sequences into the fp32 state pool.
+        It used to read the per-chunk states `h` (activation dtype, bf16), so a
+        prefix-cache hit restored a bf16-rounded state while a cache miss kept
+        fp32. `track_state` must carry the in-kernel fp32 accumulator: identical
+        to the fp32 final state of a run truncated at the boundary, and strictly
+        more precise than the bf16 `h` row for the same boundary.
+        """
+        if not torch.cuda.is_available():
+            self.skipTest("requires CUDA")
+        for backend, chunk_kda_fn in _BACKENDS.items():
+            # num_heads=2 exercises the Helion small-head track config; 16
+            # crosses _PREFILL_SMALL_HEAD_THRESHOLD (12) to exercise the
+            # large-head varlen track config that real models take.
+            for num_heads in (2, 16):
+                with self.subTest(backend=backend, num_heads=num_heads):
+                    self._check_track_state(chunk_kda_fn, num_heads)
+
+    def _check_track_state(self, chunk_kda_fn, num_heads):
+        # seq0: 100 tokens, unaligned -> snapshot at the 64-token boundary
+        # (start of chunk 1). seq1: 64 tokens, aligned -> not tracked.
+        lens = [100, 64]
+        q, k, v, gate, beta, a_log, dt_bias, state, cu_seqlens = _make_varlen_inputs(
+            0, lens, num_heads=num_heads
+        )
+        num_heads, head_dim = q.shape[2], q.shape[3]
+
+        track_state = torch.full(
+            (len(lens), num_heads, head_dim, head_dim),
+            float("nan"),
+            device="cuda",
+            dtype=torch.float32,
+        )
+        track_chunk_idx = torch.tensor([1, -1], dtype=torch.int32, device="cuda")
+        _, h = _run_chunk_kda(
+            chunk_kda_fn,
+            q,
+            k,
+            v,
+            gate,
+            beta,
+            a_log,
+            dt_bias,
+            state,
+            cu_seqlens,
+            output_intermediate_states=True,
+            track_state=track_state,
+            track_chunk_idx=track_chunk_idx,
+        )
+
+        # The untracked row must stay untouched; the tracked row must be finite.
+        self.assertTrue(torch.all(torch.isnan(track_state[1])))
+        self.assertFalse(torch.any(torch.isnan(track_state[0])))
+
+        # Reference: truncate seq0 at the boundary; the pool's fp32 row then
+        # receives the in-place final state for the same prefix — the
+        # established fp32 path the snapshot must agree with.
+        ref_state = torch.zeros(
+            1, num_heads, head_dim, head_dim, device="cuda", dtype=torch.float32
+        )
+        ref_cu_seqlens = torch.tensor([0, CHUNK_SIZE], dtype=torch.int32, device="cuda")
+        _run_chunk_kda(
+            chunk_kda_fn,
+            q[:, :CHUNK_SIZE],
+            k[:, :CHUNK_SIZE],
+            v[:, :CHUNK_SIZE],
+            gate[:, :CHUNK_SIZE],
+            beta[:, :CHUNK_SIZE],
+            a_log,
+            dt_bias,
+            ref_state,
+            ref_cu_seqlens,
+        )
+        torch.testing.assert_close(track_state[0], ref_state[0], rtol=1e-5, atol=1e-5)
+
+        # The guard: h packs one row per (seq, chunk); row 1 is seq0's state at
+        # the boundary, rounded to bf16. If the snapshot were re-routed through
+        # h, it could not match the fp32 reference above.
+        self.assertTrue(
+            torch.equal(h[0, 1].float(), track_state[0].to(torch.bfloat16).float()),
+            "h row should be exactly the bf16 rounding of the fp32 snapshot",
+        )
+        self.assertFalse(
+            torch.equal(track_state[0], track_state[0].to(torch.bfloat16).float()),
+            "test inputs must make bf16 rounding lossy",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/attention/linear/kernels/test_kda_ptx.py

@@ -0,0 +1,131 @@
+"""Unit tests for the PTX KDA prefill routing wrapper."""
+
+import unittest
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.layers.attention.linear.kernels.kda_ptx import PtxKDAKernel
+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")
+
+
+class _RejectTriton:
+    def extend(self, *args, **kwargs):
+        raise AssertionError("native-eligible batch unexpectedly fell back to Triton")
+
+
+class TestPtxKDATrackRouting(CustomTestCase):
+    """Regression: a batch carrying the fp32 track snapshot buffer must not
+    take the native PTX path — the kernel cannot write the buffer, and the
+    backend copies it into the prefix-cache track slots unconditionally, so an
+    unwritten buffer silently corrupts later cache restores.
+    """
+
+    def _make_kernel(self):
+        kernel = PtxKDAKernel()
+        kernel._ensure_loaded = lambda: None
+        kernel._fwd = Mock(side_effect=AssertionError("native path must not run"))
+        return kernel
+
+    @staticmethod
+    def _inputs(seq_lens=(64, 100)):
+        total = sum(seq_lens)
+        H, D = 2, 128
+
+        def vals(offset):
+            return torch.full((1, total, H, D), offset, dtype=torch.bfloat16)
+
+        return {
+            "q": vals(0.1),
+            "k": vals(0.2),
+            "v": vals(0.3),
+            "g": vals(0.4),
+            "beta": torch.zeros(1, total, H, dtype=torch.bfloat16),
+            "ssm_states": torch.zeros(8, H, D, D, dtype=torch.float32),
+            "cache_indices": torch.tensor([1, 3], dtype=torch.int32),
+            "query_start_loc": torch.tensor(
+                [0] + list(torch.tensor(seq_lens).cumsum(0).tolist()),
+                dtype=torch.int32,
+            ),
+            "A_log": torch.zeros(H, dtype=torch.float32),
+            "dt_bias": torch.zeros(H * D, dtype=torch.float32),
+            "extend_seq_lens_cpu": list(seq_lens),
+        }
+
+    def test_batch_with_track_state_routes_to_triton(self):
+        kernel = self._make_kernel()
+        kernel._triton.extend = Mock(return_value="triton-out")
+        x = self._inputs()
+        track_state = torch.zeros(2, 2, 128, 128, dtype=torch.float32)
+        track_chunk_idx = torch.tensor([1, -1], dtype=torch.int32)
+
+        with patch(
+            "sglang.srt.layers.attention.linear.kernels.kda_ptx.mamba_cache_chunk_size",
+            return_value=64,
+        ):
+            out = kernel.extend(
+                x["q"],
+                x["k"],
+                x["v"],
+                x["g"],
+                x["beta"],
+                ssm_states=x["ssm_states"],
+                cache_indices=x["cache_indices"],
+                query_start_loc=x["query_start_loc"],
+                A_log=x["A_log"],
+                dt_bias=x["dt_bias"],
+                return_intermediate_states=True,
+                track_ssm_h_src=torch.tensor([1], dtype=torch.long),
+                track_state=track_state,
+                track_chunk_idx=track_chunk_idx,
+                extend_seq_lens_cpu=x["extend_seq_lens_cpu"],
+            )
+
+        self.assertEqual(out, "triton-out")
+        kernel._fwd.assert_not_called()
+        kernel._triton.extend.assert_called_once()
+        forwarded = kernel._triton.extend.call_args.kwargs
+        self.assertIs(forwarded["track_state"], track_state)
+        self.assertIs(forwarded["track_chunk_idx"], track_chunk_idx)
+
+    def test_batch_without_track_state_stays_native(self):
+        kernel = self._make_kernel()
+        kernel._triton = _RejectTriton()
+        h = torch.zeros(3, 2, 128, 128, dtype=torch.float32)
+
+        def fake_fwd(*args, **kwargs):
+            return [
+                args[2].clone(),  # out == v
+                kwargs["initial_state"].clone(),  # final_state
+                *([None] * 8),
+                h,  # result[10]
+            ]
+
+        kernel._fwd = fake_fwd
+        x = self._inputs()
+
+        out, h_out = kernel.extend(
+            x["q"],
+            x["k"],
+            x["v"],
+            x["g"],
+            x["beta"],
+            ssm_states=x["ssm_states"],
+            cache_indices=x["cache_indices"],
+            query_start_loc=x["query_start_loc"],
+            A_log=x["A_log"],
+            dt_bias=x["dt_bias"],
+            return_intermediate_states=True,
+            track_ssm_h_src=torch.empty(0, dtype=torch.long),
+            extend_seq_lens_cpu=x["extend_seq_lens_cpu"],
+        )
+
+        self.assertEqual(tuple(out.shape), (1, 164, 2, 128))
+        self.assertIs(h_out, h)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/attention/test_kda_helion_dispatcher.py

@@ -194,5 +194,56 @@ class TestHelionKDADispatcher(unittest.TestCase):
         self.assertEqual(args.linear_attn_backend, "helion")
 
 
+class TestKDATrackStateSnapshotDeclaration(unittest.TestCase):
+    """Bookkeeping: every KDA prefill kernel must declare whether extend()
+    honors the fp32 track snapshot (``supports_track_state_snapshot``).
+
+    KDAAttnBackend allocates the snapshot buffer whenever a tracked batch has
+    chunk-unaligned sequences and asserts the flag before use. A kernel that
+    serves extend() without the flag must reject tracked batches loudly
+    (NotImplementedError); a missing declaration used to mean the buffer was
+    silently left unwritten and prefix-cache restores read garbage (the
+    FlashKDA fallback once dropped the track arguments exactly this way).
+    """
+
+    def test_every_kda_prefill_kernel_declares_the_contract(self):
+        from sglang.srt.layers.attention.linear.kernels.kda_cutedsl import (
+            CuteDSLKDAKernel,
+        )
+        from sglang.srt.layers.attention.linear.kernels.kda_flashinfer import (
+            FlashInferKDAKernel,
+        )
+        from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (
+            FlashKDAKernel,
+        )
+        from sglang.srt.layers.attention.linear.kernels.kda_nvidia import (
+            NvidiaKDAKernel,
+        )
+        from sglang.srt.layers.attention.linear.kernels.kda_ptx import (
+            PtxKDAKernel,
+        )
+
+        # Native support or fallback that forwards the snapshot arguments.
+        for cls in (
+            TritonKDAKernel,
+            HelionKDAKernel,
+            NvidiaKDAKernel,
+            PtxKDAKernel,
+            FlashKDAKernel,
+        ):
+            self.assertTrue(
+                cls.supports_track_state_snapshot,
+                f"{cls.__name__} must declare supports_track_state_snapshot "
+                f"(native support or a fallback that forwards track_state)",
+            )
+        # Reject tracked batches loudly instead (extend() raises).
+        for cls in (CuteDSLKDAKernel, FlashInferKDAKernel):
+            self.assertFalse(
+                cls.supports_track_state_snapshot,
+                f"{cls.__name__} rejects tracked batches; it must not claim "
+                f"snapshot support it does not have",
+            )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/layers/attention/test_mamba_track_state_dtype.py

@@ -0,0 +1,94 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    MambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
+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")
+
+# Above the fp16 midpoint 1 + 2^-11 (single-rounds up to 1 + 2^-10) but below
+# the bf16 midpoint 1 + 2^-8 (rounds to 1.0, which then stays 1.0 in fp16):
+# any fp32 -> bf16 -> fp16 double rounding loses the increment. The 2^-23 tail
+# is the last fp32 mantissa bit at 1.x, so the probe is fp32-exact (a 2^-24
+# tail would round back to the midpoint itself).
+DOUBLE_ROUND_PROBE = 1.0 + 2.0**-11 + 2.0**-23
+
+
+class TestTrackMambaStateDtype(CustomTestCase):
+    """The fp32 track snapshot is cast to the pool dtype exactly once.
+
+    ``_track_mamba_state_extend`` reads the in-kernel fp32 snapshot
+    (``h_track_buf``) and casts it to ``ssm_states.dtype`` in a single ``.to``.
+    This must hold for every ``--mamba-ssm-dtype``: fp32 keeps full precision,
+    bf16 matches the (already correct) legacy path, and fp16 must not inherit
+    the old double rounding through the bf16 per-chunk states ``h``.
+    """
+
+    @staticmethod
+    def _run_track_copy(pool_dtype, h_track_buf, dst_slots, batch_rows):
+        metadata = ForwardMetadata(
+            has_mamba_track_mask=True,
+            # Only numel() gates the copy on this path; the h-row values
+            # themselves are unused when h_track_buf is given.
+            track_ssm_h_src=torch.zeros(len(dst_slots), dtype=torch.long),
+            track_ssm_h_dst=torch.tensor(dst_slots),
+            track_ssm_h_batch_src=torch.tensor(batch_rows),
+            track_ssm_final_src=torch.empty(0, dtype=torch.long),
+            track_ssm_final_dst=torch.empty(0, dtype=torch.long),
+            # Required by the dataclass; unused on this path.
+            query_start_loc=torch.zeros(1, dtype=torch.int32),
+            mamba_cache_indices=torch.zeros(1, dtype=torch.long),
+        )
+        ssm_states = torch.zeros(8, *h_track_buf.shape[1:], dtype=pool_dtype)
+        # The method touches no `self` state; call it unbound so this stays a
+        # pure bookkeeping test.
+        MambaAttnBackendBase._track_mamba_state_extend(
+            None, None, None, ssm_states, metadata, h_track_buf=h_track_buf
+        )
+        return ssm_states
+
+    def test_snapshot_cast_once_to_pool_dtype(self):
+        torch.manual_seed(0)
+        h_track_buf = torch.randn(3, 2, 4, 4, dtype=torch.float32)
+        h_track_buf[0, 0, 0, 0] = DOUBLE_ROUND_PROBE
+        for pool_dtype in (torch.float32, torch.bfloat16, torch.float16):
+            with self.subTest(pool_dtype=pool_dtype):
+                ssm_states = self._run_track_copy(
+                    pool_dtype, h_track_buf, dst_slots=[5, 2], batch_rows=[0, 2]
+                )
+                # Single rounding of the fp32 snapshot, in batch-row order.
+                self.assertTrue(
+                    torch.equal(ssm_states[5], h_track_buf[0].to(pool_dtype))
+                )
+                self.assertTrue(
+                    torch.equal(ssm_states[2], h_track_buf[2].to(pool_dtype))
+                )
+                untouched = torch.ones(8, dtype=torch.bool)
+                untouched[[5, 2]] = False
+                self.assertTrue(torch.all(ssm_states[untouched] == 0))
+
+    def test_fp16_pool_is_not_double_rounded_through_bf16(self):
+        h_track_buf = torch.full((1, 1, 1, 1), DOUBLE_ROUND_PROBE)
+        ssm_states = self._run_track_copy(
+            torch.float16, h_track_buf, dst_slots=[3], batch_rows=[0]
+        )
+        # fp32 -> fp16 rounds the probe UP to 1 + 2^-10; the legacy path
+        # (fp32 -> bf16 h -> fp16) collapsed it to exactly 1.0.
+        self.assertEqual(ssm_states[3, 0, 0, 0].item(), 1.0 + 2.0**-10)
+
+    def test_no_unaligned_rows_leaves_pool_untouched(self):
+        # Aligned-only tracking: the h branch is gated off entirely.
+        h_track_buf = torch.randn(2, 1, 1, 1, dtype=torch.float32)
+        ssm_states = self._run_track_copy(
+            torch.float16, h_track_buf, dst_slots=[], batch_rows=[]
+        )
+        self.assertTrue(torch.all(ssm_states == 0))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/layers/test_mamba2_track_ssm_indices.py

@@ -37,8 +37,10 @@ def _split(extend_lens, prefix_lens, track_seqlens, track_mask):
     backend = _backend()
     cache_indices = torch.arange(len(extend_lens))
     (
+        _track_chunk_idx,
         h_src,
         h_dst,
+        _h_batch_src,
         _final_src,
         _final_dst,
         seq_idx,