@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 94e6b2d55f

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 14, 05:30 UTC

Allow users to specify kv cache memory size (#21489)

Signed-off-by: Boyuan Feng <boyuan@meta.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Boyuan Feng committed Sep 11, 2025 at 06:41 UTC 94e6b2d55f7fd539f01591c7521ec104b25f4e68

10 files changed +236 -47

vllm/config/cache.py

@@ -113,6 +113,15 @@ class CacheConfig:

113 necessary for implementing this optimization in some models (e.g. Gemma3n)

114 """

115

116 + kv_cache_memory_bytes: Optional[int] = None

117 + """Size of KV Cache per GPU in bytes. By default, this is set to None

118 + and vllm can automatically infer the kv cache size based on

119 + gpu_memory_utilization. However, users may want to manually specify

120 + the kv cache memory size. kv_cache_memory_bytes allows more fine-grain

121 + control of how much memory gets used when compared with using

122 + gpu_memory_memory_utilization. Note that kv_cache_memory_bytes

123 + (when not-None) ignores gpu_memory_utilization"""

124 +

125 def compute_hash(self) -> str:

126 """

127 WARNING: Whenever a new field is added to this config,

vllm/engine/arg_utils.py

+11 -1

@@ -227,8 +227,14 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, Any]:

227 elif contains_type(type_hints, int):

228 kwargs[name]["type"] = int

229 # Special case for large integers

230 - if name in {"max_model_len", "max_num_batched_tokens"}:

230 + human_readable_ints = {

231 + "max_model_len",

232 + "max_num_batched_tokens",

233 + "kv_cache_memory_bytes",

234 + }

235 + if name in human_readable_ints:

236 kwargs[name]["type"] = human_readable_int

237 + kwargs[name]["help"] += f"\n\n{human_readable_int.__doc__}"

238 elif contains_type(type_hints, float):

239 kwargs[name]["type"] = float

