[ROCm][Test] Cover MoRI graph replay and output lifetime (#58093)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Sep 23, 2026 at 16:16 UTC
8b660ce96bc2e7ef164e7c9c26b969495a88f82c
4 files changed
+201
-2
.buildkite/test-amd.yaml
+31
-1
@@ -3571,6 +3571,35 @@ steps:
3571
- pytest -v -s v1/e2e/general/test_sharded_sampling.py
3572
- pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py
3573
3574
+- label: ":amd: (MI355) MoRI EP Healthy Baseline"
3575
+ timeout_in_minutes: 10
3576
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3577
+ dind: false
3578
+ agent_pool: mi355_2
3579
+ num_gpus: 2
3580
+ working_dir: "/vllm-workspace/tests"
3581
+ source_file_dependencies:
3582
+ - csrc/moe/
3583
+ - csrc/rocm/
3584
+ - tests/kernels/moe/
3585
+ - vllm/model_executor/layers/fused_moe/
3586
+ - vllm/model_executor/layers/quantization/
3587
+ - vllm/distributed/
3588
+ - vllm/config/
3589
+ - vllm/forward_context.py
3590
+ - vllm/v1/worker/workspace.py
3591
+ - vllm/utils/import_utils.py
3592
+ - vllm/utils/math_utils.py
3593
+ - vllm/utils/torch_utils.py
3594
+ - vllm/platforms/
3595
+ - vllm/_aiter_ops.py
3596
+ optional: true
3597
+ # TODO: Add fault injection, detection and recovery once MoRI supports them.
3598
+ # This job currently checks healthy EP execution only.
3599
+ commands:
3600
+ - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
3601
+ - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
3602
+
3603
- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD
3604
timeout_in_minutes: 45
3605
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -4414,7 +4443,8 @@ steps:
4443
- vllm/distributed/device_communicators/
4444
- vllm/config
4445
commands:
4417
- - pytest -v -s kernels/moe/test_moe_layer.py
4446
+ # MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
4447
+ - pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
4448
4449
- label: ":amd: (MI355) MoRI EP Numerics"
4450
timeout_in_minutes: 20
.buildkite/test_areas/fault_tolerance.yaml
+30
@@ -24,3 +24,33 @@ steps:
24
# https://github.com/NVIDIA/nccl/issues/1838
25
- export NCCL_CUMEM_HOST_ENABLE=0
26
- pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
27
+ mirror:
28
+ amd:
29
+ label: ":amd: (MI355) MoRI EP Healthy Baseline"
30
+ dind: false
31
+ device: mi355_2
32
+ num_devices: 2
33
+ timeout_in_minutes: 10
34
+ working_dir: "/vllm-workspace/tests"
35
+ depends_on:
36
+ - image-build-amd
37
+ source_file_dependencies:
38
+ - csrc/moe/
39
+ - csrc/rocm/
40
+ - tests/kernels/moe/
41
+ - vllm/model_executor/layers/fused_moe/
42
+ - vllm/model_executor/layers/quantization/
43
+ - vllm/distributed/
44
+ - vllm/config/
45
+ - vllm/forward_context.py
46
+ - vllm/v1/worker/workspace.py
47
+ - vllm/utils/import_utils.py
48
+ - vllm/utils/math_utils.py
49
+ - vllm/utils/torch_utils.py
50
+ - vllm/platforms/
51
+ - vllm/_aiter_ops.py
52
+ # TODO: Add fault injection, detection and recovery once MoRI supports them.
53
+ # This mirror currently checks healthy EP execution only.
54
+ commands:
55
+ - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
56
+ - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]'
.buildkite/test_areas/kernels.yaml
+3
@@ -713,6 +713,9 @@ steps:
713
- vllm/utils/torch_utils.py
714
- vllm/platforms/
715
- vllm/_aiter_ops.py
716
+ commands:
717
+ # MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
718
+ - pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
719
720
- label: ":nvidia: (B200) FusedMoE Layer Kernels"
721
key: kernels-fusedmoe-layer-test-2-b200s
tests/kernels/moe/test_moe_layer.py
+137
-1
@@ -1484,6 +1484,7 @@ def _run_one_config(
1484
use_shared_experts: bool,
1485
use_gate: bool,
1486
use_routed_input_transform: bool,
1487
+ input_scale: float = 1.0,
1488
**kwargs,
1489
) -> None:
1490
"""Generic test loop that sets up environment and delegates to test_body_fn.
@@ -1530,6 +1531,8 @@ def _run_one_config(
1531
1532
# Extract data from test_data
1533
hidden_states = test_data.hidden_states
1534
+ if input_scale != 1.0:
1535
+ hidden_states = hidden_states * input_scale
1536
router_logits = test_data.router_logits
1537
w1 = test_data.w1
1538
w2 = test_data.w2
@@ -1982,7 +1985,140 @@ def _parallel_worker_rocm_deepep(
1985
os._exit(exit_code)
1986
1987
1985
-# TODO: add cudagraphs/torch.compile tests
1988
+def _test_body_mori_graph(
1989
+ moe_layer: MoERunner,
1990
+ hidden_states: torch.Tensor,
1991
+ router_logits: torch.Tensor,
1992
+ vllm_config: VllmConfig,
1993
+ num_tokens: int,
1994
+ num_tokens_across_dp: torch.Tensor,
1995
+ baseline_output: torch.Tensor,
1996
+ **kwargs,
1997
+) -> tuple[torch.Tensor, torch.Tensor]:
1998
+ from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
1999
+ AiterExperts,
2000
+ )
2001
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
2002
+ MoriPrepareAndFinalize,
2003
+ )
2004
+
2005
+ kernel = moe_layer._quant_method.moe_kernel
2006
+ assert kernel is not None
2007
+ assert isinstance(kernel.prepare_finalize, MoriPrepareAndFinalize)
2008
+ assert isinstance(kernel.fused_experts, AiterExperts)
2009
+ assert not kernel.prepare_finalize.use_fp8_dispatch
2010
+
2011
+ def check(actual: torch.Tensor, expected: torch.Tensor) -> None:
2012
+ torch.testing.assert_close(actual, expected, atol=3.5e-2, rtol=3.5e-2)
2013
+ # These inputs produce small outputs; an absolute-only budget can
2014
+ # accept all-zero output. Also bound error relative to the reference.
2015
+ ref_norm = expected.float().norm()
2016
+ assert ref_norm > 0
2017
+ assert (actual.float() - expected.float()).norm() < 3.5e-2 * ref_norm
2018
+
2019
+ # Forward-context metadata belongs on CPU, outside graph capture.
2020
+ with set_forward_context(
2021
+ None,
2022
+ vllm_config,
2023
+ num_tokens=num_tokens,
2024
+ num_tokens_across_dp=num_tokens_across_dp.cpu(),
2025
+ ):
2026
+ initial_hidden = hidden_states.clone()
2027
+ initial_router = router_logits.clone()
2028
+ for _ in range(2):
2029
+ eager_output = moe_layer(hidden_states, router_logits)
2030
+ saved_eager_output = eager_output.clone()
2031
+ torch.accelerator.synchronize()
2032
+
2033
+ graph = torch.cuda.CUDAGraph()
2034
+ with torch.cuda.graph(graph):
2035
+ captured_output = moe_layer(hidden_states, router_logits)
2036
+ graph.replay()
2037
+ original_output = captured_output.clone()
2038
+ check(original_output, baseline_output)
2039
+
2040
+ for scale, shift in [(-0.5, 1), (1.75, 3)]:
2041
+ hidden_states.copy_(initial_hidden * scale)
2042
+ router_logits.copy_(initial_router.roll(shift, dims=-1))
2043
+ expected = moe_layer(hidden_states, router_logits).clone()
2044
+ graph.replay()
2045
+ check(captured_output, expected)
2046
+ assert (expected.float() - original_output.float()).norm() > (
2047
+ 0.1 * original_output.float().norm()
2048
+ )
2049
+
2050
+ # Later calls must not overwrite an earlier eager return tensor.
2051
+ check(eager_output, saved_eager_output)
2052
+ check(original_output, baseline_output)
2053
+ return baseline_output, original_output
2054
+
2055
+
2056
+def _mori_graph_worker(
2057
+ pgi: ProcessGroupInfo,
2058
+ vllm_config: VllmConfig,
2059
+ cpu_group,
2060
+ m: int,
2061
+ n: int,
2062
+) -> None:
2063
+ from vllm._aiter_ops import rocm_aiter_ops
2064
+
2065
+ rocm_aiter_ops.refresh_env_variables()
2066
+ _run_one_config(
2067
+ vllm_config=vllm_config,
2068
+ ep_size=2,
2069
+ dp_size=2,
2070
+ tp_size=1,
2071
+ dp_rank=pgi.rank,
2072
+ tp_rank=0,
2073
+ is_sequence_parallel=False,
2074
+ m=m,
2075
+ n=n,
2076
+ k=512,
2077
+ num_experts=8,
2078
+ top_k=2,
2079
+ quantization=None,
2080
+ backend="mori_high_throughput",
2081
+ test_body_fn=_test_body_mori_graph,
2082
+ use_shared_experts=False,
2083
+ use_gate=False,
2084
+ use_routed_input_transform=False,
2085
+ input_scale=1.0 + pgi.rank,
2086
+ )
2087
+
2088
+
2089
+@pytest.mark.parametrize("m,n", [(1, 256), (32, 512), (45, 256)])
2090
+def test_moe_layer_mori_graph(m: int, n: int, monkeypatch):
2091
+ """BF16 EP dispatch/combine preserves routing and outputs across replays."""
2092
+ from vllm._aiter_ops import is_aiter_found_and_supported
2093
+
2094
+ if not current_platform.is_rocm() or not has_mori():
2095
+ pytest.skip("Requires ROCm and MoRI")
2096
+ if not is_aiter_found_and_supported():
2097
+ pytest.skip("Requires supported AITER kernels")
2098
+ if current_platform.device_count() < 2:
2099
+ pytest.skip("Requires two GPUs")
2100
+
2101
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
2102
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
2103
+ monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
2104
+ compilation_config = CompilationConfig()
2105
+ compilation_config.pass_config.fuse_allreduce_rms = False
2106
+ vllm_config = VllmConfig(
2107
+ parallel_config=ParallelConfig(
2108
+ data_parallel_size=2,
2109
+ tensor_parallel_size=1,
2110
+ enable_expert_parallel=True,
2111
+ all2all_backend="mori_high_throughput",
2112
+ ),
2113
+ compilation_config=compilation_config,
2114
+ scheduler_config=SchedulerConfig.default_factory(
2115
+ max_num_batched_tokens=64, max_num_seqs=64
2116
+ ),
2117
+ )
2118
+ parallel_launch_with_config(2, _mori_graph_worker, vllm_config, None, m, n)
2119
+
2120
+
2121
+# TODO: add torch.compile tests and graph coverage for other backends.
2122
@pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS)
2123
@pytest.mark.parametrize("backend", BACKENDS)
2124
@pytest.mark.parametrize("enable_eplb", [False, True])