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

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 2, 06:00 UTC

[MM][CG] Profile encoder CUDA graph pool memory (#41714)

Signed-off-by: JooHo Lee <jooho414@gmail.com>

JooHo Lee committed Jun 2, 2026 at 13:27 UTC a045c7425f39997f4a319d9e671c9d00590e321a

4 files changed +173 -92

tests/v1/cudagraph/test_encoder_cudagraph.py

+18 -2

@@ -109,6 +109,7 @@ def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:

109 mgr.max_batch_size = 16

110 mgr.use_dp = False

111 mgr.budget_graphs = {}

112 + mgr.graph_pool = None

113 mgr.graph_hits = 0

114 mgr.graph_misses = 0

115 mgr.log_stats_interval = 100

@@ -180,6 +181,10 @@ class TestFindBudgetGraph:

181 # Budget selection still works correctly after sorting

182 assert mgr._find_smallest_fitting_budget_given_tokens(3000) == 4096

183

184 + def test_num_graphs_to_capture_tracks_budgets(self):

185 + mgr = _make_manager_with_budgets([8192, 2048, 4096])

186 + assert mgr.get_num_graphs_to_capture() == 3

187 +

188

189 # ---------------------------------------------------------------------------

190 # get_cumulative_stats

@@ -409,6 +414,7 @@ def _make_manager_for_gpu(

414 )

415 mgr.use_dp = False

416 mgr.budget_graphs = {}

417 + mgr.graph_pool = None

418 mgr.graph_hits = 0

419 mgr.graph_misses = 0

420 mgr.log_stats_interval = 100

@@ -467,7 +473,8 @@ class TestEncoderCudaGraphCaptureReplay:

473 self.mgr = _make_manager_for_gpu(

474 self.model, _BUDGETS, _MAX_BATCH, self.device, self.dtype

475 )

470 - self.mgr.capture()

476 + self.graph_pool = current_platform.graph_pool_handle()

477 + self.mgr.capture(graph_pool=self.graph_pool)

478

479 # --- capture ---

480

@@ -475,6 +482,14 @@ class TestEncoderCudaGraphCaptureReplay:

482 assert len(self.mgr.budget_graphs) == len(_BUDGETS)

483 assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)

484

485 + def test_capture_uses_supplied_graph_pool(self):

486 + assert self.mgr.graph_pool is self.graph_pool

487 +

488 + def test_clear_releases_graphs_and_pool(self):

489 + self.mgr.clear()

490 + assert self.mgr.budget_graphs == {}

491 + assert self.mgr.graph_pool is None

492 +

493 # --- output shape ---

494

495 def test_execute_returns_one_tensor_per_image(self):

@@ -742,7 +757,8 @@ class TestEncoderCudaGraphVideoReplay:

757 self.dtype,

758 max_frames_per_batch=_VIDEO_MAX_FRAMES,

759 )

745 - self.mgr.capture()

760 + self.graph_pool = current_platform.graph_pool_handle()

761 + self.mgr.capture(graph_pool=self.graph_pool)

762

763 # --- capture ---

764

vllm/utils/import_utils.py

+1 -1

@@ -404,7 +404,7 @@ def _has_module(module_name: str) -> bool:

404 if importlib.util.find_spec(module_name) is None:

405 return False

406 importlib.import_module(module_name)

407 - except ImportError:

407 + except Exception:

408 logger.warning(

409 "Module %s was found but failed to import", module_name, exc_info=True

410 )

vllm/v1/worker/encoder_cudagraph.py

+14 -3

@@ -153,6 +153,7 @@ class EncoderCudaGraphManager:

153 )

154

155 self.budget_graphs: dict[int, BudgetGraphMetadata] = {}

156 + self.graph_pool: Any | None = None

157 self.graph_hits = 0

158 self.graph_misses = 0

159 self.log_stats_interval = 100

@@ -183,9 +184,16 @@ class EncoderCudaGraphManager:

184 """Check if a modality is supported by this manager."""

185 return modality in self.config.modalities

186

186 - def capture(self):

187 + def clear(self) -> None:

188 + """Release captured encoder CUDA graphs and the manager-local pool."""

189 + self.budget_graphs.clear()

190 + self.graph_pool = None

191 +

192 + def capture(self, graph_pool: Any):

