[Bugfix][Core] Retire Mamba states across null gaps (#55450)
Signed-off-by: Luca Motz <luca.motz@icloud.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
lucamotz committed
Sep 11, 2026 at 09:50 UTC
fadfe1c7d4df4c29b3dbc159ed790238a885b084
3 files changed
+100
-1
tests/v1/core/test_single_type_kv_cache_manager.py
+76
@@ -103,6 +103,82 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership():
103
manager._relocate_speculative_block([pinned_block], 0)
104
105
106
+def test_mamba_retirement_crosses_null_gaps():
107
+ spec = MambaSpec(
108
+ block_size=4,
109
+ shapes=((1, 1),),
110
+ dtypes=(torch.float32,),
111
+ mamba_cache_mode="align",
112
+ )
113
+ pool = BlockPool(num_gpu_blocks=8, enable_caching=False, hash_block_size=4)
114
+ manager = MambaManager(
115
+ spec,
116
+ block_pool=pool,
117
+ enable_caching=False,
118
+ kv_cache_group_id=0,
119
+ scheduler_block_size=4,
120
+ )
121
+ old, committed, in_flight = pool.get_new_blocks(3)
122
+ manager.req_to_blocks["r"] = [old, pool.null_block, committed, in_flight]
123
+
124
+ # The state at token 12 and the following in-flight state must survive.
125
+ for _ in range(2):
126
+ manager.remove_skipped_blocks("r", processed_computed_tokens=12)
127
+ assert manager.req_to_blocks["r"] == [
128
+ pool.null_block,
129
+ pool.null_block,
130
+ committed,
131
+ in_flight,
132
+ ]
133
+ assert old.ref_cnt == 0
134
+ assert committed.ref_cnt == in_flight.ref_cnt == 1
135
+ assert pool.get_num_free_blocks() == 5
136
+
137
+ manager.free("r")
138
+ manager.req_to_blocks["r"] = pool.get_new_blocks(2)
139
+ manager.remove_skipped_blocks("r", processed_computed_tokens=5)
140
+ assert manager.req_to_blocks["r"][0].is_null
141
+ assert manager.req_to_blocks["r"][1].ref_cnt == 1
142
+
143
+
144
+@pytest.mark.parametrize("block_size", [3584, 4608])
145
+@pytest.mark.parametrize("in_flight_chunks", [0, 1, 2])
146
+def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks):
147
+ spec = MambaSpec(
148
+ block_size=block_size,
149
+ shapes=((1, 1),),
150
+ dtypes=(torch.float32,),
151
+ mamba_cache_mode="align",
152
+ num_speculative_blocks=5,
153
+ )
154
+ pool = BlockPool(num_gpu_blocks=1000, enable_caching=True, hash_block_size=256)
155
+ manager = MambaManager(
156
+ spec,
157
+ block_pool=pool,
158
+ enable_caching=True,
159
+ kv_cache_group_id=0,
160
+ scheduler_block_size=block_size,
161
+ )
162
+ initial_free = pool.get_num_free_blocks()
163
+ chunk = 8192 // block_size * block_size
164
+ peak_held = 0
165
+ # Replay full prefill chunks with an async committed-token lag.
166
+ for target in range(chunk, 480031 + 1, chunk):
167
+ computed = target - chunk
168
+ committed = max(0, computed - in_flight_chunks * chunk)
169
+ manager.remove_skipped_blocks("r", committed)
170
+ manager.get_num_blocks_to_allocate(
171
+ "r", target + 5, [], computed, computed, target
172
+ )
173
+ manager.allocate_new_blocks("r", target + 5, target)
174
+ peak_held = max(peak_held, initial_free - pool.get_num_free_blocks())
175
+
176
+ # Five speculative blocks, current/previous states, and bounded in-flight states.
177
+ assert peak_held == 7 + in_flight_chunks
178
+ manager.free("r")
179
+ assert pool.get_num_free_blocks() == initial_free
180
+
181
+
182
def get_sliding_window_manager(
183
sliding_window_spec,
184
block_pool,
tests/v1/e2e/general/test_mamba_prefix_cache.py
+1
-1
@@ -837,7 +837,7 @@ def get_mamba_prefix_cache_step_configs(
837
StepAction(
838
560 * 10,
839
4,
840
- [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1]
840
+ [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
841
if a
842
else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
843
(9, 10),
vllm/v1/core/single_type_kv_cache_manager.py
+23
@@ -1460,6 +1460,7 @@ class MambaManager(SingleTypeKVCacheManager):
1460
# Mapping from request ID to the index of the block
1461
# allocated in the previous step
1462
self.last_state_block_idx: dict[str, int] = {}
1463
+ self._num_retired_blocks: dict[str, int] = {}
1464
# The set of the requests that have been allocated blocks
1465
self._allocated_block_reqs: set[str] = set()
1466
# checkpoint position and reserved block index for the current
@@ -1610,6 +1611,27 @@ class MambaManager(SingleTypeKVCacheManager):
1611
1612
return mask
1613
1614
+ def _remove_blocks_in_range(
1615
+ self, request_id: str, first_block: int, last_block: int
1616
+ ) -> None:
1617
+ if self.mamba_cache_mode != "align":
1618
+ return super()._remove_blocks_in_range(request_id, first_block, last_block)
1619
+ blocks = self.req_to_blocks.get(request_id, [])
1620
+ first_block = max(first_block, self._num_retired_blocks.get(request_id, 0))
1621
+ last_block = min(last_block, len(blocks))
1622
+ if first_block >= last_block:
1623
+ return
1624
+ freed: list[KVCacheBlock] = []
1625
+ # Mamba prefill leaves null gaps between states awaiting retirement.
1626
+ for i in range(last_block - 1, first_block - 1, -1):
1627
+ if blocks[i].is_null:
1628
+ continue
1629
+ freed.append(blocks[i])
1630
+ blocks[i] = self._null_block
1631
+ if freed:
1632
+ self.block_pool.free_blocks(freed)
1633
+ self._num_retired_blocks[request_id] = last_block
1634
+
1635
def remove_skipped_blocks(
1636
self,
1637
request_id: str,
@@ -1930,6 +1952,7 @@ class MambaManager(SingleTypeKVCacheManager):
1952
if self.mamba_cache_mode == "align":
1953
self._allocated_block_reqs.discard(request_id)
1954
self.last_state_block_idx.pop(request_id, None)
1955
+ self._num_retired_blocks.pop(request_id, None)
1956
self._checkpoints.pop(request_id, None)
1957
self._producer_partial_tail_reqs.pop(request_id, None)
1958
# An offer is only guaranteed to hold committed bytes until the end