@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 88eb946cb1

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

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

Andreas Karatzas committed Aug 21, 2026 at 15:15 UTC 88eb946cb1a0901bcbb842de463d6c1cae035fa9

2 files changed +174 -70

.buildkite/test_areas/kernels.yaml

-53

@@ -468,56 +468,3 @@ steps:

468 commands:

469 - pytest -v -s kernels/moe/test_moe_layer.py

470 - pytest -v -s kernels/moe/test_deepep_v2_moe.py

471 -

472 -- label: Kernels FusedMoE Layer Test (2xMI355)

473 - key: kernels-fusedmoe-layer-test-2-mi355

474 - depends_on:

475 - - image-build-amd

476 - timeout_in_minutes: 180

477 - dind: false

478 - device: mi355_2

479 - soft_fail: true

480 - optional: true

481 - working_dir: "/vllm-workspace/tests"

482 - source_file_dependencies:

483 - - csrc/moe/

484 - - csrc/rocm/

485 - - tests/kernels/moe

486 - - vllm/model_executor/layers/fused_moe/

487 - - vllm/model_executor/layers/quantization/

488 - - vllm/distributed/

489 - - vllm/config/

490 - - vllm/forward_context.py

491 - - vllm/v1/worker/workspace.py

492 - - vllm/utils/import_utils.py

493 - - vllm/utils/math_utils.py

494 - - vllm/utils/torch_utils.py

495 - - vllm/platforms/

496 - - vllm/_aiter_ops.py

497 - commands:

498 - - pytest -v -s kernels/moe/test_moe_layer.py

499 -

500 -- label: Kernels FP8 MoE Test (MI355)

501 - key: kernels-fp8-moe-test-mi355

502 - depends_on:

503 - - image-build-amd

504 - timeout_in_minutes: 180

505 - dind: false

506 - device: mi355_1

507 - soft_fail: true

508 - optional: true

509 - working_dir: "/vllm-workspace/tests"

510 - source_file_dependencies:

511 - - csrc/moe/

512 - - vllm/model_executor/layers/fused_moe/

513 - - tests/kernels/moe/test_deepep_moe.py

514 - - vllm/_aiter_ops.py

515 - - vllm/platforms/rocm.py

516 - - vllm/envs.py

517 - commands:

518 - - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py

519 - - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py

520 - - pytest -v -s kernels/moe/test_moe.py

521 - - pytest -v -s kernels/moe/test_block_int8.py

522 - - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py

523 - - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

tests/kernels/moe/test_moe_layer.py

+174 -17

@@ -7,12 +7,13 @@ Run `pytest tests/kernels/test_moe_layer.py`.

8 import functools

9 import os

10 +import sys

11 import tempfile

12 import traceback

13 import types

14 from collections.abc import Callable

15 from contextlib import suppress

15 -from dataclasses import astuple, dataclass, fields

16 +from dataclasses import astuple, dataclass, fields, replace

17 from itertools import product

18 from typing import get_args

@@ -28,6 +29,7 @@ from tests.kernels.moe.modular_kernel_tools.parallel_utils import (

29 from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_weights

30 from vllm.config import (

31 CompilationConfig,

32 + EPLBConfig,

33 ParallelConfig,

34 SchedulerConfig,

35 VllmConfig,

@@ -383,6 +385,56 @@ class MoETestConfig:

385 return MoETestConfig(*values)

386

387

388 +def _group_deepep_ll_configs(

389 + test_configs: list[MoETestConfig],

390 +) -> list[list[MoETestConfig]]:

391 + """Group configs that can share one DeepEP low-latency buffer."""

392 + groups: dict[tuple[int, int, int], list[MoETestConfig]] = {}

393 + for test_config in test_configs:

394 + # max_num_tokens and the remaining buffer arguments are fixed for one

395 + # outer pytest item. These are the arguments that can vary here.

396 + buffer_key = (

397 + test_config.k,

398 + test_config.ep_size,

399 + test_config.num_experts,

400 + )

401 + groups.setdefault(buffer_key, []).append(test_config)

402 + return list(groups.values())

403 +

404 +

405 +@pytest.mark.cpu_test

406 +def test_group_deepep_ll_configs_by_buffer_requirements():

407 + config_a = MoETestConfig(

408 + 1,

409 + 128,

410 + 2048,

411 + 8,

412 + 2,

413 + torch.bfloat16,

414 + None,

415 + False,

416 + False,

417 + False,

418 + backend="deepep_low_latency",

419 + ep_size=2,

420 + dp_size=2,

421 + )

422 + config_b = replace(config_a, num_experts=64)

423 + config_a_later = replace(config_a, m=32)

424 + config_c = replace(config_a, k=4096)

425 + config_d = replace(config_a, ep_size=4, dp_size=4)

426 +

427 + assert _group_deepep_ll_configs(

428 + [config_a, config_b, config_a_later, config_c, config_d]

429 + ) == [

430 + [config_a, config_a_later],

431 + [config_b],

432 + [config_c],

433 + [config_d],

434 + ]

435 + assert _group_deepep_ll_configs([]) == []

436 +

437 +

438 def generate_valid_test_configs(

439 backend: str,

440 ep_size: int,

@@ -520,6 +572,17 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:

572 f"{config.backend} does not support quantization={config.quantization}",

573 )

574

575 + if (

576 + on_gfx950()

577 + and config.backend == "deepep_low_latency"

578 + and config.quantization == "modelopt_fp4"

579 + ):

580 + return (

581 + False,

582 + "DeepEP low latency requires a batched NVFP4 MoE backend, "

583 + "which is not available on gfx950.",

584 + )

585 +

586 if config.backend in MORI_BACKENDS:

587 if os.environ.get("VLLM_TEST_ENABLE_MORI_MOE_LAYER") != "1":

588 return False, "mori MoE layer matrix is opt-in"

@@ -1719,6 +1782,7 @@ def _parallel_worker(

1782 test_configs: list[MoETestConfig],

1783 verbosity: int,

1784 failure_report_path: str | None = None,

1785 + deep_ep_handle_keepalive: list[object] | None = None,

1786 **kwargs,

1787 ) -> None:

1788 set_random_seed(7)

@@ -1776,15 +1840,30 @@ def _parallel_worker(

1840 finally:

1841 # DeepEP managers are not reliably reusable across many subtests in

1842 # a single worker process. Tear them down after each DeepEP case so

1779 - # later subtests do not inherit stale communication state.

1780 - if test_config.backend in {

1781 - "deepep_low_latency",

1782 - "deepep_high_throughput",

1783 - }:

1843 + # later subtests do not inherit stale communication state. Skip this

1844 + # on ROCm: rocSHMEM cannot reinitialize the allocator after a DeepEP

1845 + # buffer is destroyed in the same process.

1846 + if current_platform.is_cuda() and test_config.backend in DEEPEP_BACKENDS:

1847 torch.accelerator.synchronize()

1848 all2all_manager = get_ep_group().device_communicator.all2all_manager

1849 if all2all_manager is not None:

1850 all2all_manager.destroy()

1851 + elif (

1852 + deep_ep_handle_keepalive is not None

1853 + and current_platform.is_rocm()

1854 + and test_config.backend in DEEPEP_BACKENDS

1855 + ):

1856 + # The manager cache is weak. Keep its handle alive until the

1857 + # launcher hard-exits this ROCm worker, otherwise each subtest

1858 + # implicitly destroys and recreates the DeepEP buffer.

1859 + all2all_manager = get_ep_group().device_communicator.all2all_manager

1860 + if all2all_manager is not None:

1861 + handle_cache = getattr(all2all_manager, "handle_cache", None)

1862 + if handle_cache is not None and not deep_ep_handle_keepalive:

1863 + with handle_cache._lock:

1864 + cached_handles = list(handle_cache._cache.values())

1865 + if cached_handles:

1866 + deep_ep_handle_keepalive.extend(cached_handles)

1867 total = total + 1

1868 torch.distributed.barrier()

1869

@@ -1845,12 +1924,60 @@ def _parallel_worker(

1924 f"{failure_details_str}\n{report}"

1925 )

1926 if is_logging_rank and failure_report_path is not None:

1848 - with open(failure_report_path, "w", encoding="utf-8") as report_file:

1927 + with open(failure_report_path, "a", encoding="utf-8") as report_file:

1928 report_file.write(failure_report)

1929 if is_logging_rank:

1930 raise RuntimeError(failure_report)

1931

1932

1933 +def _parallel_worker_rocm_deepep(

1934 + pgi: ProcessGroupInfo,

1935 + vllm_config: VllmConfig,

1936 + cpu_group,

1937 + test_configs: list[MoETestConfig],

1938 + verbosity: int,

1939 + failure_report_path: str | None = None,

1940 + deep_ep_handle_keepalive: list[object] | None = None,

1941 + **kwargs,

1942 +) -> None:

1943 + """Run a ROCm DeepEP batch without unsafe Python/HIP teardown."""

1944 + assert current_platform.is_rocm()

1945 + assert deep_ep_handle_keepalive is not None

1946 +

1947 + exit_code = 1

1948 + try:

1949 + _parallel_worker(

1950 + pgi,

1951 + vllm_config,

1952 + cpu_group,

1953 + test_configs,

1954 + verbosity,

1955 + failure_report_path=failure_report_path,

1956 + deep_ep_handle_keepalive=deep_ep_handle_keepalive,

1957 + **kwargs,

1958 + )

1959 + exit_code = 0

1960 + except BaseException as ex:

1961 + print(ex)

1962 + traceback.print_exc()

1963 + finally:

1964 + try:

1965 + torch.accelerator.synchronize()

1966 + # Do not run vLLM cleanup: it explicitly destroys the DeepEP

1967 + # buffer. Destroy only the default group, matching the accepted

1968 + # ROCm workaround in tests/kernels/moe/parallel_utils.py.

1969 + torch.distributed.destroy_process_group()

1970 + except BaseException:

1971 + traceback.print_exc()

1972 + exit_code = 1

1973 + finally:

1974 + # Bypass the HIP atexit use-after-free fixed upstream by

1975 + # https://github.com/ROCm/rocm-systems/pull/6942.

1976 + sys.stdout.flush()

1977 + sys.stderr.flush()

1978 + os._exit(exit_code)

1979 +

1980 +

1981 # TODO: add cudagraphs/torch.compile tests

1982 @pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS)

1983 @pytest.mark.parametrize("backend", BACKENDS)

@@ -1909,6 +2036,12 @@ def test_moe_layer(

2036 # moe_backend=flashinfer_trtllm / flashinfer_cutlass / flashinfer_cutedsl

2037 # (BF16, FP8 and NVFP4 paths), and VLLM_USE_FLASHINFER_MOE_INT4=1.

2038

2039 + # Repeated NIXL memory registration in this broad layer matrix fails on

2040 + # gfx950. NIXL EPLB is covered by dedicated tests, so use Gloo here to

2041 + # preserve the layer/EPLB coverage.

2042 + eplb_config = EPLBConfig(

2043 + communicator="torch_gloo" if enable_eplb and on_gfx950() else None

2044 + )

2045 parallel_config = ParallelConfig(

2046 pipeline_parallel_size=1,

2047 data_parallel_size=dp_size,

@@ -1916,6 +2049,7 @@ def test_moe_layer(

2049 enable_expert_parallel=use_ep,

2050 all2all_backend=backend,

2051 enable_eplb=enable_eplb,

2052 + eplb_config=eplb_config,

2053 )

2054

2055 compilation_config = CompilationConfig()

@@ -1955,19 +2089,42 @@ def test_moe_layer(

2089 ) as failure_report_file:

2090 failure_report_path = failure_report_file.name

2091

2092 + test_config_batches = [test_configs]

2093 + if current_platform.is_rocm() and backend == "deepep_low_latency":

2094 + # rocSHMEM cannot destroy one low-latency buffer and initialize a

2095 + # differently sized one in the same process. Use one worker lifetime

2096 + # per compatible buffer shape while preserving every matrix case.

2097 + test_config_batches = _group_deepep_ll_configs(test_configs)

2098 +

2099 + rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS

2100 + parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker

2101 + launch_failures: list[str] = []

2102 +

2103 try:

1959 - parallel_launch_with_config(

1960 - world_size,

1961 - _parallel_worker,

1962 - vllm_config,

1963 - None,

1964 - test_configs,

1965 - verbosity,

1966 - failure_report_path=failure_report_path,

1967 - )

2104 + for test_config_batch in test_config_batches:

2105 + report_size_before = os.path.getsize(failure_report_path)

2106 + try:

2107 + parallel_launch_with_config(

2108 + world_size,

2109 + parallel_worker,

2110 + vllm_config,

2111 + None,

2112 + test_config_batch,

2113 + verbosity,

2114 + failure_report_path=failure_report_path,

2115 + deep_ep_handle_keepalive=[] if rocm_deepep else None,

2116 + )

2117 + except Exception as ex:

2118 + # Normal subtest failures are already in the shared report.

2119 + # Preserve launcher/setup errors that occur before that write.

2120 + if os.path.getsize(failure_report_path) == report_size_before:

2121 + launch_failures.append(str(ex))

2122 +

2123 if os.path.getsize(failure_report_path) > 0:

2124 with open(failure_report_path, encoding="utf-8") as report_file:

1970 - pytest.fail(report_file.read())

2125 + launch_failures.insert(0, report_file.read())

2126 + if launch_failures:

2127 + pytest.fail("\n\n".join(launch_failures))

2128 finally:

2129 with suppress(FileNotFoundError):

2130 os.remove(failure_report_path)