diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 162ed03d23b..e46ca1691c2 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1427,7 +1427,10 @@ def get_pcp_group() -> GroupCoordinator: @contextmanager -def graph_capture(device: torch.device): +def graph_capture( + device: torch.device, + graph_capture_context: GraphCaptureContext | None = None, +): """ `graph_capture` is a context manager which should surround the code that is capturing the CUDA graph. Its main purpose is to ensure that some @@ -1440,8 +1443,13 @@ def graph_capture(device: torch.device): the graph capture is running on a separate stream from the default stream, in order to explicitly distinguish the kernels to capture from other kernels possibly launched on background in the default stream. + + A caller may pass an explicit ``graph_capture_context`` to control the + stream used (e.g. to capture on the default stream). """ - context = GraphCaptureContext(torch.cuda.Stream(device=device)) + context = graph_capture_context or GraphCaptureContext( + torch.cuda.Stream(device=device) + ) with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context): yield context diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 377fb670e39..ffc596417fb 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -42,6 +42,7 @@ from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks from vllm.distributed.parallel_state import ( + GraphCaptureContext, get_dcp_group, get_pp_group, get_tp_group, @@ -6607,11 +6608,29 @@ class GPUModelRunner( per_graph_estimate = {} encoder_memory_estimate = 0 + # On ROCm, capture these throwaway profiling graphs on the current stream + # instead of the fresh side stream graph_capture() allocates by default. + # torch's allocator pools free blocks per stream, so a side-stream forward + # strands a persistent aiter scratch buffer in a separate pool, shifting + # the physical placement of the real KV cache allocated afterward and + # slowing bandwidth-bound decode ~20%. The graphs are discarded, so a + # side stream is unnecessary here. + # cap_ctx=None keeps the side-stream path on CUDA, where the current + # stream is the legacy default stream, on which capture cannot begin. + cap_ctx = ( + GraphCaptureContext(torch.cuda.current_stream(self.device)) + if current_platform.is_rocm() + else None + ) + # Cleanup-only guard: CUDA graph capture errors should still propagate # because encoder graph capture is opt-in. try: set_cudagraph_capturing_enabled(True) - with self._freeze_gc(), graph_capture(device=self.device): + with ( + self._freeze_gc(), + graph_capture(device=self.device, graph_capture_context=cap_ctx), + ): torch.accelerator.synchronize() torch.accelerator.empty_cache() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 871f2f31c00..5fb0c387737 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -492,11 +492,14 @@ class Worker(WorkerBase): ) # Profile CUDA graph memory if graphs will be captured. - # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info - # behave differently and can produce incorrect/negative estimates. + # ROCm is included: #44825 moved the profiler to + # torch.accelerator.get_memory_info (reliable on ROCm, as used by + # the AMD-CI mem tests), and graph_pool_handle resolves to the same + # torch.cuda handle the live capture path already uses on ROCm. + # XPU stays excluded (see #39977). cudagraph_memory_estimate = 0 if ( - current_platform.is_cuda() + current_platform.is_cuda_alike() and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE ): @@ -512,8 +515,7 @@ class Worker(WorkerBase): + profile_result.weights_memory ) - # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op. - # On CUDA, respect the opt-in flag as originally designed. + # Respect the opt-in flag as originally designed. cudagraph_memory_estimate_applied = ( cudagraph_memory_estimate if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS