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 15, 06:00 UTC
[Bugfix] Respect num-gpu-blocks-override in v1 (#19503)
Signed-off-by: Jon Swenson <jmswen@gmail.com>
jmswen committed
Jun 12, 2025 at 04:00 UTC
c9280e634652a8643d2d28b44bac7379238b9caf
2 files changed
+17
tests/v1/core/test_kv_cache_utils.py
+16
@@ -900,3 +900,19 @@ def test_get_kv_cache_config():
900
with pytest.raises(NotImplementedError):
901
get_kv_cache_config(vllm_config, kv_cache_specs_hybrid,
902
mem_per_block_per_layer * 2 * 32)
903
+
904
+ # Test num_gpu_blocks_override
905
+ vllm_config.cache_config.num_gpu_blocks_override = 16
906
+ kv_cache_config_override_blocks = get_kv_cache_config(
907
+ vllm_config, kv_cache_specs_full, mem_per_block_per_layer * 2 * 32)
908
+ assert kv_cache_config_override_blocks == KVCacheConfig(
909
+ num_blocks=16,
910
+ kv_cache_tensors=[
911
+ KVCacheTensor(size=mem_per_block_per_layer * 16,
912
+ shared_by=["layer_1"]),
913
+ KVCacheTensor(size=mem_per_block_per_layer * 16,
914
+ shared_by=["layer_2"]),
915
+ ],
916
+ kv_cache_groups=[
917
+ KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())
918
+ ])
\ No newline at end of file
vllm/v1/core/kv_cache_utils.py
+1
@@ -660,6 +660,7 @@ def get_num_blocks(vllm_config: VllmConfig, num_layers: int,
660
logger.info(
661
"Overriding num_gpu_blocks=%d with "
662
"num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)
663
+ num_blocks = num_gpu_blocks_override
664
return num_blocks
665
666