[Bugfix][ROCm] Capture CUDA graphs on the current stream (#53818)
Signed-off-by: fanxingran <xingran.fan@amd.com>
fanxingran committed
Aug 26, 2026 at 23:16 UTC
080a66a69c6fd1fe464756f88ab958baad66ce69
3 files changed
+11
-3
vllm/v1/spec_decode/gemma4.py
+2
-1
@@ -17,6 +17,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config, replace
17
from vllm.distributed.parallel_state import get_pp_group
18
from vllm.logger import init_logger
19
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20
+from vllm.utils.torch_utils import current_stream
21
from vllm.v1.attention.backend import CommonAttentionMetadata
22
from vllm.v1.kv_cache_interface import (
23
KVCacheConfig,
@@ -130,7 +131,7 @@ class Gemma4Proposer(SpecDecodeBaseProposer):
131
torch.accelerator.synchronize()
132
133
g = torch.cuda.CUDAGraph()
133
- with torch.cuda.graph(g):
134
+ with torch.cuda.graph(g, stream=current_stream()):
135
static_output = masked_emb.get_top_tokens(
136
static_input,
137
lm_head_weight,
vllm/v1/worker/encoder_cudagraph.py
+5
-1
@@ -20,6 +20,7 @@ from vllm.model_executor.models.interfaces import (
20
from vllm.model_executor.models.utils import scatter_output_slices
21
from vllm.model_executor.models.vision import get_load_balance_assignment
22
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
23
+from vllm.utils.torch_utils import current_stream
24
from vllm.v1.worker.encoder_cudagraph_defs import (
25
EncoderCudaGraphConfig,
26
EncoderItemSpec,
@@ -270,7 +271,10 @@ class EncoderCudaGraphManager:
271
output_buffer = torch.empty_like(output)
272
273
graph = torch.cuda.CUDAGraph()
273
- with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool):
274
+ with (
275
+ torch.inference_mode(),
276
+ torch.cuda.graph(graph, pool=self.graph_pool, stream=current_stream()),
277
+ ):
278
output = self.model.encoder_cudagraph_forward({**values}, path=path)
279
output_buffer.copy_(output)
280
vllm/v1/worker/gpu/cudagraph_utils.py
+4
-1
@@ -31,6 +31,7 @@ from vllm.model_executor.offloader.base import get_offloader
31
from vllm.platforms import current_platform
32
from vllm.sequence import IntermediateTensors
33
from vllm.utils.math_utils import round_up
34
+from vllm.utils.torch_utils import current_stream
35
from vllm.v1.kv_cache_interface import KVCacheConfig
36
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
37
from vllm.v1.worker.gpu.block_table import BlockTables
@@ -377,7 +378,9 @@ class CudaGraphManager:
378
if self._capture_mem_samples is not None:
379
torch.accelerator.synchronize()
380
free_before = torch.accelerator.get_memory_info()[0]
380
- with torch.cuda.graph(graph, self.pool):
381
+ with torch.cuda.graph(
382
+ graph, self.pool, stream=current_stream()
383
+ ):
384
forward_fn(CUDAGraphMode.NONE)
385
# Join offloader's copy stream after forward to avoid
386
# unjoined stream error. The last layer's start_prefetch