[GDN] Support FlashInfer GDN prefill with extra-buffer radix cache (#29735)

docs_new/docs/advanced_features/attention_backend.mdx

@@ -338,7 +338,7 @@ GDN (Gated Delta Network) is a linear attention mechanism with O(n) complexity,
 
 The GDN linear attention layers have their own kernel backends, selected via `--linear-attn-backend` (default: `triton`). You can override the kernel per phase with `--linear-attn-decode-backend` and `--linear-attn-prefill-backend`.
 
-On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN prefill when the per-phase override is unset, the base linear-attention backend is Triton, recurrent state is BF16, key/value head dimensions are 128, dynamic chunking and page-major KV layout are disabled, and `--chunked-prefill-size` is between 1 and 8192. Radix caching may be disabled or use the `no_buffer` strategy; extra-buffer strategies require state checkpoint support.
+On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN prefill when the per-phase override is unset, the base linear-attention backend is Triton, recurrent state is BF16, key/value head dimensions are 128, dynamic chunking and page-major KV layout are disabled, and `--chunked-prefill-size` is between 1 and 8192. Radix caching may be disabled or use `no_buffer`, `extra_buffer`, or `extra_buffer_lazy`; the extra-buffer paths use state checkpoints.
 
 <table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
   <colgroup>

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

@@ -724,16 +724,18 @@ class MambaAttnBackendBase(AttentionBackend):
     def _track_mamba_state_extend(
         self,
         forward_batch: ForwardBatch,
-        h: torch.Tensor,
+        h: Optional[torch.Tensor],
         ssm_states: torch.Tensor,
         forward_metadata: ForwardMetadata,
     ):
         """Copy extend SSM state at the last chunk boundary to track slots (source
         depends on chunk alignment; see `_init_track_ssm_indices`)."""
         if forward_metadata.has_mamba_track_mask:
-            h = h.squeeze(0)
-
+            # Triton always returns h; FlashInfer returns it only when checkpoints
+            # were requested. Aligned-only tracking reads the final state below.
             if forward_metadata.track_ssm_h_src.numel() > 0:
+                assert h is not None
+                h = h.squeeze(0)
                 ssm_states[forward_metadata.track_ssm_h_dst] = h[
                     forward_metadata.track_ssm_h_src
                 ].to(ssm_states.dtype, copy=False)

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

@@ -75,10 +75,6 @@ def maybe_set_default_flashinfer_gdn_prefill(model_runner: ModelRunner) -> None:
     ):
         return
 
-    # Extra-buffer strategies need intermediate state checkpoints.
-    if args.uses_mamba_radix_cache and args.mamba_radix_cache_strategy != "no_buffer":
-        return
-
     cuda_version = torch.version.cuda
     chunk_size = args.chunked_prefill_size
     config = hybrid_gdn_config(model_runner.model_config)
@@ -204,6 +200,10 @@ class GDNKernelDispatcher:
             f"packed_decode={self.supports_packed_decode}"
         )
 
+    @property
+    def extend_uses_state_checkpoints(self) -> bool:
+        return self.extend_kernel.uses_state_checkpoints
+
     def packed_decode(
         self,
         mixed_qkv: torch.Tensor,
@@ -362,6 +362,14 @@ class GDNAttnBackend(MambaAttnBackendBase):
                     self.forward_metadata.mamba_track_mask_indices
                 ]
             )
+            if self.kernel_dispatcher.extend_uses_state_checkpoints:
+                from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+                    maybe_build_flashinfer_checkpoint_plan,
+                )
+
+                maybe_build_flashinfer_checkpoint_plan(
+                    forward_batch, self.forward_metadata, self.device
+                )
 
     def forward_decode(
         self,
@@ -649,6 +657,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 ssm_states=ssm_states_contig,
                 cache_indices=state_cache_indices,
                 query_start_loc=query_start_loc,
+                state_checkpoint_cu_starts=(
+                    forward_metadata.state_checkpoint_cu_starts
+                ),
+                num_state_checkpoints=forward_metadata.num_state_checkpoints,
+                state_checkpoint_every_n_tokens=(
+                    forward_metadata.state_checkpoint_every_n_tokens
+                ),
             )
 
             if is_npu() and last_recurrent_state is not None:
