[ROCm] Re-enable cudagraph memory profiling, captured on the current stream (#48526)

Signed-off-by: pei.zhang <[email protected]>
Co-authored-by: Claude <[email protected]>
Co-authored-by: Andreas Karatzas <[email protected]>

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
 

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()
 

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