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 6, 23:00 UTC
[Core] Simplify core kv-cache blocks initialization logic (#36521)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 10, 2026 at 13:20 UTC
65b2f405dca824adad17a42a71c908c6ebbcfd9a
4 files changed
+28
-37
tests/models/test_initialization.py
+8
-2
@@ -88,9 +88,15 @@ def can_initialize(
88
[10 * GiB_bytes],
89
)
90
scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
91
+ vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
92
+ kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
93
+ if kv_cache_groups:
94
+ vllm_config.cache_config.block_size = min(
95
+ g.kv_cache_spec.block_size for g in kv_cache_groups
96
+ )
97
92
- # gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
93
- return 1, 0, scheduler_kv_cache_config
98
+ vllm_config.validate_block_size()
99
+ return scheduler_kv_cache_config
100
101
if model_arch == "MiniMaxVL01ForConditionalGeneration":
102
pytest.skip(
vllm/v1/engine/core.py
+11
-18
@@ -117,18 +117,7 @@ class EngineCore:
117
self._eep_scale_up_before_kv_init()
118
119
# Setup KV Caches and update CacheConfig after profiling.
120
- num_gpu_blocks, num_cpu_blocks, kv_cache_config = self._initialize_kv_caches(
121
- vllm_config
122
- )
123
- if kv_cache_config.kv_cache_groups:
124
- vllm_config.cache_config.block_size = min(
125
- g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups
126
- )
127
- vllm_config.validate_block_size()
128
- vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
129
- vllm_config.cache_config.num_cpu_blocks = num_cpu_blocks
130
- self.collective_rpc("initialize_cache", args=(num_gpu_blocks, num_cpu_blocks))
131
-
120
+ kv_cache_config = self._initialize_kv_caches(vllm_config)
121
self.structured_output_manager = StructuredOutputManager(vllm_config)
122
123
# Setup scheduler.
@@ -233,9 +222,7 @@ class EngineCore:
222
enable_envs_cache()
223
224
@instrument(span_name="Prepare model")
236
- def _initialize_kv_caches(
237
- self, vllm_config: VllmConfig
238
- ) -> tuple[int, int, KVCacheConfig]:
225
+ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
226
start = time.time()
227
228
# Get all kv cache needed by the model
@@ -276,8 +263,14 @@ class EngineCore:
263
self.collective_rpc("update_max_model_len", args=(max_model_len_after,))
264
265
scheduler_kv_cache_config = generate_scheduler_kv_cache_config(kv_cache_configs)
279
- num_gpu_blocks = scheduler_kv_cache_config.num_blocks
280
- num_cpu_blocks = 0
266
+ vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
267
+ kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
268
+ if kv_cache_groups:
269
+ vllm_config.cache_config.block_size = min(
270
+ g.kv_cache_spec.block_size for g in kv_cache_groups
271
+ )
272
+
273
+ vllm_config.validate_block_size()
274
275
# Initialize kv cache and warmup the execution
276
self.model_executor.initialize_from_config(kv_cache_configs)
@@ -288,7 +281,7 @@ class EngineCore:
281
elapsed,
282
scope="local",
283
)
291
- return num_gpu_blocks, num_cpu_blocks, scheduler_kv_cache_config
284
+ return scheduler_kv_cache_config
285
286
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
287
return self.model_executor.supported_tasks
vllm/v1/worker/gpu_worker.py
+9
-13
@@ -203,21 +203,17 @@ class Worker(WorkerBase):
203
self.model_runner.init_fp8_kv_scales()
204
205
def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
206
- if self.vllm_config.model_config.enable_sleep_mode:
207
- from vllm.device_allocator.cumem import CuMemAllocator
208
-
209
- allocator = CuMemAllocator.get_instance()
210
- if tag == "weights":
211
- assert allocator.get_current_usage() == 0, (
212
- "Sleep mode can only be used for one instance per process."
213
- )
214
- return allocator.use_memory_pool(tag=tag)
215
- else:
206
+ if not self.vllm_config.model_config.enable_sleep_mode:
207
return nullcontext()
208
218
- def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:
219
- self.cache_config.num_gpu_blocks = num_gpu_blocks
220
- self.cache_config.num_cpu_blocks = num_cpu_blocks
209
+ from vllm.device_allocator.cumem import CuMemAllocator
210
+
211
+ allocator = CuMemAllocator.get_instance()
212
+ if tag == "weights":
213
+ assert allocator.get_current_usage() == 0, (
214
+ "Sleep mode can only be used for one instance per process."
215
+ )
216
+ return allocator.use_memory_pool(tag=tag)
217
218
@instrument(span_name="Init device")
219
def init_device(self):
vllm/v1/worker/worker_base.py
-4
@@ -104,10 +104,6 @@ class WorkerBase:
104
"""
105
raise NotImplementedError
106
107
- def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None:
108
- """Initialize the KV cache with the given size in blocks."""
109
- raise NotImplementedError
110
-
107
def reset_mm_cache(self) -> None:
108
reset_fn = getattr(self.model_runner, "reset_mm_cache", None)
109
if callable(reset_fn):