| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 5 | from vllm.logger import init_logger |
| 6 | from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig |
| 7 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 8 | Fp8MoeBackend, |
| 9 | backend_to_kernel_cls, |
| 10 | ) |
| 11 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 12 | kMxfp8Dynamic, |
| 13 | kMxfp8Static, |
| 14 | ) |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | # Ordered by priority. |
| 19 | _SUPPORTED_BACKENDS = ( |
| 20 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 21 | Fp8MoeBackend.DEEPGEMM, |
| 22 | Fp8MoeBackend.MARLIN, |
| 23 | Fp8MoeBackend.XPU, |
| 24 | # AITER FlyDSL (gfx950): auto-picked by select_mxfp8_moe_backend when |
| 25 | # is_supported_config passes (gfx950 + flydsl installed + not EP). On other |
| 26 | # devices / no flydsl / EP it is skipped and native is used. |
| 27 | Fp8MoeBackend.AITER_MXFP8, |
| 28 | Fp8MoeBackend.HUMMING, |
| 29 | Fp8MoeBackend.TRITON_MXFP8, |
| 30 | Fp8MoeBackend.EMULATION, |
| 31 | ) |
| 32 | |
| 33 | _BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { |
| 34 | "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 35 | "deep_gemm": Fp8MoeBackend.DEEPGEMM, |
| 36 | "marlin": Fp8MoeBackend.MARLIN, |
| 37 | "xpu": Fp8MoeBackend.XPU, |
| 38 | "aiter": Fp8MoeBackend.AITER_MXFP8, |
| 39 | "triton": Fp8MoeBackend.TRITON_MXFP8, |
| 40 | "humming": Fp8MoeBackend.HUMMING, |
| 41 | } |
| 42 | |
| 43 | |
| 44 | def _mxfp8_backend_to_kernel_cls( |
| 45 | backend: Fp8MoeBackend, |
| 46 | ) -> list[type[mk.FusedMoEExperts]]: |
| 47 | """Resolve the MXFP8 expert classes for a backend. |
| 48 | |
| 49 | DeepGEMM resolves directly to ``DeepGemmExperts`` (not the |
| 50 | ``TritonOrDeepGemmExperts`` wrapper, whose Triton fallback cannot handle the |
| 51 | MXFP8 1x32 scheme); all other backends defer to the FP8 resolver. |
| 52 | """ |
| 53 | if backend == Fp8MoeBackend.DEEPGEMM: |
| 54 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( |
| 55 | DeepGemmExperts, |
| 56 | ) |
| 57 | |
| 58 | return [DeepGemmExperts] |
| 59 | if backend == Fp8MoeBackend.AITER_MXFP8: |
| 60 | from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp8_moe import ( |
| 61 | AiterMxfp8Experts, |
| 62 | ) |
| 63 | |
| 64 | return [AiterMxfp8Experts] |
| 65 | if backend == Fp8MoeBackend.TRITON_MXFP8: |
| 66 | from vllm.model_executor.layers.fused_moe.experts.mxfp8_native_moe import ( |
| 67 | Mxfp8NativeTritonExperts, |
| 68 | ) |
| 69 | |
| 70 | return [Mxfp8NativeTritonExperts] |
| 71 | if backend == Fp8MoeBackend.EMULATION: |
| 72 | from vllm.model_executor.layers.fused_moe.experts.mxfp8_emulation_moe import ( |
| 73 | Mxfp8EmulationTritonExperts, |
| 74 | ) |
| 75 | |
| 76 | return [Mxfp8EmulationTritonExperts] |
| 77 | return backend_to_kernel_cls(backend) |
| 78 | |
| 79 | |
| 80 | def _select_kernel_cls( |
| 81 | backend: Fp8MoeBackend, |
| 82 | config: FusedMoEConfig, |
| 83 | ) -> type[mk.FusedMoEExperts]: |
| 84 | """Select the first supported expert class for the MXFP8 config.""" |
| 85 | activation_format = ( |
| 86 | mk.FusedMoEActivationFormat.BatchedExperts |
| 87 | if config.moe_parallel_config.use_batched_activation_format |
| 88 | else mk.FusedMoEActivationFormat.Standard |
| 89 | ) |
| 90 | last_reason: str | None = None |
| 91 | for cls in _mxfp8_backend_to_kernel_cls(backend): |
| 92 | supported, reason = cls.is_supported_config( |
| 93 | cls, |
| 94 | config, |
| 95 | kMxfp8Static, |
| 96 | kMxfp8Dynamic, |
| 97 | activation_format, |
| 98 | ) |
| 99 | if supported: |
| 100 | return cls |
| 101 | last_reason = reason |
| 102 | raise ValueError( |
| 103 | f"No supported MXFP8 expert class for {backend.value}: {last_reason}" |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | def select_mxfp8_moe_backend( |
| 108 | config: FusedMoEConfig, |
| 109 | ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: |
| 110 | """Select the MXFP8 MoE backend and the best expert class. |
| 111 | |
| 112 | Returns: |
| 113 | A tuple of (fp8_backend, experts_cls). |
| 114 | |
| 115 | """ |
| 116 | runner_backend = config.moe_backend |
| 117 | if runner_backend != "auto": |
| 118 | backend = _BACKEND_NAME_MAP.get(runner_backend) |
| 119 | if backend is None: |
| 120 | raise ValueError( |
| 121 | f"moe_backend='{runner_backend}' is not supported for " |
| 122 | f"MXFP8 MoE. Expected one of " |
| 123 | f"{list(_BACKEND_NAME_MAP.keys())}." |
| 124 | ) |
| 125 | logger.info_once( |
| 126 | "Using '%s' MxFp8 MoE backend (user-requested).", |
| 127 | backend.value, |
| 128 | ) |
| 129 | return backend, _select_kernel_cls(backend, config) |
| 130 | |
| 131 | # Auto-select: pick the first supported backend. |
| 132 | for backend in _SUPPORTED_BACKENDS: |
| 133 | try: |
| 134 | experts_cls = _select_kernel_cls(backend, config) |
| 135 | except ValueError: |
| 136 | continue |
| 137 | logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) |
| 138 | return backend, experts_cls |
| 139 | |
| 140 | # TODO: add debug log with reason. |
| 141 | raise ValueError("No MXFP8 MoE backends available.") |