102 lines
4.2 KiB
Diff
102 lines
4.2 KiB
Diff
From 8cb229146ba10ea093d647deeaf9a074afaa5bee Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Tue, 22 Sep 2026 04:12:36 -0400
|
|
Subject: [PATCH 008/103] core: cap the Mamba align-mode admission estimate
|
|
(port of upstream #57050)
|
|
|
|
---
|
|
.../core/test_single_type_kv_cache_manager.py | 60 +++++++++++++++++++
|
|
vllm/v1/core/single_type_kv_cache_manager.py | 8 +--
|
|
2 files changed, 64 insertions(+), 4 deletions(-)
|
|
|
|
diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py
|
|
index 14e46d7..1598cec 100644
|
|
--- a/tests/v1/core/test_single_type_kv_cache_manager.py
|
|
+++ b/tests/v1/core/test_single_type_kv_cache_manager.py
|
|
@@ -857,3 +857,63 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
|
|
f"but allocator pulled {len(new_blocks)}"
|
|
)
|
|
total_computed = num_tokens
|
|
+
|
|
+
|
|
+@pytest.mark.parametrize("block_size", [896, 1536])
|
|
+@pytest.mark.parametrize("num_speculative_blocks", [0, 1, 4])
|
|
+@pytest.mark.parametrize("prompt_tokens", [25121, 704547])
|
|
+def test_mamba_checkpoint_admission_matches_allocation(
|
|
+ block_size, num_speculative_blocks, prompt_tokens
|
|
+):
|
|
+ """Checkpoint admission must match the subsequent physical allocation.
|
|
+
|
|
+ Ported verbatim from upstream vllm-project/vllm #57050 (merged 2026-09-16, f30a195) for v2.
|
|
+ """
|
|
+ spec = MambaSpec(
|
|
+ block_size=block_size,
|
|
+ shapes=((1, 1),),
|
|
+ dtypes=(torch.float32,),
|
|
+ mamba_cache_mode="align",
|
|
+ num_speculative_blocks=num_speculative_blocks,
|
|
+ num_prefill_checkpoint_blocks=1,
|
|
+ prefill_checkpoint_alignment=64,
|
|
+ )
|
|
+ pool = BlockPool(
|
|
+ num_gpu_blocks=2048,
|
|
+ enable_caching=True,
|
|
+ hash_block_size=128,
|
|
+ )
|
|
+ manager = MambaManager(
|
|
+ spec,
|
|
+ block_pool=pool,
|
|
+ enable_caching=True,
|
|
+ kv_cache_group_id=0,
|
|
+ scheduler_block_size=block_size,
|
|
+ )
|
|
+ request_id = "prefill"
|
|
+ computed_tokens = 23040
|
|
+
|
|
+ def estimate(num_tokens, total_computed_tokens, apply_admission_cap):
|
|
+ return manager.get_num_blocks_to_allocate(
|
|
+ request_id=request_id,
|
|
+ num_tokens=num_tokens,
|
|
+ new_computed_blocks=[],
|
|
+ total_computed_tokens=total_computed_tokens,
|
|
+ num_local_computed_tokens=total_computed_tokens,
|
|
+ num_tokens_main_model=num_tokens,
|
|
+ apply_admission_cap=apply_admission_cap,
|
|
+ )
|
|
+
|
|
+ estimate(computed_tokens, 0, False)
|
|
+ manager.allocate_new_blocks(request_id, computed_tokens, computed_tokens)
|
|
+ assert request_id in manager._allocated_block_reqs
|
|
+
|
|
+ admission_estimate = estimate(prompt_tokens, computed_tokens, True)
|
|
+ allocation_estimate = estimate(prompt_tokens, computed_tokens, False)
|
|
+ assert request_id in manager._checkpoints
|
|
+
|
|
+ free_before = pool.get_num_free_blocks()
|
|
+ manager.allocate_new_blocks(request_id, prompt_tokens, prompt_tokens)
|
|
+ allocated = free_before - pool.get_num_free_blocks()
|
|
+
|
|
+ assert admission_estimate == allocation_estimate == allocated
|
|
diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py
|
|
index 40027e7..611e1ce 100644
|
|
--- a/vllm/v1/core/single_type_kv_cache_manager.py
|
|
+++ b/vllm/v1/core/single_type_kv_cache_manager.py
|
|
@@ -1796,10 +1796,10 @@ class MambaManager(SingleTypeKVCacheManager):
|
|
self._checkpoints.pop(request_id, None)
|
|
if num_new_blocks > 0:
|
|
blocks_allocated = request_id in self._allocated_block_reqs
|
|
- if not (checkpoint_block and blocks_allocated):
|
|
- num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block
|
|
- if not blocks_allocated:
|
|
- num_new_blocks += self.num_speculative_blocks
|
|
+ physical_block_cap = 1 + int(has_partial_hit) + checkpoint_block
|
|
+ if not blocks_allocated or checkpoint_block:
|
|
+ physical_block_cap += self.num_speculative_blocks
|
|
+ num_new_blocks = min(num_new_blocks, physical_block_cap)
|
|
|
|
num_evictable_computed_blocks = self._get_num_evictable_blocks(
|
|
new_computed_blocks
|
|
--
|
|
2.43.0
|
|
|