[ROCm][CI] Stabilize MI355 FusedMoE test group (#53025)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Author: AndreasKaratzas

.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

tests/kernels/moe/test_moe_layer.py

@@ -7,12 +7,13 @@
 
 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.utils import TestMLP, make_test_weights, moe_quantize_weights
 from vllm.config import (
     CompilationConfig,
+    EPLBConfig,
     ParallelConfig,
     SchedulerConfig,
     VllmConfig,
@@ -383,6 +385,56 @@ def convert(v: str, ty):
         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,13 +2036,20 @@ 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,
         tensor_parallel_size=tp_size,
         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)