[Fix] hybrid: Preserve deterministic prefill and prefix caching

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

@@ -47,6 +47,8 @@ def chunk_gated_delta_rule_fwd(
     inplace_update: bool = True,
     chunk_offsets: torch.LongTensor | None = None,
     o_chunk_indices: torch.LongTensor | None = None,
+    track_state: Optional[torch.Tensor] = None,
+    track_chunk_idx: Optional[torch.Tensor] = None,
 ):
     g = chunk_local_cumsum(
         g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
@@ -62,6 +64,9 @@ def chunk_gated_delta_rule_fwd(
         chunk_indices=chunk_indices,
     )
 
+    snapshot_args = {}
+    if track_state is not None:
+        snapshot_args = dict(track_state=track_state, track_chunk_idx=track_chunk_idx)
     h, v_new = chunk_gated_delta_rule_fwd_h(
         k=k,
         w=w,
@@ -73,6 +78,7 @@ def chunk_gated_delta_rule_fwd(
         chunk_indices=chunk_indices,
         inplace_update=inplace_update,
         chunk_offsets=chunk_offsets,
+        **snapshot_args,
     )
     o = chunk_fwd_o(
         q=q,
@@ -109,6 +115,8 @@ def forward(
         inplace_update: bool = True,
         chunk_indices: Optional[torch.LongTensor] = None,
         chunk_offsets: Optional[torch.LongTensor] = None,
+        track_state: Optional[torch.Tensor] = None,
+        track_chunk_idx: Optional[torch.Tensor] = None,
     ):
         q_orig = q
         k_orig = k
@@ -136,6 +144,8 @@ def forward(
             inplace_update=inplace_update,
             chunk_offsets=chunk_offsets,
             o_chunk_indices=o_chunk_indices,
+            track_state=track_state,
+            track_chunk_idx=track_chunk_idx,
         )
         return o.to(q.dtype), h
 
@@ -156,6 +166,8 @@ def chunk_gated_delta_rule(
     inplace_update: bool = True,
     chunk_indices: Optional[torch.LongTensor] = None,
     chunk_offsets: Optional[torch.LongTensor] = None,
+    track_state: Optional[torch.Tensor] = None,
+    track_chunk_idx: Optional[torch.Tensor] = None,
 ):
     r"""
     Args:
@@ -184,6 +196,10 @@ def chunk_gated_delta_rule(
         head_first (Optional[bool]):
             Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
             Default: `False`.
+        track_state (Optional[torch.Tensor]):
+            FP32 buffer `[N, H, V, K]` for one selected intermediate state per sequence.
+        track_chunk_idx (Optional[torch.Tensor]):
+            Local chunk index of each selected state, or -1 for no snapshot.
         inplace_update (Optional[bool]):
             Whether to write final states back to `initial_state`. Default: `True`.
 
@@ -222,6 +238,7 @@ def chunk_gated_delta_rule(
             cu_seqlens=cu_seqlens
         )
     """
+    assert (track_state is None) == (track_chunk_idx is None)
     assert q.dtype == k.dtype == v.dtype
     assert q.dtype != torch.float32, (
         "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
@@ -275,6 +292,8 @@ def chunk_gated_delta_rule(
         inplace_update,
         chunk_indices,
         chunk_offsets,
+        track_state,
+        track_chunk_idx,
     )
     if head_first:
         o = rearrange(o, "b t h ... -> b h t ...")

python/sglang/kernels/ops/mamba/gdn_prefill_track.py

@@ -74,6 +74,7 @@ def _gdn_track_ssm_kernel(
     stride_h_row,
     numel,
     BLOCK: tl.constexpr,
+    H_PER_SEQUENCE: tl.constexpr,
 ):
     r = tl.program_id(0)
     mode = tl.load(track_mode + r)
@@ -86,7 +87,9 @@ def _gdn_track_ssm_kernel(
         src = tl.load(cache_indices + r).to(tl.int64)
         val = tl.load(ssm_pool + src * stride_ssm_slot + offs, mask=mask)
     else:
-        row = tl.load(track_h_row + r).to(tl.int64)
+        row = (
+            r.to(tl.int64) if H_PER_SEQUENCE else tl.load(track_h_row + r).to(tl.int64)
+        )
         val = tl.load(h + row * stride_h_row + offs, mask=mask)
         val = val.to(ssm_pool.dtype.element_ty)
     tl.store(ssm_pool + dst * stride_ssm_slot + offs, val, mask=mask)
@@ -130,9 +133,13 @@ def gdn_track_ssm_state(
     track_h_row: torch.Tensor,
     track_dst: torch.Tensor,
     cache_indices: torch.Tensor,
+    h_is_per_sequence: bool = False,
 ) -> None:
     """Snapshot the tracked rows' SSM state into their track slots (mode 1:
-    the row's final state, mode 2: its chunk-boundary state in ``h``)."""
+    the row's final state, mode 2: its chunk-boundary state in ``h``).
+    ``h_is_per_sequence`` selects the FP32 snapshot for each row instead of
+    the activation-dtype, per-chunk h buffer.
+    """
     rows = track_mode.shape[0]
     numel = ssm_pool[0].numel()
     assert ssm_pool[0].is_contiguous(), "per-slot SSM state must be contiguous"
@@ -150,4 +157,5 @@ def gdn_track_ssm_state(
         h.stride(0),
         numel,
         BLOCK=block,
+        H_PER_SEQUENCE=h_is_per_sequence,
     )

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

@@ -1170,6 +1170,18 @@ def _forward_extend_in_graph(
         ).transpose(0, 1)
         query, key, value = self._split_prefill_qkv(layer, mixed_qkv)
         g, beta = self._prefill_gates(layer, a, b)
+        h_track_buf = None
+        if (
+            meta.track
+            and self.kernel_dispatcher.extend_kernel.supports_track_state_snapshot
+        ):
+            # Match eager tracking: preserve the selected FP32 accumulator,
+            # rather than restoring from the BF16 h used by attention output.
+            h_track_buf = torch.empty(
+                (meta.max_seqs, *ssm_states.shape[1:]),
+                dtype=torch.float32,
+                device=ssm_states.device,
+            )
         core_attn_out, _, h = self.kernel_dispatcher.extend(
             q=query,
             k=key,
@@ -1181,15 +1193,18 @@ def _forward_extend_in_graph(
             query_start_loc=meta.cu_seqlens,
             chunk_indices=meta.chunk_indices,
             chunk_offsets=meta.chunk_offsets,
+            track_state=h_track_buf,
+            track_chunk_idx=meta.track_chunk_idx if h_track_buf is not None else None,
         )
         if meta.track:
             gdn_track_ssm_state(
                 ssm_pool=ssm_states,
-                h=h,
+                h=h_track_buf.unsqueeze(0) if h_track_buf is not None else h,
                 track_mode=meta.track_mode,
                 track_h_row=meta.track_h_row,
                 track_dst=meta.track_dst,
                 cache_indices=meta.cache_indices,
+                h_is_per_sequence=h_track_buf is not None,
             )
         return core_attn_out
 

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

@@ -78,6 +78,7 @@ class GDNPrefillGraphMetadata(msgspec.Struct, frozen=True):
     conv_block_table: torch.Tensor  # [max_conv_blocks, 2] (seq, block)
     track_mode: torch.Tensor  # [max_seqs]
     track_h_row: torch.Tensor  # [max_seqs] row of the per-chunk states h
+    track_chunk_idx: torch.Tensor  # [max_seqs] selected local chunk, or -1
     track_conv_start: torch.Tensor  # [max_seqs]
     num_logical_tokens: torch.Tensor  # [1]
     host: np.ndarray  # int32 staging for `packed`
@@ -107,6 +108,7 @@ def allocate(
             2 * max_conv_blocks,  # conv_block_table
             max_seqs,  # track_mode
             max_seqs,  # track_h_row
+            max_seqs,  # track_chunk_idx
             max_seqs,  # track_conv_start
             1,  # num_logical_tokens
         ]
@@ -125,8 +127,9 @@ def allocate(
             conv_block_table=views[3].view(max_conv_blocks, 2),
             track_mode=views[4],
             track_h_row=views[5],
-            track_conv_start=views[6],
-            num_logical_tokens=views[7],
+            track_chunk_idx=views[6],
+            track_conv_start=views[7],
+            num_logical_tokens=views[8],
             host=np.zeros(sum(sizes), dtype=np.int32),
             cache_indices=torch.full(
                 (max_seqs,), PAD_SLOT_ID, dtype=cache_index_dtype, device=device
@@ -204,18 +207,21 @@ def _write_host_tables(
 
         track_mode = np.zeros(self.max_seqs, dtype=np.int64)
         track_h_row = np.zeros(self.max_seqs, dtype=np.int64)
+        track_chunk_idx = np.full(self.max_seqs, -1, dtype=np.int64)
         track_conv_start = np.zeros(self.max_seqs, dtype=np.int64)
         if track_rows is not None:
             # Same source rows as the eager plan: h_offsets[row] + chunk.
             for row, chunk in track_rows.chunk_rows:
                 track_h_row[row] = offsets[row] + chunk
+                track_chunk_idx[row] = chunk
                 track_mode[row] = TRACK_CHUNK_STATE
             for row in track_rows.final_rows:
                 track_mode[row] = TRACK_FINAL_STATE
             for row, aligned_len in track_rows.conv_rows:
                 track_conv_start[row] = cu[row] + aligned_len - track_rows.conv_len
         o = self._write(o=o, values=track_mode)
         o = self._write(o=o, values=track_h_row)
+        o = self._write(o=o, values=track_chunk_idx)
         o = self._write(o=o, values=track_conv_start)
         o = self._write(o=o, values=np.array([total]))
         assert o == self.host.shape[0]

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

@@ -42,6 +42,8 @@ class TritonGDNKernel(LinearAttnKernelBase):
 
     supports_packed_decode: bool = not is_cpu() and not is_npu()
     supports_strided_target_verify_qkv: bool = True
+    # CPU/NPU use external chunk kernels; XPU overrides the h kernel.
+    supports_track_state_snapshot: bool = not is_cpu() and not is_npu() and not is_xpu()
 
     def packed_decode(
         self,
@@ -207,6 +209,13 @@ def extend(
             # control. Its existing behavior is equivalent to True.
             inplace_update_args = {}
 
+        snapshot_args = {}
+        if kwargs.get("track_state") is not None:
+            assert self.supports_track_state_snapshot
+            snapshot_args = dict(
+                track_state=kwargs["track_state"],
+                track_chunk_idx=kwargs["track_chunk_idx"],
+            )
         return chunk_gated_delta_rule(
             q=q,
             k=k,
@@ -220,6 +229,7 @@ def extend(
             **recurrent_state_indices_args,
             **inplace_update_args,
             **graph_table_args,
+            **snapshot_args,
         )
 
     def target_verify(

python/sglang/srt/managers/schedule_policy.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import math
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.cp.base import get_cp_strategy, is_cp_sharded_padded
@@ -838,16 +839,23 @@ def _request_slot_reserve_for_req(
     def ceil_paged_tokens(self, tokens: int) -> int:
         return -(-tokens // self.page_size) * self.page_size
 
-    def floor_chunked_tokens(self, trunc_len: int, remaining: int) -> int:
-        """Round a *truncated* chunk down to the mamba kernel chunk grid.
+    def floor_chunked_tokens(
+        self,
+        trunc_len: int,
+        remaining: int,
+        truncation_align_size: Optional[int] = None,
+    ) -> int:
+        """Align intermediate chunks; the final chunk may have any length.
 
-        The chunk that finishes a request may stop anywhere, so only shorten
-        when `trunc_len` really is a cut. Leave the cut alone when rounding
-        would empty it: making no progress is worse than losing a donation.
+        An explicit alignment is a correctness constraint: defer a cut that
+        rounds to zero. DCP's checkpoint-only grid remains best effort when
+        no explicit alignment is requested.
         """
-        if self.mamba_chunk_grid <= 1 or trunc_len >= remaining:
+        if trunc_len >= remaining:
             return trunc_len
-        return trunc_len // self.mamba_chunk_grid * self.mamba_chunk_grid or trunc_len
+        grid = math.lcm(self.mamba_chunk_grid, truncation_align_size or 1)
+        aligned = trunc_len // grid * grid
+        return aligned if truncation_align_size is not None else aligned or trunc_len
 
     def budget_state(self):
         no_token = (
@@ -1122,7 +1130,7 @@ def add_dllm_staging_req(self, req: Req):
             else AddReqResult.CONTINUE
         )
 
-    def add_chunked_req(self, req: Req):
+    def add_chunked_req(self, req: Req, truncation_align_size: Optional[int] = None):
         if self.dllm_config is not None:
             _rem_tokens = self._get_dllm_remain_tokens(req)
         else:
@@ -1169,8 +1177,12 @@ def add_chunked_req(self, req: Req):
             return req
         truncated = cand_extend_input_len > _rem_tokens
         new_len = self.floor_chunked_tokens(
-            min(cand_extend_input_len, _rem_tokens), cand_extend_input_len
+            min(cand_extend_input_len, _rem_tokens),
+            cand_extend_input_len,
+            truncation_align_size,
         )
+        if new_len <= 0:
+            return req
         request_slot_reserve = self._request_slot_reserve_for_req(
             req, extend_input_len=new_len
         )
@@ -1227,7 +1239,9 @@ def _lock_node(self, last_node: TreeNode, *, lock_host: bool = False):
             if host_lock_params is not None:
                 self.tree_cache.dec_host_lock_ref(last_node, host_lock_params)
 
-    def add_one_req_ignore_eos(self, req: Req):
+    def add_one_req_ignore_eos(
+        self, req: Req, truncation_align_size: Optional[int] = None
+    ):
         cand_extend_input_len = len(req.full_untruncated_fill_ids) - req.prefix_len
         if is_cp_sharded_padded():
             # One physical page on each rank reserves a complete logical group.
@@ -1349,9 +1363,12 @@ def add_req_state(r, insert_sort=False):
 
             # Chunked prefill
             trunc_len = self.floor_chunked_tokens(
-                self.rem_chunk_tokens, cand_extend_input_len
+                self.rem_chunk_tokens, cand_extend_input_len, truncation_align_size
             )
 
+            if trunc_len <= 0:
+                return AddReqResult.OTHER
+
             if (tile_stop := self._check_prefill_tile_budget(trunc_len)) is not None:
                 return tile_stop
 
@@ -1406,7 +1423,7 @@ def add_one_req(
             return AddReqResult.OTHER
 
         if req.sampling_params.ignore_eos and getattr(self.tree_cache, "disable", True):
-            return self.add_one_req_ignore_eos(req)
+            return self.add_one_req_ignore_eos(req, truncation_align_size)
 
         # Reserve page_size for page-alignment overhead: the paged allocator may
         # consume one extra page per request (see alloc_extend), which
@@ -1664,16 +1681,14 @@ def _select_prefill_admission(
                 # limit rather than rem_chunk_tokens so an SWA-capped chunk stays
                 # capped.
                 extend_len = chunk_tokens_limit
-                if truncation_align_size is not None:
-                    extend_len = (
-                        extend_len // truncation_align_size * truncation_align_size
-                    )
             else:
                 extend_len = chunk_tokens_limit // self.page_size * self.page_size
-                if truncation_align_size is not None:
-                    extend_len = (
-                        extend_len // truncation_align_size * truncation_align_size
-                    )
+            extend_len = self.floor_chunked_tokens(
+                extend_len,
+                len(req.full_untruncated_fill_ids) - prefix_len,
+                truncation_align_size,
+            )
+            if not self.exact_chunk_fill:
                 end = (prefix_len + extend_len) // self.page_size * self.page_size
                 extend_len = end - prefix_len
             if extend_len <= 0: