From 64590119170b30fca6c8cf6c0473165676cdb865 Mon Sep 17 00:00:00 2001
From: Xinan Miao <1403572259@qq.com>
Date: Wed, 21 Oct 2026 15:35:10 +0800
Subject: [PATCH 1/3] [Fix] scheduler: Align deterministic SSM prefill chunks

---
 python/sglang/srt/managers/schedule_policy.py | 55 ++++++++++++-------
 python/sglang/srt/managers/scheduler.py       | 13 ++++-
 .../test_scheduler_held_row_admission.py      |  2 +-
 3 files changed, 47 insertions(+), 23 deletions(-)

diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py
index e0884305d734..40aef9e354d4 100644
--- a/python/sglang/srt/managers/schedule_policy.py
+++ b/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:
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index ae83302f799b..96fa295390e8 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -54,6 +54,7 @@
     get_serving,
     get_spec,
     is_scheduler_io_rank,
+    mamba_cache_chunk_size,
     publish,
     spawn_world_rank,
 )
@@ -1816,7 +1817,7 @@ def maybe_init_ngram_embedding(self):
         self.use_ngram_embedding = self.tp_worker.model_config.use_ngram_embedding
 
     def init_deterministic_inference_config(self):
-        """Initialize deterministic inference configuration for different attention backends."""
+        """Align deterministic prefill chunks for attention and recurrent state updates."""
         if not get_exec().deterministic.enable_deterministic_inference:
             self.truncation_align_size = None
             return
@@ -1831,6 +1832,12 @@ def init_deterministic_inference_config(self):
         self.truncation_align_size = (
             get_int_env_var(env_var, default_size) if env_var else None
         )
