709 lines
23 KiB
Python
709 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import copy
|
|
import inspect
|
|
import textwrap
|
|
import traceback
|
|
from itertools import product
|
|
from typing import Any
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import AiterExperts
|
|
from vllm.platforms import current_platform
|
|
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
|
|
from vllm.utils.import_utils import has_aiter, has_deep_ep, has_deep_gemm
|
|
from vllm.utils.torch_utils import set_random_seed
|
|
from vllm.v1.worker.workspace import init_workspace_manager
|
|
|
|
from .modular_kernel_tools.common import (
|
|
Config,
|
|
RankTensors,
|
|
WeightTensors,
|
|
reference_moe_impl,
|
|
run_modular_kernel,
|
|
)
|
|
from .modular_kernel_tools.mk_objects import (
|
|
MK_FUSED_EXPERT_TYPES,
|
|
MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
|
|
MK_QUANT_CONFIGS,
|
|
MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES,
|
|
TestMoEQuantConfig,
|
|
expert_info,
|
|
)
|
|
from .modular_kernel_tools.parallel_utils import (
|
|
ProcessGroupInfo,
|
|
parallel_launch_with_config,
|
|
)
|
|
from .utils import check_accuracy
|
|
|
|
has_any_multi_gpu_package = (
|
|
has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe()
|
|
)
|
|
|
|
meets_multi_gpu_requirements = pytest.mark.skipif(
|
|
not has_any_multi_gpu_package,
|
|
reason="Requires deep_ep or deep_gemm or flashinfer packages",
|
|
)
|
|
|
|
|
|
def format_result(verbose, msg, ex=None):
|
|
if ex is not None:
|
|
x = str(ex)
|
|
newx = x.strip(" \n\t")[:16]
|
|
if len(newx) < len(x):
|
|
newx = newx + " ..."
|
|
|
|
prefix = "E\t"
|
|
print(f"{textwrap.indent(traceback.format_exc(), prefix)}")
|
|
print(f"FAILED {msg} - {newx}\n")
|
|
elif verbose:
|
|
print(f"PASSED {msg}")
|
|
else:
|
|
print(".", end="")
|
|
|
|
|
|
def assert_aiter_quant_scheme_case(config: Config) -> None:
|
|
"""Make the AITER (weight_quant_key, activation_quant_key) pair this
|
|
config exercises explicit, instead of AiterExperts being reached only
|
|
indirectly through the general quant-config sweep.
|
|
See https://github.com/vllm-project/vllm/issues/54966."""
|
|
fe_cls = config.fused_experts_type
|
|
if fe_cls is not AiterExperts:
|
|
return
|
|
|
|
if config.quant_config is None:
|
|
w_key, a_key = None, None
|
|
else:
|
|
w_key, a_key = config.fp8_quant_key_pair()
|
|
|
|
assert fe_cls._supports_quant_scheme(w_key, a_key), (
|
|
f"AITER case (weight_key={w_key}, activation_key={a_key}) reached "
|
|
"the modular-kernel harness, but AiterExperts._supports_quant_scheme "
|
|
"does not declare it supported."
|
|
)
|
|
print(f"[AITER case] weight_key={w_key}, activation_key={a_key}")
|
|
|
|
|
|
def assert_aiter_activation_case(config: Config) -> None:
|
|
"""Make the AITER activation this config exercises explicit, instead of
|
|
AiterExperts being reached only indirectly through the general
|
|
activation sweep. See https://github.com/vllm-project/vllm/issues/54966."""
|
|
fe_cls = config.fused_experts_type
|
|
if fe_cls is not AiterExperts:
|
|
return
|
|
|
|
assert fe_cls._supports_activation(config.activation), (
|
|
f"AITER case activation={config.activation} reached the "
|
|
"modular-kernel harness, but AiterExperts._supports_activation "
|
|
"does not declare it supported."
|
|
)
|
|
print(f"[AITER case] activation={config.activation}")
|
|
|
|
|
|
def rank_worker(
|
|
pgi: ProcessGroupInfo,
|
|
vllm_config: VllmConfig,
|
|
cpu_group,
|
|
base_config: Config,
|
|
weights: WeightTensors,
|
|
verbose: bool,
|
|
):
|
|
# Initialize workspace manager in child process
|
|
device = torch.device(f"cuda:{pgi.local_rank}")
|
|
init_workspace_manager(device)
|
|
|
|
set_random_seed(pgi.rank)
|
|
|
|
# get weights to this device
|
|
weights.to_current_device()
|
|
|
|
Ms = base_config.Ms
|
|
assert isinstance(Ms, list)
|
|
TOPKs = base_config.topks
|
|
assert isinstance(TOPKs, list)
|
|
|
|
exceptions = []
|
|
count = 0
|
|
|
|
for m, topk in product(Ms, TOPKs):
|
|
# override m and topk
|
|
config = copy.deepcopy(base_config)
|
|
config.Ms = m
|
|
config.topks = topk
|
|
|
|
try:
|
|
print(f"Running[{pgi.rank}]: m={m}, topk={topk} ...")
|
|
count = count + 1
|
|
|
|
# inputs for rank
|
|
rank_tensors = RankTensors.make(config, pgi)
|
|
|
|
# Skip unsupported: AITER block-scaled MoE does not
|
|
# support apply_router_weight_on_input (topk=1 path).
|
|
# https://github.com/ROCm/aiter/issues/2418
|
|
if (
|
|
topk == 1
|
|
and config.supports_apply_weight_on_input()
|
|
and config.fused_experts_type is AiterExperts
|
|
and config.quant_block_shape is not None
|
|
):
|
|
print(
|
|
f"Skipping[{pgi.rank}]: m={m}, topk={topk}"
|
|
" (AITER block-scaled + weight-on-input,"
|
|
" https://github.com/ROCm/aiter/issues/2418)"
|
|
)
|
|
count -= 1
|
|
continue
|
|
|
|
# Skip unsupported: AITER x DeepEP-HT/Mori dispatch crashes
|
|
# with an illegal memory access at world_size>1 on gfx942.
|
|
# https://github.com/vllm-project/vllm/issues/57029
|
|
if (
|
|
config.world_size > 1
|
|
and config.fused_experts_type is AiterExperts
|
|
and getattr(config.prepare_finalize_type, "__name__", "")
|
|
in ("DeepEPHTPrepareAndFinalize", "MoriPrepareAndFinalize")
|
|
):
|
|
print(
|
|
f"Skipping[{pgi.rank}]: m={m}, topk={topk}"
|
|
" (AITER x DeepEP-HT/Mori illegal memory access,"
|
|
" https://github.com/vllm-project/vllm/issues/57029)"
|
|
)
|
|
count -= 1
|
|
continue
|
|
|
|
assert_aiter_quant_scheme_case(config)
|
|
assert_aiter_activation_case(config)
|
|
|
|
# modular kernel out
|
|
mk_out = run_modular_kernel(pgi, vllm_config, config, weights, rank_tensors)
|
|
|
|
with set_current_vllm_config(vllm_config):
|
|
ref_out = reference_moe_impl(config, weights, rank_tensors)
|
|
|
|
if config.quant_dtype == "nvfp4":
|
|
atol = 1e-1 if config.K < 4096 else 2e-1
|
|
rtol = 1e-1 if config.K < 4096 else 2e-1
|
|
else:
|
|
atol = 3e-2
|
|
rtol = 3e-2
|
|
|
|
# On ROCm, AITER FP8 fused MoE uses hardware FP8
|
|
# dot-product which can produce slightly larger error
|
|
# than dequant+f32 matmul at FP8 representable-value
|
|
# boundaries. Allow a small percentage of elements to
|
|
# exceed the base tolerance by a bounded margin.
|
|
# https://github.com/ROCm/aiter/issues/2421
|
|
from vllm.platforms import current_platform as _cp
|
|
|
|
is_aiter_fp8 = (
|
|
_cp.is_rocm()
|
|
and config.fused_experts_type is AiterExperts
|
|
and config.quant_config is not None
|
|
)
|
|
if is_aiter_fp8:
|
|
check_accuracy(ref_out, mk_out, atol=atol, rtol=rtol, percent=0.9)
|
|
else:
|
|
torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)
|
|
format_result(verbose, config.describe())
|
|
except Exception as ex:
|
|
format_result(verbose, config.describe(), ex)
|
|
exceptions.append(ex)
|
|
|
|
if len(exceptions) > 0:
|
|
raise RuntimeError(
|
|
f"{len(exceptions)} of {count} tests failed in child process, "
|
|
f"rank={pgi.rank}."
|
|
)
|
|
else:
|
|
print(f"{count} of {count} tests passed in child process, rank={pgi.rank}.")
|
|
|
|
|
|
def run(config: Config, verbose: bool):
|
|
assert config.is_valid()[0]
|
|
assert not is_nyi_config(config)
|
|
|
|
weights: WeightTensors = WeightTensors.make(config)
|
|
|
|
vllm_config, env_dict = config.make_env_data()
|
|
parallel_launch_with_config(
|
|
config.world_size, rank_worker, vllm_config, env_dict, config, weights, verbose
|
|
)
|
|
|
|
|
|
Ms = [32, 64]
|
|
# hidden sizes, making this too large will cause fp4 tests to fail.
|
|
# Also needs to be a multiple of 1024 for deep_gemm.
|
|
Ks = [2048]
|
|
Ns = [1024]
|
|
TOPKs = [4, 1]
|
|
Es = [32]
|
|
DTYPEs = [torch.bfloat16]
|
|
MK_ACTIVATIONS = [
|
|
MoEActivation.SILU,
|
|
MoEActivation.GELU,
|
|
]
|
|
|
|
|
|
def is_nyi_config(config: Config) -> bool:
|
|
# We know these configs to be legitimate. but still fail.
|
|
info = expert_info(config.fused_experts_type)
|
|
if info.needs_matching_quant:
|
|
# The triton kernels expect both per-act-token-quant and
|
|
# per-out-ch-quant or neither.
|
|
unsupported_quant_config = (
|
|
config.is_per_act_token_quant + config.is_per_out_ch_quant
|
|
) == 1
|
|
if unsupported_quant_config:
|
|
return True
|
|
|
|
if config.activation != MoEActivation.SILU:
|
|
if config.fused_experts_type is not AiterExperts:
|
|
return True # AITER-only for this axis, for now
|
|
if config.quant_dtype is not None:
|
|
return True # unquantized-only for this axis, for now
|
|
|
|
return False
|
|
|
|
|
|
def generate_valid_test_cases(
|
|
world_size: int, prepare_finalize_types
|
|
) -> list[tuple[Any, ...]]:
|
|
cases = []
|
|
total = 0
|
|
|
|
for k, n, e, dtype, quant_config, activation, combination in product(
|
|
Ks,
|
|
Ns,
|
|
Es,
|
|
DTYPEs,
|
|
MK_QUANT_CONFIGS,
|
|
MK_ACTIVATIONS,
|
|
product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
|
|
):
|
|
total = total + 1
|
|
|
|
config = Config(
|
|
Ms=Ms,
|
|
K=k,
|
|
N=n,
|
|
E=e,
|
|
topks=TOPKs,
|
|
dtype=dtype,
|
|
quant_config=quant_config,
|
|
activation=activation,
|
|
prepare_finalize_type=combination[0],
|
|
fused_experts_type=combination[1],
|
|
world_size=world_size,
|
|
)
|
|
|
|
# TODO(bnell): figure out how to get verbose flag here.
|
|
verbose = False # pytestconfig.getoption('verbose') > 0
|
|
|
|
valid, reason = config.is_valid()
|
|
|
|
if not valid:
|
|
if verbose:
|
|
print(f"Test config {config} is not valid: {reason}")
|
|
continue
|
|
|
|
if is_nyi_config(config):
|
|
if verbose:
|
|
print(f"Test config {config} is nyi.")
|
|
continue
|
|
|
|
cases.append(
|
|
(
|
|
k,
|
|
n,
|
|
e,
|
|
dtype,
|
|
quant_config,
|
|
activation,
|
|
combination[0],
|
|
combination[1],
|
|
world_size,
|
|
)
|
|
)
|
|
|
|
print(f"{len(cases)} of {total} valid configs generated.")
|
|
|
|
return cases
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"k,n,e,dtype,quant_config,activation,"
|
|
"prepare_finalize_type,fused_experts_type,world_size",
|
|
generate_valid_test_cases(
|
|
world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
|
|
),
|
|
)
|
|
@meets_multi_gpu_requirements
|
|
def test_modular_kernel_combinations_multigpu(
|
|
k: int,
|
|
n: int,
|
|
e: int,
|
|
dtype: torch.dtype,
|
|
quant_config: TestMoEQuantConfig | None,
|
|
activation: MoEActivation,
|
|
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
|
|
fused_experts_type: mk.FusedMoEExperts,
|
|
world_size: int,
|
|
pytestconfig,
|
|
):
|
|
if current_platform.device_count() < world_size:
|
|
pytest.skip(
|
|
f"Not enough GPUs available to run, got "
|
|
f"{current_platform.device_count()} expected "
|
|
f"{world_size}."
|
|
)
|
|
|
|
config = Config(
|
|
Ms=Ms,
|
|
K=k,
|
|
N=n,
|
|
E=e,
|
|
topks=TOPKs,
|
|
dtype=dtype,
|
|
quant_config=quant_config,
|
|
activation=activation,
|
|
prepare_finalize_type=prepare_finalize_type,
|
|
fused_experts_type=fused_experts_type,
|
|
world_size=world_size,
|
|
)
|
|
verbosity = pytestconfig.getoption("verbose")
|
|
run(config, verbosity > 0)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"k,n,e,dtype,quant_config,activation,"
|
|
"prepare_finalize_type,fused_experts_type,world_size",
|
|
generate_valid_test_cases(
|
|
world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
|
|
),
|
|
)
|
|
def test_modular_kernel_combinations_singlegpu(
|
|
k: int,
|
|
n: int,
|
|
e: int,
|
|
dtype: torch.dtype,
|
|
quant_config: TestMoEQuantConfig | None,
|
|
activation: MoEActivation,
|
|
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
|
|
fused_experts_type: mk.FusedMoEExperts,
|
|
world_size: int,
|
|
pytestconfig,
|
|
workspace_init,
|
|
):
|
|
"""Note: float8_e4m3fn is not supported on CUDA architecture < 89,
|
|
and those tests will be skipped on unsupported hardware."""
|
|
config = Config(
|
|
Ms=Ms,
|
|
K=k,
|
|
N=n,
|
|
E=e,
|
|
topks=TOPKs,
|
|
dtype=dtype,
|
|
quant_config=quant_config,
|
|
activation=activation,
|
|
prepare_finalize_type=prepare_finalize_type,
|
|
fused_experts_type=fused_experts_type,
|
|
world_size=world_size,
|
|
)
|
|
|
|
if (
|
|
quant_config is not None and quant_config.quant_dtype == torch.float8_e4m3fn
|
|
) and not current_platform.has_device_capability(89):
|
|
pytest.skip(
|
|
"Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
|
|
)
|
|
verbosity = pytestconfig.getoption("verbose")
|
|
run(config, verbosity > 0)
|
|
|
|
|
|
# AITER sorting-backend dispatch env-var matrix (issue #54966 step 3) -------
|
|
#
|
|
# AITER_USE_CK_MOE_SORTING / AITER_USE_FLYDSL_MOE_SORTING are aiter globals
|
|
# read once at import time, so each combo needs a fresh child process (see
|
|
# parallel_launch_with_config) rather than in-process monkeypatching.
|
|
#
|
|
# AITER_MOE_SORT_BACKEND is not covered: it only matters when output_aux=True,
|
|
# which vLLM only sets on the MXFP4 path -- not yet wired into AiterExperts.
|
|
|
|
|
|
def _aiter_sort_backend_hooks_available() -> bool:
|
|
"""Whether the installed aiter build exposes the sort-backend hooks that
|
|
AITER_USE_CK_MOE_SORTING / AITER_USE_FLYDSL_MOE_SORTING select between.
|
|
"""
|
|
try:
|
|
import aiter.fused_moe as aiter_fused_moe
|
|
except ImportError:
|
|
return False
|
|
|
|
required = (
|
|
"_moe_sorting_impl",
|
|
"_flydsl_moe_sorting",
|
|
"_USE_CK_MOE_SORTING",
|
|
"_USE_FLYDSL_MOE_SORTING",
|
|
)
|
|
return all(hasattr(aiter_fused_moe, name) for name in required)
|
|
|
|
|
|
# Deliberately does NOT call _aiter_sort_backend_hooks_available() here: this
|
|
# marker is evaluated at module-import time, which also happens inside each
|
|
# spawned child (to unpickle its worker) -- before that case's env vars are
|
|
# applied, permanently freezing aiter's sort-backend globals. The hooks check
|
|
# instead runs inside the test body below (parent process only).
|
|
require_aiter_moe = pytest.mark.skipif(
|
|
not (
|
|
current_platform.is_rocm()
|
|
and has_aiter()
|
|
and rocm_aiter_ops.is_fused_moe_enabled()
|
|
),
|
|
reason=(
|
|
"AITER MoE sorting-backend dispatch needs ROCm + AITER, with "
|
|
"VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1 set before "
|
|
"the test process starts."
|
|
),
|
|
)
|
|
|
|
# (AITER_USE_CK_MOE_SORTING, AITER_USE_FLYDSL_MOE_SORTING, expected_backend)
|
|
AITER_SORTING_BACKEND_ENV_CASES = [
|
|
(1, 0, "ck"),
|
|
(1, 1, "ck"), # CK wins over FlyDSL even when both are requested.
|
|
(0, 1, "flydsl"),
|
|
(0, 0, "opus"), # default when neither flag is set.
|
|
]
|
|
|
|
|
|
def _aiter_sorting_backend_worker(
|
|
pgi: ProcessGroupInfo,
|
|
vllm_config: VllmConfig,
|
|
cpu_group,
|
|
config: Config,
|
|
weights: WeightTensors,
|
|
verbose: bool,
|
|
ck: int,
|
|
flydsl: int,
|
|
expected_backend: str,
|
|
):
|
|
# Imported lazily in the spawned child, after env vars are applied, so
|
|
# aiter's import-time globals pick up this case's env.
|
|
import aiter.fused_moe as aiter_fused_moe
|
|
|
|
# Assert the env vars actually stuck in aiter's import-time globals,
|
|
# rather than inferring it indirectly from which sorting function fires.
|
|
assert bool(ck) == aiter_fused_moe._USE_CK_MOE_SORTING, (
|
|
f"aiter.fused_moe._USE_CK_MOE_SORTING={aiter_fused_moe._USE_CK_MOE_SORTING} "
|
|
f"but AITER_USE_CK_MOE_SORTING={ck} was set before this process started -- "
|
|
"the env var never reached aiter's import-time globals."
|
|
)
|
|
assert bool(flydsl) == aiter_fused_moe._USE_FLYDSL_MOE_SORTING, (
|
|
"aiter.fused_moe._USE_FLYDSL_MOE_SORTING="
|
|
f"{aiter_fused_moe._USE_FLYDSL_MOE_SORTING} but AITER_USE_FLYDSL_MOE_SORTING="
|
|
f"{flydsl} was set before this process started -- the env var never "
|
|
"reached aiter's import-time globals."
|
|
)
|
|
|
|
# Captured before patching, since the patched attribute would resolve to
|
|
# the mock's own (*args, **kwargs) signature instead of the real one.
|
|
original_moe_sorting_impl = aiter_fused_moe._moe_sorting_impl
|
|
|
|
with (
|
|
mock.patch.object(
|
|
aiter_fused_moe,
|
|
"_moe_sorting_impl",
|
|
wraps=original_moe_sorting_impl,
|
|
) as sorting_impl_mock,
|
|
mock.patch.object(
|
|
aiter_fused_moe,
|
|
"_flydsl_moe_sorting",
|
|
wraps=aiter_fused_moe._flydsl_moe_sorting,
|
|
) as flydsl_mock,
|
|
):
|
|
# Reuses the real correctness/accuracy checking rank_worker already
|
|
# does for every other AiterExperts combo, instead of duplicating it.
|
|
rank_worker(pgi, vllm_config, cpu_group, config, weights, verbose)
|
|
|
|
if expected_backend == "flydsl":
|
|
assert flydsl_mock.call_count > 0, "Expected FlyDSL sorting to fire."
|
|
assert sorting_impl_mock.call_count == 0, (
|
|
"FlyDSL was requested and eligible, but the opus/CK sorting "
|
|
"path fired instead -- a silent fallback."
|
|
)
|
|
else:
|
|
assert flydsl_mock.call_count == 0, (
|
|
f"Expected the '{expected_backend}' sorting path, but FlyDSL "
|
|
"fired instead."
|
|
)
|
|
assert sorting_impl_mock.call_count > 0, (
|
|
f"Expected the '{expected_backend}' sorting path to fire."
|
|
)
|
|
# Bind by signature so this works regardless of whether use_opus
|
|
# is passed positionally or by keyword.
|
|
call = sorting_impl_mock.call_args
|
|
bound = inspect.signature(original_moe_sorting_impl).bind(
|
|
*call.args, **call.kwargs
|
|
)
|
|
use_opus = bound.arguments["use_opus"]
|
|
assert use_opus == (expected_backend == "opus"), (
|
|
f"Expected use_opus={expected_backend == 'opus'} for the "
|
|
f"'{expected_backend}' sorting path, got use_opus={use_opus}."
|
|
)
|
|
|
|
|
|
@require_aiter_moe
|
|
@pytest.mark.parametrize("ck,flydsl,expected_backend", AITER_SORTING_BACKEND_ENV_CASES)
|
|
def test_aiter_moe_sorting_backend_dispatch_env_matrix(
|
|
ck: int, flydsl: int, expected_backend: str
|
|
):
|
|
"""See https://github.com/vllm-project/vllm/issues/54966 ("Test AITER
|
|
backend and dispatch settings")."""
|
|
if not _aiter_sort_backend_hooks_available():
|
|
pytest.skip(
|
|
"Installed aiter build predates the sort-backend hooks "
|
|
"(_moe_sorting_impl, _flydsl_moe_sorting, _USE_CK_MOE_SORTING, "
|
|
"_USE_FLYDSL_MOE_SORTING)."
|
|
)
|
|
|
|
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
|
MoEPrepareAndFinalizeNoDPEPModular,
|
|
)
|
|
|
|
config = Config(
|
|
Ms=[32],
|
|
K=2048,
|
|
N=1024,
|
|
E=32,
|
|
topks=[4],
|
|
dtype=torch.bfloat16,
|
|
quant_config=None,
|
|
prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
|
|
fused_experts_type=AiterExperts,
|
|
world_size=1,
|
|
)
|
|
assert config.is_valid()[0]
|
|
|
|
weights = WeightTensors.make(config)
|
|
vllm_config, env_dict = config.make_env_data()
|
|
env_dict = {
|
|
**env_dict,
|
|
"AITER_USE_CK_MOE_SORTING": str(ck),
|
|
"AITER_USE_FLYDSL_MOE_SORTING": str(flydsl),
|
|
}
|
|
|
|
parallel_launch_with_config(
|
|
config.world_size,
|
|
_aiter_sorting_backend_worker,
|
|
vllm_config,
|
|
env_dict,
|
|
config,
|
|
weights,
|
|
False,
|
|
ck,
|
|
flydsl,
|
|
expected_backend,
|
|
)
|
|
|
|
|
|
def _aiter_dispatch_policy_worker(
|
|
pgi: ProcessGroupInfo,
|
|
vllm_config: VllmConfig,
|
|
cpu_group,
|
|
config: Config,
|
|
weights: WeightTensors,
|
|
verbose: bool,
|
|
dispatch_policy: int,
|
|
):
|
|
# Drives the real AiterExperts.apply() -> rocm_aiter_fused_experts() call
|
|
# through the modular kernel, so it also catches a regression in apply()'s
|
|
# own forwarding (unlike test_rocm_aiter_moe.py's version, which supplies
|
|
# the policy value directly to rocm_aiter_fused_experts()).
|
|
with mock.patch.object(
|
|
rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe
|
|
) as fused_moe_mock:
|
|
rank_worker(pgi, vllm_config, cpu_group, config, weights, verbose)
|
|
|
|
assert fused_moe_mock.call_count > 0, (
|
|
"Expected AiterExperts.apply() to call rocm_aiter_ops.fused_moe."
|
|
)
|
|
for call in fused_moe_mock.call_args_list:
|
|
forwarded = call.kwargs["moe_sorting_dispatch_policy"]
|
|
assert forwarded == dispatch_policy, (
|
|
f"AiterExperts.apply() forwarded moe_sorting_dispatch_policy="
|
|
f"{forwarded}, but VLLM_ROCM_AITER_MOE_DISPATCH_POLICY="
|
|
f"{dispatch_policy} was set before this process started."
|
|
)
|
|
|
|
|
|
@require_aiter_moe
|
|
@pytest.mark.parametrize("dispatch_policy", [0, 1, 2])
|
|
def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int):
|
|
"""See https://github.com/vllm-project/vllm/issues/54966 ("Test AITER
|
|
backend and dispatch settings").
|
|
|
|
Complements test_aiter_moe_dispatch_policy_forwarded_to_fused_moe in
|
|
test_rocm_aiter_moe.py by driving the real AiterExperts.apply() call
|
|
instead of supplying moe_sorting_dispatch_policy directly.
|
|
"""
|
|
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
|
MoEPrepareAndFinalizeNoDPEPModular,
|
|
)
|
|
|
|
config = Config(
|
|
Ms=[32],
|
|
K=2048,
|
|
N=1024,
|
|
E=32,
|
|
topks=[4],
|
|
dtype=torch.bfloat16,
|
|
quant_config=None,
|
|
prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular,
|
|
fused_experts_type=AiterExperts,
|
|
world_size=1,
|
|
)
|
|
assert config.is_valid()[0]
|
|
|
|
weights = WeightTensors.make(config)
|
|
vllm_config, env_dict = config.make_env_data()
|
|
env_dict = {
|
|
**env_dict,
|
|
"VLLM_ROCM_AITER_MOE_DISPATCH_POLICY": str(dispatch_policy),
|
|
}
|
|
|
|
parallel_launch_with_config(
|
|
config.world_size,
|
|
_aiter_dispatch_policy_worker,
|
|
vllm_config,
|
|
env_dict,
|
|
config,
|
|
weights,
|
|
False,
|
|
dispatch_policy,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Ability to test individual PrepareAndFinalize and FusedExperts combination
|
|
from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser
|
|
|
|
parser = make_config_arg_parser(
|
|
description=(
|
|
"Run single prepare-finalize & fused-experts combination test"
|
|
"Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations "
|
|
"--pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts"
|
|
)
|
|
)
|
|
args = parser.parse_args()
|
|
config = make_config(args)
|
|
|
|
run(config, True)
|