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

python/sglang/kernels/ops/attention/fla/chunk_delta_h.py

@@ -63,6 +63,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     stride_init_state,
     cu_seqlens,
     chunk_offsets,
+    track_state,
+    track_chunk_idx,
+    stride_track_state,
     T,
     H: tl.constexpr,
     Hg: tl.constexpr,
@@ -78,6 +81,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     IS_VARLEN: tl.constexpr,
     NT_BUCKET: tl.constexpr,
     USE_EXP2: tl.constexpr,
+    TRACK_STATE: tl.constexpr,
 ):
     i_v, i_nh = tl.program_id(0), tl.program_id(1)
     i_n, i_h = i_nh // H, i_nh % H
@@ -130,6 +134,15 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
     if INPLACE_UPDATE:
         ht = ht + i_h * V * K
 
+    if TRACK_STATE:
+        i_track = tl.load(track_chunk_idx + i_n).to(tl.int32)
+        p_track_base = track_state + (i_n * stride_track_state + i_h * V * K).to(
+            tl.int64
+        )
+    else:
+        i_track = -1
+        p_track_base = track_state
+
     # load initial state
     if USE_INITIAL_STATE and valid_state:
         p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0))
@@ -172,6 +185,27 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
             )
             tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
 
+        if TRACK_STATE and i_t == i_track:
+            p_t1 = tl.make_block_ptr(
+                p_track_base, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)
+            )
+            tl.store(p_t1, b_h1, boundary_check=(0, 1))
+            if K > 64:
+                p_t2 = tl.make_block_ptr(
+                    p_track_base, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
+                )
+                tl.store(p_t2, b_h2, boundary_check=(0, 1))
+            if K > 128:
+                p_t3 = tl.make_block_ptr(
+                    p_track_base, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
+                )
+                tl.store(p_t3, b_h3, boundary_check=(0, 1))
+            if K > 192:
+                p_t4 = tl.make_block_ptr(
+                    p_track_base, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
+                )
+                tl.store(p_t4, b_h4, boundary_check=(0, 1))
+
         p_w = tl.make_block_ptr(
             w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)
         )
