From 76fb6d210a57c7efbb61fa6cb029d1fa1911bfb5 Mon Sep 17 00:00:00 2001 From: 92hyungjun <78071764+92hyungjun@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:49:23 +0900 Subject: [PATCH] [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 Co-authored-by: Claude Co-authored-by: Nick Hill --- tests/v1/core/test_kv_cache_utils.py | 74 ++++++++++++++++++- tests/v1/e2e/general/test_async_scheduling.py | 11 +-- tests/v1/e2e/general/test_context_length.py | 8 +- tests/v1/engine/test_init_error_messaging.py | 5 +- tests/v1/sample/test_logprobs.py | 2 +- vllm/v1/core/kv_cache_utils.py | 25 +++++-- 6 files changed, 109 insertions(+), 16 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 829f45608a..48943aeb60 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -27,6 +27,7 @@ from vllm.v1.core.kv_cache_utils import ( BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, + check_enough_kv_cache_memory, estimate_max_model_len, generate_block_hash_extra_keys, generate_scheduler_kv_cache_config, @@ -2685,7 +2686,9 @@ def test_auto_fit_max_model_len_with_hybrid(): "layer_2": new_kv_cache_spec(), } - available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma) + # One extra block on top of what a 1024-token request needs: the pool + # reserves one block as the null block. + available_memory = mem_per_block_per_layer * (1024 // 16 + 1 + gamma + 1) _kv_cache_configs = get_kv_cache_configs( vllm_config, [kv_cache_specs], [available_memory] ) @@ -3095,3 +3098,72 @@ def test_resolve_block_hashes_rejects_mismatched_view(): mismatched = BlockHashListWithBlockSize(raw, 2, 8) with pytest.raises(AssertionError): resolve_block_hashes(mismatched, 2, 4) + + +@pytest.mark.parametrize("use_override", [True, False]) +def test_kv_cache_reserves_null_block_for_max_model_len(use_override): + """A KV cache sized to exactly the blocks a max_model_len request needs must + be rejected. BlockPool keeps one block as the null block, so only + num_blocks - 1 are usable; accepting num_blocks == needed would leave the + request unschedulable and hang the engine. Covers both the + num_gpu_blocks_override path and the memory-derived block count. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + if use_override: + # Ample raw memory; the override alone constrains the block count. + vllm_config.cache_config.num_gpu_blocks_override = 32 + available_memory = [spec.page_size_bytes * 1024] + else: + available_memory = [spec.page_size_bytes * 32] + + with pytest.raises(ValueError, match="max seq len"): + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + +def test_auto_fit_max_model_len_reserves_null_block(): + """Auto-fit (max_model_len=-1) must size max_model_len against usable + blocks, not the total pool. With memory for exactly 64 blocks, one is the + null block, so auto-fit must settle on 63 * block_size; picking the full + pool would leave a max-length request unschedulable and hang the engine. + """ + block_size = 16 + model_config = ModelConfig(max_model_len=1024) + model_config.original_max_model_len = -1 + vllm_config = VllmConfig(model_config=model_config) + spec = new_kv_cache_spec(block_size=block_size) + + # Exactly the 1024 / 16 = 64 blocks a full-length request would need. + available_memory = [spec.page_size_bytes * 64] + + get_kv_cache_configs(vllm_config, [{"layer1": spec}], available_memory) + + assert vllm_config.model_config.max_model_len == 63 * block_size + + +def test_check_enough_kv_cache_memory_reserves_null_block(): + """The public admission check must reject a KV cache sized to exactly the + blocks a max_model_len request needs. BlockPool keeps one block as the + null block, so only num_blocks - 1 are usable; accepting + num_blocks == needed would leave the request unschedulable and hang the + engine. + """ + block_size = 16 + max_model_len = 512 # needs 512 / 16 = 32 blocks + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=max_model_len)) + spec = new_kv_cache_spec(block_size=block_size) + + # 32 blocks -> only 31 usable after the null block: one short -> reject. + with pytest.raises(ValueError, match="max seq len"): + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 32 + ) + + # 33 blocks -> 32 usable after the null block -> accept. + check_enough_kv_cache_memory( + vllm_config, {"layer1": spec}, spec.page_size_bytes * 33 + ) diff --git a/tests/v1/e2e/general/test_async_scheduling.py b/tests/v1/e2e/general/test_async_scheduling.py index 7f5a115145..1024829307 100644 --- a/tests/v1/e2e/general/test_async_scheduling.py +++ b/tests/v1/e2e/general/test_async_scheduling.py @@ -332,11 +332,12 @@ def run_test( ): spec_decoding = spec_config is not None cache_arg: dict[str, Any] = ( - # Force preemptions: with 32 blocks the cache holds at most a single - # max-length request, so the ~34 concurrent prompts contend and trigger - # preemption. (Prompts here are << max_model_len, so dropping - # max_model_len from 4096 to 512 doesn't change generation behavior.) - dict(num_gpu_blocks_override=32, max_model_len=512) + # Force preemptions: with 33 blocks (one is the reserved null block) + # the cache holds at most a single max-length request, so the ~34 + # concurrent prompts contend and trigger preemption. (Prompts here are + # << max_model_len, so dropping max_model_len from 4096 to 512 doesn't + # change generation behavior.) + dict(num_gpu_blocks_override=33, max_model_len=512) if test_preemption else dict(gpu_memory_utilization=0.9, max_model_len=4096) ) diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index 955e5c5bf0..7e77b31fb1 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -77,8 +77,12 @@ def test_auto_fit_max_model_len_rejects_oversized_input( # Use a small KV cache budget to force auto-fit to a small # max_model_len. Pin block_size=16 so the budget is independent - # of the platform's default block size. - kv_cache_bytes = 1_000_000 # 1 MB + # of the platform's default block size. One block for this model is + # 2 (K/V) * 16 (block) * 12 (heads) * 64 (head_dim) * 2 (fp16) + # * 12 (layers) = 589,824 bytes, and the pool must cover at least + # two: one is reserved as the null block, so a 1 MB single-block + # pool cannot serve any tokens at all. + kv_cache_bytes = 2_000_000 # 2 MB = 3 blocks -> 2 usable with vllm_runner( model_name=model, diff --git a/tests/v1/engine/test_init_error_messaging.py b/tests/v1/engine/test_init_error_messaging.py index bc23a68f9d..5626dec170 100644 --- a/tests/v1/engine/test_init_error_messaging.py +++ b/tests/v1/engine/test_init_error_messaging.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from vllm.v1.core.kv_cache_utils import check_enough_kv_cache_memory from vllm.v1.kv_cache_interface import FullAttentionSpec @@ -18,7 +19,7 @@ def test_kv_cache_oom_no_memory(): block_size=16, num_kv_heads=8, head_size=128, - dtype="float16", + dtype=torch.float16, ) } @@ -46,7 +47,7 @@ def test_kv_cache_oom_insufficient_memory(monkeypatch): block_size=16, num_kv_heads=8, head_size=128, - dtype="float16", + dtype=torch.float16, ) } diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 05a1e93618..f1e88bec8a 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1244,7 +1244,7 @@ def test_prompt_logprobs_with_chunking_and_preemption(): max_model_len=512, enable_chunked_prefill=True, max_num_batched_tokens=48, # Force prefill chunking - num_gpu_blocks_override=32, # Force preemptions + num_gpu_blocks_override=33, # Force preemptions (32 usable + null block) disable_log_stats=False, gpu_memory_utilization=0.25, ) as vllm_model: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 4ac55dbbf7..46e9b8e96b 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -892,8 +892,17 @@ def check_enough_kv_cache_memory( # No need to check for available memory if the kv_cache_spec is empty if kv_cache_spec: + # Reserve the null block BlockPool permanently holds back, so the check + # plans against usable blocks, as in get_kv_cache_configs. Group a copy + # of the specs since grouping may unify them in-place. + groups = get_kv_cache_groups(vllm_config, dict(kv_cache_spec)) + check_memory = ( + available_memory - _pool_bytes_per_block(vllm_config, groups) + if groups + else available_memory + ) _check_enough_kv_cache_memory( - available_memory, + check_memory, lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), vllm_config.model_config.max_model_len, lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), @@ -2228,13 +2237,19 @@ def get_kv_cache_configs( adjusted_memory.append(override * bytes_per_block) available_memory = adjusted_memory + # Reserve the null block BlockPool permanently holds back, so auto-fit and + # the capacity check both plan against usable blocks. Allocation below + # still uses the full memory. + check_memory = [ + avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + for groups, avail_mem in zip(projected_groups_per_worker, available_memory) + ] + if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len( - vllm_config, projected_groups_per_worker, available_memory - ) + _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) # Check if the available memory is enough per worker. - for groups, avail_mem in zip(projected_groups_per_worker, available_memory): + for groups, avail_mem in zip(projected_groups_per_worker, check_memory): if not groups: continue _check_enough_kv_cache_memory(