[Bugfix][Core] Preserve Mamba running CoW after external hits (#51766)
Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Dao007forever committed
Aug 11, 2026 at 03:15 UTC
ce07118669a8a0a7d8ea66d2ac08787c88c25938
2 files changed
+73
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+72
@@ -333,6 +333,78 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
333
assert moved[0].block_hash_num_tokens == 6
334
335
336
+def test_external_mamba_hit_same_block_uses_running_cow_on_continue():
337
+ """An external mid-block hit must become a running request even when its
338
+ first continuation does not need another Mamba block."""
339
+ hash_block_size = 2
340
+ mamba_block_size = 4 * hash_block_size
341
+ kv_cache_config = KVCacheConfig(
342
+ num_blocks=32,
343
+ kv_cache_tensors=[],
344
+ kv_cache_groups=[
345
+ KVCacheGroupSpec(
346
+ ["full"],
347
+ FullAttentionSpec(
348
+ block_size=hash_block_size,
349
+ num_kv_heads=1,
350
+ head_size=1,
351
+ dtype=torch.float32,
352
+ ),
353
+ ),
354
+ KVCacheGroupSpec(
355
+ ["mamba"],
356
+ MambaSpec(
357
+ block_size=mamba_block_size,
358
+ shapes=(1, 1),
359
+ dtypes=(torch.float32,),
360
+ mamba_cache_mode="align",
361
+ ),
362
+ ),
363
+ ],
364
+ )
365
+ manager = make_kv_cache_manager(
366
+ kv_cache_config=kv_cache_config,
367
+ max_model_len=8192,
368
+ enable_caching=True,
369
+ hash_block_size=hash_block_size,
370
+ )
371
+
372
+ request = make_request("0", [0] * 15, hash_block_size, sha256)
373
+ loaded_blocks = manager.allocate_slots(
374
+ request,
375
+ num_new_tokens=0,
376
+ num_external_computed_tokens=10,
377
+ delay_cache_blocks=True,
378
+ )
379
+ assert loaded_blocks is not None
380
+
381
+ request.num_computed_tokens = 10
382
+ first_step_blocks = manager.allocate_slots(request, num_new_tokens=4)
383
+ assert first_step_blocks is not None
384
+
385
+ source_block_id = manager.get_blocks("0").get_block_ids()[1][1]
386
+ partial_hash = request.block_hashes[14 // hash_block_size - 1]
387
+ partial_block = manager.block_pool.get_cached_block(
388
+ partial_hash, kv_cache_group_ids=[1]
389
+ )
390
+ assert partial_block is not None
391
+ assert partial_block[0].block_id == source_block_id
392
+
393
+ request.num_computed_tokens = 14
394
+ continuation_blocks = manager.allocate_slots(request, num_new_tokens=1)
395
+ assert continuation_blocks is not None
396
+
397
+ assert continuation_blocks.get_block_ids()[1] == []
398
+ assert manager.get_blocks("0").get_block_ids()[1][1] == source_block_id
399
+ copies, _ = manager.take_kv_cache_block_copies()
400
+ cow_copy = next(c for c in copies if c.src_block_id == source_block_id)
401
+ assert cow_copy.dst_block_id != source_block_id
402
+
403
+ moved = manager.block_pool.get_cached_block(partial_hash, kv_cache_group_ids=[1])
404
+ assert moved is not None
405
+ assert moved[0].block_id == cow_copy.dst_block_id
406
+
407
+
408
def test_take_partial_tail_offloads_returns_cow_target():
409
"""The connector offload hand-off exposes the mamba CoW *target* block Y
410
(the durable boundary state), not the overwritten source X, and only at
vllm/v1/core/single_type_kv_cache_manager.py
+1
@@ -1559,6 +1559,7 @@ class MambaManager(SingleTypeKVCacheManager):
1559
# `num_required_blocks` might be less than `len(req_blocks)` if blocks are
1560
# over-allocated at last round.
1561
if num_required_blocks <= len(req_blocks) and not has_partial_hit:
1562
+ self._allocated_block_reqs.add(request_id)
1563
return []
1564
else:
1565
prev_block_len = len(req_blocks)