| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from contextlib import contextmanager |
| 5 | from typing import Any |
| 6 | |
| 7 | from vllm.model_executor.layers.fused_moe.activation import ( |
| 8 | ApplyMoEActivationConfig, |
| 9 | MoEActivation, |
| 10 | activation_without_mul, |
| 11 | apply_moe_activation, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.fused_moe.config import ( |
| 14 | FusedMoEConfig, |
| 15 | FusedMoEParallelConfig, |
| 16 | FusedMoEQuantConfig, |
| 17 | RoutingMethodType, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( |
| 20 | FusedMoEMethodBase, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.layer import ( |
| 23 | FusedMoEFactory, |
| 24 | fused_moe_make_expert_params_mapping, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 27 | FusedMoEActivationFormat, |
| 28 | FusedMoEExpertsModular, |
| 29 | FusedMoEPrepareAndFinalizeModular, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.routed_experts import ( |
| 32 | FusedMoeWeightScaleSupported, |
| 33 | RoutedExperts, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( |
| 36 | FusedMoERouter, |
| 37 | ) |
| 38 | from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear |
| 39 | from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( |
| 40 | MoERunner, |
| 41 | ) |
| 42 | from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( |
| 43 | SharedExperts, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 46 | UnquantizedFusedMoEMethod, |
| 47 | ) |
| 48 | from vllm.triton_utils import HAS_TRITON |
| 49 | |
| 50 | _config: dict[str, Any] | None = None |
| 51 | |
| 52 | |
| 53 | @contextmanager |
| 54 | def override_config(config): |
| 55 | global _config |
| 56 | old_config = _config |
| 57 | _config = config |
| 58 | yield |
| 59 | _config = old_config |
| 60 | |
| 61 | |
| 62 | def get_config() -> dict[str, Any] | None: |
| 63 | return _config |
| 64 | |
| 65 | |
| 66 | __all__ = [ |
| 67 | "ApplyMoEActivationConfig", |
| 68 | "FusedMoEFactory", |
| 69 | "FusedMoERouter", |
| 70 | "FusedMoEConfig", |
| 71 | "FusedMoEQuantConfig", |
| 72 | "FusedMoEParallelConfig", |
| 73 | "FusedMoEMethodBase", |
| 74 | "MoEActivation", |
| 75 | "UnquantizedFusedMoEMethod", |
| 76 | "FusedMoeWeightScaleSupported", |
| 77 | "FusedMoEExpertsModular", |
| 78 | "FusedMoEActivationFormat", |
| 79 | "FusedMoEPrepareAndFinalizeModular", |
| 80 | "GateLinear", |
| 81 | "MoERunner", |
| 82 | "RoutingMethodType", |
| 83 | "RoutedExperts", |
| 84 | "SharedExperts", |
| 85 | "activation_without_mul", |
| 86 | "apply_moe_activation", |
| 87 | "fused_moe_make_expert_params_mapping", |
| 88 | "override_config", |
| 89 | "get_config", |
| 90 | ] |
| 91 | |
| 92 | if HAS_TRITON: |
| 93 | # import to register the custom ops |
| 94 | from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( |
| 95 | BatchedDeepGemmExperts, |
| 96 | ) |
| 97 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 98 | CutlassBatchedExpertsFp8, |
| 99 | CutlassExpertsFp8, |
| 100 | CutlassExpertsW4A8Fp8, |
| 101 | ) |
| 102 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( |
| 103 | DeepGemmExperts, |
| 104 | ) |
| 105 | from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( |
| 106 | BatchedTritonExperts, |
| 107 | ) |
| 108 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 109 | AiterExperts, |
| 110 | ) |
| 111 | from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( |
| 112 | TritonOrDeepGemmExperts, |
| 113 | ) |
| 114 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( |
| 115 | TritonExperts, |
| 116 | TritonWNA16Experts, |
| 117 | ) |
| 118 | from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( |
| 119 | XPUExperts, |
| 120 | XPUExpertsFp8, |
| 121 | XPUExpertsMxFp4, |
| 122 | ) |
| 123 | from vllm.model_executor.layers.fused_moe.fused_moe import ( |
| 124 | fused_experts, |
| 125 | get_config_file_name, |
| 126 | ) |
| 127 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( |
| 128 | fused_topk, |
| 129 | ) |
| 130 | from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( |
| 131 | GroupedTopk, |
| 132 | ) |
| 133 | |
| 134 | __all__ += [ |
| 135 | "AiterExperts", |
| 136 | "fused_topk", |
| 137 | "fused_experts", |
| 138 | "get_config_file_name", |
| 139 | "GroupedTopk", |
| 140 | "CutlassExpertsFp8", |
| 141 | "CutlassBatchedExpertsFp8", |
| 142 | "CutlassExpertsW4A8Fp8", |
| 143 | "TritonExperts", |
| 144 | "TritonWNA16Experts", |
| 145 | "BatchedTritonExperts", |
| 146 | "DeepGemmExperts", |
| 147 | "BatchedDeepGemmExperts", |
| 148 | "TritonOrDeepGemmExperts", |
| 149 | "XPUExperts", |
| 150 | "XPUExpertsFp8", |
| 151 | "XPUExpertsBlockFp8", |
| 152 | "XPUExpertsMxFp8", |
| 153 | "XPUExpertsMxFp4", |
| 154 | ] |
| 155 | else: |
| 156 | # Some model classes directly use the custom ops. Add placeholders |
| 157 | # to avoid import errors. |
| 158 | def _raise_exception(method: str): |
| 159 | raise NotImplementedError(f"{method} is not implemented as lack of triton.") |
| 160 | |
| 161 | fused_topk = lambda *args, **kwargs: _raise_exception("fused_topk") |
| 162 | fused_experts = lambda *args, **kwargs: _raise_exception("fused_experts") |