[Bugfix] Fix incorrect Mamba block allocation estimate that prevents request admission (#57050)
Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com>
Wei Zhao committed
Sep 16, 2026 at 03:23 UTC
f30a195bbb15b920d9c2c40e6a3466d8961ab101
2 files changed
+61
-4
tests/v1/core/test_single_type_kv_cache_manager.py
+57
@@ -179,6 +179,63 @@ def test_mamba_retirement_bounds_prefill_states(block_size, in_flight_chunks):
179
assert pool.get_num_free_blocks() == initial_free
180
181
182
+@pytest.mark.parametrize("block_size", [896, 1536])
183
+@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4])
184
+@pytest.mark.parametrize("prompt_tokens", [25121, 704547])
185
+def test_mamba_checkpoint_admission_matches_allocation(
186
+ block_size, num_speculative_blocks, prompt_tokens
187
+):
188
+ """Checkpoint admission must match the subsequent physical allocation."""
189
+ spec = MambaSpec(
190
+ block_size=block_size,
191
+ shapes=((1, 1),),
192
+ dtypes=(torch.float32,),
193
+ mamba_cache_mode="align",
194
+ num_speculative_blocks=num_speculative_blocks,
195
+ num_prefill_checkpoint_blocks=1,
196
+ prefill_checkpoint_alignment=64,
197
+ )
198
+ pool = BlockPool(
199
+ num_gpu_blocks=2048,
200
+ enable_caching=True,
201
+ hash_block_size=128,
202
+ )
203
+ manager = MambaManager(
204
+ spec,
205
+ block_pool=pool,
206
+ enable_caching=True,
207
+ kv_cache_group_id=0,
208
+ scheduler_block_size=block_size,
209
+ )
210
+ request_id = "prefill"
211
+ computed_tokens = 23040
212
+
213
+ def estimate(num_tokens, total_computed_tokens, apply_admission_cap):
214
+ return manager.get_num_blocks_to_allocate(
215
+ request_id=request_id,
216
+ num_tokens=num_tokens,
217
+ new_computed_blocks=[],
218
+ total_computed_tokens=total_computed_tokens,
219
+ num_local_computed_tokens=total_computed_tokens,
220
+ num_tokens_main_model=num_tokens,
221
+ apply_admission_cap=apply_admission_cap,
222
+ )
223
+
224
+ estimate(computed_tokens, 0, False)
225
+ manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens)
226
+ assert request_id in manager._allocated_block_reqs
227
+
228
+ admission_estimate = estimate(prompt_tokens, computed_tokens, True)
229
+ allocation_estimate = estimate(prompt_tokens, computed_tokens, False)
230
+ assert request_id in manager._checkpoints
231
+
232
+ free_before = pool.get_num_free_blocks()
233
+ manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens)
234
+ allocated = free_before - pool.get_num_free_blocks()
235
+
236
+ assert admission_estimate == allocation_estimate == allocated
237
+
238
+
239
def get_sliding_window_manager(
240
sliding_window_spec,
241
block_pool,
vllm/v1/core/single_type_kv_cache_manager.py
+4
-4
@@ -1796,10 +1796,10 @@ class MambaManager(SingleTypeKVCacheManager):
1796
self._checkpoints.pop(request_id, None)
1797
if num_new_blocks > 0:
1798
blocks_allocated = request_id in self._allocated_block_reqs
1799
- if not (checkpoint_block and blocks_allocated):
1800
- num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block
1801
- if not blocks_allocated:
1802
- num_new_blocks += self.num_speculative_blocks
1799
+ physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block
1800
+ if not blocks_allocated or checkpoint_block:
1801
+ physical_block_cap += self.num_speculative_blocks
1802
+ num_new_blocks = min(num_new_blocks, physical_block_cap)
1803
1804
num_evictable_computed_blocks = self._get_num_evictable_blocks(
1805
new_computed_blocks