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

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

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Andreas Karatzas
2026-08-21 13:15:25 -07:00
committed by GitHub
parent c0ff33404b
commit 88eb946cb1
2 changed files with 174 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
-53
View File
@@ -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
+174 -17
View File
@@ -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)