[ROCm][CI] Fix cuda graph mem profile issue (#48764)
Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Charlie Fu committed
Jul 15, 2026 at 23:23 UTC
6a9f24aa8cb856235528d01a829a4ba85fc1c19d
1 file changed
+13
-10
vllm/v1/worker/gpu_model_runner.py
+13
-10
@@ -120,6 +120,7 @@ from vllm.utils.platform_utils import num_compute_units
120
from vllm.utils.torch_utils import (
121
PIN_MEMORY,
122
async_tensor_h2d,
123
+ current_stream,
124
get_dtype_size,
125
is_quantized_kv_cache,
126
kv_cache_dtype_str_to_dtype,
@@ -6608,17 +6609,19 @@ class GPUModelRunner(
6609
per_graph_estimate = {}
6610
encoder_memory_estimate = 0
6611
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.
6612
+ # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated
6613
+ # compute stream instead of the fresh side stream graph_capture()
6614
+ # allocates by default. torch's allocator pools free blocks per stream,
6615
+ # so a side-stream forward strands a persistent aiter scratch buffer in
6616
+ # a separate pool, shifting the physical placement of the real KV cache
6617
+ # allocated afterward and slowing bandwidth-bound decode ~20%. The
6618
+ # graphs are discarded, so a side stream is unnecessary here.
6619
+ # Use current_stream(), not torch.cuda.current_stream(): before vLLM
6620
+ # initializes its dedicated stream, torch returns the per-thread default
6621
+ # stream (cuda_stream=0), which cannot be used for cudagraph capture.
6622
+ # cap_ctx=None keeps the side-stream path on CUDA.
6623
cap_ctx = (
6621
- GraphCaptureContext(torch.cuda.current_stream(self.device))
6624
+ GraphCaptureContext(current_stream())
6625
if current_platform.is_rocm()
6626
else None
6627
)