[Bugfix][Core] Reserve the KV null block when validating max_model_len (#47272)
Signed-off-by: hyungjun oh <92hyungjun@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
92hyungjun committed
Aug 20, 2026 at 11:49 UTC
76fb6d210a57c7efbb61fa6cb029d1fa1911bfb5
6 files changed
+109
-16
tests/v1/core/test_kv_cache_utils.py
+73
-1
@@ -27,6 +27,7 @@ from vllm.v1.core.kv_cache_utils import (
27
BlockHash,
28
FreeKVCacheBlockQueue,
29
KVCacheBlock,
30
+ check_enough_kv_cache_memory,
31
estimate_max_model_len,
32
generate_block_hash_extra_keys,
33
generate_scheduler_kv_cache_config,
@@ -2685,7 +2686,9 @@ def test_auto_fit_max_model_len_with_hybrid():
2686
"layer_2": new_kv_cache_spec(),
2687
}
2688
2688
- available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma)
2689
+ # One extra block on top of what a 1024-token request needs: the pool
2690
+ # reserves one block as the null block.
2691
+ available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1)
2692
_kv_cache_configs = get_kv_cache_configs(
2693
vllm_config, [kv_cache_specs], [available_memory]
2694
)
@@ -3095,3 +3098,72 @@ def test_resolve_block_hashes_rejects_mismatched_view():
3098
mismatched = BlockHashListWithBlockSize(raw, 2, 8)
3099
with pytest.raises(AssertionError):
3100
resolve_block_hashes(mismatched, 2, 4)
3101
+
3102
+
3103
+@pytest.mark.parametrize("use_override", [True, False])
3104
+def test_kv_cache_reserves_null_block_for_max_model_len(use_override):
3105
+ """A KV cache sized to exactly the blocks a max_model_len request needs must
3106
+ be rejected. BlockPool keeps one block as the null block, so only
3107
+ num_blocks - 1 are usable; accepting num_blocks == needed would leave the
3108
+ request unschedulable and hang the engine. Covers both the
3109
+ num_gpu_blocks_override path and the memory-derived block count.
3110
+ """
3111
+ block_size = 16
3112
+ max_model_len = 512 # needs 512 / 16 = 32 blocks
3113
+ vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
3114
+ spec = new_kv_cache_spec(block_size=block_size)
3115
+
3116
+ # 32 blocks -> only 31 usable after the null block: one short -> reject.
3117
+ if use_override:
3118
+ # Ample raw memory; the override alone constrains the block count.
3119
+ vllm_config.cache_config.num_gpu_blocks_override = 32
3120
+ available_memory = [spec.page_size_bytes * 1024]
3121
+ else:
3122
+ available_memory = [spec.page_size_bytes * 32]
3123
+
3124
+ with pytest.raises(ValueError, match="max seq len"):
3125
+ get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory)
3126
+
3127
+
3128
+def test_auto_fit_max_model_len_reserves_null_block():
3129
+ """Auto-fit (max_model_len=-1) must size max_model_len against usable
3130
+ blocks, not the total pool. With memory for exactly 64 blocks, one is the
3131
+ null block, so auto-fit must settle on 63 * block_size; picking the full
3132
+ pool would leave a max-length request unschedulable and hang the engine.
3133
+ """
3134
+ block_size = 16
3135
+ model_config = ModelConfig(max_model_len=1024)
3136
+ model_config.original_max_model_len = -1
3137
+ vllm_config = VllmConfig(model_config=model_config)
3138
+ spec = new_kv_cache_spec(block_size=block_size)
3139
+
3140
+ # Exactly the 1024 / 16 = 64 blocks a full-length request would need.
3141
+ available_memory = [spec.page_size_bytes * 64]
3142
+
3143
+ get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory)
3144
+
3145
+ assert vllm_config.model_config.max_model_len == 63 * block_size
3146
+
3147
+
3148
+def test_check_enough_kv_cache_memory_reserves_null_block():
3149
+ """The public admission check must reject a KV cache sized to exactly the
3150
+ blocks a max_model_len request needs. BlockPool keeps one block as the
3151
+ null block, so only num_blocks - 1 are usable; accepting
3152
+ num_blocks == needed would leave the request unschedulable and hang the
3153
+ engine.
3154
+ """
3155
+ block_size = 16
3156
+ max_model_len = 512 # needs 512 / 16 = 32 blocks
3157
+ vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len))
3158
+ spec = new_kv_cache_spec(block_size=block_size)
3159
+
3160
+ # 32 blocks -> only 31 usable after the null block: one short -> reject.
3161
+ with pytest.raises(ValueError, match="max seq len"):
3162
+ check_enough_kv_cache_memory(
3163
+ vllm_config, {"layer1": spec}, spec.page_size_bytes * 32
3164
+ )
3165
+
3166
+ # 33 blocks -> 32 usable after the null block -> accept.
3167
+ check_enough_kv_cache_memory(
3168
+ vllm_config, {"layer1": spec}, spec.page_size_bytes * 33
3169
+ )
tests/v1/e2e/general/test_async_scheduling.py
+6
-5
@@ -332,11 +332,12 @@ def run_test(
332
):
333
spec_decoding = spec_config is not None
334
cache_arg: dict[str, Any] = (
335
- # Force preemptions: with 32 blocks the cache holds at most a single
336
- # max-length request, so the ~34 concurrent prompts contend and trigger
337
- # preemption. (Prompts here are << max_model_len, so dropping
338
- # max_model_len from 4096 to 512 doesn't change generation behavior.)
339
- dict(num_gpu_blocks_override=32, max_model_len=512)
335
+ # Force preemptions: with 33 blocks (one is the reserved null block)
336
+ # the cache holds at most a single max-length request, so the ~34
337
+ # concurrent prompts contend and trigger preemption. (Prompts here are
338
+ # << max_model_len, so dropping max_model_len from 4096 to 512 doesn't
339
+ # change generation behavior.)
340
+ dict(num_gpu_blocks_override=33, max_model_len=512)
341
if test_preemption
342
else dict(gpu_memory_utilization=0.9, max_model_len=4096)
343
)
tests/v1/e2e/general/test_context_length.py
+6
-2
@@ -77,8 +77,12 @@ def test_auto_fit_max_model_len_rejects_oversized_input(
77
78
# Use a small KV cache budget to force auto-fit to a small
79
# max_model_len. Pin block_size=16 so the budget is independent
80
- # of the platform's default block size.
81
- kv_cache_bytes = 1_000_000 # 1 MB
80
+ # of the platform's default block size. One block for this model is
81
+ # 2 (K/V) * 16 (block) * 12 (heads) * 64 (head_dim) * 2 (fp16)
82
+ # * 12 (layers) = 589,824 bytes, and the pool must cover at least
83
+ # two: one is reserved as the null block, so a 1 MB single-block
84
+ # pool cannot serve any tokens at all.
85
+ kv_cache_bytes = 2_000_000 # 2 MB = 3 blocks -> 2 usable
86
87
with vllm_runner(
88
model_name=model,
tests/v1/engine/test_init_error_messaging.py
+3
-2
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+import torch
6
7
from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory
8
from vllm.v1.kv_cache_interface import FullAttentionSpec
@@ -18,7 +19,7 @@ def test_kv_cache_oom_no_memory():
19
block_size=16,
20
num_kv_heads=8,
21
head_size=128,
21
- dtype="float16",
22
+ dtype=torch.float16,
23
)
24
}
25
@@ -46,7 +47,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch):
47
block_size=16,
48
num_kv_heads=8,
49
head_size=128,
49
- dtype="float16",
50
+ dtype=torch.float16,
51
)
52
}
53
tests/v1/sample/test_logprobs.py
+1
-1
@@ -1244,7 +1244,7 @@ def test_prompt_logprobs_with_chunking_and_preemption():
1244
max_model_len=512,
1245
enable_chunked_prefill=True,
1246
max_num_batched_tokens=48, # Force prefill chunking
1247
- num_gpu_blocks_override=32, # Force preemptions
1247
+ num_gpu_blocks_override=33, # Force preemptions (32 usable + null block)
1248
disable_log_stats=False,
1249
gpu_memory_utilization=0.25,
1250
) as vllm_model:
vllm/v1/core/kv_cache_utils.py
+20
-5
@@ -892,8 +892,17 @@ def check_enough_kv_cache_memory(
892
893
# No need to check for available memory if the kv_cache_spec is empty
894
if kv_cache_spec:
895
+ # Reserve the null block BlockPool permanently holds back, so the check
896
+ # plans against usable blocks, as in get_kv_cache_configs. Group a copy
897
+ # of the specs since grouping may unify them in-place.
898
+ groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec))
899
+ check_memory = (
900
+ available_memory - _pool_bytes_per_block(vllm_config, groups)
901
+ if groups
902
+ else available_memory
903
+ )
904
_check_enough_kv_cache_memory(
896
- available_memory,
905
+ check_memory,
906
lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()),
907
vllm_config.model_config.max_model_len,
908
lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am),
@@ -2228,13 +2237,19 @@ def get_kv_cache_configs(
2237
adjusted_memory.append(override * bytes_per_block)
2238
available_memory = adjusted_memory
2239
2240
+ # Reserve the null block BlockPool permanently holds back, so auto-fit and
2241
+ # the capacity check both plan against usable blocks. Allocation below
2242
+ # still uses the full memory.
2243
+ check_memory = [
2244
+ avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem
2245
+ for groups, avail_mem in zip(projected_groups_per_worker, available_memory)
2246
+ ]
2247
+
2248
if vllm_config.model_config.original_max_model_len == -1:
2232
- _auto_fit_max_model_len(
2233
- vllm_config, projected_groups_per_worker, available_memory
2234
- )
2249
+ _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory)
2250
2251
# Check if the available memory is enough per worker.
2237
- for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
2252
+ for groups, avail_mem in zip(projected_groups_per_worker, check_memory):
2253
if not groups:
2254
continue
2255
_check_enough_kv_cache_memory(