[None][fix] Use one mamba slot sentinel to save memory

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -458,6 +458,12 @@ def __init__(
         # mamba cache index, maps request_id -> state indices
         self.mamba_cache_index: Dict[int, int] = {}
 
+        # Permanent slot shared by every CUDA-graph padding sentinel id
+        # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per
+        # draft length). Pool sizing must include +1 headroom for this;
+        # see MixedMambaHybridCacheManager.
+        self._padding_slot: int = self.mamba_cache_free_blocks.pop()
+
         # save intermediate state indices for requests
         self.intermediate_state_indices = torch.arange(max_batch_size,
                                                        dtype=torch.int32,
@@ -513,36 +519,44 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
         request_ids = context_ids + generation_ids
         self._prepare_mamba_cache_blocks(request_ids)
 
+    def _is_padding_sentinel(self, request_id: int) -> bool:
+        # cuda_graph_runner caches one dummy per runtime_draft_len value
+        # (see _get_padded_batch), so any id in the range of dummy request IDs
+        # may be live concurrently.
+        from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \
+            CUDA_GRAPH_DUMMY_REQUEST_ID
+        max_dl = self.speculative_num_draft_tokens or 0
+        return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <=
+                CUDA_GRAPH_DUMMY_REQUEST_ID)
+
     def add_dummy_requests(self, request_ids: List[int], **kwargs):
-        # Allocate a permanent slot for every dummy request ID, including
-        # the CUDA-graph padding sentinel. Padding entries in a batch all
-        # reference the same dummy request ID, so they share one slot via
-        # mamba_cache_index lookup in get_state_indices. This mirrors how
-        # MTP's per-draft-len padding dummies already behave (they use
-        # CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len, which was never
-        # filtered here) and keeps padding writes off every live
-        # request's slot, even under the overlap scheduler where a prior
-        # batch's completed requests linger in mamba_cache_index until
-        # _process_previous_batch runs.
-        if request_ids:
-            for r in request_ids:
-                if r not in self.mamba_cache_index:
-                    if len(self.mamba_cache_free_blocks) == 0:
-                        raise RuntimeError("run out of mamba cache blocks")
-                    block = self.mamba_cache_free_blocks.pop()
-                    self.mamba_cache_index[r] = block
+        # Sentinels alias to the shared _padding_slot; non-sentinel
+        # dummies (warmup, attention-DP idle padding) get their own
+        # slot and are freed individually.
+        if not request_ids:
+            return
+        for r in request_ids:
+            if r in self.mamba_cache_index:
+                continue
+            if self._is_padding_sentinel(r):
+                self.mamba_cache_index[r] = self._padding_slot
+            else:
+                if len(self.mamba_cache_free_blocks) == 0:
+                    raise RuntimeError("run out of mamba cache blocks")
+                block = self.mamba_cache_free_blocks.pop()
+                self.mamba_cache_index[r] = block
 
     def free_resources(self, request: LlmRequest):
         request_id = request.py_request_id
-        if request_id in self.mamba_cache_index:
-            block = self.mamba_cache_index.pop(request_id)
+        if request_id not in self.mamba_cache_index:
+            return
+        block = self.mamba_cache_index.pop(request_id)
+        # _padding_slot stays reserved; only non-sentinel blocks return.
+        if not self._is_padding_sentinel(request_id):
             self.mamba_cache_free_blocks.append(block)
 
     def get_state_indices(self, request_ids: List[int],
                           is_padding: List[bool]) -> List[int]:
-        # Padding entries reuse the slot pre-allocated by their dummy
-        # request in add_dummy_requests; see that method for the
-        # overlap-scheduler rationale.
         return [self.mamba_cache_index[rid] for rid in request_ids]
 
     def get_conv_states(self, layer_idx: int) -> torch.Tensor:
@@ -846,12 +860,10 @@ def __init__(
         # mamba hybrid cache requires block reuse to be disabled in KV cache config
         assert not kv_cache_config.enable_block_reuse, "mamba hybrid cache requires block reuse to be disabled in KV cache config"
 
-        # Reserve one Mamba slot per possible CUDA-graph padding dummy
-        # (one per runtime_draft_len in 0..max_draft_len) so a full
-        # max_batch_size of real requests still leaves room for padding.
-        max_draft_len = (spec_config.max_draft_len
-                         if spec_config is not None else 0)
-        pool_size = max_batch_size + max_draft_len + 1
+        # +1 headroom for PythonMambaCacheManager._padding_slot, which
+        # is shared by every CUDA-graph padding sentinel regardless of
+        # max_draft_len.
+        pool_size = max_batch_size + 1
 
         MambaCacheManager.__init__(
             self,

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -19,8 +19,8 @@
 
 
 def _make_mgr(max_batch_size=4, max_draft_len=2):
-    # Pool size mirrors MambaHybridCacheManager's +max_draft_len+1 headroom.
-    pool = max_batch_size + max_draft_len + 1
+    # +1 headroom matches MixedMambaHybridCacheManager.pool_size.
+    pool = max_batch_size + 1
     return PythonMambaCacheManager(
         d_state=8,
         d_conv=4,
@@ -39,52 +39,75 @@ def _make_mgr(max_batch_size=4, max_draft_len=2):
 
 @skip_no_cuda
 def test_padding_slot_not_held_by_parked_real():
-    """get_state_indices must not hand the padding position a slot
-    owned by a live request outside the current batch. Padding entries
-    all reuse the pre-allocated slot of their dummy request (added via
-    add_dummy_requests), which is distinct from every real request's
-    slot."""
+    """Padding must not resolve to a slot owned by a parked real
+    request outside the current batch."""
     mgr = _make_mgr(max_batch_size=4, max_draft_len=2)
-    # Four real requests claim slots; pool has max_batch_size+max_draft_len+1 = 7 slots.
     mgr._prepare_mamba_cache_blocks([100, 101, 102, 103])
-    # Pre-allocate the padding dummy's slot (what _get_padded_batch does
-    # via kv_cache_manager.add_dummy_requests before get_state_indices).
     mgr.add_dummy_requests([CUDA_GRAPH_DUMMY_REQUEST_ID])
-    # Current batch has only two reals; 102 and 103 are "parked".
+    # 102 and 103 are parked outside the current batch.
     request_ids = [100, 101, CUDA_GRAPH_DUMMY_REQUEST_ID]
     indices = mgr.get_state_indices(request_ids, [False, False, True])
     real_slots = {mgr.mamba_cache_index[r] for r in [100, 101, 102, 103]}
-    assert indices[2] not in real_slots, (
-        f"padding slot {indices[2]} overlaps a real request's slot (real slots: {real_slots})"
-    )
-    # Padding should reuse the dummy's reserved slot, not allocate a new one.
+    assert indices[2] not in real_slots
     assert indices[2] == mgr.mamba_cache_index[CUDA_GRAPH_DUMMY_REQUEST_ID]
 
 
 @skip_no_cuda
 def test_padding_survives_overlap_scheduler_pressure():
-    """Regression for the overlap-scheduler + attention-dp + CUDA-graph
-    padding combo: prior-iter completions linger in mamba_cache_index
-    until _process_previous_batch runs, so get_state_indices must not
-    require N unused pool slots to serve N padding entries."""
+    """Under the overlap scheduler, prior-iter completions linger in
+    mamba_cache_index, so N padding entries must not need N free
+    slots."""
     mgr = _make_mgr(max_batch_size=4, max_draft_len=0)
-    # Fill the pool with "live" real requests (simulates completed
-    # requests from prior iter that haven't been freed yet).
     mgr._prepare_mamba_cache_blocks([100, 101, 102, 103])
-    # Pre-allocate the padding dummy's slot.
     mgr.add_dummy_requests([CUDA_GRAPH_DUMMY_REQUEST_ID])
-    # Current batch: 1 real request + 3 padding entries (attention-dp
-    # pushed padded_batch_size to 4 on this rank even though only 1 real
-    # gen request is scheduled here).
+    # 1 real + 3 padding (attention-dp padded_batch_size=4 on this rank).
     request_ids = [100] + [CUDA_GRAPH_DUMMY_REQUEST_ID] * 3
     is_padding = [False] + [True] * 3
     indices = mgr.get_state_indices(request_ids, is_padding)
-    # All padding entries share the dummy's slot.
     dummy_slot = mgr.mamba_cache_index[CUDA_GRAPH_DUMMY_REQUEST_ID]
     assert indices[0] == mgr.mamba_cache_index[100]
     assert indices[1:] == [dummy_slot] * 3
 
 
+@skip_no_cuda
+def test_all_draft_len_sentinels_share_one_slot():
+    """All per-draft-len sentinels must collapse to a single slot, so
+    the pool needs only +1 headroom regardless of max_draft_len."""
+    max_batch_size, max_draft_len = 4, 3
+    mgr = _make_mgr(max_batch_size=max_batch_size, max_draft_len=max_draft_len)
+    mgr._prepare_mamba_cache_blocks([100, 101, 102, 103])
+
+    sentinels = [CUDA_GRAPH_DUMMY_REQUEST_ID - k for k in range(max_draft_len + 1)]
+    mgr.add_dummy_requests(sentinels)
+
+    shared = mgr.mamba_cache_index[sentinels[0]]
+    real_slots = {mgr.mamba_cache_index[r] for r in [100, 101, 102, 103]}
+    assert shared not in real_slots
+    for s in sentinels:
+        assert mgr.mamba_cache_index[s] == shared
+    assert mgr.mamba_cache_free_blocks == []
+
+
+@skip_no_cuda
+def test_padding_slot_is_permanent():
+    """free_resources drops a sentinel's index entry but the shared
+    slot stays reserved for the next batch."""
+    mgr = _make_mgr(max_batch_size=4, max_draft_len=2)
+    sentinels = [CUDA_GRAPH_DUMMY_REQUEST_ID - k for k in range(3)]
+    mgr.add_dummy_requests(sentinels)
+    shared = mgr.mamba_cache_index[sentinels[0]]
+
+    def _fake(rid):
+        return SimpleNamespace(py_request_id=rid)
+
+    for s in sentinels:
+        mgr.free_resources(_fake(s))
+        assert s not in mgr.mamba_cache_index
+        assert shared not in mgr.mamba_cache_free_blocks
+
+    assert mgr._padding_slot == shared
+
+
 @skip_no_cuda
 def test_update_mamba_states_mtp_path():
     """MTP forward path: update_mamba_states must scatter using the
@@ -126,7 +149,6 @@ def test_update_mamba_states_mtp_path():
 
 @skip_no_cuda
 def test_update_mamba_states_autodeploy_path():
-    """Test update_mamba_states in AutoDeploy path."""
     mgr = _make_mgr()
     mgr._prepare_mamba_cache_blocks([200, 201, 202])
 
@@ -165,7 +187,6 @@ def test_update_mamba_states_autodeploy_path():
 
 @skip_no_cuda
 def test_non_mtp_pytorch_prepare_and_get_state_indices_flow():
-    """Test non-MTP PyTorch backend prepare and get_state_indices flow."""
     mgr = _make_mgr(max_batch_size=4, max_draft_len=0)
     # Simulate a non-MTP step: mix of context + generation requests,
     # plus a CUDA-graph padding dummy.