From 88eb946cb1a0901bcbb842de463d6c1cae035fa9 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 21 Aug 2026 15:15:25 -0500 Subject: [PATCH] [ROCm][CI] Stabilize MI355 FusedMoE test group (#53025) Signed-off-by: Andreas Karatzas --- .buildkite/test_areas/kernels.yaml | 53 -------- tests/kernels/moe/test_moe_layer.py | 191 +++++++++++++++++++++++++--- 2 files changed, 174 insertions(+), 70 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index e3ec0d17e2..6197790c03 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -468,56 +468,3 @@ steps: commands: - pytest -v -s kernels/moe/test_moe_layer.py - pytest -v -s kernels/moe/test_deepep_v2_moe.py - -- label: Kernels FusedMoE Layer Test (2xMI355) - key: kernels-fusedmoe-layer-test-2-mi355 - depends_on: - - image-build-amd - timeout_in_minutes: 180 - dind: false - device: mi355_2 - soft_fail: true - optional: true - 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 - commands: - - pytest -v -s kernels/moe/test_moe_layer.py - -- label: Kernels FP8 MoE Test (MI355) - key: kernels-fp8-moe-test-mi355 - depends_on: - - image-build-amd - timeout_in_minutes: 180 - dind: false - device: mi355_1 - soft_fail: true - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/moe/ - - vllm/model_executor/layers/fused_moe/ - - tests/kernels/moe/test_deepep_moe.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/envs.py - commands: - - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py - - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py - - pytest -v -s kernels/moe/test_moe.py - - pytest -v -s kernels/moe/test_block_int8.py - - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py - - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 8a081784fb..9d2208eb51 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -7,12 +7,13 @@ Run `pytest tests/kernels/test_moe_layer.py`. import functools import os +import sys import tempfile import traceback import types from collections.abc import Callable from contextlib import suppress -from dataclasses import astuple, dataclass, fields +from dataclasses import astuple, dataclass, fields, replace from itertools import product from typing import get_args @@ -28,6 +29,7 @@ from tests.kernels.moe.modular_kernel_tools.parallel_utils import ( from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_weights from vllm.config import ( CompilationConfig, + EPLBConfig, ParallelConfig, SchedulerConfig, VllmConfig, @@ -383,6 +385,56 @@ class MoETestConfig: return MoETestConfig(*values) +def _group_deepep_ll_configs( + test_configs: list[MoETestConfig], +) -> list[list[MoETestConfig]]: + """Group configs that can share one DeepEP low-latency buffer.""" + groups: dict[tuple[int, int, int], list[MoETestConfig]] = {} + for test_config in test_configs: + # max_num_tokens and the remaining buffer arguments are fixed for one + # outer pytest item. These are the arguments that can vary here. + buffer_key = ( + test_config.k, + test_config.ep_size, + test_config.num_experts, + ) + groups.setdefault(buffer_key, []).append(test_config) + return list(groups.values()) + + +@pytest.mark.cpu_test +def test_group_deepep_ll_configs_by_buffer_requirements(): + config_a = MoETestConfig( + 1, + 128, + 2048, + 8, + 2, + torch.bfloat16, + None, + False, + False, + False, + backend="deepep_low_latency", + ep_size=2, + dp_size=2, + ) + config_b = replace(config_a, num_experts=64) + config_a_later = replace(config_a, m=32) + config_c = replace(config_a, k=4096) + config_d = replace(config_a, ep_size=4, dp_size=4) + + assert _group_deepep_ll_configs( + [config_a, config_b, config_a_later, config_c, config_d] + ) == [ + [config_a, config_a_later], + [config_b], + [config_c], + [config_d], + ] + assert _group_deepep_ll_configs([]) == [] + + def generate_valid_test_configs( backend: str, ep_size: int, @@ -520,6 +572,17 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: f"{config.backend} does not support quantization={config.quantization}", ) + if ( + on_gfx950() + and config.backend == "deepep_low_latency" + and config.quantization == "modelopt_fp4" + ): + return ( + False, + "DeepEP low latency requires a batched NVFP4 MoE backend, " + "which is not available on gfx950.", + ) + if config.backend in MORI_BACKENDS: if os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") != "1": return False, "mori MoE layer matrix is opt-in" @@ -1719,6 +1782,7 @@ def _parallel_worker( test_configs: list[MoETestConfig], verbosity: int, failure_report_path: str | None = None, + deep_ep_handle_keepalive: list[object] | None = None, **kwargs, ) -> None: set_random_seed(7) @@ -1776,15 +1840,30 @@ def _parallel_worker( finally: # DeepEP managers are not reliably reusable across many subtests in # a single worker process. Tear them down after each DeepEP case so - # later subtests do not inherit stale communication state. - if test_config.backend in { - "deepep_low_latency", - "deepep_high_throughput", - }: + # later subtests do not inherit stale communication state. Skip this + # on ROCm: rocSHMEM cannot reinitialize the allocator after a DeepEP + # buffer is destroyed in the same process. + if current_platform.is_cuda() and test_config.backend in DEEPEP_BACKENDS: torch.accelerator.synchronize() all2all_manager = get_ep_group().device_communicator.all2all_manager if all2all_manager is not None: all2all_manager.destroy() + elif ( + deep_ep_handle_keepalive is not None + and current_platform.is_rocm() + and test_config.backend in DEEPEP_BACKENDS + ): + # The manager cache is weak. Keep its handle alive until the + # launcher hard-exits this ROCm worker, otherwise each subtest + # implicitly destroys and recreates the DeepEP buffer. + all2all_manager = get_ep_group().device_communicator.all2all_manager + if all2all_manager is not None: + handle_cache = getattr(all2all_manager, "handle_cache", None) + if handle_cache is not None and not deep_ep_handle_keepalive: + with handle_cache._lock: + cached_handles = list(handle_cache._cache.values()) + if cached_handles: + deep_ep_handle_keepalive.extend(cached_handles) total = total + 1 torch.distributed.barrier() @@ -1845,12 +1924,60 @@ def _parallel_worker( f"{failure_details_str}\n{report}" ) if is_logging_rank and failure_report_path is not None: - with open(failure_report_path, "w", encoding="utf-8") as report_file: + with open(failure_report_path, "a", encoding="utf-8") as report_file: report_file.write(failure_report) if is_logging_rank: raise RuntimeError(failure_report) +def _parallel_worker_rocm_deepep( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + test_configs: list[MoETestConfig], + verbosity: int, + failure_report_path: str | None = None, + deep_ep_handle_keepalive: list[object] | None = None, + **kwargs, +) -> None: + """Run a ROCm DeepEP batch without unsafe Python/HIP teardown.""" + assert current_platform.is_rocm() + assert deep_ep_handle_keepalive is not None + + exit_code = 1 + try: + _parallel_worker( + pgi, + vllm_config, + cpu_group, + test_configs, + verbosity, + failure_report_path=failure_report_path, + deep_ep_handle_keepalive=deep_ep_handle_keepalive, + **kwargs, + ) + exit_code = 0 + except BaseException as ex: + print(ex) + traceback.print_exc() + finally: + try: + torch.accelerator.synchronize() + # Do not run vLLM cleanup: it explicitly destroys the DeepEP + # buffer. Destroy only the default group, matching the accepted + # ROCm workaround in tests/kernels/moe/parallel_utils.py. + torch.distributed.destroy_process_group() + except BaseException: + traceback.print_exc() + exit_code = 1 + finally: + # Bypass the HIP atexit use-after-free fixed upstream by + # https://github.com/ROCm/rocm-systems/pull/6942. + sys.stdout.flush() + sys.stderr.flush() + os._exit(exit_code) + + # TODO: add cudagraphs/torch.compile tests @pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS) @pytest.mark.parametrize("backend", BACKENDS) @@ -1909,6 +2036,12 @@ def test_moe_layer( # moe_backend=flashinfer_trtllm / flashinfer_cutlass / flashinfer_cutedsl # (BF16, FP8 and NVFP4 paths), and VLLM_USE_FLASHINFER_MOE_INT4=1. + # Repeated NIXL memory registration in this broad layer matrix fails on + # gfx950. NIXL EPLB is covered by dedicated tests, so use Gloo here to + # preserve the layer/EPLB coverage. + eplb_config = EPLBConfig( + communicator="torch_gloo" if enable_eplb and on_gfx950() else None + ) parallel_config = ParallelConfig( pipeline_parallel_size=1, data_parallel_size=dp_size, @@ -1916,6 +2049,7 @@ def test_moe_layer( enable_expert_parallel=use_ep, all2all_backend=backend, enable_eplb=enable_eplb, + eplb_config=eplb_config, ) compilation_config = CompilationConfig() @@ -1955,19 +2089,42 @@ def test_moe_layer( ) as failure_report_file: failure_report_path = failure_report_file.name + test_config_batches = [test_configs] + if current_platform.is_rocm() and backend == "deepep_low_latency": + # rocSHMEM cannot destroy one low-latency buffer and initialize a + # differently sized one in the same process. Use one worker lifetime + # per compatible buffer shape while preserving every matrix case. + test_config_batches = _group_deepep_ll_configs(test_configs) + + rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS + parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker + launch_failures: list[str] = [] + try: - parallel_launch_with_config( - world_size, - _parallel_worker, - vllm_config, - None, - test_configs, - verbosity, - failure_report_path=failure_report_path, - ) + for test_config_batch in test_config_batches: + report_size_before = os.path.getsize(failure_report_path) + try: + parallel_launch_with_config( + world_size, + parallel_worker, + vllm_config, + None, + test_config_batch, + verbosity, + failure_report_path=failure_report_path, + deep_ep_handle_keepalive=[] if rocm_deepep else None, + ) + except Exception as ex: + # Normal subtest failures are already in the shared report. + # Preserve launcher/setup errors that occur before that write. + if os.path.getsize(failure_report_path) == report_size_before: + launch_failures.append(str(ex)) + if os.path.getsize(failure_report_path) > 0: with open(failure_report_path, encoding="utf-8") as report_file: - pytest.fail(report_file.read()) + launch_failures.insert(0, report_file.read()) + if launch_failures: + pytest.fail("\n\n".join(launch_failures)) finally: with suppress(FileNotFoundError): os.remove(failure_report_path)