[ROCm][Bugfix] Reduce CUDA graph divergences (#57229)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Sep 16, 2026 at 22:12 UTC
2d8a5c5741dedb3559d26ce95fa894771cffa881
5 files changed
+8
-90
tests/test_config.py
-39
@@ -477,45 +477,6 @@ def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):
477
default_breakable_cudagraph_architectures.cache_clear()
478
479
480
-@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
481
-@pytest.mark.parametrize(
482
- ("architecture", "use_v2", "mode", "expected"),
483
- [
484
- ("DeepseekV4ForCausalLM", True, None, CUDAGraphMode.NONE),
485
- ("DeepseekV4ForConditionalGeneration", True, None, CUDAGraphMode.NONE),
486
- ("DeepseekV4ForCausalLM", False, None, None),
487
- ("LlamaForCausalLM", True, None, None),
488
- (
489
- "DeepseekV4ForCausalLM",
490
- True,
491
- CUDAGraphMode.FULL_DECODE_ONLY,
492
- CUDAGraphMode.FULL_DECODE_ONLY,
493
- ),
494
- ],
495
-)
496
-def test_rocm_gfx950_deepseek_v4_cudagraph_default(
497
- monkeypatch, architecture, use_v2, mode, expected
498
-):
499
- from vllm._aiter_ops import rocm_aiter_ops
500
- from vllm.platforms import rocm
501
-
502
- monkeypatch.setattr(rocm, "on_gfx950", lambda: True)
503
- monkeypatch.setattr(rocm_aiter_ops, "is_fused_moe_enabled", lambda: False)
504
- monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: False)
505
- monkeypatch.setattr(
506
- rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", lambda: False
507
- )
508
- config = SimpleNamespace(
509
- compilation_config=CompilationConfig(cudagraph_mode=mode),
510
- model_config=SimpleNamespace(architecture=architecture),
511
- use_v2_model_runner=use_v2,
512
- )
513
-
514
- rocm.RocmPlatform.apply_config_platform_defaults(config)
515
-
516
- assert config.compilation_config.cudagraph_mode == expected
517
-
518
-
480
@pytest.mark.parametrize(
481
("model", "architecture"),
482
[
vllm/distributed/parallel_state.py
+2
-10
@@ -1632,10 +1632,7 @@ def get_pcp_group() -> GroupCoordinator:
1632
1633
1634
@contextmanager
1635
-def graph_capture(
1636
- device: torch.device,
1637
- graph_capture_context: GraphCaptureContext | None = None,
1638
-):
1635
+def graph_capture(device: torch.device):
1636
"""`graph_capture` is a context manager which should surround the code that
1637
is capturing the CUDA graph. Its main purpose is to ensure that some
1638
operations will be run after the graph is captured, before the graph
@@ -1647,13 +1644,8 @@ def graph_capture(
1644
the graph capture is running on a separate stream from the default stream,
1645
in order to explicitly distinguish the kernels to capture
1646
from other kernels possibly launched on background in the default stream.
1650
-
1651
- A caller may pass an explicit ``graph_capture_context`` to control the
1652
- stream used (e.g. to capture on the default stream).
1647
"""
1654
- context = graph_capture_context or GraphCaptureContext(
1655
- torch.cuda.Stream(device=device)
1656
- )
1648
+ context = GraphCaptureContext(torch.cuda.Stream(device=device))
1649
with (
1650
get_tp_group().graph_capture(context),
1651
get_pp_group().graph_capture(context),
vllm/platforms/rocm.py
-17
@@ -865,25 +865,8 @@ class RocmPlatform(Platform):
865
@classmethod
866
def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
867
from vllm._aiter_ops import rocm_aiter_ops
868
- from vllm.config.compilation import CUDAGraphMode
868
869
compilation_config = vllm_config.compilation_config
871
- model_config = vllm_config.model_config
872
- if (
873
- compilation_config.cudagraph_mode is None
874
- and model_config is not None
875
- and on_gfx950()
876
- and vllm_config.use_v2_model_runner
877
- and model_config.architecture
878
- in {
879
- "DeepseekV4ForCausalLM",
880
- "DeepseekV4ForConditionalGeneration",
881
- }
882
- ):
883
- # Default to eager after reported gfx950/MRV2 accuracy regressions:
884
- # https://github.com/vllm-project/vllm/issues/52644
885
- compilation_config.cudagraph_mode = CUDAGraphMode.NONE
886
-
870
use_aiter_fused_moe = rocm_aiter_ops.is_fused_moe_enabled()
871
use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enabled()
872
use_aiter_fused_se = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
vllm/v1/worker/gpu_model_runner.py
+1
-23
@@ -46,7 +46,6 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
46
copy_kv_blocks,
47
)
48
from vllm.distributed.parallel_state import (
49
- GraphCaptureContext,
49
get_dcp_group,
50
get_pp_group,
51
get_tp_group,
@@ -128,7 +127,6 @@ from vllm.utils.platform_utils import num_compute_units
127
from vllm.utils.torch_utils import (
128
PIN_MEMORY,
129
async_tensor_h2d,
131
- current_stream,
130
kv_cache_dtype_str_to_dtype,
131
)
132
from vllm.v1.attention.backend import (
@@ -6712,31 +6710,11 @@ class GPUModelRunner(
6710
per_graph_estimate = {}
6711
encoder_memory_estimate = 0
6712
6715
- # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated
6716
- # compute stream instead of the fresh side stream graph_capture()
6717
- # allocates by default. torch's allocator pools free blocks per stream,
6718
- # so a side-stream forward strands a persistent aiter scratch buffer in
6719
- # a separate pool, shifting the physical placement of the real KV cache
6720
- # allocated afterward and slowing bandwidth-bound decode ~20%. The
6721
- # graphs are discarded, so a side stream is unnecessary here.
6722
- # Use current_stream(), not torch.cuda.current_stream(): before vLLM
6723
- # initializes its dedicated stream, torch returns the per-thread default
6724
- # stream (cuda_stream=0), which cannot be used for cudagraph capture.
6725
- # cap_ctx=None keeps the side-stream path on CUDA.
6726
- cap_ctx = (
6727
- GraphCaptureContext(current_stream())
6728
- if current_platform.is_rocm()
6729
- else None
6730
- )
6731
-
6713
# Cleanup-only guard: CUDA graph capture errors should still propagate
6714
# because encoder graph capture is opt-in.
6715
try:
6716
set_cudagraph_capturing_enabled(True)
6736
- with (
6737
- self._freeze_gc(),
6738
- graph_capture(device=self.device, graph_capture_context=cap_ctx),
6739
- ):
6717
+ with self._freeze_gc(), graph_capture(device=self.device):
6718
torch.accelerator.synchronize()
6719
torch.accelerator.empty_cache()
6720
vllm/v1/worker/gpu_worker.py
+5
-1
@@ -98,6 +98,7 @@ from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase
98
from vllm.v1.worker.workspace import init_workspace_manager
99
100
from ...model_executor.model_loader import TensorizerLoader
101
+from .gpu.cudagraph_utils import has_compiled_submodule
102
from .gpu.warmup import warmup_kernels
103
from .utils import request_memory
104
@@ -772,7 +773,10 @@ class Worker(WorkerBase):
773
def compile_or_warm_up_model(self) -> CompilationTimes:
774
warmup_sizes: list[int] = []
775
775
- if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
776
+ if (
777
+ self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE
778
+ and has_compiled_submodule(self.model_runner.get_model())
779
+ ):
780
# warm up sizes that are not in cudagraph capture sizes,
781
# but users still want to compile for better performance,
782
# e.g. for the max-num-batched token size in chunked prefill.