@@ -326,10 +360,21 @@ def chunk_gated_delta_rule_fwd_h(
     cu_seqlens: Optional[torch.LongTensor] = None,
     chunk_indices: Optional[torch.LongTensor] = None,
     use_exp2: bool = False,
+    track_state: Optional[torch.Tensor] = None,
+    track_chunk_idx: Optional[torch.Tensor] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     assert not (use_exp2 and g is not None), (
         "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
     )
+    assert (track_state is None) == (track_chunk_idx is None), (
+        "track_state and track_chunk_idx must be passed together"
+    )
+    if track_state is not None:
+        # The caller rounds once to the pool dtype; a narrower buffer would
+        # silently double-round the snapshot.
+        assert track_state.dtype == torch.float32, (
+            f"track_state must be fp32, got {track_state.dtype}"
+        )
     B, T, Hg, K, V = *k.shape, u.shape[-1]
     H = u.shape[-2]
     BT = CHUNK_SIZE
@@ -369,6 +414,9 @@ def chunk_gated_delta_rule_fwd_h(
         stride_init_state=(initial_state.stride(0) if initial_state is not None else 0),
         cu_seqlens=cu_seqlens,
         chunk_offsets=chunk_offsets,
+        track_state=track_state,
+        track_chunk_idx=track_chunk_idx,
+        stride_track_state=(track_state.stride(0) if track_state is not None else 0),
         T=T,
         H=H,
         Hg=Hg,
@@ -383,5 +431,6 @@ def chunk_gated_delta_rule_fwd_h(
         IS_VARLEN=cu_seqlens is not None,
         NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
         USE_EXP2=use_exp2,
+        TRACK_STATE=track_state is not None,
     )
     return h, v_new

python/sglang/kernels/ops/attention/fla/kda.py

@@ -1094,6 +1094,8 @@ def chunk_kda_fwd(
     dt_bias: Optional[torch.Tensor] = None,
     lower_bound: Optional[float] = None,
     output_intermediate_states: bool = False,
+    track_state: Optional[torch.Tensor] = None,
+    track_chunk_idx: Optional[torch.Tensor] = None,
 ):
     chunk_size = 64
     # Pre-compute chunk indices once and thread through all downstream kernels.
@@ -1169,6 +1171,8 @@ def chunk_kda_fwd(
         cu_seqlens=cu_seqlens,
         chunk_indices=chunk_indices,
         use_exp2=True,
+        track_state=track_state,
+        track_chunk_idx=track_chunk_idx,
     )
     del w, u, kg
 
@@ -1210,6 +1214,8 @@ def chunk_kda(
     dt_bias: Optional[torch.Tensor] = None,
     lower_bound: Optional[float] = None,
     output_intermediate_states: bool = False,
+    track_state: Optional[torch.Tensor] = None,
+    track_chunk_idx: Optional[torch.Tensor] = None,
     beta_is_raw: bool = False,
     **kwargs,
 ):
@@ -1238,4 +1244,6 @@ def chunk_kda(
         dt_bias=dt_bias,
         lower_bound=lower_bound,
         output_intermediate_states=output_intermediate_states,
+        track_state=track_state,
+        track_chunk_idx=track_chunk_idx,
     )

python/sglang/kernels/ops/attention/helion/kda_prefill.py

@@ -1024,6 +1024,40 @@ _STATE_VARLEN_SMALL_HEAD_CONFIG = helion.Config(
     range_unroll_factors=[0, 0],
 )
 
+# Track variants of the two varlen configs. The fp32 track snapshot adds one
+# load + one store mid-body, which would shift the positional indexing /
+# eviction lists above; separate kernels keep the non-track configs untouched.
+# Tracked batches are rare (prefix-cache checkpointing only), so these trade
+# the hand-tuned positional lists for plainly correct settings.
+_STATE_VARLEN_TRACK_CONFIG = helion.Config(
+    atomic_indexing=[],
+    block_sizes=[64],
+    indexing="pointer",
+    l2_groupings=[4],
+    loop_orders=[[0, 2, 1]],
+    num_stages=2,
+    num_warps=4,
+    pid_type="flat",
+    range_flattens=[None, None],
+    range_multi_buffers=[None, False],
+    range_num_stages=[],
+    range_unroll_factors=[0, 2],
+)
+_STATE_VARLEN_SMALL_HEAD_TRACK_CONFIG = helion.Config(
+    atomic_indexing=[],
+    block_sizes=[32],
+    indexing="pointer",
+    l2_groupings=[1],
+    loop_orders=[[1, 2, 0]],
+    num_stages=3,
+    num_warps=8,
+    pid_type="flat",
+    range_flattens=[None, None],
+    range_multi_buffers=[None, True],
+    range_num_stages=[],
+    range_unroll_factors=[0, 0],
+)
+
 
 @helion.kernel(
     static_shapes=False,
@@ -1040,6 +1074,9 @@ def _chunk_state(
     cu_seqlens: torch.Tensor,
     chunk_indices: torch.Tensor,
     chunk_offsets: torch.Tensor,
+    track_state: torch.Tensor,
+    track_chunk_idx: torch.Tensor,
+    has_track: hl.constexpr,  # pyrefly: ignore[bad-function-definition]
     is_varlen: hl.constexpr,  # pyrefly: ignore[bad-function-definition]
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Propagate KDA state between chunks and update the state pool in place."""
@@ -1070,6 +1107,11 @@ def _chunk_state(
             initial_state.stride(2),
             initial_state.stride(3),
             initial_state_indices.stride(0),
+            track_state.stride(0),
+            track_state.stride(1),
+            track_state.stride(2),
+            track_state.stride(3),
+            track_chunk_idx.stride(0),
         )
     )
 
@@ -1108,6 +1150,9 @@ def _chunk_state(
             :,
         ].float()
 
+        if has_track:
+            i_track = track_chunk_idx[tile_sequence.id]
+
         for token_tile in hl.tile(sequence_length, block_size=64):
             global_chunk = output_offset + token_tile.id
             h_rows[
@@ -1115,6 +1160,17 @@ def _chunk_state(
                 tile_v,
                 :,
             ] = state.to(h.dtype)
+            if has_track:
+                # Snapshot the fp32 accumulator at the tracked chunk boundary
+                # (the h store above rounds to the activation dtype; -1 marks
+                # untracked sequences and never matches a chunk id).
+                if token_tile.id == i_track:
+                    track_state[
+                        tile_sequence.id,
+                        tile_h.id,
+                        tile_v.index,
+                        :,
+                    ] = state
             token = begin + token_tile.index
             valid = token < end
             row = token * H + tile_h.id
@@ -1163,13 +1219,29 @@ _chunk_state_varlen_small_head = helion.kernel(
     config=_STATE_VARLEN_SMALL_HEAD_CONFIG,
     ignore_warnings=_IGNORED_WARNINGS,
 )(_chunk_state.fn)
+_chunk_state_varlen_track = helion.kernel(
+    static_shapes=False,
+    config=_STATE_VARLEN_TRACK_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)(_chunk_state.fn)
+_chunk_state_varlen_small_head_track = helion.kernel(
+    static_shapes=False,
+    config=_STATE_VARLEN_SMALL_HEAD_TRACK_CONFIG,
+    ignore_warnings=_IGNORED_WARNINGS,
+)(_chunk_state.fn)
 
 
-def _select_state_kernel(*, is_varlen: bool, num_heads: int) -> helion.Kernel:
+def _select_state_kernel(
+    *, is_varlen: bool, num_heads: int, has_track: bool
+) -> helion.Kernel:
     if not is_varlen:
         return _chunk_state
     if num_heads <= _PREFILL_SMALL_HEAD_THRESHOLD:
+        if has_track:
+            return _chunk_state_varlen_small_head_track
         return _chunk_state_varlen_small_head
+    if has_track:
+        return _chunk_state_varlen_track
     return _chunk_state_varlen
 
 
@@ -1314,6 +1386,8 @@ def chunk_kda(
     dt_bias: torch.Tensor | None = None,
     lower_bound: float | None = None,
     output_intermediate_states: bool = False,
+    track_state: torch.Tensor | None = None,
+    track_chunk_idx: torch.Tensor | None = None,
     beta_is_raw: bool = False,
     **kwargs: object,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
@@ -1322,6 +1396,13 @@ def chunk_kda(
         scale = k.shape[-1] ** -0.5
     if initial_state is None or initial_state_indices is None:
         raise ValueError("KDA prefill requires an indexed initial-state pool")
+    assert (track_state is None) == (track_chunk_idx is None), (
+        "track_state and track_chunk_idx must be passed together"
+    )
+    if track_state is not None:
+        assert track_state.dtype == torch.float32, (
+            f"track_state must be fp32, got {track_state.dtype}"
+        )
 
     num_tokens = q.shape[1]
     if g.shape[1] < num_tokens or beta.shape[1] < num_tokens:
@@ -1349,6 +1430,8 @@ def chunk_kda(
             dt_bias=dt_bias,
             lower_bound=lower_bound,
             output_intermediate_states=output_intermediate_states,
+            track_state=track_state,
+            track_chunk_idx=track_chunk_idx,
         )
 
     q = q.contiguous()
@@ -1395,7 +1478,10 @@ def chunk_kda(
     else:
         metadata = torch.empty(0, device=q.device, dtype=torch.int32)
         chunk_offsets = torch.empty(0, device=q.device, dtype=torch.long)
-    state_kernel = _select_state_kernel(is_varlen=is_varlen, num_heads=q.size(2))
+    has_track = track_state is not None
+    state_kernel = _select_state_kernel(
+        is_varlen=is_varlen, num_heads=q.size(2), has_track=has_track
+    )
     h, v_new = state_kernel(
         kg,
         w,
@@ -1410,6 +1496,11 @@ def chunk_kda(
             else torch.empty(0, 2, device=q.device, dtype=torch.long)
         ),
         chunk_offsets,
+        # Unused when has_track is False; bind in-place tensors as stand-ins so
+        # the kernel signature always sees real tensors.
+        track_state if has_track else initial_state,
+        track_chunk_idx if has_track else initial_state_indices,
+        has_track,
         is_varlen,
     )
     if chunk_indices is None: