iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Test] Cover MoRI graph replay and output lifetime (#58093)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-09-23 14:16:51 -07:00
committed by GitHub
parent 97b1b12117
commit 8b660ce96b
4 changed files with 201 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test-amd.yaml
+31 -1
View File
@@ -3571,6 +3571,35 @@ steps:
- pytest -v -s v1/e2e/general/test_sharded_sampling.py
- pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py
- label: ":amd: (MI355) MoRI EP Healthy Baseline"
timeout_in_minutes: 10
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
dind: false
agent_pool: mi355_2
num_gpus: 2
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- csrc/moe/
- csrc/rocm/
- tests/kernels/moe/
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/
- vllm/config/
- vllm/forward_context.py
- vllm/v1/worker/workspace.py
- vllm/utils/import_utils.py
- vllm/utils/math_utils.py
- vllm/utils/torch_utils.py
- vllm/platforms/
- vllm/_aiter_ops.py
optional: true
# TODO: Add fault injection, detection and recovery once MoRI supports them.
# This job currently checks healthy EP execution only.
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
- 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]'
- label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" # TBD
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -4414,7 +4443,8 @@ steps:
- vllm/distributed/device_communicators/
- vllm/config
commands:
- pytest -v -s kernels/moe/test_moe_layer.py
# MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
- pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
- label: ":amd: (MI355) MoRI EP Numerics"
timeout_in_minutes: 20
.buildkite/test_areas/fault_tolerance.yaml
+30
View File
@@ -24,3 +24,33 @@ steps:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
- pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py
mirror:
amd:
label: ":amd: (MI355) MoRI EP Healthy Baseline"
dind: false
device: mi355_2
num_devices: 2
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
depends_on:
- image-build-amd
source_file_dependencies:
- csrc/moe/
- csrc/rocm/
- tests/kernels/moe/
- vllm/model_executor/layers/fused_moe/
- vllm/model_executor/layers/quantization/
- vllm/distributed/
- vllm/config/
- vllm/forward_context.py
- vllm/v1/worker/workspace.py
- vllm/utils/import_utils.py
- vllm/utils/math_utils.py
- vllm/utils/torch_utils.py
- vllm/platforms/
- vllm/_aiter_ops.py
# TODO: Add fault injection, detection and recovery once MoRI supports them.
# This mirror currently checks healthy EP execution only.
commands:
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph
- 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
View File
@@ -713,6 +713,9 @@ steps:
- vllm/utils/torch_utils.py
- vllm/platforms/
- vllm/_aiter_ops.py
commands:
# MoRI graph cases run in the Fault Tolerance group's healthy EP baseline.
- pytest -v -s kernels/moe/test_moe_layer.py -k "not test_moe_layer_mori_graph"
- label: ":nvidia: (B200) FusedMoE Layer Kernels"
key: kernels-fusedmoe-layer-test-2-b200s
tests/kernels/moe/test_moe_layer.py
+137 -1
View File
@@ -1484,6 +1484,7 @@ def _run_one_config(
use_shared_experts: bool,
use_gate: bool,
use_routed_input_transform: bool,
input_scale: float = 1.0,
**kwargs,
) -> None:
"""Generic test loop that sets up environment and delegates to test_body_fn.
@@ -1530,6 +1531,8 @@ def _run_one_config(
# Extract data from test_data
hidden_states = test_data.hidden_states
if input_scale != 1.0:
hidden_states = hidden_states * input_scale
router_logits = test_data.router_logits
w1 = test_data.w1
w2 = test_data.w2
@@ -1982,7 +1985,140 @@ def _parallel_worker_rocm_deepep(
os._exit(exit_code)
# TODO: add cudagraphs/torch.compile tests
def _test_body_mori_graph(
moe_layer: MoERunner,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
vllm_config: VllmConfig,
num_tokens: int,
num_tokens_across_dp: torch.Tensor,
baseline_output: torch.Tensor,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
AiterExperts,
)
from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
MoriPrepareAndFinalize,
)
kernel = moe_layer._quant_method.moe_kernel
assert kernel is not None
assert isinstance(kernel.prepare_finalize, MoriPrepareAndFinalize)
assert isinstance(kernel.fused_experts, AiterExperts)
assert not kernel.prepare_finalize.use_fp8_dispatch
def check(actual: torch.Tensor, expected: torch.Tensor) -> None:
torch.testing.assert_close(actual, expected, atol=3.5e-2, rtol=3.5e-2)
# These inputs produce small outputs; an absolute-only budget can
# accept all-zero output. Also bound error relative to the reference.
ref_norm = expected.float().norm()
assert ref_norm > 0
assert (actual.float() - expected.float()).norm() < 3.5e-2 * ref_norm
# Forward-context metadata belongs on CPU, outside graph capture.
with set_forward_context(
None,
vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp.cpu(),
):
initial_hidden = hidden_states.clone()
initial_router = router_logits.clone()
for _ in range(2):
eager_output = moe_layer(hidden_states, router_logits)
saved_eager_output = eager_output.clone()
torch.accelerator.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured_output = moe_layer(hidden_states, router_logits)
graph.replay()
original_output = captured_output.clone()
check(original_output, baseline_output)
for scale, shift in [(-0.5, 1), (1.75, 3)]:
hidden_states.copy_(initial_hidden * scale)
router_logits.copy_(initial_router.roll(shift, dims=-1))
expected = moe_layer(hidden_states, router_logits).clone()
graph.replay()
check(captured_output, expected)
assert (expected.float() - original_output.float()).norm() > (
0.1 * original_output.float().norm()
)
# Later calls must not overwrite an earlier eager return tensor.
check(eager_output, saved_eager_output)
check(original_output, baseline_output)
return baseline_output, original_output
def _mori_graph_worker(
pgi: ProcessGroupInfo,
vllm_config: VllmConfig,
cpu_group,
m: int,
n: int,
) -> None:
from vllm._aiter_ops import rocm_aiter_ops
rocm_aiter_ops.refresh_env_variables()
_run_one_config(
vllm_config=vllm_config,
ep_size=2,
dp_size=2,
tp_size=1,
dp_rank=pgi.rank,
tp_rank=0,
is_sequence_parallel=False,
m=m,
n=n,
k=512,
num_experts=8,
top_k=2,
quantization=None,
backend="mori_high_throughput",
test_body_fn=_test_body_mori_graph,
use_shared_experts=False,
use_gate=False,
use_routed_input_transform=False,
input_scale=1.0 + pgi.rank,
)
@pytest.mark.parametrize("m,n", [(1, 256), (32, 512), (45, 256)])
def test_moe_layer_mori_graph(m: int, n: int, monkeypatch):
"""BF16 EP dispatch/combine preserves routing and outputs across replays."""
from vllm._aiter_ops import is_aiter_found_and_supported
if not current_platform.is_rocm() or not has_mori():
pytest.skip("Requires ROCm and MoRI")
if not is_aiter_found_and_supported():
pytest.skip("Requires supported AITER kernels")
if current_platform.device_count() < 2:
pytest.skip("Requires two GPUs")
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "0")
compilation_config = CompilationConfig()
compilation_config.pass_config.fuse_allreduce_rms = False
vllm_config = VllmConfig(
parallel_config=ParallelConfig(
data_parallel_size=2,
tensor_parallel_size=1,
enable_expert_parallel=True,
all2all_backend="mori_high_throughput",
),
compilation_config=compilation_config,
scheduler_config=SchedulerConfig.default_factory(
max_num_batched_tokens=64, max_num_seqs=64
),
)
parallel_launch_with_config(2, _mori_graph_worker, vllm_config, None, m, n)
# TODO: add torch.compile tests and graph coverage for other backends.
@pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS)
@pytest.mark.parametrize("backend", BACKENDS)
@pytest.mark.parametrize("enable_eplb", [False, True])