diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 8bd434a355..5b630d773f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -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 diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml index 2e85a1c3f8..a578be14ed 100644 --- a/.buildkite/test_areas/fault_tolerance.yaml +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -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]' diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4fb7cfcc3c..ba683d58ad 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -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 diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 5b1af59731..5fe5ed1ab2 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -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])