wehub-resource-sync
405 lines
16 KiB
Python
405 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from tests.kernels.moe.utils import make_dummy_moe_config
|
|
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
|
|
from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
|
|
UnquantizedMoeBackend,
|
|
select_unquantized_moe_backend,
|
|
)
|
|
from vllm.platforms import current_platform
|
|
|
|
skipif_not_cuda_rocm = pytest.mark.skipif(
|
|
not (current_platform.is_cuda() or current_platform.is_rocm()),
|
|
reason="Only supported on CUDA/ROCm platforms.",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"platform_method,expected_backend",
|
|
[
|
|
("is_cuda", UnquantizedMoeBackend.TRITON), # Default CUDA without FlashInfer
|
|
("is_rocm", UnquantizedMoeBackend.TRITON), # ROCm without AITER
|
|
("is_cpu", UnquantizedMoeBackend.CPU),
|
|
("is_xpu", UnquantizedMoeBackend.XPU),
|
|
("is_tpu", UnquantizedMoeBackend.TPU),
|
|
("is_out_of_tree", UnquantizedMoeBackend.OOT),
|
|
],
|
|
)
|
|
@patch(
|
|
"vllm.utils.flashinfer.has_flashinfer",
|
|
return_value=False,
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops.is_fused_moe_enabled",
|
|
return_value=False,
|
|
)
|
|
def test_select_default_backend_by_platform(
|
|
mock_aiter_enabled,
|
|
mock_has_flashinfer,
|
|
monkeypatch,
|
|
platform_method,
|
|
expected_backend,
|
|
):
|
|
"""Test default backend selection per platform with all optional
|
|
accelerators (FlashInfer, AITER) disabled."""
|
|
with patch(
|
|
"vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
|
|
) as mock_platform:
|
|
# Set all platform checks to False
|
|
mock_platform.is_cuda.return_value = False
|
|
mock_platform.is_rocm.return_value = False
|
|
mock_platform.is_cpu.return_value = False
|
|
mock_platform.is_xpu.return_value = False
|
|
mock_platform.is_tpu.return_value = False
|
|
mock_platform.is_out_of_tree.return_value = False
|
|
|
|
# Set only the specified platform to True
|
|
getattr(mock_platform, platform_method).return_value = True
|
|
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=False),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(current_platform, platform_method, return_value=True),
|
|
):
|
|
moe_config = make_dummy_moe_config()
|
|
selected_backend, expert_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == expected_backend
|
|
if expected_backend in [
|
|
UnquantizedMoeBackend.CPU,
|
|
UnquantizedMoeBackend.OOT,
|
|
UnquantizedMoeBackend.TPU,
|
|
]:
|
|
assert expert_cls is None
|
|
else:
|
|
assert expert_cls is not None
|
|
|
|
|
|
@patch(
|
|
"vllm.utils.flashinfer.has_flashinfer",
|
|
return_value=False,
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.oracle.unquantized.rocm_aiter_ops.is_fused_moe_enabled",
|
|
return_value=True,
|
|
)
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_rocm(), reason="ROCm-specific backend selection test"
|
|
)
|
|
def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
|
|
"""Test ROCm backend selection when AITER is available."""
|
|
with patch(
|
|
"vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
|
|
) as mock_platform:
|
|
mock_platform.is_cuda.return_value = False
|
|
mock_platform.is_rocm.return_value = True
|
|
mock_platform.is_cpu.return_value = False
|
|
mock_platform.is_xpu.return_value = False
|
|
mock_platform.is_tpu.return_value = False
|
|
mock_platform.is_out_of_tree.return_value = False
|
|
|
|
moe_config = make_dummy_moe_config()
|
|
selected_backend, expert_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config,
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.AITER
|
|
assert expert_cls is not None
|
|
|
|
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
|
|
return_value=(True, None),
|
|
)
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
|
)
|
|
def test_select_cuda_flashinfer_trtllm_backend(
|
|
mock_is_supported_trtllm_monolithic,
|
|
):
|
|
"""Test CUDA backend selection when FlashInfer TRTLLM is available and enabled."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(current_platform, "has_device_capability", return_value=True),
|
|
):
|
|
moe_config = make_dummy_moe_config()
|
|
moe_config.moe_backend = "flashinfer_trtllm"
|
|
# TRTLLM requires EP and does not support DP
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.use_dp = False
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
|
|
assert experts_cls is not None
|
|
assert experts_cls.__name__ == "TrtLlmBf16ExpertsMonolithic"
|
|
|
|
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
|
|
return_value=(False, "monolithic unsupported"),
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsModular.is_supported_config",
|
|
return_value=(True, None),
|
|
)
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
|
)
|
|
def test_select_cuda_flashinfer_trtllm_modular_backend(
|
|
mock_is_supported_trtllm_modular,
|
|
mock_is_supported_trtllm_monolithic,
|
|
):
|
|
"""Test CUDA backend selection falls back to FlashInfer TRTLLM modular."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(current_platform, "has_device_capability", return_value=True),
|
|
):
|
|
moe_config = make_dummy_moe_config()
|
|
moe_config.moe_backend = "flashinfer_trtllm"
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.use_dp = False
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
|
|
assert experts_cls is not None
|
|
assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular"
|
|
|
|
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
|
|
return_value=True,
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"all2all_backend",
|
|
[
|
|
"mori_high_throughput",
|
|
"mori_low_latency",
|
|
"flashinfer_nvlink_two_sided",
|
|
"flashinfer_nvlink_one_sided",
|
|
],
|
|
)
|
|
def test_select_cuda_flashinfer_trtllm_modular_for_standard_all2all(
|
|
mock_supports_current_device,
|
|
all2all_backend,
|
|
):
|
|
"""Test non-AG/RS standard-format all2all backends select modular BF16."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(
|
|
current_platform, "is_device_capability_family", return_value=False
|
|
),
|
|
):
|
|
moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
|
|
moe_config.moe_backend = "flashinfer_trtllm"
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.dp_size = 2
|
|
moe_config.moe_parallel_config.ep_size = 2
|
|
moe_config.moe_parallel_config.all2all_backend = all2all_backend
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
|
|
assert experts_cls is not None
|
|
assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular"
|
|
|
|
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
|
|
return_value=True,
|
|
)
|
|
def test_select_cuda_deepep_ht_falls_back_from_trtllm(
|
|
mock_supports_current_device,
|
|
):
|
|
"""Test DeepEP HT avoids the unsupported BF16 TRTLLM modular path."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(
|
|
current_platform, "is_device_capability_family", return_value=False
|
|
),
|
|
):
|
|
moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
|
|
moe_config.moe_backend = "auto"
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.dp_size = 2
|
|
moe_config.moe_parallel_config.ep_size = 2
|
|
moe_config.moe_parallel_config.all2all_backend = "deepep_high_throughput"
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.TRITON
|
|
assert experts_cls is not None
|
|
assert experts_cls.__name__ == "TritonExperts"
|
|
|
|
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
|
|
return_value=True,
|
|
)
|
|
def test_select_cuda_flashinfer_trtllm_ag_rs_uses_monolithic(
|
|
mock_supports_current_device,
|
|
):
|
|
"""Test AG/RS stays on BF16 TRTLLM monolithic when TRTLLM is supported."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(
|
|
current_platform, "is_device_capability_family", return_value=False
|
|
),
|
|
):
|
|
moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
|
|
moe_config.moe_backend = "flashinfer_trtllm"
|
|
moe_config.routing_method = RoutingMethodType.Renormalize
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.dp_size = 2
|
|
moe_config.moe_parallel_config.ep_size = 2
|
|
moe_config.moe_parallel_config.all2all_backend = "allgather_reducescatter"
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
|
|
assert experts_cls is not None
|
|
assert experts_cls.__name__ == "TrtLlmBf16ExpertsMonolithic"
|
|
|
|
|
|
@patch(
|
|
"vllm.utils.flashinfer.has_flashinfer",
|
|
return_value=True,
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
|
|
return_value=(False, None),
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsModular.is_supported_config",
|
|
return_value=(False, None),
|
|
)
|
|
@patch(
|
|
"vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe.FlashInferExperts.is_supported_config",
|
|
return_value=(True, None),
|
|
)
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
|
)
|
|
def test_select_cuda_flashinfer_cutlass_backend(
|
|
mock_is_supported_cutlass,
|
|
mock_is_supported_trtllm_modular,
|
|
mock_is_supported_trtllm_monolithic,
|
|
mock_has_flashinfer,
|
|
):
|
|
"""Test CUDA backend selection when FlashInfer TRTLLM is not available
|
|
and FlashInfer CUTLASS is available."""
|
|
with (
|
|
patch.object(current_platform, "is_cuda", return_value=True),
|
|
patch.object(current_platform, "is_rocm", return_value=False),
|
|
patch.object(current_platform, "is_cpu", return_value=False),
|
|
patch.object(current_platform, "is_xpu", return_value=False),
|
|
patch.object(current_platform, "is_tpu", return_value=False),
|
|
patch.object(current_platform, "is_out_of_tree", return_value=False),
|
|
patch.object(current_platform, "has_device_capability", return_value=True),
|
|
):
|
|
moe_config = make_dummy_moe_config()
|
|
# Select FlashInfer CUTLASS explicitly
|
|
moe_config.moe_backend = "flashinfer_cutlass"
|
|
# CUTLASS requires EP and does not support DP
|
|
moe_config.moe_parallel_config.use_ep = True
|
|
moe_config.moe_parallel_config.use_dp = False
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
|
|
assert experts_cls is not None
|
|
|
|
|
|
@skipif_not_cuda_rocm
|
|
def test_select_lora_backend_prefers_triton():
|
|
"""LoRA-enabled unquantized MoE should select Triton backend."""
|
|
moe_config = make_dummy_moe_config()
|
|
moe_config.is_lora_enabled = True
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.TRITON
|
|
assert experts_cls is not None
|
|
|
|
|
|
@skipif_not_cuda_rocm
|
|
def test_select_lora_explicit_non_triton_backend():
|
|
"""LoRA should override explicit non-Triton backend to Triton."""
|
|
moe_config = make_dummy_moe_config()
|
|
moe_config.is_lora_enabled = True
|
|
|
|
# Use string from mapping in function map_unquantized_backend()
|
|
moe_config.moe_backend = "flashinfer_cutlass"
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.TRITON
|
|
assert experts_cls is not None
|
|
|
|
|
|
@skipif_not_cuda_rocm
|
|
@pytest.mark.parametrize("is_lora_enabled", [False, True])
|
|
def test_select_explicit_triton_backend(is_lora_enabled):
|
|
"""Explicit triton backend selection should return Triton."""
|
|
moe_config = make_dummy_moe_config()
|
|
moe_config.is_lora_enabled = is_lora_enabled
|
|
moe_config.moe_backend = "triton"
|
|
|
|
selected_backend, experts_cls = select_unquantized_moe_backend(
|
|
moe_config=moe_config
|
|
)
|
|
|
|
assert selected_backend == UnquantizedMoeBackend.TRITON
|
|
assert experts_cls is not None
|