@@ -663,7 +678,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
                 conv_states[cache_indices] = conv_states_contig
                 ssm_states[cache_indices] = ssm_states_contig
 
-            if h is not None:
+            if forward_metadata.has_mamba_track_mask:
                 self._track_mamba_state_extend(
                     forward_batch, h, ssm_states, forward_metadata
                 )

python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py

@@ -5,20 +5,27 @@ Both SM90 and SM100 use the same pool layout: [pool, HV, V, K] (K-last).
 SM90 (Hopper): full support — decode, prefill, MTP.  State dtype: fp32.
 SM100 (Blackwell): full support — decode, prefill, MTP.
 
-Requires flashinfer >= 0.6.7.
+Requires flashinfer >= 0.6.14.
 """
 
+from __future__ import annotations
+
 import logging
 import os
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
     LinearAttnKernelBase,
 )
+from sglang.srt.runtime_context import get_server_args
 from sglang.srt.utils import is_cuda
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.mamba.mamba2_metadata import ForwardMetadata
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
 logger = logging.getLogger(__name__)
 
 # ---------------------------------------------------------------------------
@@ -31,6 +38,47 @@ _flashinfer_gated_delta_rule_decode = None
 _flashinfer_gated_delta_rule_mtp_bf16 = None
 
 
+def maybe_build_flashinfer_checkpoint_plan(
+    forward_batch: ForwardBatch,
+    forward_metadata: ForwardMetadata,
+    device: str,
+) -> None:
+    """Populate packed FlashInfer checkpoint metadata when tracking requires it."""
+    if (
+        forward_metadata.track_ssm_h_src is None
+        or forward_metadata.track_ssm_h_src.numel() == 0
+    ):
+        return
+
+    checkpoint_every_n_tokens = get_server_args().mamba_cache_chunk_size
+    extend_seq_lens = forward_batch.extend_seq_lens.to(device="cpu", dtype=torch.int64)
+    track_mask = forward_batch.mamba_track_mask.to(device="cpu", dtype=torch.bool)
+    relative_track_lens = forward_batch.mamba_track_seqlens.to(
+        device="cpu", dtype=torch.int64
+    ) - forward_batch.extend_prefix_lens.to(device="cpu", dtype=torch.int64)
+
+    checkpoint_counts = extend_seq_lens // checkpoint_every_n_tokens
+    checkpoint_cu_starts = torch.zeros(checkpoint_counts.numel() + 1, dtype=torch.int64)
+    checkpoint_cu_starts[1:] = torch.cumsum(checkpoint_counts, dim=0)
+
+    use_checkpoint = track_mask & (relative_track_lens % checkpoint_every_n_tokens != 0)
+    track_checkpoint_src = checkpoint_cu_starts[:-1][use_checkpoint] + (
+        relative_track_lens[use_checkpoint] // checkpoint_every_n_tokens - 1
+    )
+    if track_checkpoint_src.numel() and track_checkpoint_src.min() < 0:
+        raise ValueError("Tracked GDN state precedes the first FlashInfer checkpoint.")
+    assert track_checkpoint_src.numel() == forward_metadata.track_ssm_h_dst.numel()
+
+    forward_metadata.track_ssm_h_src = track_checkpoint_src.to(
+        device, non_blocking=True
+    )
+    forward_metadata.state_checkpoint_cu_starts = checkpoint_cu_starts.to(
+        device, non_blocking=True
+    )
+    forward_metadata.num_state_checkpoints = int(checkpoint_cu_starts[-1])
+    forward_metadata.state_checkpoint_every_n_tokens = checkpoint_every_n_tokens
+
+
 def _get_flashinfer_gdn_kernels():
     """Lazy import for FlashInfer GDN prefill, decode and verify (MTP) kernels.
 
@@ -89,9 +137,11 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
     SM90 (Hopper): decode uses gather/scatter; prefill and MTP verify supported.
     SM100 (Blackwell): decode uses gather/scatter; prefill and MTP verify supported.
 
