[ROCm] Re-enable cudagraph memory profiling, captured on the current stream (#48526)
Signed-off-by: pei.zhang <pei.zhang@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
peizhang56 committed
Jul 15, 2026 at 08:03 UTC
05eed72aec6c05e6d500c7276b47f7652bb37af6
3 files changed
+37
-8
vllm/distributed/parallel_state.py
+10
-2
@@ -1427,7 +1427,10 @@ def get_pcp_group() -> GroupCoordinator:
1427
1428
1429
@contextmanager
1430
-def graph_capture(device: torch.device):
1430
+def graph_capture(
1431
+ device: torch.device,
1432
+ graph_capture_context: GraphCaptureContext | None = None,
1433
+):
1434
"""
1435
`graph_capture` is a context manager which should surround the code that
1436
is capturing the CUDA graph. Its main purpose is to ensure that some
@@ -1440,8 +1443,13 @@ def graph_capture(device: torch.device):
1443
the graph capture is running on a separate stream from the default stream,
1444
in order to explicitly distinguish the kernels to capture
1445
from other kernels possibly launched on background in the default stream.
1446
+
1447
+ A caller may pass an explicit ``graph_capture_context`` to control the
1448
+ stream used (e.g. to capture on the default stream).
1449
"""
1444
- context = GraphCaptureContext(torch.cuda.Stream(device=device))
1450
+ context = graph_capture_context or GraphCaptureContext(
1451
+ torch.cuda.Stream(device=device)
1452
+ )
1453
with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context):
1454
yield context
1455
vllm/v1/worker/gpu_model_runner.py
+20
-1
@@ -42,6 +42,7 @@ from vllm.distributed.eplb.eplb_state import EplbState
42
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
43
from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
44
from vllm.distributed.parallel_state import (
45
+ GraphCaptureContext,
46
get_dcp_group,
47
get_pp_group,
48
get_tp_group,
@@ -6607,11 +6608,29 @@ class GPUModelRunner(
6608
per_graph_estimate = {}
6609
encoder_memory_estimate = 0
6610
6611
+ # On ROCm, capture these throwaway profiling graphs on the current stream
6612
+ # instead of the fresh side stream graph_capture() allocates by default.
6613
+ # torch's allocator pools free blocks per stream, so a side-stream forward
6614
+ # strands a persistent aiter scratch buffer in a separate pool, shifting
6615
+ # the physical placement of the real KV cache allocated afterward and
6616
+ # slowing bandwidth-bound decode ~20%. The graphs are discarded, so a
6617
+ # side stream is unnecessary here.
6618
+ # cap_ctx=None keeps the side-stream path on CUDA, where the current
6619
+ # stream is the legacy default stream, on which capture cannot begin.
6620
+ cap_ctx = (
6621
+ GraphCaptureContext(torch.cuda.current_stream(self.device))
6622
+ if current_platform.is_rocm()
6623
+ else None
6624
+ )
6625
+
6626
# Cleanup-only guard: CUDA graph capture errors should still propagate
6627
# because encoder graph capture is opt-in.
6628
try:
6629
set_cudagraph_capturing_enabled(True)
6614
- with self._freeze_gc(), graph_capture(device=self.device):
6630
+ with (
6631
+ self._freeze_gc(),
6632
+ graph_capture(device=self.device, graph_capture_context=cap_ctx),
6633
+ ):
6634
torch.accelerator.synchronize()
6635
torch.accelerator.empty_cache()
6636
vllm/v1/worker/gpu_worker.py
+7
-5
@@ -492,11 +492,14 @@ class Worker(WorkerBase):
492
)
493
494
# Profile CUDA graph memory if graphs will be captured.
495
- # Skip on ROCm/HIP/XPU as graph pool handles and get_memory_info
496
- # behave differently and can produce incorrect/negative estimates.
495
+ # ROCm is included: #44825 moved the profiler to
496
+ # torch.accelerator.get_memory_info (reliable on ROCm, as used by
497
+ # the AMD-CI mem tests), and graph_pool_handle resolves to the same
498
+ # torch.cuda handle the live capture path already uses on ROCm.
499
+ # XPU stays excluded (see #39977).
500
cudagraph_memory_estimate = 0
501
if (
499
- current_platform.is_cuda()
502
+ current_platform.is_cuda_alike()
503
and self.vllm_config.compilation_config.cudagraph_mode
504
!= CUDAGraphMode.NONE
505
):
@@ -512,8 +515,7 @@ class Worker(WorkerBase):
515
+ profile_result.weights_memory
516
)
517
515
- # On ROCm, cudagraph_memory_estimate is always 0 so this is a no-op.
516
- # On CUDA, respect the opt-in flag as originally designed.
518
+ # Respect the opt-in flag as originally designed.
519
cudagraph_memory_estimate_applied = (
520
cudagraph_memory_estimate
521
if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS