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

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

[UX][Startup] Account for CUDA graphs during memory profiling (#30515)

Matthew Bonanni committed Mar 7, 2026 at 16:49 UTC ebb9cc5f2b26d73222c08e42b32fcf59e831386c

6 files changed +360 -61

vllm/compilation/cuda_graph.py

+19 -1

@@ -2,10 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import dataclasses

5 +import weakref

6 from collections import Counter

7 from collections.abc import Callable

8 from contextlib import ExitStack

8 -from typing import Any

9 +from typing import Any, ClassVar

10 from unittest.mock import patch

12 import torch

@@ -162,6 +163,14 @@ class CUDAGraphWrapper:

163 guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".

164 """

165

166 + _all_instances: ClassVar[weakref.WeakSet["CUDAGraphWrapper"]] = weakref.WeakSet()

167 +

168 + @classmethod

169 + def clear_all_graphs(cls) -> None:

170 + """Clear captured graphs from all CUDAGraphWrapper instances."""

171 + for instance in list(cls._all_instances):

172 + instance.clear_graphs()

173 +

174 def __init__(

175 self,

176 runnable: Callable[..., Any],

@@ -192,6 +201,8 @@ class CUDAGraphWrapper:

201 # cudagraphs for.

202 self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {}

203

204 + CUDAGraphWrapper._all_instances.add(self)

205 +

206 def __getattr__(self, key: str) -> Any:

207 # allow accessing the attributes of the runnable.

208 if hasattr(self.runnable, key):

@@ -205,6 +216,13 @@ class CUDAGraphWrapper:

216 # in case we need to access the original runnable.

217 return self.runnable

218

219 + @property

220 + def cudagraph_wrapper(self) -> "CUDAGraphWrapper":

221 + return self

222 +

223 + def clear_graphs(self) -> None:

224 + self.concrete_cudagraph_entries.clear()

225 +

226 def __call__(self, *args: Any, **kwargs: Any) -> Any | None:

227 forward_context = get_forward_context()

228 batch_descriptor = forward_context.batch_descriptor

vllm/envs.py

@@ -244,6 +244,7 @@ if TYPE_CHECKING:

244 VLLM_CUDA_COMPATIBILITY_PATH: str | None = None

245 VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False

246 VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False

247 + VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = False

248

249

250 def get_default_cache_root():

@@ -1628,6 +1629,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

1629 "VLLM_ELASTIC_EP_DRAIN_REQUESTS": lambda: bool(

1630 int(os.getenv("VLLM_ELASTIC_EP_DRAIN_REQUESTS", "0"))

1631 ),

1632 + # If set to 1, enable CUDA graph memory estimation during memory profiling.

1633 + # This profiles CUDA graph memory usage to provide more accurate KV cache

1634 + # memory allocation. Disabled by default to preserve existing behavior.

1635 + "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool(

1636 + int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "0"))

1637 + ),

1638 }

1639

1640

vllm/v1/cudagraph_dispatcher.py

+5 -2

@@ -334,8 +334,11 @@ class CudagraphDispatcher:

334 for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]:

335 descs = list(self.cudagraph_keys[mode])

336 if descs:

337 - # Sort by num_tokens descending (largest first)

338 - descs.sort(key=lambda d: d.num_tokens, reverse=True)

337 + # Sort by (num_tokens, num_active_loras) descending

338 + descs.sort(

339 + key=lambda d: (d.num_tokens, d.num_active_loras),

340 + reverse=True,

341 + )

342 result.append((mode, descs))

343

344 return result

vllm/v1/worker/gpu_model_runner.py

+227 -52

@@ -29,6 +29,7 @@ from vllm.config import (

29 CUDAGraphMode,

30 VllmConfig,

31 get_layers_from_vllm_config,

32 + set_current_vllm_config,

33 update_config,

34 )

35 from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer

@@ -94,6 +95,7 @@ from vllm.multimodal.inputs import (

95 PlaceholderRange,

96 )

97 from vllm.multimodal.utils import group_and_batch_mm_kwargs

98 +from vllm.platforms import current_platform

99 from vllm.pooling_params import PoolingParams

100 from vllm.sampling_params import SamplingType

101 from vllm.sequence import IntermediateTensors

@@ -596,6 +598,17 @@ class GPUModelRunner(

598 self.async_output_copy_stream = torch.cuda.Stream()

599 self.prepare_inputs_event = torch.Event()

600

601 + # self.cudagraph_batch_sizes sorts in ascending order.

602 + if (

603 + self.compilation_config.cudagraph_capture_sizes

604 + and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

605 + ):

606 + self.cudagraph_batch_sizes = sorted(

607 + self.compilation_config.cudagraph_capture_sizes

608 + )

609 + else:

610 + self.cudagraph_batch_sizes = []

611 +

612 # Cache the device properties.

613 self._init_device_properties()

614

@@ -4727,6 +4740,7 @@ class GPUModelRunner(

4740 remove_lora: bool = True,

4741 is_graph_capturing: bool = False,

4742 num_active_loras: int = 0,

4743 + profile_seq_lens: int | None = None,

4744 ) -> tuple[torch.Tensor, torch.Tensor]:

4745 """

4746 Run a dummy forward pass to warm up/profile run or capture the

@@ -4751,6 +4765,9 @@ class GPUModelRunner(

4765 remove_lora: If False, dummy LoRAs are not destroyed after the run

4766 num_active_loras: Number of distinct active LoRAs to capture for.

4767 LoRA is activated when num_active_loras > 0.

4768 + profile_seq_lens: If provided, use this value for seq_lens instead

4769 + of max_query_len. Used to profile attention workspace that

4770 + scales with context length.

4771 """

4772 mm_config = self.vllm_config.model_config.multimodal_config

4773 if mm_config and mm_config.mm_encoder_only:

@@ -4881,11 +4898,13 @@ class GPUModelRunner(

4898 # If force_attention is True, we always capture attention.

4899 # Otherwise, it only happens for cudagraph_runtime_mode=FULL.

4900 if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:

4884 - if create_mixed_batch:

4901 + if profile_seq_lens is not None:

4902 + seq_lens = profile_seq_lens # type: ignore[assignment]

4903 + elif create_mixed_batch:

4904 # In the mixed batch mode (used for FI warmup), we use

4905 # shorter sequence lengths to run faster.

4906 # TODO(luka) better system for describing dummy batches

4888 - seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]

4907 + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] # type: ignore[assignment]

4908 else:

4909 seq_lens = max_query_len # type: ignore[assignment]

4910 self.seq_lens.np[:num_reqs] = seq_lens

@@ -5298,6 +5317,167 @@ class GPUModelRunner(

5317 self.encoder_cache.clear()

5318 gc.collect()

5319

5320 + def _init_minimal_kv_cache_for_profiling(self) -> None:

5321 + from vllm.v1.core.kv_cache_utils import (

5322 + get_kv_cache_config_from_groups,

5323 + get_kv_cache_groups,

5324 + )

5325 +

5326 + kv_cache_spec = self.get_kv_cache_spec()

5327 + kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec)

5328 + min_blocks = self.compilation_config.max_cudagraph_capture_size or 1

5329 + if kv_cache_groups:

5330 + page_size = kv_cache_groups[0].kv_cache_spec.page_size_bytes

5331 + group_size = max(len(g.layer_names) for g in kv_cache_groups)

5332 + available_memory = min_blocks * page_size * group_size

5333 + else:

5334 + available_memory = 1 # Attention-free model

5335 +

5336 + minimal_config = get_kv_cache_config_from_groups(

5337 + self.vllm_config, kv_cache_groups, available_memory=available_memory

5338 + )

5339 +

5340 + self.initialize_kv_cache(minimal_config)

5341 + self.cache_config.num_gpu_blocks = minimal_config.num_blocks

5342 +

5343 + logger.debug("Initialized minimal KV cache for CUDA graph profiling")

5344 +

5345 + @staticmethod

5346 + @contextmanager

5347 + def _freeze_gc():

5348 + gc.collect()

5349 + should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC

5350 + if should_freeze:

5351 + gc.freeze()

5352 + try:

5353 + yield

5354 + finally:

5355 + if should_freeze:

5356 + gc.unfreeze()

5357 + gc.collect()

5358 +

5359 + def _cleanup_profiling_kv_cache(self) -> None:

5360 + torch.accelerator.synchronize()

5361 + if hasattr(self, "kv_caches") and self.kv_caches:

5362 + for i in range(len(self.kv_caches)):

5363 + self.kv_caches[i] = None # type: ignore

5364 + self.kv_caches.clear()

5365 + if hasattr(self, "cross_layers_kv_cache"):

5366 + self.cross_layers_kv_cache = None

5367 + self.cross_layers_attn_backend = None

5368 + if hasattr(self, "attn_groups"):

5369 + self.attn_groups.clear()

5370 + if hasattr(self, "kv_cache_config"):

5371 + delattr(self, "kv_cache_config")

5372 + self.cache_config.num_gpu_blocks = None

5373 +

5374 + for layer in self.compilation_config.static_forward_context.values():

5375 + if hasattr(layer, "kv_cache"):

5376 + layer.kv_cache = []

5377 +

5378 + gc.collect()

5379 + torch.accelerator.empty_cache()

5380 +

5381 + logger.debug("Cleaned up profiling KV cache and CUDA graphs")

5382 +

5383 + @torch.inference_mode()

5384 + def profile_cudagraph_memory(self) -> int:

5385 + with set_current_vllm_config(self.vllm_config):

5386 + self._init_minimal_kv_cache_for_profiling()

5387 +

5388 + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured

5389 +

5390 + capture_descs = self.cudagraph_dispatcher.get_capture_descs()

5391 +

5392 + total_graphs = sum(len(descs) for _, descs in capture_descs)

5393 + if total_graphs == 0:

5394 + logger.debug("No CUDA graphs will be captured, skipping profiling")

5395 + self._cleanup_profiling_kv_cache()

5396 + return 0

5397 +

5398 + logger.info(

5399 + "Profiling CUDA graph memory: %s",

5400 + ", ".join(

5401 + f"{mode.name}={len(descs)} (largest={descs[0].num_tokens})"

5402 + for mode, descs in capture_descs

5403 + if descs

5404 + ),

5405 + )

5406 +

5407 + # Use a temporary pool for profiling to avoid fragmentation in the main pool.

5408 + profiling_pool = current_platform.graph_pool_handle()

5409 + original_pools: dict[int, Any] = {}

5410 + for instance in list(CUDAGraphWrapper._all_instances):

5411 + original_pools[id(instance)] = instance.graph_pool

5412 + instance.graph_pool = profiling_pool

5413 +

5414 + set_cudagraph_capturing_enabled(True)

5415 + with self._freeze_gc(), graph_capture(device=self.device):

5416 + shared_memory_estimate = {}

5417 + per_graph_estimate = {}

5418 + torch.accelerator.synchronize()

5419 + torch.accelerator.empty_cache()

5420 +

5421 + for mode, descs in capture_descs:

5422 + profile_descs = descs[:2]

5423 + mem_samples: list[int] = []

5424 +

5425 + for i, desc in enumerate(profile_descs):

5426 + mem_before = torch.cuda.mem_get_info()[0]

5427 + self._warmup_and_capture(

5428 + desc,

5429 + cudagraph_runtime_mode=mode,

5430 + profile_seq_lens=(

5431 + min(

5432 + self.max_model_len,

5433 + self.max_num_tokens // desc.num_tokens,

5434 + )

5435 + if mode == CUDAGraphMode.FULL and i == 0

5436 + else None

5437 + ),

5438 + )

5439 + torch.accelerator.synchronize()

5440 + free_after = torch.cuda.mem_get_info()[0]

5441 + mem_samples.append(mem_before - free_after)

5442 +

5443 + first_capture = mem_samples[0]

5444 + # Use at least 1 MiB per graph for driver overhead

5445 + per_graph = max(mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20)

5446 +

5447 + shared_memory_estimate[mode] = first_capture

5448 + per_graph_estimate[mode] = per_graph * (len(descs) - 1)

5449 +

5450 + logger.debug(

5451 + "Estimated %s CUDA graph memory: "

5452 + "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph",

5453 + mode.name,

5454 + first_capture / (1 << 20),

5455 + len(descs),

5456 + per_graph / (1 << 20),

5457 + )

5458 +

5459 + set_cudagraph_capturing_enabled(False)

5460 + CUDAGraphWrapper.clear_all_graphs()

5461 + for instance in list(CUDAGraphWrapper._all_instances):

5462 + if id(instance) in original_pools:

5463 + instance.graph_pool = original_pools[id(instance)]

5464 + self.maybe_remove_all_loras(self.lora_config)

5465 + self._cleanup_profiling_kv_cache()

5466 + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

5467 +

5468 + # FULL and PIECEWISE graphs share the global pool at runtime and are

5469 + # never replayed concurrently, so the pool overlays their memory.

5470 + # Take the max to avoid double-counting the overlap.

5471 + total_estimate = max(shared_memory_estimate.values()) + sum(

5472 + per_graph_estimate.values()

5473 + )

5474 + logger.info(

5475 + "Estimated CUDA graph memory: %.2f GiB total",

5476 + total_estimate / (1 << 30),

5477 + )

5478 +

5479 + return int(total_estimate)

5480 +

5481 @instrument(span_name="Capture model")

5482 def capture_model(self) -> int:

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

@@ -5311,27 +5491,13 @@ class GPUModelRunner(

5491

5492 start_time = time.perf_counter()

5493

5314 - @contextmanager

5315 - def freeze_gc():

5316 - # Optimize garbage collection during CUDA graph capture.

5317 - # Clean up, then freeze all remaining objects from being included

5318 - # in future collections.

5319 - gc.collect()

5320 - should_freeze = not envs.VLLM_ENABLE_CUDAGRAPH_GC

5321 - if should_freeze:

5322 - gc.freeze()

5323 - try:

5324 - yield

5325 - finally:

5326 - if should_freeze:

5327 - gc.unfreeze()

5328 - gc.collect()

5329 -

5494 # Trigger CUDA graph capture for specific shapes.

5495 # Capture the large shapes first so that the smaller shapes

5496 # can reuse the memory pool allocated for the large shapes.

5497 set_cudagraph_capturing_enabled(True)

5334 - with freeze_gc(), graph_capture(device=self.device):

5498 + with self._freeze_gc(), graph_capture(device=self.device):

5499 + torch.accelerator.synchronize()

5500 + torch.accelerator.empty_cache()

5501 start_free_gpu_memory = torch.cuda.mem_get_info()[0]

5502

5503 for (

@@ -5342,6 +5508,7 @@ class GPUModelRunner(

5508 batch_descriptors=batch_descs,

5509 cudagraph_runtime_mode=runtime_mode,

5510 )

5511 + torch.accelerator.synchronize()

5512

5513 torch.accelerator.synchronize()

5514 end_free_gpu_memory = torch.cuda.mem_get_info()[0]

@@ -5353,6 +5520,9 @@ class GPUModelRunner(

5520 # after here.

5521 set_cudagraph_capturing_enabled(False)

5522

5523 + torch.accelerator.synchronize()

5524 + torch.accelerator.empty_cache()

5525 +

5526 # Lock workspace to prevent resizing during execution.

5527 # Max workspace sizes should have been captured during warmup/profiling.

5528 lock_workspace()

@@ -5369,6 +5539,40 @@ class GPUModelRunner(

5539 )

5540 return cuda_graph_size

5541

5542 + def _warmup_and_capture(

5543 + self,

5544 + desc: BatchDescriptor,

5545 + cudagraph_runtime_mode: CUDAGraphMode,

5546 + profile_seq_lens: int | None = None,

5547 + allow_microbatching: bool = False,

5548 + num_warmups: int | None = None,

5549 + ):

5550 + if num_warmups is None:

5551 + num_warmups = self.compilation_config.cudagraph_num_of_warmups

5552 + force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL

5553 + for _ in range(num_warmups):

5554 + self._dummy_run(

5555 + desc.num_tokens,

5556 + cudagraph_runtime_mode=CUDAGraphMode.NONE,

5557 + force_attention=force_attention,

5558 + uniform_decode=desc.uniform,

5559 + allow_microbatching=allow_microbatching,

5560 + skip_eplb=True,

5561 + remove_lora=False,

5562 + num_active_loras=desc.num_active_loras,

5563 + )

5564 + self._dummy_run(

5565 + desc.num_tokens,

5566 + cudagraph_runtime_mode=cudagraph_runtime_mode,

5567 + uniform_decode=desc.uniform,

5568 + allow_microbatching=allow_microbatching,

5569 + skip_eplb=True,

5570 + remove_lora=False,

5571 + num_active_loras=desc.num_active_loras,

5572 + is_graph_capturing=True,

5573 + profile_seq_lens=profile_seq_lens,

5574 + )

5575 +

5576 def _capture_cudagraphs(

5577 self,

5578 batch_descriptors: list[BatchDescriptor],

@@ -5383,15 +5587,6 @@ class GPUModelRunner(

5587 return

5588

5589 uniform_decode = batch_descriptors[0].uniform

5386 - force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL

5387 -

5388 - dummy_run = functools.partial(

5389 - self._dummy_run,

5390 - uniform_decode=uniform_decode,

5391 - skip_eplb=True,

5392 - remove_lora=False,

5393 - force_attention=force_attention,

5394 - )

5590

5591 # Only rank 0 should print progress bar during capture

5592 if is_global_first_rank():

@@ -5406,9 +5601,6 @@ class GPUModelRunner(

5601

5602 # We skip EPLB here since we don't want to record dummy metrics

5603 for batch_desc in batch_descriptors:

5409 - num_tokens = batch_desc.num_tokens

5410 - num_active_loras = batch_desc.num_active_loras

5411 -

5604 # We currently only capture ubatched graphs when its a FULL

5605 # cudagraph, a uniform decode batch, and the number of tokens

5606 # is above the threshold. Otherwise we just capture a non-ubatched

@@ -5419,33 +5611,16 @@ class GPUModelRunner(

5611 and uniform_decode

5612 and check_ubatch_thresholds(

5613 config=self.vllm_config.parallel_config,

5422 - num_tokens=num_tokens,

5614 + num_tokens=batch_desc.num_tokens,

5615 uniform_decode=uniform_decode,

5616 )

5617 )

5426 -

5427 - for _ in range(self.compilation_config.cudagraph_num_of_warmups):

5428 - # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.

5429 - # But be careful, warm up with `NONE` is orthogonal to

5430 - # if we want to warm up attention or not. This is

5431 - # different from the case where `FULL` implies capture

5432 - # attention while `PIECEWISE` implies no attention.

5433 -

5434 - dummy_run(

5435 - num_tokens,

5436 - cudagraph_runtime_mode=CUDAGraphMode.NONE,

5437 - allow_microbatching=allow_microbatching,

5438 - num_active_loras=num_active_loras,

5439 - )

5440 -

5441 - # Capture run

5442 - dummy_run(

5443 - num_tokens,

5618 + self._warmup_and_capture(

5619 + batch_desc,

5620 cudagraph_runtime_mode=cudagraph_runtime_mode,

5621 allow_microbatching=allow_microbatching,

5446 - num_active_loras=num_active_loras,

5447 - is_graph_capturing=True,

5622 )

5623 + torch.accelerator.synchronize()

5624 self.maybe_remove_all_loras(self.lora_config)

5625

5626 def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:

vllm/v1/worker/gpu_ubatch_wrapper.py

+11 -2

@@ -112,16 +112,25 @@ class UBatchWrapper:

112 self.cudagraphs: dict[int, CUDAGraphMetaData] = {}

113

114 self.cudagraph_wrapper = None

115 - self.graph_pool = None

115 if runtime_mode is not CUDAGraphMode.NONE:

116 self.cudagraph_wrapper = CUDAGraphWrapper(

117 runnable, vllm_config, runtime_mode=runtime_mode

118 )

120 - self.graph_pool = current_platform.get_global_graph_pool()

119

120 self.sm_control = self._create_sm_control_context(vllm_config)

121 self.device = device

122

123 + @property

124 + def graph_pool(self):

125 + if self.cudagraph_wrapper is not None:

126 + return self.cudagraph_wrapper.graph_pool

127 + return None

128 +

129 + def clear_graphs(self) -> None:

130 + self.cudagraphs.clear()

131 + if self.cudagraph_wrapper is not None:

132 + self.cudagraph_wrapper.clear_graphs()

133 +

134 @staticmethod

135 def _create_sm_control_context(vllm_config: VllmConfig):

136 comm_sms: int = envs.VLLM_DBO_COMM_SMS

vllm/v1/worker/gpu_worker.py

+91 -4

@@ -44,6 +44,7 @@ from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper

44 from vllm.sequence import IntermediateTensors

45 from vllm.tasks import SupportedTask

46 from vllm.tracing import instrument

47 +from vllm.utils.mem_constants import GiB_bytes

48 from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling

49 from vllm.utils.torch_utils import set_random_seed

50 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

@@ -390,8 +391,36 @@ class Worker(WorkerBase):

391 ) as profile_result:

392 self.model_runner.profile_run()

393

394 + profile_torch_peak = current_platform.memory_stats(self.device).get(

395 + "allocated_bytes.all.peak", 0

396 + )

397 +

398 + # Profile CUDA graph memory if graphs will be captured.

399 + cudagraph_memory_estimate = 0

400 + if not self.model_config.enforce_eager:

401 + cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()

402 +

403 + # Use the pre-cudagraph torch peak to avoid double-counting.

404 + profile_result.torch_peak_increase = (

405 + profile_torch_peak - profile_result.before_profile.torch_peak

406 + )

407 + profile_result.non_kv_cache_memory = (

408 + profile_result.non_torch_increase

409 + + profile_result.torch_peak_increase

410 + + profile_result.weights_memory

411 + )

412 +

413 + cudagraph_memory_estimate_applied = (

414 + cudagraph_memory_estimate

415 + if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS

416 + else 0

417 + )

418 +

419 self.non_torch_memory = profile_result.non_torch_increase

394 - self.peak_activation_memory = profile_result.torch_peak_increase

420 + self.peak_activation_memory = (

421 + profile_result.torch_peak_increase + cudagraph_memory_estimate_applied

422 + )

423 + self.cudagraph_memory_estimate = cudagraph_memory_estimate

424

425 free_gpu_memory = profile_result.after_profile.free_memory

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

@@ -406,7 +435,9 @@ class Worker(WorkerBase):

435 "isolate vLLM in its own container."

436 )

437 self.available_kv_cache_memory_bytes = (

409 - self.requested_memory - profile_result.non_kv_cache_memory

438 + self.requested_memory

439 + - profile_result.non_kv_cache_memory

440 + - cudagraph_memory_estimate_applied

441 )

442

443 unrequested_memory = self.init_snapshot.free_memory - self.requested_memory

@@ -428,6 +459,46 @@ class Worker(WorkerBase):

459 scope="local",

460 )

461

462 + if cudagraph_memory_estimate > 0:

463 + total_mem = self.init_snapshot.total_memory

464 + current_util = self.cache_config.gpu_memory_utilization

465 + cg_util_delta = cudagraph_memory_estimate / total_mem

466 + if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:

467 + equiv_util = round(current_util - cg_util_delta, 4)

468 + suggested_util = min(

469 + round(current_util + cg_util_delta, 4),

470 + 1.0,

471 + )

472 + logger.info(

473 + "CUDA graph memory profiling is enabled "

474 + "(VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1). "

475 + "This will become the default in v0.19. "

476 + "The current --gpu-memory-utilization=%.4f is equivalent "

477 + "to --gpu-memory-utilization=%.4f without CUDA graph "

478 + "memory profiling. To maintain the same effective KV "

479 + "cache size as before, increase "

480 + "--gpu-memory-utilization to %.4f.",

481 + current_util,

482 + equiv_util,

483 + suggested_util,

484 + )

485 + else:

486 + suggested_util = min(

487 + round(current_util + cg_util_delta, 4),

488 + 1.0,

489 + )

490 + logger.info(

491 + "In v0.19, CUDA graph memory profiling will be enabled "

492 + "by default (VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1), "

493 + "which more accurately accounts for CUDA graph memory "

494 + "during KV cache allocation. To try it now, set "

495 + "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1 and increase "

496 + "--gpu-memory-utilization from %.4f to %.4f to maintain "

497 + "the same effective KV cache size.",

498 + current_util,

499 + suggested_util,

500 + )

501 +

502 return int(self.available_kv_cache_memory_bytes)

503

504 def get_kv_connector_handshake_metadata(self) -> dict | None:

@@ -487,14 +558,14 @@ class Worker(WorkerBase):

558

559 @instrument(span_name="Warmup (GPU)")

560 def compile_or_warm_up_model(self) -> float:

490 - warmup_sizes = []

561 + warmup_sizes: list[int] = []

562

563 if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:

564 # warm up sizes that are not in cudagraph capture sizes,

565 # but users still want to compile for better performance,

566 # e.g. for the max-num-batched token size in chunked prefill.

567 compile_sizes = self.vllm_config.compilation_config.compile_sizes

497 - warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []

568 + warmup_sizes = compile_sizes.copy() if compile_sizes is not None else [] # type: ignore[assignment]

569 cg_capture_sizes: list[int] = []

570

571 if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

@@ -526,6 +597,22 @@ class Worker(WorkerBase):

597 if not self.model_config.enforce_eager:

598 cuda_graph_memory_bytes = self.model_runner.capture_model()

599

600 + # Compare actual vs estimated CUDA graph memory (if we did profiling)

601 + if (

602 + hasattr(self, "cudagraph_memory_estimate")

603 + and self.cudagraph_memory_estimate > 0

604 + ):

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

606 + diff = abs(cuda_graph_memory_bytes - self.cudagraph_memory_estimate)

607 + logger.info(

608 + "CUDA graph pool memory: %s GiB (actual), %s GiB (estimated), "

609 + "difference: %s GiB (%.1f%%).",

610 + GiB(cuda_graph_memory_bytes),

611 + GiB(self.cudagraph_memory_estimate),

612 + GiB(diff),

613 + 100 * diff / max(cuda_graph_memory_bytes, 1),

614 + )

615 +

616 if self.cache_config.kv_cache_memory_bytes is None and hasattr(

617 self, "peak_activation_memory"

618 ):