-    Requires flashinfer >= 0.6.7.
+    Requires flashinfer >= 0.6.14.
     """
 
+    uses_state_checkpoints = True
+
     def __init__(self):
         (
             available,
@@ -231,6 +281,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         ssm_states: torch.Tensor,
         cache_indices: torch.Tensor,
         query_start_loc: torch.Tensor,
+        state_checkpoint_cu_starts: Optional[torch.Tensor] = None,
+        num_state_checkpoints: int = 0,
+        state_checkpoint_every_n_tokens: int = 0,
         **kwargs,
     ) -> tuple:
         from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
@@ -253,23 +306,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
             # slot) so the FlashInfer kernel never reads out-of-bounds state.
             ssm_cache_indices = cache_indices.clamp(min=0).to(torch.int64)
             initial_state_fi = ssm_states[ssm_cache_indices].contiguous()
-            # Pre-allocate bf16 output_state so the kernel compiles and writes the
-            # bf16 state path directly, avoiding a fp32 allocation and a subsequent
-            # fp32->bf16 conversion in the scatter step.
-            output_state_fi = torch.empty_like(initial_state_fi)
-            output_fi, output_state_fi = self._prefill_fn(
-                q=q_fi,
-                k=k_fi,
-                v=v_fi,
-                g=alpha_fi,
-                beta=beta_fi,
-                scale=None,
-                initial_state=initial_state_fi,
-                output_final_state=True,
-                cu_seqlens=query_start_loc,  # already int32
-                use_qk_l2norm_in_kernel=False,
-                output_state=output_state_fi,
-            )
+            cu_seqlens = query_start_loc  # already int32
         else:
             # SM90: preserve original negative-index handling (remap to last slot).
             ssm_cache_indices = torch.where(
@@ -279,18 +316,33 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
             ).to(torch.int64)
             # State must be float32; kernel requires int64 cu_seqlens.
             initial_state_fi = ssm_states[ssm_cache_indices].to(torch.float32)
-            output_fi, output_state_fi = self._prefill_fn(
-                q=q_fi,
-                k=k_fi,
-                v=v_fi,
-                g=alpha_fi,
-                beta=beta_fi,
-                scale=None,
-                initial_state=initial_state_fi,
-                output_final_state=True,
-                cu_seqlens=query_start_loc.to(torch.int64),
-                use_qk_l2norm_in_kernel=False,
+            cu_seqlens = query_start_loc.to(torch.int64)
+
+        # Keep final state and checkpoints in the same kernel state dtype.
+        output_state_fi = torch.empty_like(initial_state_fi)
+        state_checkpoints = (
+            initial_state_fi.new_empty(
+                (num_state_checkpoints, *initial_state_fi.shape[1:])
             )
+            if num_state_checkpoints > 0
+            else None
+        )
+        output_fi, output_state_fi = self._prefill_fn(
+            q=q_fi,
+            k=k_fi,
+            v=v_fi,
+            g=alpha_fi,
+            beta=beta_fi,
+            scale=None,
+            initial_state=initial_state_fi,
+            output_final_state=True,
+            cu_seqlens=cu_seqlens,
+            use_qk_l2norm_in_kernel=False,
+            output_state=output_state_fi,
+            state_checkpoints=state_checkpoints,
+            checkpoint_cu_starts=state_checkpoint_cu_starts,
+            checkpoint_every_n_tokens=state_checkpoint_every_n_tokens,
+        )
 
         # Write back state to pool
         ssm_states.index_copy_(
@@ -302,9 +354,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         # Output: [seq, HV, V] -> [1, seq, HV, V]
         core_attn_out = output_fi.view(1, total_seq_len, num_v_heads, head_v_dim)
 
-        # Return (output, last_recurrent_state, h) to match Triton kernel interface.
-        # h=None since FlashInfer doesn't provide intermediate states.
-        return core_attn_out, None, None
+        # Match Triton's [1, checkpoints, H, V, K] intermediate-state layout.
+        h = state_checkpoints.unsqueeze(0) if state_checkpoints is not None else None
+        return core_attn_out, None, h
 
     # ---- target_verify (MTP) ----
 

python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py

@@ -10,6 +10,8 @@ class LinearAttnKernelBase(ABC):
     and provides decode/extend/target_verify methods with a unified interface.
     """
 