240 elif (contains_type(type_hints, dict)

@@ -335,6 +341,7 @@ class EngineArgs:

341 swap_space: float = CacheConfig.swap_space

342 cpu_offload_gb: float = CacheConfig.cpu_offload_gb

343 gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization

344 + kv_cache_memory_bytes: Optional[int] = CacheConfig.kv_cache_memory_bytes

345 max_num_batched_tokens: Optional[

346 int] = SchedulerConfig.max_num_batched_tokens

347 max_num_partial_prefills: int = SchedulerConfig.max_num_partial_prefills

@@ -734,6 +741,8 @@ class EngineArgs:

741 cache_group.add_argument("--block-size", **cache_kwargs["block_size"])

742 cache_group.add_argument("--gpu-memory-utilization",

743 **cache_kwargs["gpu_memory_utilization"])

744 + cache_group.add_argument("--kv-cache-memory-bytes",

745 + **cache_kwargs["kv_cache_memory_bytes"])

746 cache_group.add_argument("--swap-space", **cache_kwargs["swap_space"])

747 cache_group.add_argument("--kv-cache-dtype",

748 **cache_kwargs["cache_dtype"])

@@ -1174,6 +1183,7 @@ class EngineArgs:

1183 cache_config = CacheConfig(

1184 block_size=self.block_size,

1185 gpu_memory_utilization=self.gpu_memory_utilization,

1186 + kv_cache_memory_bytes=self.kv_cache_memory_bytes,

1187 swap_space=self.swap_space,

1188 cache_dtype=self.kv_cache_dtype,

1189 is_attention_free=model_config.is_attention_free,

vllm/engine/llm_engine.py

+2 -1

@@ -278,7 +278,8 @@ class LLMEngine:

278 self.cache_config.block_size,

279 "gpu_memory_utilization":

280 self.cache_config.gpu_memory_utilization,

281 -

281 + "kv_cache_memory_bytes":

282 + self.cache_config.kv_cache_memory_bytes,

283 # Quantization

284 "quantization":

285 self.model_config.quantization,

vllm/entrypoints/llm.py

+10

@@ -110,6 +110,14 @@ class LLM:

110 values will increase the KV cache size and thus improve the model's

111 throughput. However, if the value is too high, it may cause out-of-

112 memory (OOM) errors.

113 + kv_cache_memory_bytes: Size of KV Cache per GPU in bytes. By default,

114 + this is set to None and vllm can automatically infer the kv cache

115 + size based on gpu_memory_utilization. However, users may want to

116 + manually specify the kv cache memory size. kv_cache_memory_bytes

117 + allows more fine-grain control of how much memory gets used when

118 + compared with using gpu_memory_memory_utilization. Note that

119 + kv_cache_memory_bytes (when not-None) ignores

120 + gpu_memory_utilization

121 swap_space: The size (GiB) of CPU memory per GPU to use as swap space.

122 This can be used for temporarily storing the states of the requests

123 when their `best_of` sampling parameters are larger than 1. If all

@@ -184,6 +192,7 @@ class LLM:

192 hf_overrides: Optional[HfOverrides] = None,

193 mm_processor_kwargs: Optional[dict[str, Any]] = None,

194 override_pooler_config: Optional[PoolerConfig] = None,

195 + kv_cache_memory_bytes: Optional[int] = None,

196 compilation_config: Optional[Union[int, dict[str, Any],

197 CompilationConfig]] = None,

198 logits_processors: Optional[list[Union[str,

@@ -251,6 +260,7 @@ class LLM:

260 tokenizer_revision=tokenizer_revision,

261 seed=seed,

262 gpu_memory_utilization=gpu_memory_utilization,

263 + kv_cache_memory_bytes=kv_cache_memory_bytes,

264 swap_space=swap_space,

265 cpu_offload_gb=cpu_offload_gb,

266 enforce_eager=enforce_eager,

vllm/utils/__init__.py

+4 -1

@@ -2791,7 +2791,10 @@ def memory_profiling(

2791 result.torch_peak_increase = diff_profile.torch_peak

2792 result.non_torch_increase = diff_from_create.non_torch_memory

2793 result.profile_time = diff_profile.timestamp

2794 - result.non_kv_cache_memory = result.non_torch_increase + result.torch_peak_increase + result.weights_memory # noqa

2794 +

2795 + non_torch_memory = result.non_torch_increase

2796 + peak_activation_memory = result.torch_peak_increase

2797 + result.non_kv_cache_memory = non_torch_memory + peak_activation_memory + result.weights_memory # noqa

2798

2799

2800 # Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501

vllm/v1/utils.py

+2 -1

@@ -355,7 +355,8 @@ def report_usage_stats(

355 vllm_config.cache_config.block_size,

356 "gpu_memory_utilization":

357 vllm_config.cache_config.gpu_memory_utilization,

358 -

358 + "kv_cache_memory_bytes":

359 + vllm_config.cache_config.kv_cache_memory_bytes,

360 # Quantization

361 "quantization":

362 vllm_config.model_config.quantization,

vllm/v1/worker/gpu_model_runner.py

+3 -2

@@ -3041,12 +3041,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3041 self.encoder_cache.clear()

3042 gc.collect()

3043

3044 - def capture_model(self) -> None:

3044 + def capture_model(self) -> int:

3045 if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:

3046 logger.warning(

3047 "Skipping CUDA graph capture. To turn on CUDA graph capture, "

3048 "ensure `cudagraph_mode` was not manually set to `NONE`")

3049 - return

3049 + return 0

3050 else:

3051 self.initialize_cudagraph_capture()

3052

@@ -3117,6 +3117,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3117 # This usually takes 5~20 seconds.

3118 logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",

3119 elapsed_time, cuda_graph_size / (1 << 30))

3120 + return cuda_graph_size

3121

3122 def _capture_cudagraphs(self, compilation_cases: list[int],

3123 cudagraph_runtime_mode: CUDAGraphMode,

vllm/v1/worker/gpu_worker.py

+77 -7

@@ -231,18 +231,40 @@ class Worker(WorkerBase):

231 You may limit the usage of GPU memory

232 by adjusting the `gpu_memory_utilization` parameter.

233 """

234 + GiB = lambda b: b / GiB_bytes

235 + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:

236 + # still need a profile run which compiles the model for

237 + # max_num_batched_tokens

238 + self.model_runner.profile_run()

239 +

240 + msg = (

241 + f"Initial free memory {GiB(self.init_snapshot.free_memory)} "

242 + f"GiB, reserved {GiB(kv_cache_memory_bytes):.2f}GiB memory for "

243 + "KV Cache as specified by kv_cache_memory_bytes config and "

244 + "skipped memory profiling. This does does not respect the "

245 + "gpu_memory_utilization config. Only use kv_cache_memory_bytes "

246 + "config when you want manual control of KV cache memory "

247 + "size. If OOM'ed, check the difference of initial free "

248 + "memory between the current run and the previous run "

249 + "where kv_cache_memory_bytes is suggested and update it "

250 + "correspondingly.")

251 + logger.info(msg)

252 + return kv_cache_memory_bytes

253 +

254 torch.cuda.empty_cache()

255 torch.cuda.reset_peak_memory_stats()

236 - GiB = lambda b: b / GiB_bytes

256

257 # Execute a forward pass with dummy inputs to profile the memory usage

258 # of the model.

259 with memory_profiling(

260 self.init_snapshot,

242 - weights_memory=int(

243 - self.model_runner.model_memory_usage)) as profile_result:

261 + weights_memory=int(self.model_runner.model_memory_usage),

262 + ) as profile_result:

263 self.model_runner.profile_run()

264

265 + self.non_torch_memory = profile_result.non_torch_increase

266 + self.peak_activation_memory = profile_result.torch_peak_increase

267 +

268 free_gpu_memory = profile_result.after_profile.free_memory

269 # NOTE(woosuk): Here we assume that the other processes using the same

270 # GPU did not change their memory usage during the profiling.

@@ -254,7 +276,7 @@ class Worker(WorkerBase):

276 "release GPU memory while vLLM is profiling during initialization. "

277 "To fix this, ensure consistent GPU memory allocation or "

278 "isolate vLLM in its own container.")

257 - available_kv_cache_memory = self.requested_memory \

279 + self.available_kv_cache_memory_bytes = self.requested_memory \

280 - profile_result.non_kv_cache_memory

281

282 unrequested_memory = self.init_snapshot.free_memory \

@@ -274,10 +296,10 @@ class Worker(WorkerBase):

296 )

297 logger.debug(profile_result)

298 logger.info("Available KV cache memory: %.2f GiB",

277 - GiB(available_kv_cache_memory))

299 + GiB(self.available_kv_cache_memory_bytes))

300 gc.collect()

301

280 - return int(available_kv_cache_memory)

302 + return int(self.available_kv_cache_memory_bytes)

303

304 def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:

305 return self.model_runner.get_kv_cache_spec()

@@ -317,8 +339,56 @@ class Worker(WorkerBase):

339 # cuda graph capture.

340 kernel_warmup(self)

341

342 + cuda_graph_memory_bytes = 0

343 if not self.model_config.enforce_eager:

321 - self.model_runner.capture_model()

344 + cuda_graph_memory_bytes = self.model_runner.capture_model()

345 +

346 + if (self.cache_config.kv_cache_memory_bytes is None

347 + and hasattr(self, "peak_activation_memory")):

348 + # Suggests optimal kv cache memory size if we rely on

349 + # memory_profiling to guess the kv cache memory size which

350 + # provides peak_activation_memory and a few other memory

351 + # consumption. `memory_profiling` does not consider

352 + # CUDAGraph memory size and may not utilize all gpu memory.

353 + # Users may want fine-grained control to specify kv cache

354 + # memory size.

355 + GiB = lambda b: round(b / GiB_bytes, 2)

356 +

357 + # empirically observed that the memory profiling may

358 + # slightly underestimate the memory consumption.

359 + # So leave a small buffer (=150MiB) to avoid OOM.

360 + redundancy_buffer_memory = 150 * (1 << 20)

361 + non_kv_cache_memory = (self.model_runner.model_memory_usage +

362 + self.peak_activation_memory +

363 + self.non_torch_memory +

364 + cuda_graph_memory_bytes)

365 + kv_cache_memory_bytes_to_gpu_limit = (

366 + self.init_snapshot.free_memory - non_kv_cache_memory -

367 + redundancy_buffer_memory)

368 + kv_cache_memory_bytes_to_requested_limit = (

369 + int(self.requested_memory) - non_kv_cache_memory -

370 + redundancy_buffer_memory)

371 +

372 + msg = (

373 + f"Free memory on device "

374 + f"({GiB(self.init_snapshot.free_memory)}/"

375 + f"{GiB(self.init_snapshot.total_memory)} GiB) on startup. "

376 + f"Desired GPU memory utilization is "

377 + f"({self.cache_config.gpu_memory_utilization}, "

378 + f"{GiB(self.requested_memory)} GiB). "

379 + f"Actual usage is {GiB(self.model_runner.model_memory_usage)} "

380 + f"GiB for weight, {GiB(self.peak_activation_memory)} GiB "

381 + f"for peak activation, {GiB(self.non_torch_memory)} GiB "

382 + f"for non-torch memory, and {GiB(cuda_graph_memory_bytes)} "

383 + f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "

384 + f"config with `--kv-cache-memory="

385 + f"{kv_cache_memory_bytes_to_requested_limit}` to fit into "

386 + f"requested memory, or `--kv-cache-memory="

387 + f"{kv_cache_memory_bytes_to_gpu_limit}` to fully "

388 + f"utilize gpu memory. Current kv cache memory in use is "

389 + f"{int(self.available_kv_cache_memory_bytes)} bytes.")

390 +

391 + logger.info(msg)

392

393 # Warm up sampler and preallocate memory buffer for logits and other

394 # sampling related tensors of max possible shape to avoid memory

vllm/worker/model_runner.py

+4 -2

@@ -1337,8 +1337,9 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):

1337 return self.lora_manager.list_adapters()

1338

1339 @torch.inference_mode()

1340 - def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> None:

1341 - """Cuda graph capture a model.

1340 + def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> int:

1341 + """Cuda graph capture a model and return cudagraph memory

1342 + consumption in bytes.

1343

1344 Note that CUDA graph's performance gain is negligible if number

1345 of batched tokens are larger than 200. And since CUDA graph

@@ -1505,6 +1506,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):

1506 # This usually takes < 10 seconds.

1507 logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",

1508 elapsed_time, cuda_graph_size / GiB_bytes)

1509 + return cuda_graph_size

1510

1511 def _update_inputs_to_capture_for_enc_dec_model(self,

1512 capture_inputs: Dict[str,

vllm/worker/worker.py

+114 -32

@@ -229,6 +229,67 @@ class Worker(LocalOrDistributedWorkerBase):

229 self.model_runner.save_tensorized_model(

230 tensorizer_config=tensorizer_config, )

231

232 + @torch.inference_mode()

233 + def determine_available_kv_cache_memory(self,

234 + total_gpu_memory: int) -> float:

235 + if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:

236 + # still need a profile run which compiles the model for

237 + # max_num_batched_tokens

238 + self.model_runner.profile_run()

239 +

240 + GiB = lambda b: b / GiB_bytes

241 + msg = (

242 + f"Initial free memory "

243 + f"{GiB(self.baseline_snapshot.free_memory):.2f} "

244 + f"GiB, reserved {GiB(kv_cache_memory_bytes):.2f}GiB memory for "

245 + "KV Cache as specified by kv_cache_memory_bytes config and "

246 + "skipped memory profiling. This does does not respect the "

247 + "gpu_memory_utilization config. Only use kv_cache_memory_bytes "

248 + "config when you want manual control of KV cache memory "

249 + "size. If OOM'ed, check the difference of initial free "

250 + "memory between the current run and the previous run "

251 + "where kv_cache_memory_bytes is suggested and update it "

252 + "correspondingly.")

253 + logger.info(msg)

254 + return self.cache_config.kv_cache_memory_bytes

255 +

256 + # Execute a forward pass with dummy inputs to profile the memory usage

257 + # of the model.

258 + with memory_profiling(

259 + self.baseline_snapshot,

260 + weights_memory=self.model_runner.model_memory_usage) as result:

261 + self.model_runner.profile_run()

262 +

263 + self.non_torch_memory = result.non_torch_increase

264 + self.peak_activation_memory = result.torch_peak_increase

265 +

266 + self._assert_memory_footprint_increased_during_profiling()

267 +

268 + self.requested_memory = total_gpu_memory * \

269 + self.cache_config.gpu_memory_utilization

270 +

271 + self.available_kv_cache_memory = (self.requested_memory -

272 + result.non_kv_cache_memory)

273 +

274 + msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"

275 + "the current vLLM instance can use "

276 + "total_gpu_memory "

277 + f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"

278 + " x gpu_memory_utilization "

279 + f"({self.cache_config.gpu_memory_utilization:.2f})"

280 + f" = {(self.requested_memory / GiB_bytes):.2f}GiB\n"

281 + "model weights take "

282 + f"{(result.weights_memory / GiB_bytes):.2f}GiB;"

283 + " non_torch_memory takes "

284 + f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"

285 + " PyTorch activation peak memory takes "

286 + f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"

287 + " the rest of the memory reserved for KV Cache is "

288 + f"{(self.available_kv_cache_memory / GiB_bytes):.2f}GiB.")

289 +

290 + logger.info(msg)

291 + return self.available_kv_cache_memory

292 +

293 @torch.inference_mode()

294 def determine_num_available_blocks(self) -> Tuple[int, int]:

295 """Profiles the peak memory usage of the model to determine how many

@@ -248,20 +309,8 @@ class Worker(LocalOrDistributedWorkerBase):

309 torch.cuda.reset_peak_memory_stats()

310

311 free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()

251 -

252 - # Execute a forward pass with dummy inputs to profile the memory usage

253 - # of the model.

254 - with memory_profiling(

255 - self.baseline_snapshot,

256 - weights_memory=self.model_runner.model_memory_usage) as result:

257 - self.model_runner.profile_run()

258 -

259 - self._assert_memory_footprint_increased_during_profiling()

260 -

261 - memory_for_current_instance = total_gpu_memory * \

262 - self.cache_config.gpu_memory_utilization

263 - available_kv_cache_memory = (memory_for_current_instance -

264 - result.non_kv_cache_memory)

312 + available_kv_cache_memory = self.determine_available_kv_cache_memory(

313 + total_gpu_memory)

314

315 # Calculate the number of blocks that can be allocated with the

316 # profiled peak memory.

@@ -276,23 +325,6 @@ class Worker(LocalOrDistributedWorkerBase):

325 num_gpu_blocks = max(num_gpu_blocks, 0)

326 num_cpu_blocks = max(num_cpu_blocks, 0)

327

279 - msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"

280 - "the current vLLM instance can use "

281 - "total_gpu_memory "

282 - f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"

283 - " x gpu_memory_utilization "

284 - f"({self.cache_config.gpu_memory_utilization:.2f})"

285 - f" = {(memory_for_current_instance / GiB_bytes):.2f}GiB\n"

286 - "model weights take "

287 - f"{(result.weights_memory / GiB_bytes):.2f}GiB;"

288 - " non_torch_memory takes "

289 - f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"

290 - " PyTorch activation peak memory takes "

291 - f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"

292 - " the rest of the memory reserved for KV Cache is "

293 - f"{(available_kv_cache_memory / GiB_bytes):.2f}GiB.")

294 -

295 - logger.info(msg)

328 # Final cleanup

329 gc.collect()

330

@@ -382,8 +414,58 @@ class Worker(LocalOrDistributedWorkerBase):

414 for size in sorted(warmup_sizes, reverse=True):

415 logger.info("Compile and warming up model for size %d", size)

416 self.model_runner._dummy_run(size)

417 +

418 + cuda_graph_memory_bytes = 0

419 if not self.model_config.enforce_eager:

386 - self.model_runner.capture_model(self.gpu_cache)

420 + cuda_graph_memory_bytes = self.model_runner.capture_model(

421 + self.gpu_cache)

422 +

423 + if (self.cache_config.kv_cache_memory_bytes is None

424 + and hasattr(self, "peak_activation_memory")):

425 + # Suggests optimal kv cache memory size if we rely on

426 + # memory_profiling to guess the kv cache memory size which

427 + # provides peak_activation_memory and a few other memory

428 + # consumption. `memory_profiling` does not consider

429 + # CUDAGraph memory size and may not utilize all gpu memory.

430 + # Users may want fine-grained control to specify kv cache

431 + # memory size.

432 + GiB = lambda b: round(b / GiB_bytes, 2)

433 + non_kv_cache_memory = (self.model_runner.model_memory_usage +

434 + self.peak_activation_memory +

435 + self.non_torch_memory +

436 + cuda_graph_memory_bytes)

437 +

438 + # empirically observed that the memory profiling may

439 + # slightly underestimate the memory consumption.

440 + # So leave a small buffer (=150MiB) to avoid OOM.

441 + redundancy_buffer_memory = 150 * (1 << 20)

442 + kv_cache_memory_bytes_to_gpu_limit = (

443 + self.baseline_snapshot.free_memory - non_kv_cache_memory -

444 + redundancy_buffer_memory)

445 + kv_cache_memory_bytes_to_requested_limit = (

446 + int(self.requested_memory) - non_kv_cache_memory -

447 + redundancy_buffer_memory)

448 +

449 + msg = (

450 + f"Free memory on device "

451 + f"({GiB(self.baseline_snapshot.free_memory)}/"

452 + f"{GiB(self.baseline_snapshot.total_memory)} GiB) on startup. "

453 + f"Desired GPU memory utilization is "

454 + f"({self.cache_config.gpu_memory_utilization}, "

455 + f"{GiB(self.requested_memory)} GiB). "

456 + f"Actual usage is {GiB(self.model_runner.model_memory_usage)} "

457 + f"GiB for weight, {GiB(self.peak_activation_memory)} GiB "

458 + f"for peak activation, {GiB(self.non_torch_memory)} GiB "

459 + f"for non-torch memory, and {GiB(cuda_graph_memory_bytes)} "

460 + f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "

461 + f"config with `--kv-cache-memory="

462 + f"{kv_cache_memory_bytes_to_requested_limit}` to fit into "

463 + f"requested memory, or `--kv-cache-memory="

464 + f"{kv_cache_memory_bytes_to_gpu_limit}` to fully "

465 + f"utilize gpu memory. Current kv cache memory in use is "

466 + f"{int(self.available_kv_cache_memory)} bytes.")

467 + logger.info(msg)

468 +

469 # Reset the seed to ensure that the random state is not affected by

470 # the model initialization and profiling.

471 set_random_seed(self.model_config.seed)