obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-08 23:08:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

[MXFP8][ROCm] Fix MXFP8 MoE backend selection (#49747)

Browse Source 
Signed-off-by: Felix Marty <[email protected]>
This commit is contained in:
fxmarty-amd
2026-07-29 07:57:20 +08:00
committed by GitHub
parent e7f6a39db8
commit 5369f7b7b8
6 changed files with 244 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_mxfp8_aiter_backend_selection.py
+37 -7
View File
@@ -19,9 +19,17 @@ if not current_platform.is_rocm():
pytest.skip("This test can only run on ROCm.", allow_module_level=True)
from tests.kernels.moe.utils import make_dummy_moe_config # noqa: E402
from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
MoEActivation,
)
from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( # noqa: E402
_AITER_SWIGLU_ALPHA,
_AITER_SWIGLU_BETA,
AiterMxfp8Experts,
)
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( # noqa: E402
Mxfp8NativeTritonExperts,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import ( # noqa: E402
FusedMoEActivationFormat,
)
@@ -33,6 +41,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( # noqa: E402
_SUPPORTED_BACKENDS,
_mxfp8_backend_to_kernel_cls,
_select_kernel_cls,
select_mxfp8_moe_backend,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
kMxfp8Dynamic,
@@ -43,7 +52,17 @@ _AITER_MOD = "vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe"
def _config(ep_size: int = 1):
cfg = make_dummy_moe_config(num_experts=128, experts_per_token=4, hidden_dim=6144)
# AiterMxfp8Experts hardcodes SwiGLU-OAI: match its required activation and
# alpha/beta so is_supported_config doesn't reject the config on those grounds.
cfg = make_dummy_moe_config(
num_experts=128,
experts_per_token=4,
hidden_dim=6144,
activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
)
cfg = dataclasses.replace(
cfg, swiglu_alpha=_AITER_SWIGLU_ALPHA, swiglu_beta=_AITER_SWIGLU_BETA
)
if ep_size != 1:
cfg = dataclasses.replace(
cfg,
@@ -76,12 +95,6 @@ def test_aiter_mxfp8_registered():
]
def test_triton_selectable():
assert _BACKEND_NAME_MAP["triton"] is Fp8MoeBackend.TRITON_MXFP8
# Not auto-selected (only reachable explicitly), so FlyDSL still wins auto.
assert Fp8MoeBackend.TRITON_MXFP8 not in _SUPPORTED_BACKENDS
@pytest.mark.parametrize("ep_size", [1, 2])
def test_ep_supported(ep_size):
"""FlyDSL accepts both TP and EP: apply() forwards expert_map as expert_mask."""
@@ -133,3 +146,20 @@ def test_explicit_moe_backend_aiter():
pytest.raises(ValueError, match="flydsl package"),
):
_select_kernel_cls(Fp8MoeBackend.AITER_MXFP8, _config(1))
def test_gfx950_picks_aiter():
"""Auto-select on real ROCm hardware with flydsl usable -> FlyDSL wins."""
with _flydsl_installed(True):
backend, experts_cls = select_mxfp8_moe_backend(_config())
assert backend is Fp8MoeBackend.AITER_MXFP8
assert experts_cls is AiterMxfp8Experts
def test_gfx942_picks_triton():
"""flydsl unusable (e.g. gfx942, no FlyDSL support) -> native Triton
dot_scaled backend wins instead."""
with _flydsl_installed(False):
backend, experts_cls = select_mxfp8_moe_backend(_config())
assert backend is Fp8MoeBackend.TRITON_MXFP8
assert experts_cls is Mxfp8NativeTritonExperts
tests/models/quantization/test_mxfp8.py
+166
View File
@@ -17,8 +17,10 @@ diverse prompts from ``tests/prompts/example.txt``.
"""
import pytest
import torch
from tests.quantization.utils import is_quant_method_supported
from vllm.platforms import current_platform
from ..utils import check_logprobs_close
@@ -81,6 +83,170 @@ def test_mxfp8_logprobs(
)
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
)
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="AITER MXFP8 MoE backend is ROCm-only.",
)
@pytest.mark.quant_model
def test_mxfp8_aiter_requires_swigluoai_activation(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
select_mxfp8_moe_backend,
)
monkeypatch.setattr(
aiter_mxfp8_moe.AiterMxfp8Experts,
"_supports_current_device",
staticmethod(lambda: True),
)
monkeypatch.setattr(
aiter_mxfp8_moe,
"is_aiter_mxfp8_moe_available",
lambda: True,
)
config = FusedMoEConfig(
num_experts=8,
experts_per_token=2,
hidden_dim=256,
intermediate_size=256,
num_local_experts=8,
num_logical_experts=8,
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
activation=MoEActivation.SILU,
in_dtype=torch.bfloat16,
device="cuda",
routing_method=RoutingMethodType.Renormalize,
moe_backend="aiter",
)
with pytest.raises(ValueError, match="requires activation=swigluoai_uninterleave"):
select_mxfp8_moe_backend(config)
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
)
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="AITER MXFP8 MoE backend is ROCm-only.",
)
@pytest.mark.quant_model
def test_mxfp8_aiter_requires_swigluoai_params(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
select_mxfp8_moe_backend,
)
monkeypatch.setattr(
aiter_mxfp8_moe.AiterMxfp8Experts,
"_supports_current_device",
staticmethod(lambda: True),
)
monkeypatch.setattr(
aiter_mxfp8_moe,
"is_aiter_mxfp8_moe_available",
lambda: True,
)
config = FusedMoEConfig(
num_experts=8,
experts_per_token=2,
hidden_dim=256,
intermediate_size=256,
num_local_experts=8,
num_logical_experts=8,
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
in_dtype=torch.bfloat16,
device="cuda",
routing_method=RoutingMethodType.Renormalize,
moe_backend="aiter",
)
with pytest.raises(ValueError, match="hardcodes SwiGLU-OAI"):
select_mxfp8_moe_backend(config)
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
)
@pytest.mark.skipif(
not current_platform.is_rocm(),
reason="AITER MXFP8 MoE backend is ROCm-only.",
)
@pytest.mark.quant_model
def test_mxfp8_aiter_accepts_swigluoai_params(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.fused_moe.experts import aiter_mxfp8_moe
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
select_mxfp8_moe_backend,
)
monkeypatch.setattr(
aiter_mxfp8_moe.AiterMxfp8Experts,
"_supports_current_device",
staticmethod(lambda: True),
)
monkeypatch.setattr(
aiter_mxfp8_moe,
"is_aiter_mxfp8_moe_available",
lambda: True,
)
config = FusedMoEConfig(
num_experts=8,
experts_per_token=2,
hidden_dim=256,
intermediate_size=256,
num_local_experts=8,
num_logical_experts=8,
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
in_dtype=torch.bfloat16,
device="cuda",
routing_method=RoutingMethodType.Renormalize,
moe_backend="aiter",
swiglu_alpha=aiter_mxfp8_moe._AITER_SWIGLU_ALPHA,
swiglu_beta=aiter_mxfp8_moe._AITER_SWIGLU_BETA,
)
backend, experts_cls = select_mxfp8_moe_backend(config)
assert backend == Fp8MoeBackend.AITER_MXFP8
assert experts_cls is aiter_mxfp8_moe.AiterMxfp8Experts
@pytest.mark.skipif(
not is_quant_method_supported("mxfp8"),
reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp8_moe.py
+27
View File
@@ -6,10 +6,13 @@
``convert_to_fp8_moe_kernel_format``.
"""
import math
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
Mxfp8TritonExpertsBase,
)
@@ -17,6 +20,9 @@ from vllm.platforms import current_platform
logger = init_logger(__name__)
_AITER_SWIGLU_ALPHA = 1.702
_AITER_SWIGLU_BETA = 1.0
def is_aiter_mxfp8_moe_available() -> bool:
"""True when the FlyDSL MXFP8 MoE can run here: gfx950, the ``flydsl``
@@ -93,6 +99,27 @@ class AiterMxfp8Experts(Mxfp8TritonExpertsBase):
return False, (
"kernel requires the aiter flydsl package, which is not installed"
)
if (
is_supported
and moe_config.activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE
):
return False, (
"kernel hardcodes SwiGLU-OAI activation and requires "
f"activation={MoEActivation.SWIGLUOAI_UNINTERLEAVE.value}; "
f"got activation={moe_config.activation.value}"
)
if is_supported and (
moe_config.swiglu_alpha is None
or not math.isclose(float(moe_config.swiglu_alpha), _AITER_SWIGLU_ALPHA)
or moe_config.swiglu_beta is None
or not math.isclose(float(moe_config.swiglu_beta), _AITER_SWIGLU_BETA)
):
return False, (
"kernel hardcodes SwiGLU-OAI with "
f"alpha={_AITER_SWIGLU_ALPHA} and beta={_AITER_SWIGLU_BETA}; "
f"got swiglu_alpha={moe_config.swiglu_alpha} and "
f"swiglu_beta={moe_config.swiglu_beta}"
)
return is_supported, reason
def apply(
vllm/model_executor/layers/fused_moe/experts/mxfp8_emulation_moe.py
+2 -6
View File
@@ -107,17 +107,13 @@ class Mxfp8EmulationTritonExperts(Mxfp8TritonExpertsBase):
limit = self.quant_config.gemm1_clamp_limit
if limit is None:
raise ValueError("SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit")
alpha = self.quant_config.gemm1_alpha
alpha = 1.702 if alpha is None else float(alpha)
beta = self.quant_config.gemm1_beta
beta = 1.0 if beta is None else float(beta)
apply_moe_activation(
activation,
output,
input,
clamp_limit=float(limit),
alpha=alpha,
beta=beta,
alpha=self.gemm1_alpha,
beta=self.gemm1_beta,
)
return
super().activation(activation, output, input)
vllm/model_executor/layers/fused_moe/experts/mxfp8_native_moe.py
+3 -6
View File
@@ -362,10 +362,7 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
alpha = self.quant_config.gemm1_alpha
alpha = 1.702 if alpha is None else float(alpha)
beta = self.quant_config.gemm1_beta
beta = 1.0 if beta is None else float(beta)
# `self.gemm1_alpha` and `self.gemm1_beta`` are set by `TritonExperts.__init__`.
limit = self.quant_config.gemm1_clamp_limit
limit = None if limit is None else float(limit)
out = fused_moe_mxfp8_native(
@@ -376,8 +373,8 @@ class Mxfp8NativeTritonExperts(Mxfp8TritonExpertsBase):
self.w2_scale_val,
topk_weights,
topk_ids,
alpha=alpha,
beta=beta,
alpha=self.gemm1_alpha,
beta=self.gemm1_beta,
limit=limit,
global_num_experts=global_num_experts,
expert_map=expert_map,
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+9 -42
View File
@@ -12,10 +12,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kMxfp8Dynamic,
kMxfp8Static,
)
from vllm.platforms import current_platform
logger = init_logger(__name__)
# Ordered by priority.
_SUPPORTED_BACKENDS = (
Fp8MoeBackend.FLASHINFER_TRTLLM,
Fp8MoeBackend.DEEPGEMM,
@@ -26,6 +26,8 @@ _SUPPORTED_BACKENDS = (
# devices / no flydsl / EP it is skipped and native is used.
Fp8MoeBackend.AITER_MXFP8,
Fp8MoeBackend.HUMMING,
Fp8MoeBackend.TRITON_MXFP8,
Fp8MoeBackend.EMULATION,
)
_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
@@ -61,15 +63,12 @@ def _mxfp8_backend_to_kernel_cls(
return [AiterMxfp8Experts]
if backend == Fp8MoeBackend.TRITON_MXFP8:
# Explicit ``--moe-backend triton``: the Triton mxfp8 path, i.e.
# dot_scaled on MX-capable HW (gfx950) and BF16 emulation otherwise.
# Mirrors the ROCm auto-fallback in ``_select_rocm_mxfp8_backend``.
if current_platform.supports_mx():
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
Mxfp8NativeTritonExperts,
)
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
Mxfp8NativeTritonExperts,
)
return [Mxfp8NativeTritonExperts]
return [Mxfp8NativeTritonExperts]
if backend == Fp8MoeBackend.EMULATION:
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
Mxfp8EmulationTritonExperts,
)
@@ -105,35 +104,6 @@ def _select_kernel_cls(
)
def _select_rocm_mxfp8_backend() -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
"""ROCm fallback when no auto-selected MXFP8 backend is available.
The aiter FlyDSL backend (``AITER_MXFP8``) is auto-picked earlier by
``select_mxfp8_moe_backend`` via ``_SUPPORTED_BACKENDS`` when usable, or
explicitly via ``--moe-backend aiter``; this fallback handles the rest
(native dot_scaled on gfx950, else BF16 emulation).
"""
if current_platform.supports_mx():
from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import (
Mxfp8NativeTritonExperts,
)
logger.info_once("Using native CDNA4 (gfx950) MXFP8 dot_scaled MoE backend.")
return Fp8MoeBackend.TRITON_MXFP8, Mxfp8NativeTritonExperts
from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import (
Mxfp8EmulationTritonExperts,
)
logger.info_once(
"No native MXFP8 MoE backend available on this device; "
"MXFP8 weights will be dequantized to BF16 once at load time and the "
"MoE will run in BF16 (no per-step dequant)."
)
return Fp8MoeBackend.EMULATION, Mxfp8EmulationTritonExperts
def select_mxfp8_moe_backend(
config: FusedMoEConfig,
) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:
@@ -167,8 +137,5 @@ def select_mxfp8_moe_backend(
logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
return backend, experts_cls
# simplify the logic for rocm, refactor later when more backends are supported
if current_platform.is_rocm():
return _select_rocm_mxfp8_backend()
# TODO: add debug log with reason.
raise ValueError("No MXFP8 MoE backends available.")