+    uses_state_checkpoints: bool = False
+
     @abstractmethod
     def decode(
         self,

python/sglang/srt/layers/attention/mamba/mamba2_metadata.py

@@ -55,6 +55,9 @@ class ForwardMetadata:
     track_ssm_h_dst: Optional[torch.Tensor] = None
     track_ssm_final_src: Optional[torch.Tensor] = None
     track_ssm_final_dst: Optional[torch.Tensor] = None
+    state_checkpoint_cu_starts: Optional[torch.Tensor] = None
+    num_state_checkpoints: int = 0
+    state_checkpoint_every_n_tokens: int = 0
 
     is_target_verify: bool = False
     draft_token_num: int = 1

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -55,6 +55,7 @@ class GDNAttentionCase:
     page_size: int
     prefix_lens: tuple[int, ...]
     extend_lens: tuple[int, ...] = ()
+    linear_attn_prefill_backend: str | None = None
 
     @property
     def batch_size(self) -> int:
@@ -245,7 +246,7 @@ class MockGDNModelRunner(ModelRunner):
             enable_mis=False,
             linear_attn_backend="triton",
             linear_attn_decode_backend=None,
-            linear_attn_prefill_backend=None,
+            linear_attn_prefill_backend=case.linear_attn_prefill_backend,
             max_running_requests=None,
             revision=None,
             speculative_algorithm=None,
@@ -268,10 +269,16 @@ class MockGDNModelRunner(ModelRunner):
             state_size=head_k_dim,
             conv_kernel=2,
         )
