From 9260454564698e057bf3a376bab9cb42a95d3462 Mon Sep 17 00:00:00 2001 From: Tero Karttunen Date: Tue, 18 Aug 2026 11:08:04 +0000 Subject: [PATCH] feat(vllm): instrument the V2 model runner for graph-capture traces The capture profiler is added to capture_model() in the legacy GPUModelRunner only. Since vLLM 0.22, VllmConfig.use_v2_model_runner selects the rewritten runner in vllm/v1/worker/gpu/ by default for many configurations, and that runner contains no reference to capture_torch_profiler at all. The result is a silent no-op: the patch applies, the server starts and serves, CUDA graphs are captured normally, capture_traces/ stays empty and neither of the patch's own log lines is emitted. Nothing reports an error, so the only way to notice is to go looking for traces that were never written. The workaround until now has been to pin VLLM_USE_V2_MODEL_RUNNER=0, which forces users onto a runner vLLM no longer selects by default. Instrument CudaGraphManager.capture() in vllm/v1/worker/gpu/cudagraph_utils.py, which owns the per-shape capture loop for both PIECEWISE and FULL modes. The manager already holds vllm_config, so the profiler is built where it is used and nothing has to be threaded through model_runner.py. Behaviour matches the legacy path exactly: same rank-0 guard, same trace directory (torch_profiler_dir/capture_traces), same worker_name, and the same capture_{num_tokens}_{mode} annotation that CaptureAnnotation's capture_(\d+)_(.*) pattern parses. Warmup passes stay outside the profiled region, as they are in the legacy runner. Verified on MI350X for 0.22 and 0.25, the two ends of the affected range, with no VLLM_USE_V2_MODEL_RUNNER set anywhere: INFO [gpu_worker.py] Using V2 Model Runner INFO [cudagraph_utils.py] Rank 0: Torch profiler enabled for CUDA graph capture, traces will be saved to: /tmp/traces/capture_traces Both produced 102 capture shards, the same count the legacy runner produces for the same model, carrying annotations such as capture_512_PIECEWISE. All four patches apply to their base images with strict git apply and every touched file byte-compiles. Indenting the capture body under the profiler pushed two existing comments past the 88-column limit, so they are rewrapped; no logic in them changed. --- .../vllm_patches/config_vllm_v0.22.0.patch | 124 +++++++++++++++++ .../vllm_patches/config_vllm_v0.23.0.patch | 130 ++++++++++++++++++ .../vllm_patches/config_vllm_v0.24.0.patch | 130 ++++++++++++++++++ .../vllm_patches/config_vllm_v0.25.0.patch | 130 ++++++++++++++++++ 4 files changed, 514 insertions(+) diff --git a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.22.0.patch b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.22.0.patch index 5a8025a31..e5cbfcb26 100644 --- a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.22.0.patch +++ b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.22.0.patch @@ -31,6 +31,130 @@ index 68fa788..f0e29d0 100644 + ) + return self +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index c7a7ffe..5a3b231 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from collections import defaultdict + from collections.abc import Callable ++from contextlib import AbstractContextManager, nullcontext + from dataclasses import dataclass + from typing import Any, NamedTuple + +@@ -14,6 +15,7 @@ from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode + from vllm.distributed.parallel_state import ( + get_pp_group, ++ get_world_group, + graph_capture, + is_global_first_rank, + ) +@@ -199,6 +201,38 @@ class CudaGraphManager: + create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and + returns a tuple of (forward_fn, captured_attn_state). + """ ++ # Setup torch profiler for graph capture traces (conditional) ++ profiler_config = self.vllm_config.profiler_config ++ local_rank = get_world_group().local_rank ++ profiler: AbstractContextManager[Any] ++ if local_rank == 0 and profiler_config.capture_torch_profiler: ++ trace_dir = profiler_config.torch_profiler_dir + "/capture_traces" ++ profiler = torch.profiler.profile( ++ activities=[ ++ torch.profiler.ProfilerActivity.CPU, ++ torch.profiler.ProfilerActivity.CUDA, ++ ], ++ record_shapes=True, ++ profile_memory=True, ++ with_stack=True, ++ on_trace_ready=torch.profiler.tensorboard_trace_handler( ++ trace_dir, ++ worker_name=f"graph_capture_rank_{local_rank}", ++ use_gzip=True, ++ ), ++ ) ++ logger.info( ++ "Rank %d: Torch profiler enabled for CUDA graph capture, " ++ "traces will be saved to: %s", ++ local_rank, ++ trace_dir, ++ ) ++ else: ++ profiler = nullcontext() ++ logger.info( ++ "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank ++ ) ++ + captured_attn_states: dict[ + BatchExecutionDescriptor, CapturedAttentionState + ] = {} +@@ -224,32 +258,38 @@ class CudaGraphManager: + logger.debug( + "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + ) +- if desc.cg_mode == CUDAGraphMode.PIECEWISE: +- captured_attn_states[desc] = attn_state +- forward_fn(CUDAGraphMode.PIECEWISE) +- else: +- # Capture with fresh attention state. The warmup +- # attention state is discarded because some backends +- # (e.g. FlashMLA) perform lazy initializations that +- # must be captured in the graph. +- forward_fn, attn_state = create_forward_fn(desc) +- captured_attn_states[desc] = attn_state +- assert desc not in self.graphs, ( +- f"Graph already captured for {desc}" +- ) +- graph = torch.cuda.CUDAGraph() +- # Sync offloader's copy stream before capture. +- # Ensure any pre-capture prefetches from offloader are complete. +- get_offloader().sync_prev_onload() +- with torch.cuda.graph(graph, self.pool): +- forward_fn(CUDAGraphMode.NONE) +- # Join offloader's copy stream after forward to avoid +- # unjoined stream error. The last layer's start_prefetch +- # forks copy_stream, but wait_prefetch only happens in +- # the next forward pass. +- get_offloader().join_after_forward() +- self.graphs[desc] = graph +- compilation_counter.num_cudagraph_captured += 1 ++ with ( ++ profiler, ++ torch.profiler.record_function( ++ f"capture_{desc.num_tokens}_{desc.cg_mode.name}" ++ ), ++ ): ++ if desc.cg_mode == CUDAGraphMode.PIECEWISE: ++ captured_attn_states[desc] = attn_state ++ forward_fn(CUDAGraphMode.PIECEWISE) ++ else: ++ # Capture with fresh attention state. The warmup ++ # attention state is discarded because some backends ++ # (e.g. FlashMLA) perform lazy initializations that ++ # must be captured in the graph. ++ forward_fn, attn_state = create_forward_fn(desc) ++ captured_attn_states[desc] = attn_state ++ assert desc not in self.graphs, ( ++ f"Graph already captured for {desc}" ++ ) ++ graph = torch.cuda.CUDAGraph() ++ # Sync offloader's copy stream before capture. Ensure ++ # any pre-capture prefetches from offloader completed. ++ get_offloader().sync_prev_onload() ++ with torch.cuda.graph(graph, self.pool): ++ forward_fn(CUDAGraphMode.NONE) ++ # Join offloader's copy stream after forward to ++ # avoid unjoined stream error. The last layer's ++ # start_prefetch forks copy_stream, but ++ # wait_prefetch only happens in the next forward. ++ get_offloader().join_after_forward() ++ self.graphs[desc] = graph ++ compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True + return captured_attn_states + diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d51bf22..d03ca66 100644 --- a/vllm/v1/worker/gpu_model_runner.py diff --git a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.23.0.patch b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.23.0.patch index 3e87b75c8..873ca8c18 100644 --- a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.23.0.patch +++ b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.23.0.patch @@ -31,6 +31,136 @@ index 68fa788..f0e29d0 100644 + ) + return self +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index dff6047..2ada6c5 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from collections import defaultdict + from collections.abc import Callable ++from contextlib import AbstractContextManager, nullcontext + from dataclasses import dataclass + from typing import Any, NamedTuple, Protocol + +@@ -18,6 +19,7 @@ from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode + from vllm.distributed.parallel_state import ( + get_pp_group, ++ get_world_group, + graph_capture, + is_global_first_rank, + ) +@@ -229,6 +231,38 @@ class CudaGraphManager: + FULL cudagraph capture requires distinct metadatas for warmup and + capture. + """ ++ # Setup torch profiler for graph capture traces (conditional) ++ profiler_config = self.vllm_config.profiler_config ++ local_rank = get_world_group().local_rank ++ profiler: AbstractContextManager[Any] ++ if local_rank == 0 and profiler_config.capture_torch_profiler: ++ trace_dir = profiler_config.torch_profiler_dir + "/capture_traces" ++ profiler = torch.profiler.profile( ++ activities=[ ++ torch.profiler.ProfilerActivity.CPU, ++ torch.profiler.ProfilerActivity.CUDA, ++ ], ++ record_shapes=True, ++ profile_memory=True, ++ with_stack=True, ++ on_trace_ready=torch.profiler.tensorboard_trace_handler( ++ trace_dir, ++ worker_name=f"graph_capture_rank_{local_rank}", ++ use_gzip=True, ++ ), ++ ) ++ logger.info( ++ "Rank %d: Torch profiler enabled for CUDA graph capture, " ++ "traces will be saved to: %s", ++ local_rank, ++ trace_dir, ++ ) ++ else: ++ profiler = nullcontext() ++ logger.info( ++ "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank ++ ) ++ + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {} + with graph_capture(device=self.device): + # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger +@@ -252,35 +286,41 @@ class CudaGraphManager: + logger.debug( + "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + ) +- if desc.cg_mode == CUDAGraphMode.PIECEWISE: +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, warmup_attn_state +- ) +- forward_fn(CUDAGraphMode.PIECEWISE) +- else: +- # Capture with fresh attention state. +- forward_fn, capture_attn_state = create_forward_fn( +- desc, warmup=False +- ) +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, capture_attn_state +- ) +- assert desc not in self.graphs, ( +- f"Graph already captured for {desc}" +- ) +- graph = torch.cuda.CUDAGraph() +- # Sync offloader's copy stream before capture. +- # Ensure any pre-capture prefetches from offloader are complete. +- get_offloader().sync_prev_onload() +- with torch.cuda.graph(graph, self.pool): +- forward_fn(CUDAGraphMode.NONE) +- # Join offloader's copy stream after forward to avoid +- # unjoined stream error. The last layer's start_prefetch +- # forks copy_stream, but wait_prefetch only happens in +- # the next forward pass. +- get_offloader().join_after_forward() +- self.graphs[desc] = graph +- compilation_counter.num_cudagraph_captured += 1 ++ with ( ++ profiler, ++ torch.profiler.record_function( ++ f"capture_{desc.num_tokens}_{desc.cg_mode.name}" ++ ), ++ ): ++ if desc.cg_mode == CUDAGraphMode.PIECEWISE: ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, warmup_attn_state ++ ) ++ forward_fn(CUDAGraphMode.PIECEWISE) ++ else: ++ # Capture with fresh attention state. ++ forward_fn, capture_attn_state = create_forward_fn( ++ desc, warmup=False ++ ) ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, capture_attn_state ++ ) ++ assert desc not in self.graphs, ( ++ f"Graph already captured for {desc}" ++ ) ++ graph = torch.cuda.CUDAGraph() ++ # Sync offloader's copy stream before capture. Ensure ++ # any pre-capture prefetches from offloader completed. ++ get_offloader().sync_prev_onload() ++ with torch.cuda.graph(graph, self.pool): ++ forward_fn(CUDAGraphMode.NONE) ++ # Join offloader's copy stream after forward to ++ # avoid unjoined stream error. The last layer's ++ # start_prefetch forks copy_stream, but ++ # wait_prefetch only happens in the next forward. ++ get_offloader().join_after_forward() ++ self.graphs[desc] = graph ++ compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True + return attn_states + diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 801a857..91266be 100644 --- a/vllm/v1/worker/gpu_model_runner.py diff --git a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.24.0.patch b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.24.0.patch index 38ea4e2ff..349b1728a 100644 --- a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.24.0.patch +++ b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.24.0.patch @@ -31,6 +31,136 @@ index 68fa788..f0e29d0 100644 + ) + return self +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index dad1777..6051d7e 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from collections import defaultdict + from collections.abc import Callable ++from contextlib import AbstractContextManager, nullcontext + from dataclasses import dataclass + from itertools import product + from typing import Any, NamedTuple, Protocol +@@ -19,6 +20,7 @@ from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode + from vllm.distributed.parallel_state import ( + get_pp_group, ++ get_world_group, + graph_capture, + is_global_first_rank, + ) +@@ -274,6 +276,38 @@ class CudaGraphManager: + FULL cudagraph capture requires distinct metadatas for warmup and + capture. + """ ++ # Setup torch profiler for graph capture traces (conditional) ++ profiler_config = self.vllm_config.profiler_config ++ local_rank = get_world_group().local_rank ++ profiler: AbstractContextManager[Any] ++ if local_rank == 0 and profiler_config.capture_torch_profiler: ++ trace_dir = profiler_config.torch_profiler_dir + "/capture_traces" ++ profiler = torch.profiler.profile( ++ activities=[ ++ torch.profiler.ProfilerActivity.CPU, ++ torch.profiler.ProfilerActivity.CUDA, ++ ], ++ record_shapes=True, ++ profile_memory=True, ++ with_stack=True, ++ on_trace_ready=torch.profiler.tensorboard_trace_handler( ++ trace_dir, ++ worker_name=f"graph_capture_rank_{local_rank}", ++ use_gzip=True, ++ ), ++ ) ++ logger.info( ++ "Rank %d: Torch profiler enabled for CUDA graph capture, " ++ "traces will be saved to: %s", ++ local_rank, ++ trace_dir, ++ ) ++ else: ++ profiler = nullcontext() ++ logger.info( ++ "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank ++ ) ++ + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {} + with graph_capture(device=self.device): + # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger +@@ -297,35 +331,41 @@ class CudaGraphManager: + logger.debug( + "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + ) +- if desc.cg_mode == CUDAGraphMode.PIECEWISE: +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, warmup_attn_state +- ) +- forward_fn(CUDAGraphMode.PIECEWISE) +- else: +- # Capture with fresh attention state. +- forward_fn, capture_attn_state = create_forward_fn( +- desc, warmup=False +- ) +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, capture_attn_state +- ) +- assert desc not in self.graphs, ( +- f"Graph already captured for {desc}" +- ) +- graph = torch.cuda.CUDAGraph() +- # Sync offloader's copy stream before capture. +- # Ensure any pre-capture prefetches from offloader are complete. +- get_offloader().sync_prev_onload() +- with torch.cuda.graph(graph, self.pool): +- forward_fn(CUDAGraphMode.NONE) +- # Join offloader's copy stream after forward to avoid +- # unjoined stream error. The last layer's start_prefetch +- # forks copy_stream, but wait_prefetch only happens in +- # the next forward pass. +- get_offloader().join_after_forward() +- self.graphs[desc] = graph +- compilation_counter.num_cudagraph_captured += 1 ++ with ( ++ profiler, ++ torch.profiler.record_function( ++ f"capture_{desc.num_tokens}_{desc.cg_mode.name}" ++ ), ++ ): ++ if desc.cg_mode == CUDAGraphMode.PIECEWISE: ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, warmup_attn_state ++ ) ++ forward_fn(CUDAGraphMode.PIECEWISE) ++ else: ++ # Capture with fresh attention state. ++ forward_fn, capture_attn_state = create_forward_fn( ++ desc, warmup=False ++ ) ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, capture_attn_state ++ ) ++ assert desc not in self.graphs, ( ++ f"Graph already captured for {desc}" ++ ) ++ graph = torch.cuda.CUDAGraph() ++ # Sync offloader's copy stream before capture. Ensure ++ # any pre-capture prefetches from offloader completed. ++ get_offloader().sync_prev_onload() ++ with torch.cuda.graph(graph, self.pool): ++ forward_fn(CUDAGraphMode.NONE) ++ # Join offloader's copy stream after forward to ++ # avoid unjoined stream error. The last layer's ++ # start_prefetch forks copy_stream, but ++ # wait_prefetch only happens in the next forward. ++ get_offloader().join_after_forward() ++ self.graphs[desc] = graph ++ compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True + return attn_states + diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a8..cfffdab 100644 --- a/vllm/v1/worker/gpu_model_runner.py diff --git a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.25.0.patch b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.25.0.patch index 33c2bc134..da89acf5b 100644 --- a/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.25.0.patch +++ b/examples/custom_workflows/inference_analysis/vllm_patches/config_vllm_v0.25.0.patch @@ -31,6 +31,136 @@ index 68fa788..f0e29d0 100644 + ) + return self +diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py +index 9b93d17..9fae073 100644 +--- a/vllm/v1/worker/gpu/cudagraph_utils.py ++++ b/vllm/v1/worker/gpu/cudagraph_utils.py +@@ -2,6 +2,7 @@ + # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + from collections import defaultdict + from collections.abc import Callable ++from contextlib import AbstractContextManager, nullcontext + from dataclasses import dataclass + from itertools import product + from typing import Any, NamedTuple, Protocol +@@ -19,6 +20,7 @@ from vllm.config import VllmConfig + from vllm.config.compilation import CUDAGraphMode + from vllm.distributed.parallel_state import ( + get_pp_group, ++ get_world_group, + graph_capture, + is_global_first_rank, + ) +@@ -313,6 +315,38 @@ class CudaGraphManager: + FULL cudagraph capture requires distinct metadatas for warmup and + capture. + """ ++ # Setup torch profiler for graph capture traces (conditional) ++ profiler_config = self.vllm_config.profiler_config ++ local_rank = get_world_group().local_rank ++ profiler: AbstractContextManager[Any] ++ if local_rank == 0 and profiler_config.capture_torch_profiler: ++ trace_dir = profiler_config.torch_profiler_dir + "/capture_traces" ++ profiler = torch.profiler.profile( ++ activities=[ ++ torch.profiler.ProfilerActivity.CPU, ++ torch.profiler.ProfilerActivity.CUDA, ++ ], ++ record_shapes=True, ++ profile_memory=True, ++ with_stack=True, ++ on_trace_ready=torch.profiler.tensorboard_trace_handler( ++ trace_dir, ++ worker_name=f"graph_capture_rank_{local_rank}", ++ use_gzip=True, ++ ), ++ ) ++ logger.info( ++ "Rank %d: Torch profiler enabled for CUDA graph capture, " ++ "traces will be saved to: %s", ++ local_rank, ++ trace_dir, ++ ) ++ else: ++ profiler = nullcontext() ++ logger.info( ++ "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank ++ ) ++ + attn_states: dict[BatchExecutionDescriptor, AttentionStatePair] = {} + with graph_capture(device=self.device): + # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger +@@ -336,35 +370,41 @@ class CudaGraphManager: + logger.debug( + "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc + ) +- if desc.cg_mode == CUDAGraphMode.PIECEWISE: +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, warmup_attn_state +- ) +- forward_fn(CUDAGraphMode.PIECEWISE) +- else: +- # Capture with fresh attention state. +- forward_fn, capture_attn_state = create_forward_fn( +- desc, warmup=False +- ) +- attn_states[desc] = AttentionStatePair( +- warmup_attn_state, capture_attn_state +- ) +- assert desc not in self.graphs, ( +- f"Graph already captured for {desc}" +- ) +- graph = torch.cuda.CUDAGraph() +- # Sync offloader's copy stream before capture. +- # Ensure any pre-capture prefetches from offloader are complete. +- get_offloader().sync_prev_onload() +- with torch.cuda.graph(graph, self.pool): +- forward_fn(CUDAGraphMode.NONE) +- # Join offloader's copy stream after forward to avoid +- # unjoined stream error. The last layer's start_prefetch +- # forks copy_stream, but wait_prefetch only happens in +- # the next forward pass. +- get_offloader().join_after_forward() +- self.graphs[desc] = graph +- compilation_counter.num_cudagraph_captured += 1 ++ with ( ++ profiler, ++ torch.profiler.record_function( ++ f"capture_{desc.num_tokens}_{desc.cg_mode.name}" ++ ), ++ ): ++ if desc.cg_mode == CUDAGraphMode.PIECEWISE: ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, warmup_attn_state ++ ) ++ forward_fn(CUDAGraphMode.PIECEWISE) ++ else: ++ # Capture with fresh attention state. ++ forward_fn, capture_attn_state = create_forward_fn( ++ desc, warmup=False ++ ) ++ attn_states[desc] = AttentionStatePair( ++ warmup_attn_state, capture_attn_state ++ ) ++ assert desc not in self.graphs, ( ++ f"Graph already captured for {desc}" ++ ) ++ graph = torch.cuda.CUDAGraph() ++ # Sync offloader's copy stream before capture. Ensure ++ # any pre-capture prefetches from offloader completed. ++ get_offloader().sync_prev_onload() ++ with torch.cuda.graph(graph, self.pool): ++ forward_fn(CUDAGraphMode.NONE) ++ # Join offloader's copy stream after forward to ++ # avoid unjoined stream error. The last layer's ++ # start_prefetch forks copy_stream, but ++ # wait_prefetch only happens in the next forward. ++ get_offloader().join_after_forward() ++ self.graphs[desc] = graph ++ compilation_counter.num_cudagraph_captured += 1 + self._graphs_captured = True + return attn_states + diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e9b23f1..2805319 100644 --- a/vllm/v1/worker/gpu_model_runner.py