| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 6 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 7 | from vllm.model_executor.layers.fused_moe.config import ( |
| 8 | FusedMoEConfig, |
| 9 | FusedMoEParallelConfig, |
| 10 | FusedMoEQuantConfig, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( |
| 13 | TopKWeightAndReduceNoOP, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 16 | QuantKey, |
| 17 | kFp8Dynamic128Sym, |
| 18 | kFp8DynamicTensorSym, |
| 19 | kFp8Static128BlockSym, |
| 20 | kFp8StaticTensorSym, |
| 21 | kInt4Static, |
| 22 | kInt4Static32, |
| 23 | kMxfp4Dynamic, |
| 24 | kMxfp4Static, |
| 25 | kMxfp8Dynamic, |
| 26 | kMxfp8Static, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | |
| 30 | if current_platform.is_xpu(): |
| 31 | from vllm_xpu_kernels.fused_moe_interface import XpuFusedMoe |
| 32 | |
| 33 | |
| 34 | def prepare_fp8_moe_layer_for_xpu( |
| 35 | w13: torch.Tensor, |
| 36 | w13_scale: torch.Tensor, |
| 37 | w2: torch.Tensor, |
| 38 | w2_scale: torch.Tensor, |
| 39 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 40 | if w13_scale is not None and w13_scale.ndim == 3: |
| 41 | w13_scale = w13_scale.transpose(-1, -2).contiguous() |
| 42 | if w2_scale is not None and w2_scale.ndim == 3: |
| 43 | w2_scale = w2_scale.transpose(-1, -2).contiguous() |
| 44 | return ( |
| 45 | w13.transpose(-1, -2).contiguous(), |
| 46 | w13_scale, |
| 47 | w2.transpose(-1, -2).contiguous(), |
| 48 | w2_scale, |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | class XPUExperts(mk.FusedMoEExpertsModular): |
| 53 | def __init__( |
| 54 | self, |
| 55 | moe_config: FusedMoEConfig, |
| 56 | quant_config: FusedMoEQuantConfig, |
| 57 | max_num_tokens: int | None = None, |
| 58 | num_dispatchers: int | None = None, |
| 59 | ): |
| 60 | super().__init__( |
| 61 | moe_config, |
| 62 | quant_config, |
| 63 | max_num_tokens, |
| 64 | num_dispatchers, |
| 65 | ) |
| 66 | self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit |
| 67 | self.fused_moe_impl: XpuFusedMoe | None = None |
| 68 | is_xe2_or_xe3 = torch.ops._xpu_C.is_xe2_arch() or torch.ops._xpu_C.is_xe3_arch() |
| 69 | if not is_xe2_or_xe3: |
| 70 | raise NotImplementedError( |
| 71 | "XPUExperts is only supported on Intel Xe2/Xe3 GPUs" |
| 72 | ) |
| 73 | self._expects_unquantized_inputs = is_xe2_or_xe3 |
| 74 | |
| 75 | @property |
| 76 | def expects_unquantized_inputs(self) -> bool: |
| 77 | return self._expects_unquantized_inputs |
| 78 | |
| 79 | @staticmethod |
| 80 | def activation_format() -> mk.FusedMoEActivationFormat: |
| 81 | return mk.FusedMoEActivationFormat.Standard |
| 82 | |
| 83 | @staticmethod |
| 84 | def _supports_current_device() -> bool: |
| 85 | return current_platform.is_xpu() |
| 86 | |
| 87 | @staticmethod |
| 88 | def _supports_no_act_and_mul() -> bool: |
| 89 | return True |
| 90 | |
| 91 | @staticmethod |
| 92 | def _supports_activation(activation: MoEActivation) -> bool: |
| 93 | return activation in [ |
| 94 | MoEActivation.SILU, |
| 95 | MoEActivation.GELU, |
| 96 | MoEActivation.GELU_TANH, |
| 97 | MoEActivation.SWIGLUOAI, |
| 98 | MoEActivation.RELU2_NO_MUL, |
| 99 | ] |
| 100 | |
| 101 | @staticmethod |
| 102 | def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: |
| 103 | return True |
| 104 | |
| 105 | @staticmethod |
| 106 | def _supports_quant_scheme( |
| 107 | weight_key: QuantKey | None, |
| 108 | activation_key: QuantKey | None, |
| 109 | ) -> bool: |
| 110 | SUPPORTED_W_A = [ |
| 111 | (None, None), |
| 112 | (kFp8StaticTensorSym, None), |
| 113 | (kFp8StaticTensorSym, kFp8DynamicTensorSym), |
| 114 | ] |
| 115 | return (weight_key, activation_key) in SUPPORTED_W_A |
| 116 | |
| 117 | def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: |
| 118 | return TopKWeightAndReduceNoOP() |
| 119 | |
| 120 | def workspace_shapes( |
| 121 | self, |
| 122 | M: int, |
| 123 | N: int, |
| 124 | K: int, |
| 125 | topk: int, |
| 126 | global_num_experts: int, |
| 127 | local_num_experts: int, |
| 128 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 129 | activation: MoEActivation, |
| 130 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 131 | workspace1 = (0,) |
| 132 | workspace2 = (0,) |
| 133 | output = (M, K) |
| 134 | return (workspace1, workspace2, output) |
| 135 | |
| 136 | def apply( |
| 137 | self, |
| 138 | output: torch.Tensor, |
| 139 | hidden_states: torch.Tensor, |
| 140 | w1: torch.Tensor, |
| 141 | w2: torch.Tensor, |
| 142 | topk_weights: torch.Tensor, |
| 143 | topk_ids: torch.Tensor, |
| 144 | activation: MoEActivation, |
| 145 | global_num_experts: int, |
| 146 | expert_map: torch.Tensor | None, |
| 147 | a1q_scale: torch.Tensor | None, |
| 148 | a2_scale: torch.Tensor | None, |
| 149 | workspace13: torch.Tensor, |
| 150 | workspace2: torch.Tensor, |
| 151 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 152 | apply_router_weight_on_input: bool, |
| 153 | ): |
| 154 | if self.fused_moe_impl is None: |
| 155 | topk = topk_ids.size(-1) |
| 156 | if ( |
| 157 | self.quant_config is not None |
| 158 | and self.quant_config.weight_quant_dtype == "mxfp4" |
| 159 | ): |
| 160 | w1 = w1.view(torch.float4_e2m1fn_x2) |
| 161 | w2 = w2.view(torch.float4_e2m1fn_x2) |
| 162 | self.fused_moe_impl = XpuFusedMoe( |
| 163 | w13=w1, |
| 164 | w13_scales=self.w1_scale, |
| 165 | w13_bias=self.w1_bias, |
| 166 | w2=w2, |
| 167 | w2_scales=self.w2_scale, |
| 168 | w2_bias=self.w2_bias, |
| 169 | n_experts_per_token=topk, |
| 170 | activation=activation.value, |
| 171 | num_experts=self.moe_config.num_local_experts, |
| 172 | ep_rank=self.moe_config.ep_rank, |
| 173 | ep_size=self.moe_config.ep_size, |
| 174 | expert_map=expert_map, |
| 175 | gemm1_clamp_limit=self.gemm1_clamp_limit, |
| 176 | ) |
| 177 | assert self.fused_moe_impl is not None |
| 178 | self.fused_moe_impl.apply( |
| 179 | output=output, |
| 180 | hidden_states=hidden_states, |
| 181 | topk_weights=topk_weights, |
| 182 | topk_ids=topk_ids, |
| 183 | a1q_scale=a1q_scale, |
| 184 | ) |
| 185 | |
| 186 | |
| 187 | class XPUExpertsFp8(XPUExperts): |
| 188 | def __init__( |
| 189 | self, |
| 190 | moe_config: FusedMoEConfig, |
| 191 | quant_config: FusedMoEQuantConfig, |
| 192 | max_num_tokens: int | None = None, |
| 193 | num_dispatchers: int | None = None, |
| 194 | ): |
| 195 | super().__init__( |
| 196 | moe_config, |
| 197 | quant_config, |
| 198 | max_num_tokens, |
| 199 | num_dispatchers, |
| 200 | ) |
| 201 | |
| 202 | @staticmethod |
| 203 | def _supports_quant_scheme( |
| 204 | weight_key: QuantKey | None, |
| 205 | activation_key: QuantKey | None, |
| 206 | ) -> bool: |
| 207 | SUPPORTED_W_A = [ |
| 208 | (kFp8StaticTensorSym, None), |
| 209 | (kFp8StaticTensorSym, kFp8DynamicTensorSym), |
| 210 | ] |
| 211 | return (weight_key, activation_key) in SUPPORTED_W_A |
| 212 | |
| 213 | |
| 214 | class XPUExpertsMxFp8(XPUExpertsFp8): |
| 215 | def __init__( |
| 216 | self, |
| 217 | moe_config: FusedMoEConfig, |
| 218 | quant_config: FusedMoEQuantConfig, |
| 219 | max_num_tokens: int | None = None, |
| 220 | num_dispatchers: int | None = None, |
| 221 | ): |
| 222 | super().__init__( |
| 223 | moe_config, |
| 224 | quant_config, |
| 225 | max_num_tokens, |
| 226 | num_dispatchers, |
| 227 | ) |
| 228 | assert quant_config.quant_dtype == "mxfp8" |
| 229 | |
| 230 | @staticmethod |
| 231 | def _supports_quant_scheme( |
| 232 | weight_key: QuantKey | None, |
| 233 | activation_key: QuantKey | None, |
| 234 | ) -> bool: |
| 235 | SUPPORTED_W_A = [ |
| 236 | (kMxfp8Static, None), |
| 237 | (kMxfp8Static, kMxfp8Dynamic), |
| 238 | ] |
| 239 | return (weight_key, activation_key) in SUPPORTED_W_A |
| 240 | |
| 241 | |
| 242 | class XPUExpertsBlockFp8(XPUExperts): |
| 243 | def __init__( |
| 244 | self, |
| 245 | moe_config: FusedMoEConfig, |
| 246 | quant_config: FusedMoEQuantConfig, |
| 247 | max_num_tokens: int | None = None, |
| 248 | num_dispatchers: int | None = None, |
| 249 | ): |
| 250 | super().__init__( |
| 251 | moe_config, |
| 252 | quant_config, |
| 253 | max_num_tokens, |
| 254 | num_dispatchers, |
| 255 | ) |
| 256 | |
| 257 | @staticmethod |
| 258 | def _supports_quant_scheme( |
| 259 | weight_key: QuantKey | None, |
| 260 | activation_key: QuantKey | None, |
| 261 | ) -> bool: |
| 262 | SUPPORTED_W_A = [ |
| 263 | (kFp8Static128BlockSym, kFp8Dynamic128Sym), |
| 264 | ] |
| 265 | return (weight_key, activation_key) in SUPPORTED_W_A |
| 266 | |
| 267 | |
| 268 | class XPUExpertsWNA16(XPUExperts): |
| 269 | """W4A16 INT4-symmetric MoE backed by `xpu_fused_moe(is_int4=True)`. |
| 270 | |
| 271 | Weight layout when `is_int4=True` (per `xpu_fused_moe` docstring): |
| 272 | w13: [num_experts, 2*inter_size, hidden_size] contiguous int4-packed |
| 273 | w13_scales: [num_experts, 2*inter_size, hidden_size // group_size] |
| 274 | w2: [num_experts, hidden_size, inter_size] contiguous int4-packed |
| 275 | w2_scales: [num_experts, hidden_size, inter_size // group_size] |
| 276 | |
| 277 | Pairs with `INCXPULinearMethod` for the linear layers; together they |
| 278 | cover full-attn + MoE on Intel XPU end-to-end without IPEX. |
| 279 | """ |
| 280 | |
| 281 | def __init__( |
| 282 | self, |
| 283 | moe_config: FusedMoEConfig, |
| 284 | quant_config: FusedMoEQuantConfig, |
| 285 | max_num_tokens: int | None = None, |
| 286 | num_dispatchers: int | None = None, |
| 287 | ): |
| 288 | super().__init__( |
| 289 | moe_config, |
| 290 | quant_config, |
| 291 | max_num_tokens, |
| 292 | num_dispatchers, |
| 293 | ) |
| 294 | |
| 295 | @staticmethod |
| 296 | def _supports_quant_scheme( |
| 297 | weight_key: QuantKey | None, |
| 298 | activation_key: QuantKey | None, |
| 299 | ) -> bool: |
| 300 | return (weight_key, activation_key) in ( |
| 301 | (kInt4Static, None), |
| 302 | (kInt4Static32, None), |
| 303 | ) |
| 304 | |
| 305 | |
| 306 | class XPUExpertsMxFp4(XPUExperts): |
| 307 | def __init__( |
| 308 | self, |
| 309 | moe_config: FusedMoEConfig, |
| 310 | quant_config: FusedMoEQuantConfig, |
| 311 | max_num_tokens: int | None = None, |
| 312 | num_dispatchers: int | None = None, |
| 313 | ): |
| 314 | super().__init__( |
| 315 | moe_config, |
| 316 | quant_config, |
| 317 | max_num_tokens, |
| 318 | num_dispatchers, |
| 319 | ) |
| 320 | |
| 321 | def workspace_shapes( |
| 322 | self, |
| 323 | M: int, |
| 324 | N: int, |
| 325 | K: int, |
| 326 | topk: int, |
| 327 | global_num_experts: int, |
| 328 | local_num_experts: int, |
| 329 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 330 | activation: MoEActivation, |
| 331 | ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: |
| 332 | # K = a1q.size(-1). When activations are pre-quantized packed mxfp4, |
| 333 | # K is the packed hidden_size (= logical / 2); the kernel output is at |
| 334 | # logical hidden_size (2 * K). When unquantized (bf16), K is already |
| 335 | # the logical size. |
| 336 | logical_K = K if self.expects_unquantized_inputs else 2 * K |
| 337 | return (0,), (0,), (M, logical_K) |
| 338 | |
| 339 | @staticmethod |
| 340 | def _supports_quant_scheme( |
| 341 | weight_key: QuantKey | None, |
| 342 | activation_key: QuantKey | None, |
| 343 | ) -> bool: |
| 344 | SUPPORTED_W_A = [ |
| 345 | (kMxfp4Static, None), |
| 346 | (kMxfp4Static, kMxfp4Dynamic), |
| 347 | ] |
| 348 | return (weight_key, activation_key) in SUPPORTED_W_A |