A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 11:15 UTC
[Bugfix] Fix failure to allocate KV blocks error (#41282)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Apr 29, 2026 at 21:44 UTC
0ff1bf9bb1ee31ba1f416a4688e705be92643711
4 files changed
+69
-2
tests/v1/core/test_single_type_kv_cache_manager.py
+49
@@ -432,3 +432,52 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
432
)
433
== 15
434
)
435
+
436
+
437
+def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
438
+ """In forward steps, `get_num_blocks_to_allocate` must return exactly what
439
+ `allocate_new_blocks` will pull; otherwise `block_pool.get_new_blocks`
440
+ raises `ValueError: Cannot get N free blocks from the pool`.
441
+ """
442
+ block_size = 2
443
+ sliding_window = 8 # 4-block live window
444
+ cap = sliding_window // block_size
445
+
446
+ spec = SlidingWindowSpec(
447
+ block_size=block_size,
448
+ num_kv_heads=1,
449
+ head_size=1,
450
+ dtype=torch.float32,
451
+ sliding_window=sliding_window,
452
+ )
453
+ block_pool = BlockPool(
454
+ num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
455
+ )
456
+ manager = SlidingWindowManager(
457
+ spec,
458
+ block_pool=block_pool,
459
+ enable_caching=False,
460
+ kv_cache_group_id=0,
461
+ max_admission_blocks_per_request=cap,
462
+ )
463
+
464
+ request_id = "req"
465
+ total_computed = 0
466
+ # Walk through request forward steps. Check num_blocks returned by
467
+ # `get_num_blocks_to_allocate` matches what `allocate_new_blocks` pulls
468
+ for num_tokens in (4, 8, 12, 16):
469
+ predicted = manager.get_num_blocks_to_allocate(
470
+ request_id=request_id,
471
+ num_tokens=num_tokens,
472
+ new_computed_blocks=[],
473
+ total_computed_tokens=total_computed,
474
+ num_tokens_main_model=num_tokens,
475
+ )
476
+ new_blocks = manager.allocate_new_blocks(
477
+ request_id, num_tokens=num_tokens, num_tokens_main_model=num_tokens
478
+ )
479
+ assert predicted == len(new_blocks), (
480
+ f"num_tokens={num_tokens}: predictor returned {predicted} "
481
+ f"but allocator pulled {len(new_blocks)}"
482
+ )
483
+ total_computed = num_tokens
vllm/v1/core/kv_cache_coordinator.py
+12
-1
@@ -85,6 +85,7 @@ class KVCacheCoordinator(ABC):
85
num_encoder_tokens: int,
86
total_computed_tokens: int,
87
num_tokens_main_model: int,
88
+ apply_admission_cap: bool = False,
89
) -> int:
90
"""
91
Get the number of blocks needed to be allocated for the request.
@@ -101,6 +102,10 @@ class KVCacheCoordinator(ABC):
102
num_tokens_main_model: The number of tokens for the main model (aka target
103
model in spec decode). w/o spec decode, it is num_tokens;
104
with spec decode, it is num_tokens - num_lookahead_tokens.
105
+ apply_admission_cap: If True, apply the recycling-aware
106
+ per-request admission cap (SWA / chunked-local). Set only by
107
+ the full-sequence admission gate; per-step allocation must
108
+ leave it False so the predictor matches `allocate_new_blocks`.
109
110
Returns:
111
The number of blocks to allocate.
@@ -111,7 +116,12 @@ class KVCacheCoordinator(ABC):
116
# For cross-attention, we issue a single static allocation
117
# of blocks based on the number of encoder input tokens.
118
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
114
- request_id, num_encoder_tokens, [], 0, num_encoder_tokens
119
+ request_id,
120
+ num_encoder_tokens,
121
+ [],
122
+ 0,
123
+ num_encoder_tokens,
124
+ apply_admission_cap=apply_admission_cap,
125
)
126
else:
127
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
@@ -120,6 +130,7 @@ class KVCacheCoordinator(ABC):
130
new_computed_blocks[i],
131
total_computed_tokens,
132
num_tokens_main_model,
133
+ apply_admission_cap=apply_admission_cap,
134
)
135
return num_blocks_to_allocate
136
vllm/v1/core/kv_cache_manager.py
+1
@@ -257,6 +257,7 @@ class KVCacheManager:
257
num_encoder_tokens=num_encoder_tokens,
258
total_computed_tokens=total_computed_tokens,
259
num_tokens_main_model=full_num_tokens,
260
+ apply_admission_cap=True,
261
)
262
263
return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
vllm/v1/core/single_type_kv_cache_manager.py
+7
-1
@@ -92,6 +92,7 @@ class SingleTypeKVCacheManager(ABC):
92
new_computed_blocks: Sequence[KVCacheBlock],
93
total_computed_tokens: int,
94
num_tokens_main_model: int,
95
+ apply_admission_cap: bool = False,
96
) -> int:
97
"""
98
Get the number of blocks needed to be allocated for the request.
@@ -107,13 +108,16 @@ class SingleTypeKVCacheManager(ABC):
108
num_tokens_main_model: The number of tokens for the main model (aka target
109
model in spec decode). w/o spec decode, it is num_tokens;
110
with spec decode, it is num_tokens - num_lookahead_tokens.
111
+ apply_admission_cap: If True, clamp by `num_required_blocks` by
112
+ `_max_admission_blocks_per_request`for recycling-aware specs
113
+ (SWA, chunked-local).
114
115
Returns:
116
The number of blocks to allocate.
117
"""
118
119
num_required_blocks = cdiv(num_tokens, self.block_size)
116
- if self._max_admission_blocks_per_request is not None:
120
+ if apply_admission_cap and self._max_admission_blocks_per_request is not None:
121
# Recycling-aware specs (SWA, chunked-local) cap the per-request
122
# reservation here so admission matches the startup pool sizer
123
# (`SlidingWindowSpec.max_admission_blocks_per_request` / its
@@ -893,6 +897,7 @@ class MambaManager(SingleTypeKVCacheManager):
897
new_computed_blocks: Sequence[KVCacheBlock],
898
total_computed_tokens: int,
899
num_tokens_main_model: int,
900
+ apply_admission_cap: bool = False,
901
) -> int:
902
assert isinstance(self.kv_cache_spec, MambaSpec)
903
if (
@@ -917,6 +922,7 @@ class MambaManager(SingleTypeKVCacheManager):
922
new_computed_blocks,
923
total_computed_tokens,
924
num_tokens_main_model,
925
+ apply_admission_cap=apply_admission_cap,
926
)
927
else:
928
# We don't allocate blocks for lookahead tokens in align mode, because if