| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from enum import Enum |
| 4 | from typing import TYPE_CHECKING, Literal, Union |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 9 | from vllm import envs |
| 10 | from vllm.config import get_current_vllm_config |
| 11 | from vllm.config.kernel import MoEBackend |
| 12 | from vllm.config.quantization import QuantizationConfigArgs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.fused_moe import ( |
| 15 | FusedMoEConfig, |
| 16 | RoutedExperts, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 19 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 20 | maybe_make_prepare_finalize, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.config import ( |
| 23 | FusedMoEQuantConfig, |
| 24 | FusedMoEQuantDesc, |
| 25 | RoutingMethodType, |
| 26 | mxfp4_mxfp8_moe_quant_config, |
| 27 | mxfp4_w4a8_moe_quant_config, |
| 28 | mxfp4_w4a16_moe_quant_config, |
| 29 | ocp_mx_moe_quant_config, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 32 | swap_w13_to_w31, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 35 | _swizzle_mxfp4, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 38 | mx_scale_kwargs as _mx_scale_kwargs, |
| 39 | ) |
| 40 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( |
| 41 | OCP_MX_BLOCK_SIZE, |
| 42 | ) |
| 43 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 44 | QuantKey, |
| 45 | kFp8Dynamic128Sym, |
| 46 | kFp8StaticTensorSym, |
| 47 | kMxfp4Dynamic, |
| 48 | kMxfp4Static, |
| 49 | kMxfp8Dynamic, |
| 50 | ) |
| 51 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import all_close_1d |
| 52 | from vllm.platforms import current_platform |
| 53 | from vllm.utils.import_utils import get_triton_kernels_version |
| 54 | from vllm.utils.math_utils import round_up |
| 55 | |
| 56 | if TYPE_CHECKING: |
| 57 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 58 | |
| 59 | |
| 60 | logger = init_logger(__name__) |
| 61 | |
| 62 | triton_kernels_version = get_triton_kernels_version() |
| 63 | |
| 64 | if triton_kernels_version is not None: |
| 65 | try: |
| 66 | if triton_kernels_version == "3.8": |
| 67 | from triton_kernels.matmul import FlexCtx, PrecisionConfig |
| 68 | else: |
| 69 | from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig |
| 70 | except (ImportError, AttributeError) as e: |
| 71 | logger.error( |
| 72 | "Failed to import Triton kernels. Please make sure your triton " |
| 73 | "version is compatible. Error: %s", |
| 74 | e, |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def _pack_deepgemm_mxfp4_scales( |
| 79 | w13_weight: torch.Tensor, |
| 80 | w2_weight: torch.Tensor, |
| 81 | w13_weight_scale: torch.Tensor, |
| 82 | w2_weight_scale: torch.Tensor, |
| 83 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 84 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 85 | deepgemm_post_process_weight_scale_block, |
| 86 | ) |
| 87 | |
| 88 | num_experts = w13_weight.shape[0] |
| 89 | intermediate_size_2 = w13_weight.shape[1] # = intermediate*2 |
| 90 | hidden_size = w13_weight.shape[2] * 2 # weight is FP4-packed |
| 91 | intermediate_size = w2_weight.shape[2] * 2 # weight is FP4-packed |
| 92 | block_shape = (1, 32) # MXFP4 block (per-row, K=32) |
| 93 | |
| 94 | return ( |
| 95 | deepgemm_post_process_weight_scale_block( |
| 96 | ws=w13_weight_scale.data, |
| 97 | mn=intermediate_size_2, |
| 98 | k=hidden_size, |
| 99 | quant_block_shape=block_shape, |
| 100 | num_groups=num_experts, |
| 101 | ), |
| 102 | deepgemm_post_process_weight_scale_block( |
| 103 | ws=w2_weight_scale.data, |
| 104 | mn=hidden_size, |
| 105 | k=intermediate_size, |
| 106 | quant_block_shape=block_shape, |
| 107 | num_groups=num_experts, |
| 108 | ), |
| 109 | ) |
| 110 | |
| 111 | |
| 112 | class Mxfp4MoeBackend(Enum): |
| 113 | NONE = "None" |
| 114 | # b12x backends |
| 115 | B12X_MXFP4_MXFP8 = "B12X_MXFP4_MXFP8" |
| 116 | B12X_MXFP4_BF16 = "B12X_MXFP4_BF16" |
| 117 | # DeepGEMM FP8xFP4 backend (SM100+) |
| 118 | DEEPGEMM_MXFP4 = "DEEPGEMM_MXFP4" |
| 119 | # FlashInfer TRTLLM backends |
| 120 | FLASHINFER_TRTLLM_MXFP4_MXFP8 = "FLASHINFER_TRTLLM_MXFP4_MXFP8" |
| 121 | FLASHINFER_TRTLLM_MXFP4_BF16 = "FLASHINFER_TRTLLM_MXFP4_BF16" |
| 122 | # FlashInfer CUTLASS backends |
| 123 | FLASHINFER_CUTLASS_MXFP4_MXFP8 = "FLASHINFER_CUTLASS_MXFP4_MXFP8" |
| 124 | FLASHINFER_CUTLASS_MXFP4_BF16 = "FLASHINFER_CUTLASS_MXFP4_BF16" |
| 125 | # Marlin |
| 126 | BATCHED_MARLIN = "BATCHED_MARLIN" |
| 127 | MARLIN = "MARLIN" |
| 128 | # ROCm AITER backends |
| 129 | AITER_MXFP4_BF16 = "AITER_MXFP4_BF16" # W4A16: CK kernel (gfx950) |
| 130 | # Legacy alias, resolves to the CK kernel. New code should name either |
| 131 | # AITER_MXFP4_BF16 (CK) or AITER_TRITON_MXFP4_BF16 (Triton) explicitly. |
| 132 | AITER = "AITER_MXFP4_BF16" |
| 133 | # W4A16: aiter Triton moe_gemm_a16w4 kernel (gfx942/gfx950/gfx1250) |
| 134 | AITER_TRITON_MXFP4_BF16 = "AITER_TRITON_MXFP4_BF16" |
| 135 | AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel |
| 136 | AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel |
| 137 | # Triton |
| 138 | TRITON = "TRITON" |
| 139 | TRITON_UNFUSED = "TRITON_UNFUSED" |
| 140 | # XPU |
| 141 | XPU = "XPU" |
| 142 | # CPU |
| 143 | CPU = "CPU" |
| 144 | # Emulation |
| 145 | EMULATION = "EMULATION" |
| 146 | # Humming |
| 147 | HUMMING = "HUMMING" |
| 148 | |
| 149 | |
| 150 | # Backends that share the same TRTLLM weight format |
| 151 | TRTLLM_BACKENDS = ( |
| 152 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, |
| 153 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 154 | ) |
| 155 | |
| 156 | TRITON_BACKENDS = ( |
| 157 | Mxfp4MoeBackend.TRITON, |
| 158 | Mxfp4MoeBackend.TRITON_UNFUSED, |
| 159 | # aiter Triton W4A16 shares the triton_kernels weight format |
| 160 | Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16, |
| 161 | ) |
| 162 | |
| 163 | B12X_BACKENDS = ( |
| 164 | Mxfp4MoeBackend.B12X_MXFP4_MXFP8, |
| 165 | Mxfp4MoeBackend.B12X_MXFP4_BF16, |
| 166 | ) |
| 167 | |
| 168 | |
| 169 | def backend_to_kernel_cls( |
| 170 | backend: Mxfp4MoeBackend, |
| 171 | ) -> list[type[mk.FusedMoEExperts]]: |
| 172 | if backend in B12X_BACKENDS: |
| 173 | from vllm.model_executor.layers.fused_moe.b12x import B12xExperts |
| 174 | |
| 175 | return [B12xExperts] |
| 176 | |
| 177 | elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 178 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( |
| 179 | DeepGemmFP4Experts, |
| 180 | ) |
| 181 | |
| 182 | return [DeepGemmFP4Experts] |
| 183 | |
| 184 | elif backend in ( |
| 185 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, |
| 186 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 187 | ): |
| 188 | from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( |
| 189 | TrtLlmMxfp4ExpertsModular, |
| 190 | TrtLlmMxfp4ExpertsMonolithic, |
| 191 | ) |
| 192 | |
| 193 | # NOTE: prefer Monolithic > Modular, so return Monolithic first. |
| 194 | return [TrtLlmMxfp4ExpertsMonolithic, TrtLlmMxfp4ExpertsModular] |
| 195 | |
| 196 | elif backend in ( |
| 197 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 198 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 199 | ): |
| 200 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501 |
| 201 | FlashInferExperts, |
| 202 | ) |
| 203 | |
| 204 | return [FlashInferExperts] |
| 205 | |
| 206 | elif backend == Mxfp4MoeBackend.TRITON: |
| 207 | from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 |
| 208 | OAITritonExperts, |
| 209 | OAITritonMxfp4ExpertsMonolithic, |
| 210 | ) |
| 211 | |
| 212 | # NOTE: prefer Monolithic > Modular, so return Monolithic first. |
| 213 | return [OAITritonMxfp4ExpertsMonolithic, OAITritonExperts] |
| 214 | |
| 215 | elif backend == Mxfp4MoeBackend.TRITON_UNFUSED: |
| 216 | from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 |
| 217 | UnfusedOAITritonExperts, |
| 218 | ) |
| 219 | |
| 220 | return [UnfusedOAITritonExperts] |
| 221 | |
| 222 | elif backend == Mxfp4MoeBackend.HUMMING: |
| 223 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 224 | BatchedHummingGroupedExperts, |
| 225 | HummingGroupedExperts, |
| 226 | HummingIndexedExperts, |
| 227 | ) |
| 228 | |
| 229 | return [ |
| 230 | BatchedHummingGroupedExperts, |
| 231 | HummingGroupedExperts, |
| 232 | HummingIndexedExperts, |
| 233 | ] |
| 234 | |
| 235 | elif backend == Mxfp4MoeBackend.MARLIN: |
| 236 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 237 | MarlinExperts, |
| 238 | ) |
| 239 | |
| 240 | return [MarlinExperts] |
| 241 | |
| 242 | elif backend == Mxfp4MoeBackend.BATCHED_MARLIN: |
| 243 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 244 | BatchedMarlinExperts, |
| 245 | ) |
| 246 | |
| 247 | return [BatchedMarlinExperts] |
| 248 | |
| 249 | elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: |
| 250 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 251 | AiterExperts, |
| 252 | ) |
| 253 | |
| 254 | return [AiterExperts] |
| 255 | |
| 256 | elif backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: |
| 257 | from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( |
| 258 | AiterW4A16ExpertsMonolithic, |
| 259 | ) |
| 260 | |
| 261 | return [AiterW4A16ExpertsMonolithic] |
| 262 | |
| 263 | elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: |
| 264 | from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( |
| 265 | AiterW4A8ExpertsMonolithic, |
| 266 | ) |
| 267 | |
| 268 | return [AiterW4A8ExpertsMonolithic] |
| 269 | |
| 270 | elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: |
| 271 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 272 | AiterExperts, |
| 273 | ) |
| 274 | |
| 275 | return [AiterExperts] |
| 276 | |
| 277 | elif backend == Mxfp4MoeBackend.XPU: |
| 278 | from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMxFp4 |
| 279 | |
| 280 | return [XPUExpertsMxFp4] |
| 281 | |
| 282 | elif backend == Mxfp4MoeBackend.CPU: |
| 283 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import CPUExpertsMxfp4 |
| 284 | |
| 285 | return [CPUExpertsMxfp4] |
| 286 | |
| 287 | elif backend == Mxfp4MoeBackend.EMULATION: |
| 288 | from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import ( |
| 289 | OCP_MXQuantizationEmulationTritonExperts, |
| 290 | ) |
| 291 | |
| 292 | return [OCP_MXQuantizationEmulationTritonExperts] |
| 293 | |
| 294 | else: |
| 295 | raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}") |
| 296 | |
| 297 | |
| 298 | def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: |
| 299 | """Map a moe_backend string to its candidate Mxfp4MoeBackends. |
| 300 | |
| 301 | Vendor families return all activation variants; the caller picks one |
| 302 | via ``activation_key`` and ``is_supported_config``. |
| 303 | """ |
| 304 | mapping: dict[str, list[Mxfp4MoeBackend]] = { |
| 305 | "b12x": list(B12X_BACKENDS), |
| 306 | "deep_gemm": [Mxfp4MoeBackend.DEEPGEMM_MXFP4], |
| 307 | "flashinfer_trtllm": [ |
| 308 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, |
| 309 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 310 | ], |
| 311 | "flashinfer_trtllm_afp8": [Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8], |
| 312 | "flashinfer_cutlass": [ |
| 313 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 314 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 315 | ], |
| 316 | "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8], |
| 317 | "triton": [Mxfp4MoeBackend.TRITON], |
| 318 | "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED], |
| 319 | "humming": [Mxfp4MoeBackend.HUMMING], |
| 320 | "marlin": [Mxfp4MoeBackend.MARLIN], |
| 321 | "aiter": [ |
| 322 | Mxfp4MoeBackend.AITER_MXFP4_BF16, |
| 323 | Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16, |
| 324 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 325 | Mxfp4MoeBackend.AITER_MXFP4_MXFP4, |
| 326 | ], |
| 327 | "aiter_triton_mxfp4_bf16": [Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16], |
| 328 | "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8], |
| 329 | "aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4], |
| 330 | "xpu": [Mxfp4MoeBackend.XPU], |
| 331 | "cpu": [Mxfp4MoeBackend.CPU], |
| 332 | "emulation": [Mxfp4MoeBackend.EMULATION], |
| 333 | } |
| 334 | if backends := mapping.get(runner_backend): |
| 335 | return backends |
| 336 | raise ValueError( |
| 337 | f"moe_backend='{runner_backend}' is not supported for MXFP4 MoE. " |
| 338 | f"Expected one of {list(mapping.keys())}." |
| 339 | ) |
| 340 | |
| 341 | |
| 342 | def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: |
| 343 | """Available backends in priority order, BF16-act variant before |
| 344 | activation-quantized variant within each vendor family.""" |
| 345 | _AVAILABLE_BACKENDS = [ |
| 346 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, |
| 347 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 348 | Mxfp4MoeBackend.AITER_MXFP4_BF16, |
| 349 | Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16, |
| 350 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 351 | Mxfp4MoeBackend.AITER_MXFP4_MXFP4, |
| 352 | Mxfp4MoeBackend.TRITON, |
| 353 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 354 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 355 | # TRITON_UNFUSED has bug with MTP support |
| 356 | # TODO re-enable after kernel is fixed |
| 357 | # TRITON_UNFUSED |
| 358 | Mxfp4MoeBackend.MARLIN, |
| 359 | Mxfp4MoeBackend.BATCHED_MARLIN, |
| 360 | Mxfp4MoeBackend.XPU, |
| 361 | Mxfp4MoeBackend.CPU, |
| 362 | Mxfp4MoeBackend.EMULATION, |
| 363 | ] |
| 364 | return _AVAILABLE_BACKENDS |
| 365 | |
| 366 | |
| 367 | def _get_priority_backends() -> list[Mxfp4MoeBackend]: |
| 368 | """Get available backends in priority order. SM100+ prefers DeepGEMM FP4 / |
| 369 | TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the |
| 370 | backend-level ``is_supported_config`` check filters by device capability). |
| 371 | """ |
| 372 | if current_platform.is_rocm(): |
| 373 | return [ |
| 374 | Mxfp4MoeBackend.AITER_MXFP4_BF16, |
| 375 | Mxfp4MoeBackend.EMULATION, |
| 376 | ] |
| 377 | if current_platform.is_xpu(): |
| 378 | return [Mxfp4MoeBackend.XPU] |
| 379 | if current_platform.is_cpu(): |
| 380 | return [Mxfp4MoeBackend.CPU] |
| 381 | _AVAILABLE_BACKENDS = [ |
| 382 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 383 | Mxfp4MoeBackend.DEEPGEMM_MXFP4, |
| 384 | # TRITON_UNFUSED has bug with MTP support |
| 385 | # TODO re-enable after kernel is fixed |
| 386 | # TRITON_UNFUSED |
| 387 | Mxfp4MoeBackend.MARLIN, |
| 388 | Mxfp4MoeBackend.BATCHED_MARLIN, |
| 389 | ] |
| 390 | return _AVAILABLE_BACKENDS |
| 391 | |
| 392 | |
| 393 | def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: |
| 394 | """Map backend to its activation key (FP8, MXFP8, or None for BF16).""" |
| 395 | if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 396 | return kFp8Dynamic128Sym |
| 397 | if backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: |
| 398 | return kMxfp8Dynamic |
| 399 | if backend in ( |
| 400 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, |
| 401 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 402 | ): |
| 403 | return kMxfp8Dynamic |
| 404 | if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: |
| 405 | return kFp8StaticTensorSym |
| 406 | if backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: |
| 407 | return kMxfp4Dynamic |
| 408 | return None # BF16 activation |
| 409 | |
| 410 | |
| 411 | def _user_moe_activation_override() -> QuantKey | None: |
| 412 | """User's MoE activation override from quantization_config, or None.""" |
| 413 | args = get_current_vllm_config().model_config.quantization_config |
| 414 | if not isinstance(args, QuantizationConfigArgs) or args.moe is None: |
| 415 | return None |
| 416 | return args.moe.activation |
| 417 | |
| 418 | |
| 419 | def _resolve_activation_key( |
| 420 | model_activation_key: QuantKey | None, |
| 421 | ) -> QuantKey | None: |
| 422 | """Combine the model-supplied activation key with the user override. |
| 423 | Raises on conflict (both set and disagreeing).""" |
| 424 | user_override = _user_moe_activation_override() |
| 425 | if user_override is None: |
| 426 | return model_activation_key |
| 427 | if model_activation_key is None or model_activation_key == user_override: |
| 428 | return user_override |
| 429 | raise ValueError( |
| 430 | f"checkpoint declares MoE activation={model_activation_key} but " |
| 431 | f"quantization_config.moe.activation={user_override}; remove the " |
| 432 | f"override or align it with the checkpoint." |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | def _make_log_backend(backend: Mxfp4MoeBackend) -> str: |
| 437 | return f"Using '{backend.value}' Mxfp4 MoE backend." |
| 438 | |
| 439 | |
| 440 | def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str: |
| 441 | base = ( |
| 442 | f"Mxfp4 MoE backend '{backend.value}' does not support the " |
| 443 | f"deployment configuration" |
| 444 | ) |
| 445 | return f"{base} since {reason}." if reason else f"{base}." |
| 446 | |
| 447 | |
| 448 | def _return_or_raise( |
| 449 | backend: Mxfp4MoeBackend, |
| 450 | config: FusedMoEConfig, |
| 451 | weight_key: QuantKey | None, |
| 452 | activation_key: QuantKey | None, |
| 453 | activation_format: mk.FusedMoEActivationFormat, |
| 454 | scope: Literal["process", "global", "local"] = "local", |
| 455 | ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]: |
| 456 | reason: str | None = None |
| 457 | for k_cls in backend_to_kernel_cls(backend): |
| 458 | supported, reason = k_cls.is_supported_config( |
| 459 | k_cls, config, weight_key, activation_key, activation_format |
| 460 | ) |
| 461 | if supported: |
| 462 | logger.info_once(_make_log_backend(backend), scope=scope) |
| 463 | return backend, k_cls |
| 464 | raise ValueError(_make_log_unsupported(backend, reason)) |
| 465 | |
| 466 | |
| 467 | def _filter_by_activation( |
| 468 | backends: list[Mxfp4MoeBackend], |
| 469 | requested_activation_key: QuantKey | None, |
| 470 | ) -> list[Mxfp4MoeBackend]: |
| 471 | """Pick variants matching ``requested_activation_key``; without one, |
| 472 | prefer BF16 if the list has any, else keep the list as-is so explicit |
| 473 | non-BF16 picks (e.g. the ``_afp8`` aliases) still land.""" |
| 474 | if requested_activation_key is not None: |
| 475 | return [ |
| 476 | b |
| 477 | for b in backends |
| 478 | if _backend_activation_key(b) == requested_activation_key |
| 479 | or b == Mxfp4MoeBackend.EMULATION |
| 480 | ] |
| 481 | bf16 = [b for b in backends if _backend_activation_key(b) is None] |
| 482 | return bf16 if bf16 else backends |
| 483 | |
| 484 | |
| 485 | def _get_requested_backends( |
| 486 | runner_backend: MoEBackend, |
| 487 | requested_activation_key: QuantKey | None, |
| 488 | ) -> list[Mxfp4MoeBackend]: |
| 489 | backends = map_mxfp4_backend(runner_backend) |
| 490 | if runner_backend == "b12x": |
| 491 | if envs.VLLM_B12X_MOE_FP4_FORCE_A16: |
| 492 | return [Mxfp4MoeBackend.B12X_MXFP4_BF16] |
| 493 | # W4A8 is the high-throughput b12x path and is preferred when the |
| 494 | # model does not request an activation format. Keep this as backend |
| 495 | # policy instead of exposing a separate backend name per precision. |
| 496 | if requested_activation_key is None: |
| 497 | return backends |
| 498 | return _filter_by_activation(backends, requested_activation_key) |
| 499 | |
| 500 | |
| 501 | def _requires_qwen38_tep8_emulation( |
| 502 | config: FusedMoEConfig, |
| 503 | activation_key: QuantKey | None, |
| 504 | ) -> bool: |
| 505 | """Avoid the inaccurate native gfx950 kernel for Qwen3.8 TEP8. |
| 506 | |
| 507 | Qwen3.8 Flash Next's routed experts use the distinctive |
| 508 | ``E=512, H=2560, N=640`` W4A4 shape. With eight-way expert parallelism, |
| 509 | AITER operates on 64 local experts and pads ``N`` to 768. That native path |
| 510 | is not numerically reliable on gfx950, while OCP MX emulation preserves |
| 511 | model accuracy. Keep this guard exact so other MXFP4 shapes and smaller EP |
| 512 | configurations retain the native backend. |
| 513 | """ |
| 514 | parallel = config.moe_parallel_config |
| 515 | if not ( |
| 516 | activation_key == kMxfp4Dynamic |
| 517 | and parallel.use_ep |
| 518 | and parallel.ep_size == 8 |
| 519 | and config.num_experts == 512 |
| 520 | and config.num_local_experts == 64 |
| 521 | and config.hidden_dim == 2560 |
| 522 | and config.intermediate_size == 640 |
| 523 | ): |
| 524 | return False |
| 525 | |
| 526 | if not current_platform.is_rocm(): |
| 527 | return False |
| 528 | |
| 529 | from vllm.platforms.rocm import on_gfx950 |
| 530 | |
| 531 | return on_gfx950() |
| 532 | |
| 533 | |
| 534 | def select_mxfp4_moe_backend( |
| 535 | config: FusedMoEConfig, |
| 536 | activation_key: QuantKey | None = None, |
| 537 | ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: |
| 538 | """Select the primary MXFP4 MoE backend. |
| 539 | |
| 540 | Args: |
| 541 | config: MoE configuration |
| 542 | activation_key: Optional activation quantization key. If provided, |
| 543 | overrides the default activation key for backend selection. |
| 544 | Use kFp8StaticTensorSym for W4A8 scheme. |
| 545 | |
| 546 | Note: Shape-specific fallbacks may still occur at runtime. |
| 547 | |
| 548 | """ |
| 549 | runner_backend = config.moe_backend |
| 550 | requested_activation_key = _resolve_activation_key(activation_key) |
| 551 | |
| 552 | activation_format = ( |
| 553 | mk.FusedMoEActivationFormat.BatchedExperts |
| 554 | if config.moe_parallel_config.use_batched_activation_format |
| 555 | else mk.FusedMoEActivationFormat.Standard |
| 556 | ) |
| 557 | |
| 558 | if runner_backend != "auto": |
| 559 | requested_backends = _get_requested_backends( |
| 560 | runner_backend, requested_activation_key |
| 561 | ) |
| 562 | if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 563 | requested_backends = [ |
| 564 | Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b |
| 565 | for b in requested_backends |
| 566 | ] |
| 567 | if not requested_backends: |
| 568 | raise ValueError( |
| 569 | f"moe_backend={runner_backend!r} does not support " |
| 570 | f"activation={requested_activation_key}" |
| 571 | ) |
| 572 | last_error: Exception | None = None |
| 573 | for requested_backend in requested_backends: |
| 574 | act_key = ( |
| 575 | requested_activation_key |
| 576 | if requested_backend == Mxfp4MoeBackend.EMULATION |
| 577 | else _backend_activation_key(requested_backend) |
| 578 | ) |
| 579 | try: |
| 580 | return _return_or_raise( |
| 581 | requested_backend, |
| 582 | config, |
| 583 | kMxfp4Static, |
| 584 | act_key, |
| 585 | activation_format, |
| 586 | ) |
| 587 | except ValueError as e: |
| 588 | last_error = e |
| 589 | assert last_error is not None |
| 590 | raise last_error |
| 591 | |
| 592 | if _requires_qwen38_tep8_emulation(config, requested_activation_key): |
| 593 | backend = Mxfp4MoeBackend.EMULATION |
| 594 | logger.warning_once( |
| 595 | "Using OCP MX emulation for the Qwen3.8 Flash Next TEP8 routed " |
| 596 | "experts on gfx950 because the native AITER W4A4 kernel is not " |
| 597 | "numerically reliable for this shape. Performance will be lower." |
| 598 | ) |
| 599 | return _return_or_raise( |
| 600 | backend, |
| 601 | config, |
| 602 | kMxfp4Static, |
| 603 | requested_activation_key, |
| 604 | activation_format, |
| 605 | ) |
| 606 | |
| 607 | # Select kernels in order of backend. |
| 608 | AVAILABLE_BACKENDS = _filter_by_activation( |
| 609 | _get_priority_backends_for_gpt_oss(), requested_activation_key |
| 610 | ) |
| 611 | |
| 612 | unsupported_reasons = [] |
| 613 | for backend in AVAILABLE_BACKENDS: |
| 614 | # Use requested_activation_key if provided, otherwise use backend default |
| 615 | act_key = ( |
| 616 | requested_activation_key |
| 617 | if requested_activation_key is not None |
| 618 | else _backend_activation_key(backend) |
| 619 | ) |
| 620 | for k_cls in backend_to_kernel_cls(backend): |
| 621 | supported, reason = k_cls.is_supported_config( |
| 622 | k_cls, config, kMxfp4Static, act_key, activation_format |
| 623 | ) |
| 624 | if supported: |
| 625 | logger.info_once(_make_log_backend(backend)) |
| 626 | return backend, k_cls |
| 627 | else: |
| 628 | logger.debug_once(_make_log_unsupported(backend, reason)) |
| 629 | unsupported_reasons.append((backend, reason)) |
| 630 | |
| 631 | if current_platform.is_xpu(): |
| 632 | backend = Mxfp4MoeBackend.XPU |
| 633 | logger.info_once(_make_log_backend(backend)) |
| 634 | return _return_or_raise( |
| 635 | Mxfp4MoeBackend.XPU, |
| 636 | config, |
| 637 | kMxfp4Static, |
| 638 | None, |
| 639 | activation_format, |
| 640 | ) |
| 641 | |
| 642 | unsupported_log = "; ".join( |
| 643 | [ |
| 644 | f"backend: {backend.value}, reason: {reason}" |
| 645 | for backend, reason in unsupported_reasons |
| 646 | ] |
| 647 | ) |
| 648 | raise NotImplementedError( |
| 649 | "No MXFP4 MoE backend supports the deployment configuration. " |
| 650 | f"weight_key=kMxfp4Static, activation_key={activation_key}. " |
| 651 | f"Candidate backends were: " |
| 652 | f"{[backend.value for backend in AVAILABLE_BACKENDS]}. " |
| 653 | f"Unsupported reasons: {unsupported_log}. " |
| 654 | ) |
| 655 | |
| 656 | |
| 657 | def select_deepseek_v4_mxfp4_moe_backend( |
| 658 | config: FusedMoEConfig, |
| 659 | ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: |
| 660 | """Select the MXFP4 MoE backend with MXFP8 activation as top priority. |
| 661 | Falls back through BF16 and other backends. |
| 662 | """ |
| 663 | activation_format = ( |
| 664 | mk.FusedMoEActivationFormat.BatchedExperts |
| 665 | if config.moe_parallel_config.use_batched_activation_format |
| 666 | else mk.FusedMoEActivationFormat.Standard |
| 667 | ) |
| 668 | |
| 669 | # Honor explicit moe_backend (e.g. "marlin", "triton_unfused") before |
| 670 | # falling back to the auto priority list. |
| 671 | runner_backend = config.moe_backend |
| 672 | if runner_backend != "auto": |
| 673 | if runner_backend == "b12x": |
| 674 | requested_backends = _get_requested_backends(runner_backend, None) |
| 675 | else: |
| 676 | # Try every variant of the alias in priority order. Narrowing to |
| 677 | # the BF16 variant would drop SM100+ W4A8 variants on devices where |
| 678 | # the BF16 variant is gated to SM90. |
| 679 | requested_backends = map_mxfp4_backend(runner_backend) |
| 680 | if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 681 | requested_backends = [ |
| 682 | Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b |
| 683 | for b in requested_backends |
| 684 | ] |
| 685 | last_error: Exception | None = None |
| 686 | for requested_backend in requested_backends: |
| 687 | try: |
| 688 | return _return_or_raise( |
| 689 | requested_backend, |
| 690 | config, |
| 691 | kMxfp4Static, |
| 692 | _backend_activation_key(requested_backend), |
| 693 | activation_format, |
| 694 | ) |
| 695 | except ValueError as e: |
| 696 | last_error = e |
| 697 | assert last_error is not None |
| 698 | raise last_error |
| 699 | |
| 700 | # DeepSeek-V4 on ROCm: prefer AITER FlyDSL MoE (better perf + accuracy |
| 701 | # after shuffle/TP-offset fixes), with Triton-unfused as fallback. |
| 702 | if ( |
| 703 | current_platform.is_rocm() |
| 704 | and config.routing_method == RoutingMethodType.DeepseekV4 |
| 705 | ): |
| 706 | priority_backends = [ |
| 707 | Mxfp4MoeBackend.AITER_MXFP4_BF16, |
| 708 | Mxfp4MoeBackend.TRITON_UNFUSED, |
| 709 | ] |
| 710 | else: |
| 711 | priority_backends = _get_priority_backends() |
| 712 | |
| 713 | # Iterate priority backends: TRTLLM MXFP8, then Triton. |
| 714 | for backend in priority_backends: |
| 715 | activation_key = _backend_activation_key(backend) |
| 716 | for k_cls in backend_to_kernel_cls(backend): |
| 717 | supported, reason = k_cls.is_supported_config( |
| 718 | k_cls, config, kMxfp4Static, activation_key, activation_format |
| 719 | ) |
| 720 | if supported: |
| 721 | logger.info_once(_make_log_backend(backend), scope="local") |
| 722 | return backend, k_cls |
| 723 | else: |
| 724 | logger.debug_once(_make_log_unsupported(backend, reason), scope="local") |
| 725 | |
| 726 | raise NotImplementedError( |
| 727 | "No MXFP4 MoE backend supports the deployment configuration." |
| 728 | ) |
| 729 | |
| 730 | |
| 731 | def mxfp4_round_up_hidden_size_and_intermediate_size( |
| 732 | backend: Mxfp4MoeBackend, |
| 733 | hidden_size: int, |
| 734 | intermediate_size: int, |
| 735 | activation: MoEActivation | None = None, |
| 736 | ) -> tuple[int, int]: |
| 737 | """Round up hidden_size and intermediate_size based on backend requirements.""" |
| 738 | if backend in B12X_BACKENDS: |
| 739 | # b12x plans for the exact model dimensions. B12xExperts validates the |
| 740 | # required MXFP4 block alignment before selecting the backend. |
| 741 | return hidden_size, intermediate_size |
| 742 | if backend == Mxfp4MoeBackend.EMULATION: |
| 743 | # Emulation has no kernel tile; it only needs OCP MX block alignment so the |
| 744 | # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated |
| 745 | # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720). |
| 746 | intermediate_size = round_up(intermediate_size, OCP_MX_BLOCK_SIZE) |
| 747 | hidden_size = round_up(hidden_size, OCP_MX_BLOCK_SIZE) |
| 748 | elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 749 | # DeepGEMM requires M/N/K alignment |
| 750 | intermediate_size = round_up(intermediate_size, 128) |
| 751 | hidden_size = round_up(hidden_size, 128) |
| 752 | elif backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): |
| 753 | intermediate_size = round_up(intermediate_size, 128) |
| 754 | if current_platform.is_xpu(): |
| 755 | hidden_size = round_up(hidden_size, 128) |
| 756 | else: |
| 757 | hidden_size = round_up(hidden_size, 256) |
| 758 | elif backend in TRTLLM_BACKENDS: |
| 759 | intermediate_size = round_up(intermediate_size, 128) |
| 760 | hidden_size = round_up(hidden_size, 256) |
| 761 | elif backend in ( |
| 762 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 763 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 764 | ): |
| 765 | intermediate_size = round_up(intermediate_size, 128) |
| 766 | hidden_size = round_up(hidden_size, 128) |
| 767 | elif current_platform.is_rocm(): |
| 768 | from vllm.platforms.rocm import get_cdna_version |
| 769 | |
| 770 | is_situ_or_silu = activation in ( |
| 771 | MoEActivation.SITU, |
| 772 | MoEActivation.SILU, |
| 773 | ) |
| 774 | |
| 775 | # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size |
| 776 | # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 |
| 777 | # rather than the generic ROCm 256 round-up, which would inflate |
| 778 | # weights and OOM. |
| 779 | aiter_uses_128 = backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 |
| 780 | |
| 781 | # matmul_ogs uses block_k=128 for MXFP4 on pre-CDNA4 GPUs. |
| 782 | # CDNA4's F16xMXFP4 configuration uses block_k=256. |
| 783 | triton_uses_128 = ( |
| 784 | backend == Mxfp4MoeBackend.TRITON_UNFUSED and get_cdna_version() != 4 |
| 785 | ) |
| 786 | |
| 787 | alignment = ( |
| 788 | 128 if is_situ_or_silu and (aiter_uses_128 or triton_uses_128) else 256 |
| 789 | ) |
| 790 | |
| 791 | intermediate_size = round_up(intermediate_size, alignment) |
| 792 | hidden_size = round_up(hidden_size, alignment) |
| 793 | elif backend == Mxfp4MoeBackend.CPU: |
| 794 | # CPU AMX kernel uses BLOCK_N=32, align to 32 |
| 795 | intermediate_size = round_up(intermediate_size, 32) |
| 796 | hidden_size = round_up(hidden_size, 32) |
| 797 | else: |
| 798 | intermediate_size = round_up(intermediate_size, 64) |
| 799 | return hidden_size, intermediate_size |
| 800 | |
| 801 | |
| 802 | def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( |
| 803 | mxfp4_backend: Mxfp4MoeBackend, |
| 804 | layer: torch.nn.Module, |
| 805 | w13_weight: torch.Tensor, |
| 806 | w2_weight: torch.Tensor, |
| 807 | w13_weight_scale: torch.Tensor, |
| 808 | w2_weight_scale: torch.Tensor, |
| 809 | w13_bias: torch.Tensor | None = None, |
| 810 | w2_bias: torch.Tensor | None = None, |
| 811 | w13_input_scale: torch.Tensor | None = None, |
| 812 | w2_input_scale: torch.Tensor | None = None, |
| 813 | _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, |
| 814 | ) -> tuple[ |
| 815 | torch.Tensor, |
| 816 | torch.Tensor, |
| 817 | Union[torch.Tensor, "PrecisionConfig"], |
| 818 | Union[torch.Tensor, "PrecisionConfig"], |
| 819 | torch.Tensor | None, |
| 820 | torch.Tensor | None, |
| 821 | ]: |
| 822 | """Convert loaded weights into backend-specific kernel format.""" |
| 823 | if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 824 | w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( |
| 825 | w13_weight, |
| 826 | w2_weight, |
| 827 | w13_weight_scale, |
| 828 | w2_weight_scale, |
| 829 | ) |
| 830 | |
| 831 | return ( |
| 832 | w13_weight.data, |
| 833 | w2_weight.data, |
| 834 | w13_weight_scale, |
| 835 | w2_weight_scale, |
| 836 | w13_bias, |
| 837 | w2_bias, |
| 838 | ) |
| 839 | |
| 840 | num_experts = w13_weight.shape[0] |
| 841 | intermediate_size = w13_weight.shape[1] // 2 |
| 842 | hidden_size = w13_weight.shape[2] * 2 |
| 843 | |
| 844 | sf_block_size = 32 # mxfp4 block size |
| 845 | |
| 846 | if mxfp4_backend == Mxfp4MoeBackend.HUMMING: |
| 847 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 848 | convert_to_humming_moe_kernel_format, |
| 849 | ) |
| 850 | |
| 851 | convert_to_humming_moe_kernel_format( |
| 852 | layer, quant_config={"quant_method": "gpt_oss_mxfp4"} |
| 853 | ) |
| 854 | return ( |
| 855 | layer.w13_weight, |
| 856 | layer.w2_weight, |
| 857 | layer.w13_weight_scale, |
| 858 | layer.w2_weight_scale, |
| 859 | getattr(layer, "w13_bias", None), |
| 860 | getattr(layer, "w2_bias", None), |
| 861 | ) |
| 862 | elif mxfp4_backend in ( |
| 863 | Mxfp4MoeBackend.MARLIN, |
| 864 | Mxfp4MoeBackend.BATCHED_MARLIN, |
| 865 | ): |
| 866 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 867 | prepare_moe_mxfp4_layer_for_marlin, |
| 868 | ) |
| 869 | |
| 870 | return prepare_moe_mxfp4_layer_for_marlin( |
| 871 | layer, |
| 872 | w13_weight, |
| 873 | w2_weight, |
| 874 | w13_weight_scale, |
| 875 | w2_weight_scale, |
| 876 | w13_bias, |
| 877 | w2_bias, |
| 878 | ) |
| 879 | |
| 880 | elif mxfp4_backend in TRTLLM_BACKENDS: |
| 881 | assert _cache_permute_indices is not None |
| 882 | from flashinfer.fp4_quantization import nvfp4_block_scale_interleave |
| 883 | from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache |
| 884 | |
| 885 | # gemm1_alpha/beta/clamp_limit are created by the expert class |
| 886 | # (TrtLlmMxfp4ExpertsBase), not on the layer. |
| 887 | |
| 888 | w13_weight = w13_weight.data |
| 889 | w2_weight = w2_weight.data |
| 890 | w13_weight_scale = w13_weight_scale.data |
| 891 | w2_weight_scale = w2_weight_scale.data |
| 892 | assert w13_bias is not None and w2_bias is not None |
| 893 | w13_bias = w13_bias.data.to(torch.float32) |
| 894 | w2_bias = w2_bias.data.to(torch.float32) |
| 895 | |
| 896 | # Swap w1 and w3 as the definition of swiglu is different in trtllm-gen |
| 897 | def swap_every_two_rows(x, axis=-1): |
| 898 | shape = x.shape |
| 899 | if axis < 0: |
| 900 | axis = len(shape) + axis |
| 901 | new_shape = list(shape) |
| 902 | new_shape[axis] = shape[axis] // 2 |
| 903 | new_shape.insert(axis + 1, 2) |
| 904 | x = x.reshape(*new_shape) |
| 905 | x = x.flip(axis + 1) |
| 906 | new_shape = list(shape) |
| 907 | return x.reshape(*new_shape) |
| 908 | |
| 909 | w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2) |
| 910 | w13_weight = swap_every_two_rows(w13_weight, -2) |
| 911 | w13_bias = swap_every_two_rows(w13_bias, -1) |
| 912 | |
| 913 | # Shuffle weights and scaling factors for transposed mma output |
| 914 | gemm1_weights_shuffled = [] |
| 915 | gemm1_scales_shuffled = [] |
| 916 | gemm2_weights_shuffled = [] |
| 917 | gemm2_scales_shuffled = [] |
| 918 | gemm1_bias_shuffled = [] |
| 919 | gemm2_bias_shuffled = [] |
| 920 | epilogue_tile_m = 128 |
| 921 | for i in range(num_experts): |
| 922 | # w13 weight |
| 923 | permute_indices = get_w2_permute_indices_with_cache( |
| 924 | _cache_permute_indices, |
| 925 | w13_weight[i].view(torch.uint8), |
| 926 | epilogue_tile_m, |
| 927 | ) |
| 928 | gemm1_weights_shuffled.append( |
| 929 | w13_weight[i] |
| 930 | .view(torch.uint8)[permute_indices.to(w13_weight.device)] |
| 931 | .contiguous() |
| 932 | ) |
| 933 | # w13 scale |
| 934 | permute_sf_indices = get_w2_permute_indices_with_cache( |
| 935 | _cache_permute_indices, |
| 936 | w13_weight_scale[i].view(torch.uint8), |
| 937 | epilogue_tile_m, |
| 938 | num_elts_per_sf=16, |
| 939 | ) |
| 940 | gemm1_scales_shuffled.append( |
| 941 | nvfp4_block_scale_interleave( |
| 942 | w13_weight_scale[i] |
| 943 | .view(torch.uint8)[permute_sf_indices.to(w13_weight_scale.device)] |
| 944 | .contiguous() |
| 945 | ) |
| 946 | ) |
| 947 | # w13 bias |
| 948 | permute_bias_indices = get_w2_permute_indices_with_cache( |
| 949 | _cache_permute_indices, |
| 950 | w13_bias[i].clone().reshape(-1, 1), |
| 951 | epilogue_tile_m, |
| 952 | ) |
| 953 | gemm1_bias_shuffled.append( |
| 954 | w13_bias[i] |
| 955 | .clone() |
| 956 | .reshape(-1, 1)[permute_bias_indices.to(w13_bias.device)] |
| 957 | .contiguous() |
| 958 | ) |
| 959 | # w2 weight |
| 960 | permute_indices = get_w2_permute_indices_with_cache( |
| 961 | _cache_permute_indices, |
| 962 | w2_weight[i].view(torch.uint8), |
| 963 | epilogue_tile_m, |
| 964 | ) |
| 965 | gemm2_weights_shuffled.append( |
| 966 | w2_weight[i] |
| 967 | .view(torch.uint8)[permute_indices.to(w2_weight.device)] |
| 968 | .contiguous() |
| 969 | ) |
| 970 | # w2 scale |
| 971 | permute_sf_indices = get_w2_permute_indices_with_cache( |
| 972 | _cache_permute_indices, |
| 973 | w2_weight_scale[i].view(torch.uint8), |
| 974 | epilogue_tile_m, |
| 975 | num_elts_per_sf=16, |
| 976 | ) |
| 977 | gemm2_scales_shuffled.append( |
| 978 | nvfp4_block_scale_interleave( |
| 979 | w2_weight_scale[i] |
| 980 | .view(torch.uint8)[permute_sf_indices.to(w2_weight_scale.device)] |
| 981 | .contiguous() |
| 982 | ) |
| 983 | ) |
| 984 | # w2 bias |
| 985 | permute_indices = get_w2_permute_indices_with_cache( |
| 986 | _cache_permute_indices, |
| 987 | w2_bias[i].clone().reshape(-1, 1), |
| 988 | epilogue_tile_m, |
| 989 | ) |
| 990 | gemm2_bias_shuffled.append( |
| 991 | w2_bias[i] |
| 992 | .clone() |
| 993 | .reshape(-1, 1)[permute_indices.to(w2_bias.device)] |
| 994 | .contiguous() |
| 995 | ) |
| 996 | |
| 997 | w13_weight = torch.stack(gemm1_weights_shuffled) |
| 998 | w13_weight_scale = ( |
| 999 | torch.stack(gemm1_scales_shuffled) |
| 1000 | .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size) |
| 1001 | .view(torch.float8_e4m3fn) |
| 1002 | ) |
| 1003 | w2_weight = torch.stack(gemm2_weights_shuffled) |
| 1004 | w2_weight_scale = ( |
| 1005 | torch.stack(gemm2_scales_shuffled) |
| 1006 | .reshape(num_experts, hidden_size, intermediate_size // sf_block_size) |
| 1007 | .view(torch.float8_e4m3fn) |
| 1008 | ) |
| 1009 | w13_bias = torch.stack(gemm1_bias_shuffled).reshape(num_experts, -1) |
| 1010 | w2_bias = torch.stack(gemm2_bias_shuffled).reshape(num_experts, -1) |
| 1011 | |
| 1012 | return ( |
| 1013 | w13_weight, |
| 1014 | w2_weight, |
| 1015 | w13_weight_scale, |
| 1016 | w2_weight_scale, |
| 1017 | w13_bias, |
| 1018 | w2_bias, |
| 1019 | ) |
| 1020 | |
| 1021 | elif mxfp4_backend in ( |
| 1022 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 1023 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 1024 | ): |
| 1025 | # De-interleave and swap for w13 weight, bias, and scales |
| 1026 | w13_w = w13_weight.data |
| 1027 | gate_w, up_w = w13_w[:, ::2, :], w13_w[:, 1::2, :] |
| 1028 | deinterleaved_w13_w = torch.cat([gate_w, up_w], dim=1) |
| 1029 | w1_w, w3_w = torch.chunk(deinterleaved_w13_w, 2, dim=1) |
| 1030 | w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1) |
| 1031 | |
| 1032 | assert w13_bias is not None and w2_bias is not None |
| 1033 | w13_b = w13_bias.data.to(torch.float32) |
| 1034 | gate_b, up_b = w13_b[:, ::2], w13_b[:, 1::2] |
| 1035 | deinterleaved_w13_b = torch.cat([gate_b, up_b], dim=1) |
| 1036 | b1, b3 = torch.chunk(deinterleaved_w13_b, 2, dim=-1) |
| 1037 | w13_bias_swapped = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) |
| 1038 | |
| 1039 | w13_s = w13_weight_scale.data |
| 1040 | gate_s, up_s = w13_s[:, ::2, :], w13_s[:, 1::2, :] |
| 1041 | deinterleaved_w13_s = torch.cat([gate_s, up_s], dim=1) |
| 1042 | s1, s3 = torch.chunk(deinterleaved_w13_s, 2, dim=1) |
| 1043 | w13_scale_swapped = torch.cat([s3, s1], dim=1) |
| 1044 | |
| 1045 | if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8: |
| 1046 | from flashinfer import block_scale_interleave |
| 1047 | |
| 1048 | orig_shape = w13_scale_swapped.shape |
| 1049 | w13_scale_interleaved = block_scale_interleave( |
| 1050 | w13_scale_swapped.view(torch.uint8) |
| 1051 | ).reshape(orig_shape) |
| 1052 | |
| 1053 | w2_s = w2_weight_scale.data |
| 1054 | orig_shape = w2_s.shape |
| 1055 | w2_scale_interleaved = block_scale_interleave( |
| 1056 | w2_s.view(torch.uint8) |
| 1057 | ).reshape(orig_shape) |
| 1058 | |
| 1059 | return ( |
| 1060 | w13_weight_swapped, |
| 1061 | w2_weight, |
| 1062 | w13_scale_interleaved, |
| 1063 | w2_scale_interleaved, |
| 1064 | w13_bias_swapped, |
| 1065 | w2_bias, |
| 1066 | ) |
| 1067 | |
| 1068 | else: |
| 1069 | assert mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16 |
| 1070 | |
| 1071 | from flashinfer.fused_moe import ( |
| 1072 | interleave_moe_scales_for_sm90_mixed_gemm, |
| 1073 | interleave_moe_weights_for_sm90_mixed_gemm, |
| 1074 | ) |
| 1075 | |
| 1076 | w13_weight_interleaved = interleave_moe_weights_for_sm90_mixed_gemm( |
| 1077 | w13_weight_swapped.contiguous(), "fp4" |
| 1078 | ) |
| 1079 | w2_weight_interleaved = interleave_moe_weights_for_sm90_mixed_gemm( |
| 1080 | w2_weight.contiguous(), "fp4" |
| 1081 | ) |
| 1082 | w31_scales_interleaved = interleave_moe_scales_for_sm90_mixed_gemm( |
| 1083 | w13_scale_swapped.to(torch.uint8) |
| 1084 | ) |
| 1085 | w2_scale_interleaved = interleave_moe_scales_for_sm90_mixed_gemm( |
| 1086 | w2_weight_scale.data.to(torch.uint8) |
| 1087 | ) |
| 1088 | |
| 1089 | return ( |
| 1090 | w13_weight_interleaved, |
| 1091 | w2_weight_interleaved, |
| 1092 | w31_scales_interleaved, |
| 1093 | w2_scale_interleaved, |
| 1094 | w13_bias_swapped, |
| 1095 | w2_bias, |
| 1096 | ) |
| 1097 | |
| 1098 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: |
| 1099 | from vllm._aiter_ops import rocm_aiter_ops |
| 1100 | |
| 1101 | if w13_bias is not None: |
| 1102 | w13_bias = w13_bias.data.to(torch.float32) |
| 1103 | if w2_bias is not None: |
| 1104 | w2_bias = w2_bias.data.to(torch.float32) |
| 1105 | |
| 1106 | # e8m0_shuffle on weight scales (GFX950 swizzle layout) |
| 1107 | from aiter.utility.fp4_utils import e8m0_shuffle |
| 1108 | |
| 1109 | s0, s1, _ = w13_weight_scale.shape |
| 1110 | w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view( |
| 1111 | s0, s1, -1 |
| 1112 | ) |
| 1113 | |
| 1114 | s0, s1, _ = w2_weight_scale.shape |
| 1115 | w2_weight_scale.data = e8m0_shuffle(w2_weight_scale.view(s0 * s1, -1)).view( |
| 1116 | s0, s1, -1 |
| 1117 | ) |
| 1118 | |
| 1119 | # View as native FP4 dtype |
| 1120 | fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) |
| 1121 | if fp4_dtype is not None: |
| 1122 | w13_weight.data = w13_weight.data.view(fp4_dtype) |
| 1123 | w2_weight.data = w2_weight.data.view(fp4_dtype) |
| 1124 | |
| 1125 | # Shuffle weights for AITER CK kernel |
| 1126 | shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights( |
| 1127 | w13_weight, w2_weight |
| 1128 | ) |
| 1129 | shuffled_w13.is_shuffled = True |
| 1130 | shuffled_w2.is_shuffled = True |
| 1131 | |
| 1132 | return ( |
| 1133 | shuffled_w13, |
| 1134 | shuffled_w2, |
| 1135 | w13_weight_scale, |
| 1136 | w2_weight_scale, |
| 1137 | w13_bias, |
| 1138 | w2_bias, |
| 1139 | ) |
| 1140 | |
| 1141 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: |
| 1142 | from vllm._aiter_ops import rocm_aiter_ops |
| 1143 | |
| 1144 | if w13_bias is not None: |
| 1145 | w13_bias = w13_bias.data.to(torch.float32) |
| 1146 | if w2_bias is not None: |
| 1147 | w2_bias = w2_bias.data.to(torch.float32) |
| 1148 | |
| 1149 | e, n, k = w13_weight.shape |
| 1150 | |
| 1151 | # De-interleave w13 rows: gate/up pairs -> contiguous gate, up blocks |
| 1152 | w13_weight.view(torch.uint8).copy_( |
| 1153 | w13_weight.data.view(torch.uint8) |
| 1154 | .view(e, n // 2, 2, k) |
| 1155 | .permute(0, 2, 1, 3) |
| 1156 | .contiguous() |
| 1157 | .view(e, n, k) |
| 1158 | ) |
| 1159 | w13_weight_scale.data = ( |
| 1160 | w13_weight_scale.data.view(e, n // 2, 2, -1) |
| 1161 | .permute(0, 2, 1, 3) |
| 1162 | .contiguous() |
| 1163 | .view(e, n, -1) |
| 1164 | ) |
| 1165 | |
| 1166 | # View as native FP4 dtype for AITER shuffle |
| 1167 | w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2) |
| 1168 | w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2) |
| 1169 | |
| 1170 | # Shuffle weights and scales for AITER CK kernel layout |
| 1171 | w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True) |
| 1172 | shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( |
| 1173 | w13_weight_scale.view(-1, w13_weight_scale.shape[-1]), |
| 1174 | num_experts, |
| 1175 | True, |
| 1176 | ) |
| 1177 | |
| 1178 | w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False) |
| 1179 | shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4( |
| 1180 | w2_weight_scale.view(-1, w2_weight_scale.shape[-1]), |
| 1181 | num_experts, |
| 1182 | False, |
| 1183 | ) |
| 1184 | |
| 1185 | # Permute bias to match de-interleaved weight layout |
| 1186 | if w13_bias is not None: |
| 1187 | w13_bias = ( |
| 1188 | w13_bias.data.view(-1, n // 2, 2) |
| 1189 | .permute(0, 2, 1) |
| 1190 | .contiguous() |
| 1191 | .view(-1, n) |
| 1192 | ) |
| 1193 | |
| 1194 | w13_weight.is_shuffled = True |
| 1195 | w2_weight.is_shuffled = True |
| 1196 | |
| 1197 | return ( |
| 1198 | w13_weight, |
| 1199 | w2_weight, |
| 1200 | shuffled_w13_scale, |
| 1201 | shuffled_w2_scale, |
| 1202 | w13_bias, |
| 1203 | w2_bias, |
| 1204 | ) |
| 1205 | |
| 1206 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: |
| 1207 | # W4A8: MXFP4 weights + static FP8 activations (triton kernel) |
| 1208 | from triton_kernels.numerics import InFlexData |
| 1209 | |
| 1210 | if w13_bias is not None: |
| 1211 | w13_bias = w13_bias.to(torch.float32) |
| 1212 | if w2_bias is not None: |
| 1213 | w2_bias = w2_bias.to(torch.float32) |
| 1214 | |
| 1215 | # Process static FP8 input scales (reduce to scalar, warn if not uniform) |
| 1216 | w13_input_scale = layer.w13_input_scale |
| 1217 | w2_input_scale = layer.w2_input_scale |
| 1218 | if w13_input_scale is None or w2_input_scale is None: |
| 1219 | raise ValueError( |
| 1220 | "W4A8 (AITER_MXFP4_FP8) requires static input scales, but found " |
| 1221 | "w13_input_scale or w2_input_scale is None." |
| 1222 | ) |
| 1223 | if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale): |
| 1224 | logger.warning_once( |
| 1225 | "Found input_scales that are not equal for " |
| 1226 | "fp8 MoE layer. Using the maximum across experts " |
| 1227 | "for each layer." |
| 1228 | ) |
| 1229 | w13_input_scale = w13_input_scale.max().to(torch.float32) |
| 1230 | w2_input_scale = w2_input_scale.max().to(torch.float32) |
| 1231 | |
| 1232 | # Swizzle weights for GFX950 |
| 1233 | w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(w13_weight, w13_weight_scale) |
| 1234 | w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(w2_weight, w2_weight_scale) |
| 1235 | |
| 1236 | # Create InFlexData for activation scales |
| 1237 | lhs_data13 = InFlexData(scale=w13_input_scale) |
| 1238 | lhs_data2 = InFlexData(scale=w2_input_scale) |
| 1239 | |
| 1240 | # Create PrecisionConfig with both weight and activation info |
| 1241 | w13_precision_config = PrecisionConfig( |
| 1242 | **_mx_scale_kwargs(w13_scale), |
| 1243 | flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13), |
| 1244 | ) |
| 1245 | w2_precision_config = PrecisionConfig( |
| 1246 | **_mx_scale_kwargs(w2_scale), |
| 1247 | flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2), |
| 1248 | ) |
| 1249 | |
| 1250 | del layer.w13_weight |
| 1251 | del layer.w2_weight |
| 1252 | |
| 1253 | return ( |
| 1254 | w13_weight, |
| 1255 | w2_weight, |
| 1256 | w13_precision_config, |
| 1257 | w2_precision_config, |
| 1258 | w13_bias, |
| 1259 | w2_bias, |
| 1260 | ) |
| 1261 | |
| 1262 | elif mxfp4_backend in TRITON_BACKENDS: |
| 1263 | if w13_bias is not None: |
| 1264 | w13_bias = w13_bias.to(torch.float32) |
| 1265 | if w2_bias is not None: |
| 1266 | w2_bias = w2_bias.to(torch.float32) |
| 1267 | |
| 1268 | w13_weight, w13_flex, w13_scale = _swizzle_mxfp4( |
| 1269 | w13_weight, |
| 1270 | w13_weight_scale, |
| 1271 | ) |
| 1272 | w2_weight, w2_flex, w2_scale = _swizzle_mxfp4( |
| 1273 | w2_weight, |
| 1274 | w2_weight_scale, |
| 1275 | ) |
| 1276 | |
| 1277 | w13_precision_config = PrecisionConfig( |
| 1278 | **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) |
| 1279 | ) |
| 1280 | w2_precision_config = PrecisionConfig( |
| 1281 | **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) |
| 1282 | ) |
| 1283 | |
| 1284 | # The original mxfp4 block scales have been swizzled into the |
| 1285 | # precision configs above and are no longer read by the kernel, so |
| 1286 | # drop the now-dead weight/scale Parameters to free their memory. |
| 1287 | del layer.w13_weight |
| 1288 | del layer.w2_weight |
| 1289 | del layer.w13_weight_scale |
| 1290 | del layer.w2_weight_scale |
| 1291 | |
| 1292 | return ( |
| 1293 | w13_weight, |
| 1294 | w2_weight, |
| 1295 | w13_precision_config, |
| 1296 | w2_precision_config, |
| 1297 | w13_bias, |
| 1298 | w2_bias, |
| 1299 | ) |
| 1300 | elif mxfp4_backend == Mxfp4MoeBackend.XPU: |
| 1301 | # No additional transformation needed for XPU backend |
| 1302 | return ( |
| 1303 | w13_weight, |
| 1304 | w2_weight, |
| 1305 | w13_weight_scale, |
| 1306 | w2_weight_scale, |
| 1307 | w13_bias, |
| 1308 | w2_bias, |
| 1309 | ) |
| 1310 | elif mxfp4_backend == Mxfp4MoeBackend.CPU: |
| 1311 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 1312 | prepare_mxfp4_moe_layer_for_cpu, |
| 1313 | ) |
| 1314 | |
| 1315 | packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = ( |
| 1316 | prepare_mxfp4_moe_layer_for_cpu( |
| 1317 | w13_weight.data, |
| 1318 | w2_weight.data, |
| 1319 | w13_weight_scale.data, |
| 1320 | w2_weight_scale.data, |
| 1321 | ) |
| 1322 | ) |
| 1323 | if w13_bias is not None: |
| 1324 | w13_bias = w13_bias.data.to(torch.float32) |
| 1325 | if w2_bias is not None: |
| 1326 | w2_bias = w2_bias.data.to(torch.float32) |
| 1327 | return ( |
| 1328 | packed_w13, |
| 1329 | packed_w2, |
| 1330 | packed_w13_scale, |
| 1331 | packed_w2_scale, |
| 1332 | w13_bias, |
| 1333 | w2_bias, |
| 1334 | ) |
| 1335 | elif mxfp4_backend == Mxfp4MoeBackend.EMULATION: |
| 1336 | w13_has_per_expert_scale = ( |
| 1337 | w13_input_scale is not None |
| 1338 | and w13_input_scale.ndim == 1 |
| 1339 | and not all_close_1d(w13_input_scale) |
| 1340 | ) |
| 1341 | w2_has_per_expert_scale = ( |
| 1342 | w2_input_scale is not None |
| 1343 | and w2_input_scale.ndim == 1 |
| 1344 | and not all_close_1d(w2_input_scale) |
| 1345 | ) |
| 1346 | if w13_has_per_expert_scale or w2_has_per_expert_scale: |
| 1347 | logger.warning_once( |
| 1348 | "Found input_scales that are not equal for OCP MX MoE " |
| 1349 | "emulation. Using the maximum across experts for each layer." |
| 1350 | ) |
| 1351 | if w13_input_scale is not None: |
| 1352 | layer.w13_input_scale = torch.nn.Parameter( |
| 1353 | w13_input_scale.max().to(torch.float32), requires_grad=False |
| 1354 | ) |
| 1355 | if w2_input_scale is not None: |
| 1356 | layer.w2_input_scale = torch.nn.Parameter( |
| 1357 | w2_input_scale.max().to(torch.float32), requires_grad=False |
| 1358 | ) |
| 1359 | return ( |
| 1360 | w13_weight, |
| 1361 | w2_weight, |
| 1362 | w13_weight_scale, |
| 1363 | w2_weight_scale, |
| 1364 | w13_bias, |
| 1365 | w2_bias, |
| 1366 | ) |
| 1367 | else: |
| 1368 | raise ValueError( |
| 1369 | f"Unsupported mxfp4_backend: {mxfp4_backend}: " |
| 1370 | f"should be one of: {list(Mxfp4MoeBackend)}." |
| 1371 | ) |
| 1372 | |
| 1373 | |
| 1374 | def convert_weight_to_mxfp4_moe_kernel_format( |
| 1375 | mxfp4_backend: Mxfp4MoeBackend, |
| 1376 | layer: torch.nn.Module, |
| 1377 | w13_weight: torch.Tensor, |
| 1378 | w2_weight: torch.Tensor, |
| 1379 | w13_weight_scale: torch.Tensor, |
| 1380 | w2_weight_scale: torch.Tensor, |
| 1381 | w13_bias: torch.Tensor | None = None, |
| 1382 | w2_bias: torch.Tensor | None = None, |
| 1383 | _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, |
| 1384 | activation: MoEActivation | None = None, |
| 1385 | use_separated_a4w4: bool = False, |
| 1386 | ) -> tuple[ |
| 1387 | torch.Tensor, |
| 1388 | torch.Tensor, |
| 1389 | Union[torch.Tensor, "PrecisionConfig"], |
| 1390 | Union[torch.Tensor, "PrecisionConfig"], |
| 1391 | torch.Tensor | None, |
| 1392 | torch.Tensor | None, |
| 1393 | ]: |
| 1394 | """Convert loaded weights into backend-specific kernel format. |
| 1395 | |
| 1396 | Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends. |
| 1397 | """ |
| 1398 | is_gfx1250 = False |
| 1399 | if current_platform.is_rocm(): |
| 1400 | from vllm.platforms.rocm import on_gfx1250 |
| 1401 | |
| 1402 | is_gfx1250 = on_gfx1250() |
| 1403 | |
| 1404 | if mxfp4_backend in B12X_BACKENDS: |
| 1405 | return ( |
| 1406 | w13_weight.data, |
| 1407 | w2_weight.data, |
| 1408 | w13_weight_scale.data, |
| 1409 | w2_weight_scale.data, |
| 1410 | w13_bias, |
| 1411 | w2_bias, |
| 1412 | ) |
| 1413 | |
| 1414 | if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 1415 | w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( |
| 1416 | w13_weight, |
| 1417 | w2_weight, |
| 1418 | w13_weight_scale, |
| 1419 | w2_weight_scale, |
| 1420 | ) |
| 1421 | |
| 1422 | return ( |
| 1423 | w13_weight.data, |
| 1424 | w2_weight.data, |
| 1425 | w13_weight_scale, |
| 1426 | w2_weight_scale, |
| 1427 | w13_bias, |
| 1428 | w2_bias, |
| 1429 | ) |
| 1430 | |
| 1431 | if mxfp4_backend == Mxfp4MoeBackend.HUMMING: |
| 1432 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 1433 | convert_to_humming_moe_kernel_format, |
| 1434 | ) |
| 1435 | |
| 1436 | convert_to_humming_moe_kernel_format( |
| 1437 | layer, quant_config={"quant_method": "mxfp4"} |
| 1438 | ) |
| 1439 | return ( |
| 1440 | layer.w13_weight, |
| 1441 | layer.w2_weight, |
| 1442 | layer.w13_weight_scale, |
| 1443 | layer.w2_weight_scale, |
| 1444 | getattr(layer, "w13_bias", None), |
| 1445 | getattr(layer, "w2_bias", None), |
| 1446 | ) |
| 1447 | |
| 1448 | if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): |
| 1449 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 1450 | prepare_moe_mxfp4_layer_for_marlin, |
| 1451 | ) |
| 1452 | |
| 1453 | return prepare_moe_mxfp4_layer_for_marlin( |
| 1454 | layer, |
| 1455 | w13_weight, |
| 1456 | w2_weight, |
| 1457 | w13_weight_scale, |
| 1458 | w2_weight_scale, |
| 1459 | w13_bias, |
| 1460 | w2_bias, |
| 1461 | ) |
| 1462 | |
| 1463 | num_experts = w13_weight.shape[0] |
| 1464 | intermediate_size = w13_weight.shape[1] // 2 |
| 1465 | hidden_size = w13_weight.shape[2] * 2 |
| 1466 | |
| 1467 | sf_block_size = 32 # mxfp4 block size |
| 1468 | |
| 1469 | if mxfp4_backend in TRTLLM_BACKENDS: |
| 1470 | assert _cache_permute_indices is not None |
| 1471 | from flashinfer.fp4_quantization import nvfp4_block_scale_interleave |
| 1472 | from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache |
| 1473 | |
| 1474 | w13_weight = w13_weight.data |
| 1475 | w2_weight = w2_weight.data |
| 1476 | w13_weight_scale = w13_weight_scale.data |
| 1477 | w2_weight_scale = w2_weight_scale.data |
| 1478 | if w13_bias is not None: |
| 1479 | w13_bias = w13_bias.data.to(torch.float32) |
| 1480 | if w2_bias is not None: |
| 1481 | w2_bias = w2_bias.data.to(torch.float32) |
| 1482 | |
| 1483 | # Swap w1/w3 and interleave to match TRTLLM SwiGLU convention. |
| 1484 | # Standard loading gives contiguous [w1/gate, w3/up]. |
| 1485 | # TRTLLM kernel expects interleaved [w3_0, w1_0, w3_1, w1_1, ...]. |
| 1486 | w1_weight = w13_weight[:, :intermediate_size, :] |
| 1487 | w3_weight = w13_weight[:, intermediate_size:, :] |
| 1488 | w13_weight = torch.stack([w3_weight, w1_weight], dim=2).reshape( |
| 1489 | w13_weight.shape |
| 1490 | ) |
| 1491 | |
| 1492 | w1_scale = w13_weight_scale[:, :intermediate_size, :] |
| 1493 | w3_scale = w13_weight_scale[:, intermediate_size:, :] |
| 1494 | w13_weight_scale = torch.stack([w3_scale, w1_scale], dim=2).reshape( |
| 1495 | w13_weight_scale.shape |
| 1496 | ) |
| 1497 | |
| 1498 | if w13_bias is not None: |
| 1499 | b1 = w13_bias[:, :intermediate_size] |
| 1500 | b3 = w13_bias[:, intermediate_size:] |
| 1501 | w13_bias = torch.stack([b3, b1], dim=2).reshape(w13_bias.shape) |
| 1502 | |
| 1503 | # Shuffle weights and scaling factors for transposed mma output. |
| 1504 | # Permute indices depend only on shape (cached by torch.Size), |
| 1505 | # so compute once and apply to all experts via batched indexing. |
| 1506 | epilogue_tile_m = 128 |
| 1507 | |
| 1508 | # w13 weight permute |
| 1509 | w13_perm = get_w2_permute_indices_with_cache( |
| 1510 | _cache_permute_indices, |
| 1511 | w13_weight[0].view(torch.uint8), |
| 1512 | epilogue_tile_m, |
| 1513 | ).to(w13_weight.device) |
| 1514 | w13_weight = w13_weight.view(torch.uint8)[:, w13_perm].contiguous() |
| 1515 | |
| 1516 | # w13 scale permute + interleave |
| 1517 | w13_sf_perm = get_w2_permute_indices_with_cache( |
| 1518 | _cache_permute_indices, |
| 1519 | w13_weight_scale[0].view(torch.uint8), |
| 1520 | epilogue_tile_m, |
| 1521 | num_elts_per_sf=16, |
| 1522 | ).to(w13_weight_scale.device) |
| 1523 | w13_s = w13_weight_scale.view(torch.uint8)[:, w13_sf_perm].contiguous() |
| 1524 | E, N_s, K_s = w13_s.shape |
| 1525 | w13_weight_scale = ( |
| 1526 | nvfp4_block_scale_interleave(w13_s.reshape(E * N_s, K_s)) |
| 1527 | .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size) |
| 1528 | .view(torch.float8_e4m3fn) |
| 1529 | ) |
| 1530 | |
| 1531 | # w2 weight permute |
| 1532 | w2_perm = get_w2_permute_indices_with_cache( |
| 1533 | _cache_permute_indices, |
| 1534 | w2_weight[0].view(torch.uint8), |
| 1535 | epilogue_tile_m, |
| 1536 | ).to(w2_weight.device) |
| 1537 | w2_weight = w2_weight.view(torch.uint8)[:, w2_perm].contiguous() |
| 1538 | |
| 1539 | # w2 scale permute + interleave |
| 1540 | w2_sf_perm = get_w2_permute_indices_with_cache( |
| 1541 | _cache_permute_indices, |
| 1542 | w2_weight_scale[0].view(torch.uint8), |
| 1543 | epilogue_tile_m, |
| 1544 | num_elts_per_sf=16, |
| 1545 | ).to(w2_weight_scale.device) |
| 1546 | w2_s = w2_weight_scale.view(torch.uint8)[:, w2_sf_perm].contiguous() |
| 1547 | E2, N2_s, K2_s = w2_s.shape |
| 1548 | w2_weight_scale = ( |
| 1549 | nvfp4_block_scale_interleave(w2_s.reshape(E2 * N2_s, K2_s)) |
| 1550 | .reshape(num_experts, hidden_size, intermediate_size // sf_block_size) |
| 1551 | .view(torch.float8_e4m3fn) |
| 1552 | ) |
| 1553 | |
| 1554 | # w13 bias permute |
| 1555 | if w13_bias is not None: |
| 1556 | w13_b_perm = get_w2_permute_indices_with_cache( |
| 1557 | _cache_permute_indices, |
| 1558 | w13_bias[0].reshape(-1, 1), |
| 1559 | epilogue_tile_m, |
| 1560 | ).to(w13_bias.device) |
| 1561 | w13_bias = w13_bias.reshape(num_experts, -1, 1)[:, w13_b_perm].reshape( |
| 1562 | num_experts, -1 |
| 1563 | ) |
| 1564 | |
| 1565 | # w2 bias permute |
| 1566 | if w2_bias is not None: |
| 1567 | w2_b_perm = get_w2_permute_indices_with_cache( |
| 1568 | _cache_permute_indices, |
| 1569 | w2_bias[0].reshape(-1, 1), |
| 1570 | epilogue_tile_m, |
| 1571 | ).to(w2_bias.device) |
| 1572 | w2_bias = w2_bias.reshape(num_experts, -1, 1)[:, w2_b_perm].reshape( |
| 1573 | num_experts, -1 |
| 1574 | ) |
| 1575 | |
| 1576 | return ( |
| 1577 | w13_weight, |
| 1578 | w2_weight, |
| 1579 | w13_weight_scale, |
| 1580 | w2_weight_scale, |
| 1581 | w13_bias, |
| 1582 | w2_bias, |
| 1583 | ) |
| 1584 | |
| 1585 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250: |
| 1586 | # Initially introduced for DeepSeekV4 |
| 1587 | |
| 1588 | if w13_bias is not None: |
| 1589 | w13_bias = w13_bias.data.to(torch.float32) |
| 1590 | if w2_bias is not None: |
| 1591 | w2_bias = w2_bias.data.to(torch.float32) |
| 1592 | |
| 1593 | if activation == MoEActivation.SITU: |
| 1594 | from aiter.utility.fp4_utils import e8m0_shuffle |
| 1595 | |
| 1596 | from vllm._aiter_ops import rocm_aiter_ops |
| 1597 | |
| 1598 | fp4_dtype = torch.float4_e2m1fn_x2 |
| 1599 | e8m0_dtype = torch.float8_e8m0fnu |
| 1600 | guinterleave = rocm_aiter_ops.is_fused_moe_situv2_gate_up_interleaved() |
| 1601 | w13 = rocm_aiter_ops.shuffle_weight_a16w4( |
| 1602 | w13_weight.data.view(fp4_dtype), 16, guinterleave |
| 1603 | ) |
| 1604 | w2 = rocm_aiter_ops.shuffle_weight_a16w4( |
| 1605 | w2_weight.data.view(fp4_dtype), 16, False |
| 1606 | ) |
| 1607 | w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype) |
| 1608 | w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype) |
| 1609 | w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( |
| 1610 | w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), |
| 1611 | num_experts, |
| 1612 | guinterleave, |
| 1613 | ) |
| 1614 | w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) |
| 1615 | w13.is_shuffled = True |
| 1616 | w2.is_shuffled = True |
| 1617 | return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) |
| 1618 | |
| 1619 | import os |
| 1620 | |
| 1621 | # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to |
| 1622 | # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2 |
| 1623 | # a4w4 is separated and selects q_dtype_a independently, so the bound |
| 1624 | # is unused on that path. |
| 1625 | os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" |
| 1626 | |
| 1627 | from aiter.ops.shuffle import shuffle_scale as _shuf_s |
| 1628 | from aiter.ops.shuffle import shuffle_weight as _shuf_w |
| 1629 | |
| 1630 | # DeepSeek V4.1 a4w4 uses ATOM's SEPARATED gate/up layout instead of |
| 1631 | # the default INTERLEAVE shuffle (INTERLEAVE + fp4x2 has no tuned |
| 1632 | # kernel and produces garbage output). Must match GateMode.SEPARATED |
| 1633 | # in rocm_aiter_moe.py. |
| 1634 | is_guinterleave = not use_separated_a4w4 |
| 1635 | |
| 1636 | w13_weight = torch.nn.Parameter( |
| 1637 | _shuf_w( |
| 1638 | w13_weight.data.view(torch.float4_e2m1fn_x2), |
| 1639 | is_guinterleave=is_guinterleave, |
| 1640 | gate_up=True, |
| 1641 | ), |
| 1642 | requires_grad=False, |
| 1643 | ) |
| 1644 | shuffled_w13_scale = _shuf_s( |
| 1645 | w13_weight_scale.reshape(-1, w13_weight_scale.shape[-1]), |
| 1646 | num_experts, |
| 1647 | is_guinterleave, |
| 1648 | True, |
| 1649 | ) |
| 1650 | |
| 1651 | w2_weight = torch.nn.Parameter( |
| 1652 | _shuf_w( |
| 1653 | w2_weight.data.view(torch.float4_e2m1fn_x2), |
| 1654 | is_guinterleave=is_guinterleave, |
| 1655 | gate_up=False, |
| 1656 | ), |
| 1657 | requires_grad=False, |
| 1658 | ) |
| 1659 | # use_gu_interleave |
| 1660 | shuffled_w2_scale = _shuf_s( |
| 1661 | w2_weight_scale.reshape(-1, w2_weight_scale.shape[-1]), |
| 1662 | num_experts, |
| 1663 | is_guinterleave, |
| 1664 | False, |
| 1665 | ) |
| 1666 | |
| 1667 | w13_weight.is_shuffled = True |
| 1668 | w2_weight.is_shuffled = True |
| 1669 | |
| 1670 | return ( |
| 1671 | w13_weight, |
| 1672 | w2_weight, |
| 1673 | shuffled_w13_scale, |
| 1674 | shuffled_w2_scale, |
| 1675 | w13_bias, |
| 1676 | w2_bias, |
| 1677 | ) |
| 1678 | |
| 1679 | elif mxfp4_backend in TRITON_BACKENDS or ( |
| 1680 | mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250 |
| 1681 | ): |
| 1682 | if mxfp4_backend == Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16: |
| 1683 | # AITER moe_gemm_a16w4 needs gate/up interleaved |
| 1684 | def interleave_gate_up(w: torch.Tensor) -> torch.Tensor: |
| 1685 | gate, up = w.chunk(2, dim=1) |
| 1686 | return torch.stack((gate, up), dim=2).reshape(w.shape) |
| 1687 | |
| 1688 | w13_weight = interleave_gate_up(w13_weight) |
| 1689 | w13_weight_scale = interleave_gate_up(w13_weight_scale) |
| 1690 | |
| 1691 | if w13_bias is not None: |
| 1692 | w13_bias = interleave_gate_up(w13_bias.to(torch.float32)) |
| 1693 | elif mxfp4_backend == Mxfp4MoeBackend.TRITON: |
| 1694 | |
| 1695 | def shuffle_weight(w: torch.Tensor) -> torch.Tensor: |
| 1696 | shape = w.shape |
| 1697 | n = shape[-1] |
| 1698 | first = w[..., : n // 2] |
| 1699 | second = w[..., n // 2 :] |
| 1700 | stacked = torch.stack((first, second), dim=-1) |
| 1701 | return stacked.reshape(shape) |
| 1702 | |
| 1703 | w13_weight = shuffle_weight(w13_weight) |
| 1704 | w13_weight_scale = shuffle_weight(w13_weight_scale) |
| 1705 | |
| 1706 | if w13_bias is not None: |
| 1707 | w13_bias = shuffle_weight(w13_bias.to(torch.float32)) |
| 1708 | else: |
| 1709 | if w13_bias is not None: |
| 1710 | w13_bias = w13_bias.to(torch.float32) |
| 1711 | |
| 1712 | if w2_bias is not None: |
| 1713 | w2_bias = w2_bias.to(torch.float32) |
| 1714 | |
| 1715 | w13_weight, w13_flex, w13_scale = _swizzle_mxfp4( |
| 1716 | w13_weight, |
| 1717 | w13_weight_scale, |
| 1718 | ) |
| 1719 | w2_weight, w2_flex, w2_scale = _swizzle_mxfp4( |
| 1720 | w2_weight, |
| 1721 | w2_weight_scale, |
| 1722 | ) |
| 1723 | |
| 1724 | w13_precision_config = PrecisionConfig( |
| 1725 | **_mx_scale_kwargs(w13_scale), flex_ctx=FlexCtx(rhs_data=w13_flex) |
| 1726 | ) |
| 1727 | w2_precision_config = PrecisionConfig( |
| 1728 | **_mx_scale_kwargs(w2_scale), flex_ctx=FlexCtx(rhs_data=w2_flex) |
| 1729 | ) |
| 1730 | |
| 1731 | # The original mxfp4 block scales have been swizzled into the |
| 1732 | # precision configs above and are no longer read by the kernel, so |
| 1733 | # drop the now-dead weight/scale Parameters to free their memory. |
| 1734 | del layer.w13_weight |
| 1735 | del layer.w2_weight |
| 1736 | del layer.w13_weight_scale |
| 1737 | del layer.w2_weight_scale |
| 1738 | |
| 1739 | return ( |
| 1740 | w13_weight, |
| 1741 | w2_weight, |
| 1742 | w13_precision_config, |
| 1743 | w2_precision_config, |
| 1744 | w13_bias, |
| 1745 | w2_bias, |
| 1746 | ) |
| 1747 | elif mxfp4_backend in ( |
| 1748 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 1749 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 1750 | ): |
| 1751 | # Standard checkpoints store fused gate/up tensors as [w1; w3], while |
| 1752 | # FlashInfer CUTLASS consumes [w3; w1]. Keep weights, scales, and bias |
| 1753 | # in the same order before applying the backend-specific interleave. |
| 1754 | w13_weight = swap_w13_to_w31(w13_weight.data) |
| 1755 | w13_weight_scale = swap_w13_to_w31(w13_weight_scale.data) |
| 1756 | if w13_bias is not None: |
| 1757 | b1, b3 = torch.chunk(w13_bias.data, 2, dim=-1) |
| 1758 | w13_bias = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) |
| 1759 | if w2_bias is not None: |
| 1760 | w2_bias = w2_bias.data.to(torch.bfloat16) |
| 1761 | |
| 1762 | if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8: |
| 1763 | from flashinfer import block_scale_interleave |
| 1764 | |
| 1765 | w13_scale_shape = w13_weight_scale.shape |
| 1766 | w13_weight_scale = block_scale_interleave( |
| 1767 | w13_weight_scale.view(torch.uint8) |
| 1768 | ).reshape(w13_scale_shape) |
| 1769 | |
| 1770 | w2_scale_shape = w2_weight_scale.shape |
| 1771 | w2_weight_scale = block_scale_interleave( |
| 1772 | w2_weight_scale.data.view(torch.uint8) |
| 1773 | ).reshape(w2_scale_shape) |
| 1774 | else: |
| 1775 | from flashinfer.fused_moe import ( |
| 1776 | interleave_moe_scales_for_sm90_mixed_gemm, |
| 1777 | interleave_moe_weights_for_sm90_mixed_gemm, |
| 1778 | ) |
| 1779 | |
| 1780 | w13_weight = interleave_moe_weights_for_sm90_mixed_gemm( |
| 1781 | w13_weight.contiguous(), "fp4" |
| 1782 | ) |
| 1783 | w2_weight = interleave_moe_weights_for_sm90_mixed_gemm( |
| 1784 | w2_weight.data.contiguous(), "fp4" |
| 1785 | ) |
| 1786 | w13_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm( |
| 1787 | w13_weight_scale.to(torch.uint8) |
| 1788 | ) |
| 1789 | w2_weight_scale = interleave_moe_scales_for_sm90_mixed_gemm( |
| 1790 | w2_weight_scale.data.to(torch.uint8) |
| 1791 | ) |
| 1792 | |
| 1793 | return ( |
| 1794 | w13_weight, |
| 1795 | w2_weight, |
| 1796 | w13_weight_scale, |
| 1797 | w2_weight_scale, |
| 1798 | w13_bias, |
| 1799 | w2_bias, |
| 1800 | ) |
| 1801 | elif mxfp4_backend in ( |
| 1802 | Mxfp4MoeBackend.XPU, |
| 1803 | Mxfp4MoeBackend.EMULATION, |
| 1804 | ): |
| 1805 | # No additional transformation is needed: XPU consumes the checkpoint |
| 1806 | # layout directly, while emulation dequantizes that layout at runtime. |
| 1807 | return ( |
| 1808 | w13_weight, |
| 1809 | w2_weight, |
| 1810 | w13_weight_scale, |
| 1811 | w2_weight_scale, |
| 1812 | w13_bias, |
| 1813 | w2_bias, |
| 1814 | ) |
| 1815 | elif mxfp4_backend in ( |
| 1816 | Mxfp4MoeBackend.AITER_MXFP4_MXFP4, |
| 1817 | Mxfp4MoeBackend.AITER_MXFP4_FP8, |
| 1818 | ): |
| 1819 | return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( |
| 1820 | mxfp4_backend=mxfp4_backend, |
| 1821 | layer=layer, |
| 1822 | w13_weight=w13_weight, |
| 1823 | w2_weight=w2_weight, |
| 1824 | w13_weight_scale=w13_weight_scale, |
| 1825 | w2_weight_scale=w2_weight_scale, |
| 1826 | w13_bias=w13_bias, |
| 1827 | w2_bias=w2_bias, |
| 1828 | _cache_permute_indices=_cache_permute_indices, |
| 1829 | ) |
| 1830 | elif mxfp4_backend == Mxfp4MoeBackend.CPU: |
| 1831 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 1832 | prepare_mxfp4_moe_layer_for_cpu, |
| 1833 | ) |
| 1834 | |
| 1835 | packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = ( |
| 1836 | prepare_mxfp4_moe_layer_for_cpu( |
| 1837 | w13_weight.data, |
| 1838 | w2_weight.data, |
| 1839 | w13_weight_scale.data, |
| 1840 | w2_weight_scale.data, |
| 1841 | ) |
| 1842 | ) |
| 1843 | if w13_bias is not None: |
| 1844 | w13_bias = w13_bias.data.to(torch.float32) |
| 1845 | if w2_bias is not None: |
| 1846 | w2_bias = w2_bias.data.to(torch.float32) |
| 1847 | return ( |
| 1848 | packed_w13, |
| 1849 | packed_w2, |
| 1850 | packed_w13_scale, |
| 1851 | packed_w2_scale, |
| 1852 | w13_bias, |
| 1853 | w2_bias, |
| 1854 | ) |
| 1855 | else: |
| 1856 | raise ValueError( |
| 1857 | f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. " |
| 1858 | "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, " |
| 1859 | "CPU, or emulation backend." |
| 1860 | ) |
| 1861 | |
| 1862 | |
| 1863 | def make_mxfp4_moe_quant_config( |
| 1864 | mxfp4_backend: Mxfp4MoeBackend, |
| 1865 | w1_scale: Union[torch.Tensor, "PrecisionConfig"], |
| 1866 | w2_scale: Union[torch.Tensor, "PrecisionConfig"], |
| 1867 | gemm1_alpha: float | None = None, |
| 1868 | gemm1_beta: float | None = None, |
| 1869 | swiglu_limit: float | None = None, |
| 1870 | w1_bias: torch.Tensor | None = None, |
| 1871 | w2_bias: torch.Tensor | None = None, |
| 1872 | a1_scale: torch.Tensor | None = None, |
| 1873 | a2_scale: torch.Tensor | None = None, |
| 1874 | layer: "RoutedExperts | None" = None, |
| 1875 | ) -> FusedMoEQuantConfig | None: |
| 1876 | """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" |
| 1877 | if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: |
| 1878 | return mxfp4_mxfp8_moe_quant_config( |
| 1879 | w1_bias=w1_bias, |
| 1880 | w2_bias=w2_bias, |
| 1881 | w1_scale=w1_scale, |
| 1882 | w2_scale=w2_scale, |
| 1883 | gemm1_alpha=gemm1_alpha, |
| 1884 | gemm1_beta=gemm1_beta, |
| 1885 | gemm1_clamp_limit=swiglu_limit, |
| 1886 | ) |
| 1887 | if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: |
| 1888 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 1889 | GroupShape, |
| 1890 | ) |
| 1891 | |
| 1892 | # DeepGEMM FP4 uses FP8 per-token-group activation quantization |
| 1893 | # with block 128, matching the FP8 DeepGEMM path. |
| 1894 | _fp8_dtype = current_platform.fp8_dtype() |
| 1895 | _block_shape = GroupShape(128, 128) |
| 1896 | return FusedMoEQuantConfig( |
| 1897 | _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None), |
| 1898 | _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None), |
| 1899 | _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), |
| 1900 | _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), |
| 1901 | gemm1_alpha=gemm1_alpha, |
| 1902 | gemm1_beta=gemm1_beta, |
| 1903 | gemm1_clamp_limit=swiglu_limit, |
| 1904 | ) |
| 1905 | elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8: |
| 1906 | # TRTLLM kernel expects non-swizzled mxfp8 activation scales. |
| 1907 | return mxfp4_mxfp8_moe_quant_config( |
| 1908 | w1_bias=w1_bias, |
| 1909 | w2_bias=w2_bias, |
| 1910 | w1_scale=w1_scale, |
| 1911 | w2_scale=w2_scale, |
| 1912 | gemm1_alpha=gemm1_alpha, |
| 1913 | gemm1_beta=gemm1_beta, |
| 1914 | gemm1_clamp_limit=swiglu_limit, |
| 1915 | mx_alignment=256, |
| 1916 | is_scale_swizzled=False, |
| 1917 | ) |
| 1918 | elif mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8: |
| 1919 | # CUTLASS kernel expects swizzled mxfp8 activation scales. |
| 1920 | return mxfp4_mxfp8_moe_quant_config( |
| 1921 | w1_bias=w1_bias, |
| 1922 | w2_bias=w2_bias, |
| 1923 | w1_scale=w1_scale, |
| 1924 | w2_scale=w2_scale, |
| 1925 | gemm1_alpha=gemm1_alpha, |
| 1926 | gemm1_beta=gemm1_beta, |
| 1927 | gemm1_clamp_limit=swiglu_limit, |
| 1928 | is_scale_swizzled=True, |
| 1929 | ) |
| 1930 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: |
| 1931 | # W4A8: MXFP4 weights + static FP8 activations |
| 1932 | return mxfp4_w4a8_moe_quant_config( |
| 1933 | w1_scale=w1_scale, |
| 1934 | w2_scale=w2_scale, |
| 1935 | a1_scale=a1_scale, |
| 1936 | a2_scale=a2_scale, |
| 1937 | w1_bias=w1_bias, |
| 1938 | w2_bias=w2_bias, |
| 1939 | block_shape=None, |
| 1940 | gemm1_clamp_limit=swiglu_limit, |
| 1941 | ) |
| 1942 | elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4: |
| 1943 | return ocp_mx_moe_quant_config( |
| 1944 | quant_dtype="mxfp4", |
| 1945 | w1_bias=w1_bias, |
| 1946 | w2_bias=w2_bias, |
| 1947 | w1_scale=w1_scale, |
| 1948 | w2_scale=w2_scale, |
| 1949 | gemm1_alpha=gemm1_alpha, |
| 1950 | gemm1_beta=gemm1_beta, |
| 1951 | gemm1_clamp_limit=swiglu_limit, |
| 1952 | ) |
| 1953 | elif mxfp4_backend in ( |
| 1954 | Mxfp4MoeBackend.B12X_MXFP4_BF16, |
| 1955 | Mxfp4MoeBackend.MARLIN, |
| 1956 | Mxfp4MoeBackend.BATCHED_MARLIN, |
| 1957 | Mxfp4MoeBackend.TRITON, |
| 1958 | Mxfp4MoeBackend.TRITON_UNFUSED, |
| 1959 | Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, |
| 1960 | Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, |
| 1961 | Mxfp4MoeBackend.AITER_MXFP4_BF16, |
| 1962 | Mxfp4MoeBackend.AITER_TRITON_MXFP4_BF16, |
| 1963 | Mxfp4MoeBackend.CPU, |
| 1964 | ): |
| 1965 | return mxfp4_w4a16_moe_quant_config( |
| 1966 | w1_bias=w1_bias, |
| 1967 | w2_bias=w2_bias, |
| 1968 | w1_scale=w1_scale, |
| 1969 | w2_scale=w2_scale, |
| 1970 | gemm1_alpha=gemm1_alpha, |
| 1971 | gemm1_beta=gemm1_beta, |
| 1972 | gemm1_clamp_limit=swiglu_limit, |
| 1973 | ) |
| 1974 | elif mxfp4_backend == Mxfp4MoeBackend.HUMMING: |
| 1975 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 1976 | get_humming_moe_quant_config, |
| 1977 | ) |
| 1978 | |
| 1979 | assert layer is not None |
| 1980 | return get_humming_moe_quant_config( |
| 1981 | layer, |
| 1982 | gemm1_alpha=gemm1_alpha, |
| 1983 | gemm1_beta=gemm1_beta, |
| 1984 | gemm1_clamp_limit=swiglu_limit, |
| 1985 | ) |
| 1986 | else: |
| 1987 | return ocp_mx_moe_quant_config( |
| 1988 | quant_dtype="mxfp4", |
| 1989 | w1_bias=w1_bias, |
| 1990 | w2_bias=w2_bias, |
| 1991 | w1_scale=w1_scale, |
| 1992 | w2_scale=w2_scale, |
| 1993 | gemm1_alpha=gemm1_alpha, |
| 1994 | gemm1_beta=gemm1_beta, |
| 1995 | gemm1_clamp_limit=swiglu_limit, |
| 1996 | ) |
| 1997 | |
| 1998 | |
| 1999 | def make_mxfp4_moe_kernel( |
| 2000 | moe_quant_config: FusedMoEQuantConfig, |
| 2001 | moe_config: FusedMoEConfig, |
| 2002 | experts_cls: type[mk.FusedMoEExperts], |
| 2003 | mxfp4_backend: Mxfp4MoeBackend, |
| 2004 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 2005 | ) -> mk.FusedMoEKernel: |
| 2006 | """Create a FusedMoEKernel for the given MXFP4 backend.""" |
| 2007 | is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) |
| 2008 | |
| 2009 | prepare_finalize = maybe_make_prepare_finalize( |
| 2010 | moe=moe_config, |
| 2011 | quant_config=moe_quant_config, |
| 2012 | routing_tables=routing_tables, |
| 2013 | allow_new_interface=True, |
| 2014 | use_monolithic=is_monolithic, |
| 2015 | ) |
| 2016 | assert prepare_finalize is not None |
| 2017 | |
| 2018 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 2019 | logger.info_once("Using %s", experts_cls.__name__) |
| 2020 | |
| 2021 | # Create Experts. |
| 2022 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 2023 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 2024 | assert max_num_tokens is not None |
| 2025 | experts = experts_cls( |
| 2026 | moe_config=moe_config, |
| 2027 | quant_config=moe_quant_config, |
| 2028 | max_num_tokens=max_num_tokens, |
| 2029 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 2030 | ) |
| 2031 | else: |
| 2032 | experts = experts_cls( |
| 2033 | moe_config=moe_config, |
| 2034 | quant_config=moe_quant_config, |
| 2035 | ) |
| 2036 | |
| 2037 | kernel = mk.FusedMoEKernel( |
| 2038 | prepare_finalize, |
| 2039 | experts, |
| 2040 | ) |
| 2041 | |
| 2042 | return kernel |