+        temporal_state_dtype = (
+            dtype
+            if case.linear_attn_prefill_backend == "flashinfer"
+            and torch.cuda.get_device_capability()[0] >= 10
+            else torch.float32
+        )
         cache_params = Mamba2CacheParams(
             shape=cache_shape,
             layers=[0],
-            dtype=Mamba2StateDType(conv=dtype, temporal=torch.float32),
+            dtype=Mamba2StateDType(conv=dtype, temporal=temporal_state_dtype),
         )
         self.req_to_token_pool = HybridReqToTokenPool(
             size=pool_batch_size,
@@ -591,6 +598,14 @@ def build_gdn_attention_fixture(
 
     initialize_linear_attn_config(runner.server_args)
     linear_backend = GDNAttnBackend(runner)
+    if case.linear_attn_prefill_backend == "flashinfer":
+        from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+            FlashInferGDNKernel,
+        )
+
+        testcase.assertIsInstance(
+            linear_backend.kernel_dispatcher.extend_kernel, FlashInferGDNKernel
+        )
     backend = HybridLinearAttnBackend(full_backend, linear_backend, full_attn_layers=[])
     actual_module = ProjectedGDNAttention(
         num_k_heads=case.num_k_heads,

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

@@ -10,11 +10,14 @@ from sglang.test.test_utils import CustomTestCase
 
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 
+from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.attention_unittest.attention_methods.gdn_attention import (
     GDNAttentionCase,
+    build_gdn_attention_fixture,
     make_gdn_cases,
     run_gdn_attention_case,
+    run_gdn_fixture_eager,
 )
 from sglang.test.kits.attention_unittest.runner_modes.cuda_graph_decode_runner import (
     run_gdn_cuda_graph_decode_case,
@@ -30,6 +33,12 @@ from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import (
 register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200")
 register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large")
 
+_cuda_major = int(torch.version.cuda.split(".")[0]) if torch.version.cuda else 0
+_sm_major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0
+_supports_flashinfer_linear_gdn = _sm_major == 9 or (
+    _sm_major == 10 and _cuda_major >= 13
+)
+
 
 @unittest.skipIf(
     not torch.cuda.is_available() or not is_flashinfer_available(),
@@ -322,5 +331,71 @@ class TestFlashInferGDNBackendCorrectness(CustomTestCase):
                 )
 
 
+@unittest.skipUnless(
+    torch.cuda.is_available()
+    and is_flashinfer_available()
+    and _supports_flashinfer_linear_gdn,
+    "FlashInfer linear GDN requires SM90 or SM100/SM103 with CUDA 13+",
+)
+class TestFlashInferLinearGDNBackendCorrectness(CustomTestCase):
+    # FlashInfer's DSL prefill kernels require head size 128 on SM90 and SM100.
+    HEAD_DIM = 128
+    CHECKPOINT_CASE = GDNAttentionCase(
+        name="flashinfer_gdn_prefill_state_checkpoints",
+        backend="triton",
+        linear_attn_prefill_backend="flashinfer",
+        forward_mode=ForwardMode.EXTEND,
+        num_k_heads=2,
+        num_v_heads=4,
+        page_size=16,
+        prefix_lens=(0, 64, 128),
+        extend_lens=(64, 65, 129),
+    )
+
+    def test_prefill_tracked_state_checkpoints(self):
+        fixture = build_gdn_attention_fixture(
+            self,
+            self.CHECKPOINT_CASE,
+            head_k_dim=self.HEAD_DIM,
+            head_v_dim=self.HEAD_DIM,
+            max_context_len=320,
+            runner_batch_size=6,
+        )
+        batch = fixture.forward_batch
+        # Simulate the tracking metadata produced by the extra-buffer scheduler.
+        # This test covers checkpoint mapping and state copies, not scheduler setup.
+        batch.mamba_track_mask = torch.ones(3, dtype=torch.bool, device="cuda")
+        batch.mamba_track_indices = torch.tensor(
+            [4, 5, 6], dtype=torch.int64, device="cuda"
+        )
+        batch.mamba_track_seqlens = torch.tensor(
+            # The final entry selects the second checkpoint at absolute S256.
+            [64, 129, 257],
+            dtype=torch.int64,
+            device="cuda",
+        )
+
+        cache = fixture.runner.req_to_token_pool.mamba2_layer_cache(0)
+        initial_conv = cache.conv[0].clone()
+        initial_ssm = cache.temporal.clone()
+        flashinfer_output = run_gdn_fixture_eager(fixture)
+        flashinfer_tracked = cache.temporal[batch.mamba_track_indices].clone()
+
+        cache.conv[0].copy_(initial_conv)
+        cache.temporal.copy_(initial_ssm)
+        fixture.backend.linear_attn_backend.kernel_dispatcher.extend_kernel = (
+            TritonGDNKernel()
+        )
+        triton_output = run_gdn_fixture_eager(fixture)
+        triton_tracked = cache.temporal[batch.mamba_track_indices]
+
+        torch.testing.assert_close(
+            flashinfer_output, triton_output, atol=3e-2, rtol=3e-2
+        )
+        torch.testing.assert_close(
+            flashinfer_tracked, triton_tracked, atol=3e-2, rtol=3e-2
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

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

@@ -4,11 +4,18 @@ from unittest.mock import MagicMock, patch, sentinel
 
 import torch
 
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    MambaAttnBackendBase,
+)
 from sglang.srt.layers.attention.linear import gdn_backend
 from sglang.srt.layers.attention.linear.gdn_backend import (
+    GDNAttnBackend,
     GDNKernelDispatcher,
     maybe_set_default_flashinfer_gdn_prefill,
 )
+from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
+    maybe_build_flashinfer_checkpoint_plan,
+)
 from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
 from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -32,16 +39,13 @@ def make_runner(
         enable_dynamic_chunking=False,
         chunked_prefill_size=8192,
     )
-    for name, value in arg_overrides.items():
-        setattr(args, name, value)
-
     # The policy routes its load-time default through the audited mutation entry
     # (server_args.override); mirror that on the stub so the write lands.
-    def _override(source, **fields):
-        for _field, _value in fields.items():
-            setattr(args, _field, _value)
-
-    args.override = _override
+    args.override = MagicMock(
+        side_effect=lambda _source, **fields: vars(args).update(fields)
+    )
+    for name, value in arg_overrides.items():
+        setattr(args, name, value)
 
     return SimpleNamespace(
         server_args=args,
@@ -87,14 +91,21 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
         return runner.server_args.linear_attn_prefill_backend
 
     def test_selects_flashinfer_for_supported_sm100_gdn(self):
-        self.assertEqual(self.apply_policy(make_runner()), "flashinfer")
-
-    def test_selects_flashinfer_for_no_buffer_radix_cache(self):
-        runner = make_runner(
-            uses_mamba_radix_cache=True,
-            mamba_radix_cache_strategy="no_buffer",
-        )
+        runner = make_runner()
         self.assertEqual(self.apply_policy(runner), "flashinfer")
+        runner.server_args.override.assert_called_once_with(
+            "gdn_backend.sm100_flashinfer_default",
+            linear_attn_prefill_backend="flashinfer",
+        )
+
+    def test_selects_flashinfer_for_radix_cache_strategies(self):
+        for strategy in ("no_buffer", "extra_buffer", "extra_buffer_lazy"):
+            with self.subTest(strategy=strategy):
+                runner = make_runner(
+                    uses_mamba_radix_cache=True,
+                    mamba_radix_cache_strategy=strategy,
+                )
+                self.assertEqual(self.apply_policy(runner), "flashinfer")
 
     def test_preserves_explicit_prefill_override(self):
         for backend in ("triton", "flashinfer", "cutedsl"):
@@ -128,20 +139,6 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
         cases = (
             ("non_triton_base", {"linear_attn_backend": "cutedsl"}),
             ("page_major_kv", {"enable_page_major_kv_layout": True}),
-            (
-                "extra_buffer",
-                {
-                    "uses_mamba_radix_cache": True,
-                    "mamba_radix_cache_strategy": "extra_buffer",
-                },
-            ),
-            (
-                "extra_buffer_lazy",
-                {
-                    "uses_mamba_radix_cache": True,
-                    "mamba_radix_cache_strategy": "extra_buffer_lazy",
-                },
-            ),
             ("dynamic_chunk", {"enable_dynamic_chunking": True}),
             ("unchunked", {"chunked_prefill_size": -1}),
             ("unknown_chunk", {"chunked_prefill_size": None}),
@@ -151,6 +148,53 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
             with self.subTest(name=name):
                 self.assertIsNone(self.apply_policy(make_runner(**runner_args)))
 
+    def test_builds_compact_checkpoint_plan_for_packed_sequences(self):
+        forward_batch = SimpleNamespace(
+            extend_seq_lens=torch.tensor([63, 64, 65, 127, 128, 129]),
+            mamba_track_mask=torch.tensor([False, True, True, True, True, True]),
+            # 65 on the 128-token sequence represents an interior S64
+            # boundary encoded as S64 + 1 by the scheduler.
+            mamba_track_seqlens=torch.tensor([63, 64, 65, 127, 65, 129]),
+            extend_prefix_lens=torch.zeros(6, dtype=torch.int64),
+        )
+        metadata = SimpleNamespace(
+            track_ssm_h_src=torch.empty(4),
+            track_ssm_h_dst=torch.empty(4),
+        )
+
+        with patch(
+            "sglang.srt.layers.attention.linear.kernels.gdn_flashinfer."
+            "get_server_args",
+            return_value=SimpleNamespace(mamba_cache_chunk_size=64),
+        ):
+            maybe_build_flashinfer_checkpoint_plan(forward_batch, metadata, "cpu")
+
+        torch.testing.assert_close(
+            metadata.state_checkpoint_cu_starts,
+            torch.tensor([0, 0, 1, 2, 3, 5, 7]),
+        )
+        torch.testing.assert_close(metadata.track_ssm_h_src, torch.tensor([1, 2, 3, 6]))
+        self.assertEqual(metadata.num_state_checkpoints, 7)
+        self.assertEqual(metadata.state_checkpoint_every_n_tokens, 64)
+
+    def test_decode_tracking_without_h_source_skips_checkpoint_plan(self):
+        backend = object.__new__(GDNAttnBackend)
+        backend.device = "cpu"
+        backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True)
+        metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None)
+        forward_batch = SimpleNamespace(
+            mamba_track_mask=torch.tensor([True]),
+            mamba_track_indices=torch.tensor([7]),
+        )
+
+        def init_base(instance, _forward_batch):
+            instance.forward_metadata = metadata
+
+        with patch.object(MambaAttnBackendBase, "init_forward_metadata", init_base):
+            backend.init_forward_metadata(forward_batch)
+
+        torch.testing.assert_close(metadata.conv_states_mask_indices, torch.tensor([7]))
+
     def test_tree_verify_uses_triton_kernel(self):
         flashinfer_kernel = MagicMock(supports_target_verify=True)
         with (