+        if self.is_hybrid_ssm:
+            # Splitting a linear-attention chunk changes its recurrent update
+            # and floating-point rounding, even with batch-invariant GEMMs.
+            self.truncation_align_size = math.lcm(
+                self.truncation_align_size or 1, mamba_cache_chunk_size()
+            )
 
     def init_dsa_kpool_truncation_align(self):
         """Kpool compress-write asserts chunked extends start on pool boundaries.
@@ -4418,7 +4425,9 @@ def _get_new_batch_prefill_raw(
                 adder.rem_chunk_tokens or 0,
                 self.page_size,
             )
-            self.chunked_req = adder.add_chunked_req(self.chunked_req)
+            self.chunked_req = adder.add_chunked_req(
+                self.chunked_req, truncation_align_size=self.truncation_align_size
+            )
 
         if self.enable_lora:
             running_reqs = running_batch.reqs
diff --git a/test/registered/unit/managers/test_scheduler_held_row_admission.py b/test/registered/unit/managers/test_scheduler_held_row_admission.py
index c220e9d5d333..63b87af3acd9 100644
--- a/test/registered/unit/managers/test_scheduler_held_row_admission.py
+++ b/test/registered/unit/managers/test_scheduler_held_row_admission.py
@@ -44,7 +44,7 @@ def __init__(self, *args, **kwargs):
         self.new_chunked_req = None
         self.rem_chunk_tokens = 0
 
-    def add_chunked_req(self, req):
+    def add_chunked_req(self, req, truncation_align_size=None):
         self.can_run_list.append(req)
         return None
 

From 278792bb7f8aae5a12d1252584c8e6d1f68492bb Mon Sep 17 00:00:00 2001
From: Southwest <1403572259@qq.com>
Date: Thu, 5 Nov 2026 04:52:54 +0800
Subject: [PATCH 2/3] fix: preserve FP32 GDN prefix checkpoints across eager
 and prefill graphs

---
 .../sglang/kernels/ops/attention/fla/chunk.py | 19 +++++++++++++++++
 .../kernels/ops/mamba/gdn_prefill_track.py    | 12 +++++++++--
 .../layers/attention/linear/gdn_backend.py    | 21 +++++++++++++++----
 .../attention/linear/gdn_prefill_graph.py     | 10 +++++++--
 .../attention/linear/kernels/gdn_triton.py    | 10 +++++++++
 5 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/python/sglang/kernels/ops/attention/fla/chunk.py b/python/sglang/kernels/ops/attention/fla/chunk.py
index 918d3bcd8d8e..0ce94c2d9c1a 100644
--- a/python/sglang/kernels/ops/attention/fla/chunk.py
+++ b/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 ...")
diff --git a/python/sglang/kernels/ops/mamba/gdn_prefill_track.py b/python/sglang/kernels/ops/mamba/gdn_prefill_track.py
index b1db469d75e0..0e6f1a788015 100644
--- a/python/sglang/kernels/ops/mamba/gdn_prefill_track.py
+++ b/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,
     )
diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py
index 5e827ff2e1e5..480a37ec6b86 100644
--- a/python/sglang/srt/layers/attention/linear/gdn_backend.py
+++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py
@@ -11,9 +11,7 @@
 from sglang.kernels.ops.mamba.causal_conv1d_triton import (
     causal_conv1d_fn as causal_conv1d_fn_triton,
 )
-from sglang.kernels.ops.mamba.causal_conv1d_triton import (
-    causal_conv1d_update,
-)
+from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update
 from sglang.kernels.ops.mamba.gdn_prefill_track import (
     gdn_track_conv_window,
     gdn_track_ssm_state,
@@ -1170,6 +1168,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 +1191,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
 
diff --git a/python/sglang/srt/layers/attention/linear/gdn_prefill_graph.py b/python/sglang/srt/layers/attention/linear/gdn_prefill_graph.py
index 17bd633bd794..68d784bd07ab 100644
--- a/python/sglang/srt/layers/attention/linear/gdn_prefill_graph.py
+++ b/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,11 +207,13 @@ 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
@@ -216,6 +221,7 @@ def _write_host_tables(
                 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]
diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py
index 433e2704b442..f90e1d2a2cea 100644
--- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py
+++ b/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(

From 13771a98127524e7967d395a267b222306e7714b Mon Sep 17 00:00:00 2001
From: Southwest <1403572259@qq.com>
Date: Thu, 5 Nov 2026 05:07:57 +0800
Subject: [PATCH 3/3] fix: keep decode checkpoints out of deterministic hybrid
 prefix caching

---
 python/sglang/srt/layers/attention/linear/gdn_backend.py     | 4 +++-
 .../sglang/srt/mem_cache/unified_cache/components/mamba.py   | 5 +++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py
index 480a37ec6b86..d63a657ac532 100644
--- a/python/sglang/srt/layers/attention/linear/gdn_backend.py
+++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py
@@ -11,7 +11,9 @@
 from sglang.kernels.ops.mamba.causal_conv1d_triton import (
     causal_conv1d_fn as causal_conv1d_fn_triton,
 )
-from sglang.kernels.ops.mamba.causal_conv1d_triton import causal_conv1d_update
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_update,
+)
 from sglang.kernels.ops.mamba.gdn_prefill_track import (
     gdn_track_conv_window,
     gdn_track_ssm_state,
diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py
index 1a12f03e832d..361f5746a075 100644
--- a/python/sglang/srt/mem_cache/unified_cache/components/mamba.py
+++ b/python/sglang/srt/mem_cache/unified_cache/components/mamba.py
@@ -748,6 +748,11 @@ def _free_mamba_value(self, mamba_value: torch.Tensor) -> None:
     def _select_finished_checkpoint(
         self, req: Req, token_ids_len: int
     ) -> Optional[tuple[int, int]]:
+        # Decode and prefill need not produce bitwise-identical states/KV.
+        # Retaining decode KV in the radix tree can poison a later prefill's
+        # prefix through deduplication, even when its SSM snapshot is exact.
+        if get_exec().deterministic.enable_deterministic_inference:
+            token_ids_len = min(token_ids_len, len(req.origin_input_ids))
         # None means donate nothing, not "no slot found".
         pool = self.cache.req_to_token_pool
         keep_idx = pool.get_mamba_ping_pong_keep_idx(req)