@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e68fa1b90a
Fork 0

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, 13:45 UTC

[Core] Account for `num_gpu_blocks_override` in `max_model_len` checks (#41069)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Apr 28, 2026 at 15:44 UTC e68fa1b90a7bc52510c11fe2edeae11db15f98fc
5 files changed +128 -60
tests/compile/h100/test_startup.py
+5 -2
@@ -34,7 +34,10 @@ def _run_vllm(vllm_runner):
34 mode=CompilationMode.VLLM_COMPILE,
35 cudagraph_mode=CUDAGraphMode.NONE,
36 ),
37 - num_gpu_blocks_override=8,
37 + # Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
38 + # at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
39 + # 32 for headroom.
40 + num_gpu_blocks_override=32,
41 ):
42 pass
43
@@ -190,7 +193,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
193 cudagraph_mode=CUDAGraphMode.NONE,
194 pass_config=PassConfig(fuse_allreduce_rms=False),
195 ),
193 - num_gpu_blocks_override=8,
196 + num_gpu_blocks_override=16,
197 ):
198 pass
199
tests/v1/core/test_kv_cache_utils.py
+48
@@ -2074,6 +2074,54 @@ def test_auto_fit_max_model_len_not_triggered():
2074 assert vllm_config.model_config.max_model_len == 16
2075
2076
2077 +def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
2078 + """Auto-fit must size max_model_len against the override-clamped pool, not
2079 + the raw `available_memory`. Without this, auto-fit could pick a
2080 + max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
2081 + """
2082 + model_config = ModelConfig(max_model_len=16384)
2083 + model_config.original_max_model_len = -1 # request auto-fit
2084 + vllm_config = VllmConfig(model_config=model_config)
2085 + # Cap the cache to 32 blocks regardless of available memory.
2086 + vllm_config.cache_config.num_gpu_blocks_override = 32
2087 +
2088 + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
2089 + kv_cache_specs = {
2090 + "layer_1": new_kv_cache_spec(), # block_size=16
2091 + "layer_2": new_kv_cache_spec(),
2092 + }
2093 + # Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
2094 + large_available_memory = mem_per_block_per_layer * 2 * 1024
2095 +
2096 + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
2097 +
2098 + # 32 blocks * block_size 16 = 512 token slots, so max_model_len must
2099 + # auto-fit at or below that.
2100 + assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
2101 +
2102 +
2103 +def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
2104 + """Admission check must use the override-clamped pool size, not raw
2105 + `available_memory`. Without this, startup could accept a max_model_len
2106 + that does not actually fit in `num_gpu_blocks_override` blocks.
2107 + """
2108 + model_config = ModelConfig(max_model_len=16384)
2109 + vllm_config = VllmConfig(model_config=model_config)
2110 + # 32 blocks is far too small for max_model_len=16384 (would need 1024).
2111 + vllm_config.cache_config.num_gpu_blocks_override = 32
2112 +
2113 + mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
2114 + kv_cache_specs = {
2115 + "layer_1": new_kv_cache_spec(),
2116 + "layer_2": new_kv_cache_spec(),
2117 + }
2118 + # Plenty of raw memory: a bytes-only check against this would pass.
2119 + large_available_memory = mem_per_block_per_layer * 2 * 1024
2120 +
2121 + with pytest.raises(ValueError, match="max seq len"):
2122 + get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
2123 +
2124 +
2125 def test_unify_hybrid_kv_cache_specs():
2126 # 1. has_full_attention and has_sliding_window
2127 before_spec_1 = new_kv_cache_spec()
tests/v1/e2e/general/test_async_scheduling.py
+6 -4
@@ -324,10 +324,13 @@ def run_test(
324 ):
325 spec_decoding = spec_config is not None
326 cache_arg: dict[str, Any] = (
327 - # Force preemptions
328 - dict(num_gpu_blocks_override=32)
327 + # Force preemptions: with 32 blocks the cache holds at most a single
328 + # max-length request, so the ~34 concurrent prompts contend and trigger
329 + # preemption. (Prompts here are << max_model_len, so dropping
330 + # max_model_len from 4096 to 512 doesn't change generation behavior.)
331 + dict(num_gpu_blocks_override=32, max_model_len=512)
332 if test_preemption
330 - else dict(gpu_memory_utilization=0.9)
333 + else dict(gpu_memory_utilization=0.9, max_model_len=4096)
334 )
335 spec_mml = (spec_config or {}).get("max_model_len")
336 spec_method = (spec_config or {}).get("method", "none")
@@ -343,7 +346,6 @@ def run_test(
346
347 with VllmRunner(
348 model,
346 - max_model_len=4096,
349 enable_chunked_prefill=test_prefill_chunking,
350 # Force prefill chunking
351 max_num_batched_tokens=48 if test_prefill_chunking else None,
vllm/v1/core/kv_cache_utils.py
+68 -53
@@ -890,31 +890,48 @@ def get_max_concurrency_for_kv_cache_config(
890 return max_concurrency
891
892
893 -def may_override_num_blocks(
894 - vllm_config: VllmConfig, num_blocks: int, suppress_log: bool = False
895 -) -> int:
893 +def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int:
894 """
895 Override the number of kv cache blocks if `num_gpu_blocks_override` is set.
896 + The override is logged once, at the call site in `get_kv_cache_configs`.
897 """
898 if vllm_config.cache_config.num_gpu_blocks_override is not None:
900 - num_gpu_blocks_override = vllm_config.cache_config.num_gpu_blocks_override
901 - if not suppress_log:
902 - logger.info(
903 - "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
904 - num_blocks,
905 - num_gpu_blocks_override,
906 - )
907 - num_blocks = num_gpu_blocks_override
908 -
899 + num_blocks = vllm_config.cache_config.num_gpu_blocks_override
900 return num_blocks
901
902
903 +def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
904 + """
905 + Bytes consumed by one block in the worker's shared KV cache pool, mirroring
906 + the divisor used by `get_kv_cache_config_from_groups` to convert
907 + `available_memory` into `num_blocks`. Used to compute the effective KV cache
908 + capacity once `num_gpu_blocks_override` is applied.
909 + """
910 + if len(kv_cache_groups) == 1 and isinstance(
911 + kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
912 + ):
913 + return kv_cache_groups[0].kv_cache_spec.page_size_bytes
914 + if all(
915 + isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups
916 + ):
917 + # DeepseekV4: shared layout sized by the largest per-page-size bucket.
918 + full_mla_spec = cast(UniformTypeKVCacheSpecs, kv_cache_groups[0].kv_cache_spec)
919 + layer_tuple_page_bytes = sum(full_mla_spec.get_page_sizes())
920 + num_layer_tuples = max(
921 + cast(UniformTypeKVCacheSpecs, g.kv_cache_spec).get_num_layer_tuples()
922 + for g in kv_cache_groups
923 + )
924 + return layer_tuple_page_bytes * num_layer_tuples
925 + group_size = max(len(g.layer_names) for g in kv_cache_groups)
926 + page_size = get_uniform_page_size([g.kv_cache_spec for g in kv_cache_groups])
927 + return page_size * group_size
928 +
929 +
930 def get_num_blocks(
931 vllm_config: VllmConfig,
932 num_layers: int,
933 available_memory: int,
934 page_size: int,
917 - suppress_log: bool = False,
935 ) -> int:
936 """
937 Get the number of kv cache blocks.
@@ -924,15 +941,10 @@ def get_num_blocks(
941 num_layers: The number of layers
942 available_memory: Memory available for KV cache in bytes.
943 page_size: The page size of the KV cache.
927 - suppress_log: Whether to suppress override log messages. Used when creating a
928 - temporary/dummy KV cache config, e.g. during CG memory profiling
944 """
945 num_blocks = int(available_memory // page_size // num_layers)
946 num_blocks = max(num_blocks, 0)
932 - num_blocks = may_override_num_blocks(
933 - vllm_config, num_blocks, suppress_log=suppress_log
934 - )
935 - return num_blocks
947 + return may_override_num_blocks(vllm_config, num_blocks)
948
949
950 def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int:
@@ -1220,7 +1232,6 @@ def get_kv_cache_config_from_groups(
1232 vllm_config: VllmConfig,
1233 kv_cache_groups: list[KVCacheGroupSpec],
1234 available_memory: int,
1223 - suppress_log: bool = False,
1235 ) -> KVCacheConfig:
1236 """
1237 Generate the KV cache configuration from the KV cache groups and spec
@@ -1252,9 +1263,7 @@ def get_kv_cache_config_from_groups(
1263 num_blocks = (
1264 available_memory // kv_cache_groups[0].kv_cache_spec.page_size_bytes
1265 )
1255 - num_blocks = may_override_num_blocks(
1256 - vllm_config, num_blocks, suppress_log=suppress_log
1257 - )
1266 + num_blocks = may_override_num_blocks(vllm_config, num_blocks)
1267 per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
1268 kv_cache_tensors = [
1269 KVCacheTensor(
@@ -1288,11 +1297,7 @@ def get_kv_cache_config_from_groups(
1297 )
1298 assert group_size > 0, "group_size must be greater than 0"
1299 num_blocks = get_num_blocks(
1291 - vllm_config,
1292 - group_size,
1293 - available_memory,
1294 - page_size,
1295 - suppress_log=suppress_log,
1300 + vllm_config, group_size, available_memory, page_size
1301 )
1302 kv_cache_tensors = []
1303 for i in range(group_size):
@@ -1688,36 +1693,24 @@ def _report_kv_cache_config(
1693 vllm_config: The global VllmConfig
1694 kv_cache_config: The resolved KV cache configuration
1695 """
1691 - min_block_size = min(
1692 - [group.kv_cache_spec.block_size for group in kv_cache_config.kv_cache_groups]
1693 - )
1694 -
1695 - # Log the KV cache size and maximum concurrency.
1696 - num_tokens = (
1697 - kv_cache_config.num_blocks
1698 - // len(kv_cache_config.kv_cache_groups)
1699 - * min_block_size
1700 - )
1701 - dcp_size = vllm_config.parallel_config.decode_context_parallel_size
1702 - pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
1703 - if pcp_size * dcp_size > 1:
1704 - num_tokens *= pcp_size * dcp_size
1705 - logger.info(
1706 - "Multiplying the GPU KV cache size by the cp_world_size %d "
1707 - "(pcp_world_size %d * dcp_world_size %d).",
1708 - pcp_size * dcp_size,
1709 - pcp_size,
1710 - dcp_size,
1711 - )
1712 - num_tokens_str = f"{num_tokens:,}"
1713 - logger.info_once("GPU KV cache size: %s tokens", num_tokens_str)
1714 - max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"
1696 + max_model_len = vllm_config.model_config.max_model_len
1697 max_concurrency = get_max_concurrency_for_kv_cache_config(
1698 vllm_config, kv_cache_config
1699 )
1700 +
1701 + # GPU KV cache size in tokens = max_concurrency * max_model_len: the total
1702 + # tokens of context the pool can hold at peak utilization. Sourcing this
1703 + # from the concurrency calculation handles hybrid layouts correctly: SWA /
1704 + # chunked-local groups have a per-request block count that's capped by
1705 + # their window, so a naive `num_blocks // num_groups * block_size` formula
1706 + # underestimates capacity for these models. DCP/PCP sharding is already
1707 + # accounted for in each spec's `max_memory_usage_bytes`.
1708 + num_tokens = int(max_concurrency * max_model_len)
1709 +
1710 + logger.info_once("GPU KV cache size: %s tokens", f"{num_tokens:,}")
1711 logger.info_once(
1712 "Maximum concurrency for %s tokens per request: %.2fx",
1720 - max_model_len_str,
1713 + f"{max_model_len:,}",
1714 max_concurrency,
1715 )
1716
@@ -1988,6 +1981,28 @@ def get_kv_cache_configs(
1981 for worker_spec in kv_cache_specs
1982 ]
1983
1984 + # If `num_gpu_blocks_override` is set, the cache size that will actually
1985 + # be allocated is decoupled from the profiled `available_memory`:
1986 + # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps
1987 + # `num_blocks` to the override. Reflect that in `available_memory` here so
1988 + # auto-fit, the admission check, and the per-worker config builder all
1989 + # plan against the same effective capacity.
1990 + override = vllm_config.cache_config.num_gpu_blocks_override
1991 + if override is not None:
1992 + adjusted_memory: list[int] = []
1993 + for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
1994 + if not groups:
1995 + adjusted_memory.append(avail_mem)
1996 + continue
1997 + bytes_per_block = _pool_bytes_per_block(groups)
1998 + logger.info(
1999 + "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d",
2000 + avail_mem // bytes_per_block,
2001 + override,
2002 + )
2003 + adjusted_memory.append(override * bytes_per_block)
2004 + available_memory = adjusted_memory
2005 +
2006 if vllm_config.model_config.original_max_model_len == -1:
2007 _auto_fit_max_model_len(
2008 vllm_config, projected_groups_per_worker, available_memory
vllm/v1/worker/gpu_model_runner.py
+1 -1
@@ -5874,7 +5874,7 @@ class GPUModelRunner(
5874 saved_override = self.cache_config.num_gpu_blocks_override
5875 self.cache_config.num_gpu_blocks_override = min_blocks
5876 minimal_config = get_kv_cache_config_from_groups(
5877 - self.vllm_config, kv_cache_groups, available_memory=0, suppress_log=True
5877 + self.vllm_config, kv_cache_groups, available_memory=0
5878 )
5879 self.cache_config.num_gpu_blocks_override = saved_override
5880