193 """Capture CUDA graphs for all token budgets."""

188 - for token_budget in self.token_budgets:

194 + self.graph_pool = graph_pool

195 +

196 + for token_budget in sorted(self.token_budgets, reverse=True):

197 self._capture_budget_graph(token_budget)

198

199 logger.info(

@@ -193,6 +201,9 @@ class EncoderCudaGraphManager:

201 len(self.budget_graphs),

202 )

203

204 + def get_num_graphs_to_capture(self) -> int:

205 + return len(self.token_budgets)

206 +

207 def _capture_budget_graph(self, token_budget: int):

208 """Capture CUDA graph for a single token budget."""

209 logger.debug(

@@ -218,7 +229,7 @@ class EncoderCudaGraphManager:

229 output_buffer = torch.empty_like(output)

230

231 graph = torch.cuda.CUDAGraph()

221 - with torch.inference_mode(), torch.cuda.graph(graph):

232 + with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool):

233 output = self.model.encoder_cudagraph_forward({**values})

234 output_buffer.copy_(output)

235

vllm/v1/worker/gpu_model_runner.py

+140 -86

@@ -6304,6 +6304,42 @@ class GPUModelRunner(

6304

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

6306

6307 + @torch.inference_mode()

6308 + def _create_encoder_cudagraph_manager(self) -> "EncoderCudaGraphManager | None":

6309 + if not (

6310 + self.compilation_config.cudagraph_mm_encoder and self.supports_mm_inputs

6311 + ):

6312 + return None

6313 +

6314 + # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper, because

6315 + # @runtime_checkable Protocol isinstance() checks do not work through

6316 + # __getattr__ forwarding.

6317 + from vllm.model_executor.models.interfaces import (

6318 + SupportsEncoderCudaGraph,

6319 + supports_encoder_cudagraph,

6320 + )

6321 + from vllm.v1.worker.encoder_cudagraph import (

6322 + EncoderCudaGraphManager,

6323 + )

6324 +

6325 + raw_model = self.get_model()

6326 + if not supports_encoder_cudagraph(raw_model):

6327 + return None

6328 +

6329 + return EncoderCudaGraphManager(

6330 + vllm_config=self.vllm_config,

6331 + device=self.device,

6332 + dtype=self.dtype,

6333 + model=cast(SupportsEncoderCudaGraph, raw_model),

6334 + )

6335 +

6336 + @torch.inference_mode()

6337 + def _maybe_init_encoder_cudagraph_manager(self) -> None:

6338 + if self.encoder_cudagraph_manager is None:

6339 + self.encoder_cudagraph_manager = self._create_encoder_cudagraph_manager()

6340 + if self.encoder_cudagraph_manager is not None:

6341 + logger.info("Initialized EncoderCudaGraphManager for vision encoder")

6342 +

6343 @torch.inference_mode()

6344 def profile_cudagraph_memory(self) -> int:

6345 with set_current_vllm_config(self.vllm_config):

@@ -6312,24 +6348,40 @@ class GPUModelRunner(

6348 saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured

6349

6350 capture_descs = self.cudagraph_dispatcher.get_capture_descs()

6351 + # Use a temporary manager for memory profiling. The persistent manager

6352 + # is initialized later so it does not keep profiling-only graph state.

6353 + encoder_cudagraph_manager = self._create_encoder_cudagraph_manager()

6354

6316 - total_graphs = sum(len(descs) for _, descs in capture_descs)

6355 + decoder_graphs = sum(len(descs) for _, descs in capture_descs)

6356 + encoder_graphs = (

6357 + encoder_cudagraph_manager.get_num_graphs_to_capture()

6358 + if encoder_cudagraph_manager is not None

6359 + else 0

6360 + )

6361 + total_graphs = decoder_graphs + encoder_graphs

6362 if total_graphs == 0:

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

6364 self._cleanup_profiling_kv_cache()

6365 return 0

6366

6322 - logger.info(

6323 - "Profiling CUDA graph memory: %s",

6324 - ", ".join(

6367 + graph_groups = [

6368 + *(

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

6370 for mode, descs in capture_descs

6371 if descs

6372 ),

6329 - )

6373 + ]

6374 + if encoder_graphs > 0:

6375 + graph_groups.append(

6376 + f"ENCODER={encoder_graphs} "

6377 + f"(largest={encoder_cudagraph_manager.token_budgets[-1]})"

6378 + )

6379 +

6380 + logger.info("Profiling CUDA graph memory: %s", ", ".join(graph_groups))

6381

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

6383 profiling_pool = current_platform.graph_pool_handle()

6384 + encoder_profiling_pool = current_platform.graph_pool_handle()

6385 original_pools: dict[int, Any] = {}

6386 all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

6387 BreakableCUDAGraphWrapper._all_instances

@@ -6338,73 +6390,98 @@ class GPUModelRunner(

6390 original_pools[id(instance)] = instance.graph_pool

6391 instance.graph_pool = profiling_pool

6392

6341 - set_cudagraph_capturing_enabled(True)

6342 - with self._freeze_gc(), graph_capture(device=self.device):

6343 - shared_memory_estimate = {}

6344 - per_graph_estimate = {}

6345 - torch.accelerator.synchronize()

6346 - torch.accelerator.empty_cache()

6393 + shared_memory_estimate = {}

6394 + per_graph_estimate = {}

6395 + encoder_memory_estimate = 0

6396

6348 - for mode, descs in capture_descs:

6349 - profile_descs = descs[:2]

6350 - mem_samples: list[int] = []

6397 + # Cleanup-only guard: CUDA graph capture errors should still propagate

6398 + # because encoder graph capture is opt-in.

6399 + try:

6400 + set_cudagraph_capturing_enabled(True)

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

6402 + torch.accelerator.synchronize()

6403 + torch.accelerator.empty_cache()

6404 +

6405 + for mode, descs in capture_descs:

6406 + profile_descs = descs[:2]

6407 + mem_samples: list[int] = []

6408 +

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

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

6411 + self._warmup_and_capture(

6412 + desc,

6413 + cudagraph_runtime_mode=mode,

6414 + profile_seq_lens=(

6415 + min(

6416 + self.max_model_len,

6417 + self.max_num_tokens // desc.num_tokens,

6418 + )

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

6420 + else None

6421 + ),

6422 + )

6423 + torch.accelerator.synchronize()

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

6425 + mem_samples.append(mem_before - free_after)

6426 +

6427 + first_capture = mem_samples[0]

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

6429 + per_graph = max(

6430 + mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20

6431 + )

6432

6352 - for i, desc in enumerate(profile_descs):

6353 - mem_before = torch.cuda.mem_get_info()[0]

6354 - self._warmup_and_capture(

6355 - desc,

6356 - cudagraph_runtime_mode=mode,

6357 - profile_seq_lens=(

6358 - min(

6359 - self.max_model_len,

6360 - self.max_num_tokens // desc.num_tokens,

6361 - )

6362 - if mode == CUDAGraphMode.FULL and i == 0

6363 - else None

6364 - ),

6433 + shared_memory_estimate[mode] = first_capture

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

6435 +

6436 + logger.debug(

6437 + "Estimated %s CUDA graph memory: "

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

6439 + mode.name,

6440 + first_capture / (1 << 20),

6441 + len(descs),

6442 + per_graph / (1 << 20),

6443 )

6444 +

6445 + if encoder_cudagraph_manager is not None:

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

6447 + encoder_cudagraph_manager.capture(graph_pool=encoder_profiling_pool)

6448 torch.accelerator.synchronize()

6449 free_after = torch.cuda.mem_get_info()[0]

6368 - mem_samples.append(mem_before - free_after)

6369 -

6370 - first_capture = mem_samples[0]

6371 - # Use at least 1 MiB per graph for driver overhead

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

6373 -

6374 - shared_memory_estimate[mode] = first_capture

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

6376 -

6377 - logger.debug(

6378 - "Estimated %s CUDA graph memory: "

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

6380 - mode.name,

6381 - first_capture / (1 << 20),

6382 - len(descs),

6383 - per_graph / (1 << 20),

6384 - )

6450 + encoder_memory_estimate = max(mem_before - free_after, 0)

6451

6386 - set_cudagraph_capturing_enabled(False)

6387 - CUDAGraphWrapper.clear_all_graphs()

6388 - BreakableCUDAGraphWrapper.clear_all_graphs()

6389 - all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

6390 - BreakableCUDAGraphWrapper._all_instances

6391 - )

6392 - for instance in all_wrappers:

6393 - if id(instance) in original_pools:

6394 - instance.graph_pool = original_pools[id(instance)]

6395 - for key_set in self.cudagraph_dispatcher.cudagraph_keys.values():

6396 - key_set.clear()

6397 - self.cudagraph_dispatcher.keys_initialized = False

6398 - self.maybe_remove_all_loras(self.lora_config)

6399 - self._cleanup_profiling_kv_cache()

6400 - compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

6452 + logger.debug(

6453 + "Estimated encoder CUDA graph memory: %.2f MiB for %d graphs",

6454 + encoder_memory_estimate / (1 << 20),

6455 + encoder_graphs,

6456 + )

6457 + finally:

6458 + set_cudagraph_capturing_enabled(False)

6459 + CUDAGraphWrapper.clear_all_graphs()

6460 + BreakableCUDAGraphWrapper.clear_all_graphs()

6461 + if encoder_cudagraph_manager is not None:

6462 + encoder_cudagraph_manager.clear()

6463 + all_wrappers = list(CUDAGraphWrapper._all_instances) + list(

6464 + BreakableCUDAGraphWrapper._all_instances

6465 + )

6466 + for instance in all_wrappers:

6467 + if id(instance) in original_pools:

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

6469 + for key_set in self.cudagraph_dispatcher.cudagraph_keys.values():

6470 + key_set.clear()

6471 + self.cudagraph_dispatcher.keys_initialized = False

6472 + self.maybe_remove_all_loras(self.lora_config)

6473 + self._cleanup_profiling_kv_cache()

6474 + compilation_counter.num_cudagraph_captured = saved_num_cudagraph_captured

6475

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

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

6478 # Take the max to avoid double-counting the overlap.

6405 - total_estimate = max(shared_memory_estimate.values()) + sum(

6479 + decoder_estimate = max(shared_memory_estimate.values(), default=0) + sum(

6480 per_graph_estimate.values()

6481 )

6482 + # Encoder graphs use a manager-local pool at runtime, separate from the

6483 + # decoder pool, so add their estimate instead of overlaying it.

6484 + total_estimate = decoder_estimate + encoder_memory_estimate

6485 logger.info(

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

6487 total_estimate / (1 << 30),

@@ -6422,31 +6499,7 @@ class GPUModelRunner(

6499 return 0

6500

6501 # Initialize encoder CUDA graph manager if enabled.

6425 - # Use get_model() to unwrap CUDAGraphWrapper/UBatchWrapper,

6426 - # because @runtime_checkable Protocol isinstance() checks do not

6427 - # work through __getattr__ forwarding.

6428 - if (

6429 - self.compilation_config.cudagraph_mm_encoder

6430 - and self.supports_mm_inputs

6431 - and self.encoder_cudagraph_manager is None

6432 - ):

6433 - from vllm.model_executor.models.interfaces import (

6434 - SupportsEncoderCudaGraph,

6435 - supports_encoder_cudagraph,

6436 - )

6437 - from vllm.v1.worker.encoder_cudagraph import (

6438 - EncoderCudaGraphManager,

6439 - )

6440 -

6441 - raw_model = self.get_model()

6442 - if supports_encoder_cudagraph(raw_model):

6443 - self.encoder_cudagraph_manager = EncoderCudaGraphManager(

6444 - vllm_config=self.vllm_config,

6445 - device=self.device,

6446 - dtype=self.dtype,

6447 - model=cast(SupportsEncoderCudaGraph, raw_model),

6448 - )

6449 - logger.info("Initialized EncoderCudaGraphManager for vision encoder")

6502 + self._maybe_init_encoder_cudagraph_manager()

6503

6504 compilation_counter.num_gpu_runner_capture_triggers += 1

6505

@@ -6473,7 +6526,8 @@ class GPUModelRunner(

6526

6527 # Capture encoder CUDA graphs if enabled

6528 if self.encoder_cudagraph_manager is not None:

6476 - self.encoder_cudagraph_manager.capture()

6529 + encoder_graph_pool = current_platform.graph_pool_handle()

6530 + self.encoder_cudagraph_manager.capture(graph_pool=encoder_graph_pool)

6531

6532 torch.accelerator.synchronize()

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