forked from Karylab-cklius/vllm
91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950).
|
|
|
|
These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.model_executor.layers.fused_moe.config import (
|
|
FusedMoEConfig,
|
|
FusedMoEParallelConfig,
|
|
RoutingMethodType,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
|
|
Mxfp4MoeBackend,
|
|
select_mxfp4_moe_backend,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
kMxfp4Dynamic,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
|
|
ROCM_AVAILABLE = current_platform.is_rocm()
|
|
ROCM_GFX950 = False
|
|
ROCM_AITER_AVAILABLE = False
|
|
|
|
if ROCM_AVAILABLE:
|
|
from vllm._aiter_ops import rocm_aiter_ops
|
|
from vllm.platforms.rocm import on_gfx950
|
|
|
|
ROCM_GFX950 = on_gfx950()
|
|
ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()
|
|
|
|
|
|
def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
|
|
return FusedMoEConfig(
|
|
num_experts=8,
|
|
experts_per_token=2,
|
|
hidden_dim=256,
|
|
intermediate_size_per_partition=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=moe_backend,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
|
|
@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")
|
|
def test_w4a4_dispatches_to_aiter():
|
|
"""With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""
|
|
config = _make_w4a4_moe_config()
|
|
backend, experts_cls = select_mxfp4_moe_backend(
|
|
config, activation_key=kMxfp4Dynamic
|
|
)
|
|
assert backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4
|
|
assert experts_cls is not None
|
|
|
|
|
|
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
|
|
@pytest.mark.skipif(
|
|
ROCM_AITER_AVAILABLE,
|
|
reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)",
|
|
)
|
|
def test_w4a4_raises_without_aiter_and_no_moe_backend():
|
|
"""Without AITER and no --moe-backend, raises NotImplementedError
|
|
with hint to use --moe-backend emulation."""
|
|
config = _make_w4a4_moe_config()
|
|
with pytest.raises(NotImplementedError, match="--moe-backend emulation"):
|
|
select_mxfp4_moe_backend(config, activation_key=kMxfp4Dynamic)
|
|
|
|
|
|
@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")
|
|
def test_w4a4_dispatches_to_emulation_with_moe_backend():
|
|
"""With --moe-backend emulation, W4A4 selects EMULATION."""
|
|
config = _make_w4a4_moe_config(moe_backend="emulation")
|
|
backend, experts_cls = select_mxfp4_moe_backend(
|
|
config, activation_key=kMxfp4Dynamic
|
|
)
|
|
assert backend == Mxfp4MoeBackend.EMULATION
|
|
assert experts_cls is not None
|