| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from enum import Enum |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 10 | from vllm import envs |
| 11 | from vllm._aiter_ops import rocm_aiter_ops |
| 12 | from vllm.config.kernel import MoEBackend |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 15 | maybe_make_prepare_finalize, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.fused_moe.config import ( |
| 18 | FusedMoEConfig, |
| 19 | FusedMoEQuantConfig, |
| 20 | fp8_w8a8_moe_quant_config, |
| 21 | fp8_w8a16_moe_quant_config, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 24 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 25 | prepare_fp8_moe_layer_for_fi, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 28 | prepare_fp8_moe_layer_for_deepgemm, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 31 | prepare_fp8_moe_layer_for_marlin, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 34 | GroupShape, |
| 35 | QuantKey, |
| 36 | create_fp8_quant_key, |
| 37 | kFp8Dynamic128Sym, |
| 38 | kFp8Static128BlockSym, |
| 39 | ) |
| 40 | from vllm.platforms import current_platform |
| 41 | from vllm.utils.math_utils import round_up |
| 42 | |
| 43 | logger = init_logger(__name__) |
| 44 | |
| 45 | |
| 46 | class Fp8MoeBackend(Enum): |
| 47 | NONE = "NONE" |
| 48 | FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" |
| 49 | FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" |
| 50 | DEEPGEMM = "DEEPGEMM" |
| 51 | BATCHED_DEEPGEMM = "BATCHED_DEEPGEMM" |
| 52 | MARLIN = "MARLIN" |
| 53 | HUMMING = "HUMMING" |
| 54 | TRITON = "TRITON" |
| 55 | BATCHED_TRITON = "BATCHED_TRITON" |
| 56 | AITER = "AITER" |
| 57 | VLLM_CUTLASS = "VLLM_CUTLASS" |
| 58 | BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS" |
| 59 | XPU = "XPU" |
| 60 | # CPU FP8 W8A16 (BF16 activations); CPU_W8A8 needs native AMX-FP8. |
| 61 | CPU = "CPU" |
| 62 | CPU_W8A8 = "CPU_W8A8" |
| 63 | HPC = "HPC" |
| 64 | # Dequantize-to-BF16 emulation for MXFP8 on devices without a native |
| 65 | # MXFP8 MoE kernel (e.g. ROCm). Weights pass through unchanged here. |
| 66 | EMULATION = "EMULATION" |
| 67 | # MXFP8 MoE via a Triton ``dot_scaled`` kernel that lowers to CDNA4 |
| 68 | # (gfx950) native MX matrix-core ops. Weights stay in MXFP8 (no load-time |
| 69 | # format conversion); the FP8 values + E8M0 scales are consumed directly. |
| 70 | TRITON_MXFP8 = "TRITON_MXFP8" |
| 71 | # MXFP8 MoE via AITER (FlyDSL two-stage grouped GEMM) on gfx950. |
| 72 | AITER_MXFP8 = "AITER_MXFP8" |
| 73 | |
| 74 | |
| 75 | def _get_priority_backends( |
| 76 | moe_config: FusedMoEConfig, |
| 77 | weight_key: QuantKey | None, |
| 78 | activation_key: QuantKey | None, |
| 79 | ) -> list[Fp8MoeBackend]: |
| 80 | """Get available backends in priority order based on platform and config. |
| 81 | |
| 82 | This function can be extended to become more complex as needed. |
| 83 | """ |
| 84 | _AVAILABLE_BACKENDS = [ |
| 85 | Fp8MoeBackend.AITER, |
| 86 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 87 | Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 88 | Fp8MoeBackend.DEEPGEMM, |
| 89 | Fp8MoeBackend.VLLM_CUTLASS, |
| 90 | Fp8MoeBackend.TRITON, |
| 91 | Fp8MoeBackend.MARLIN, |
| 92 | Fp8MoeBackend.HUMMING, |
| 93 | Fp8MoeBackend.BATCHED_DEEPGEMM, |
| 94 | Fp8MoeBackend.BATCHED_VLLM_CUTLASS, |
| 95 | Fp8MoeBackend.BATCHED_TRITON, |
| 96 | Fp8MoeBackend.XPU, |
| 97 | Fp8MoeBackend.CPU_W8A8, |
| 98 | Fp8MoeBackend.CPU, |
| 99 | Fp8MoeBackend.HPC, |
| 100 | ] |
| 101 | |
| 102 | def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None: |
| 103 | backends.insert(0, backends.pop(backends.index(backend))) |
| 104 | |
| 105 | # With DeepEP v2 contiguous layout (do_expand=False), tensors are |
| 106 | # worst-case allocated with padding. TrtLLM's tile-level skipping |
| 107 | # avoids wasted compute on padding rows; other backends process all rows. |
| 108 | if ( |
| 109 | current_platform.is_cuda() |
| 110 | and current_platform.is_device_capability_family(100) |
| 111 | and moe_config.moe_parallel_config.use_deepep_v2_kernels |
| 112 | and activation_key == kFp8Dynamic128Sym |
| 113 | and weight_key == kFp8Static128BlockSym |
| 114 | ): |
| 115 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_TRTLLM) |
| 116 | |
| 117 | # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP. |
| 118 | if ( |
| 119 | current_platform.is_cuda() |
| 120 | and current_platform.is_device_capability(90) |
| 121 | and activation_key == kFp8Dynamic128Sym |
| 122 | and weight_key == kFp8Static128BlockSym |
| 123 | ): |
| 124 | if moe_config.moe_parallel_config.ep_size > 1: |
| 125 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_CUTLASS) |
| 126 | else: |
| 127 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.TRITON) |
| 128 | |
| 129 | if current_platform.is_xpu(): |
| 130 | # XPU platform supports TritonExperts and XPUExpertsFp8, |
| 131 | # move XPU backend to the front. |
| 132 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU) |
| 133 | |
| 134 | if current_platform.is_cpu(): |
| 135 | # W8A8 first: it falls through to the W8A16 backend whenever the |
| 136 | # hardware (AMX-FP8) or the config isn't supported. |
| 137 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU) |
| 138 | _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU_W8A8) |
| 139 | |
| 140 | return _AVAILABLE_BACKENDS |
| 141 | |
| 142 | |
| 143 | def backend_to_kernel_cls( |
| 144 | backend: Fp8MoeBackend, |
| 145 | ) -> list[type[mk.FusedMoEExperts]]: |
| 146 | if backend == Fp8MoeBackend.FLASHINFER_TRTLLM: |
| 147 | from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( # noqa: E501 |
| 148 | TrtLlmFp8ExpertsModular, |
| 149 | TrtLlmFp8ExpertsMonolithic, |
| 150 | ) |
| 151 | |
| 152 | return [TrtLlmFp8ExpertsMonolithic, TrtLlmFp8ExpertsModular] |
| 153 | |
| 154 | elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS: |
| 155 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501 |
| 156 | FlashInferExperts, |
| 157 | ) |
| 158 | |
| 159 | return [FlashInferExperts] |
| 160 | |
| 161 | elif backend == Fp8MoeBackend.DEEPGEMM: |
| 162 | from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( |
| 163 | TritonOrDeepGemmExperts, |
| 164 | ) |
| 165 | |
| 166 | return [TritonOrDeepGemmExperts] |
| 167 | |
| 168 | elif backend == Fp8MoeBackend.BATCHED_DEEPGEMM: |
| 169 | from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import ( |
| 170 | BatchedDeepGemmExperts, |
| 171 | ) |
| 172 | |
| 173 | return [BatchedDeepGemmExperts] |
| 174 | |
| 175 | elif backend == Fp8MoeBackend.HUMMING: |
| 176 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 177 | BatchedHummingGroupedExperts, |
| 178 | HummingGroupedExperts, |
| 179 | HummingIndexedExperts, |
| 180 | ) |
| 181 | |
| 182 | return [ |
| 183 | BatchedHummingGroupedExperts, |
| 184 | HummingGroupedExperts, |
| 185 | HummingIndexedExperts, |
| 186 | ] |
| 187 | |
| 188 | elif backend == Fp8MoeBackend.MARLIN: |
| 189 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 190 | MarlinExperts, |
| 191 | ) |
| 192 | |
| 193 | return [MarlinExperts] |
| 194 | |
| 195 | elif backend == Fp8MoeBackend.TRITON: |
| 196 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( |
| 197 | TritonExperts, |
| 198 | ) |
| 199 | |
| 200 | return [TritonExperts] |
| 201 | |
| 202 | elif backend == Fp8MoeBackend.BATCHED_TRITON: |
| 203 | from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( |
| 204 | BatchedTritonExperts, |
| 205 | ) |
| 206 | |
| 207 | return [BatchedTritonExperts] |
| 208 | |
| 209 | elif backend == Fp8MoeBackend.AITER: |
| 210 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 211 | AiterExperts, |
| 212 | ) |
| 213 | |
| 214 | return [AiterExperts] |
| 215 | |
| 216 | elif backend == Fp8MoeBackend.VLLM_CUTLASS: |
| 217 | from vllm.model_executor.layers.fused_moe.experts.triton_cutlass_moe import ( |
| 218 | TritonOrCutlassExperts, |
| 219 | ) |
| 220 | |
| 221 | return [TritonOrCutlassExperts] |
| 222 | |
| 223 | elif backend == Fp8MoeBackend.BATCHED_VLLM_CUTLASS: |
| 224 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 225 | CutlassBatchedExpertsFp8, |
| 226 | ) |
| 227 | |
| 228 | return [CutlassBatchedExpertsFp8] |
| 229 | |
| 230 | elif backend == Fp8MoeBackend.XPU: |
| 231 | from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( |
| 232 | XPUExpertsBlockFp8, |
| 233 | XPUExpertsFp8, |
| 234 | XPUExpertsMxFp8, |
| 235 | ) |
| 236 | |
| 237 | return [XPUExpertsFp8, XPUExpertsMxFp8, XPUExpertsBlockFp8] |
| 238 | |
| 239 | elif backend == Fp8MoeBackend.CPU: |
| 240 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 241 | CPUExpertsFp8, |
| 242 | ) |
| 243 | |
| 244 | return [CPUExpertsFp8] |
| 245 | |
| 246 | elif backend == Fp8MoeBackend.CPU_W8A8: |
| 247 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 248 | CPUExpertsFp8W8A8, |
| 249 | ) |
| 250 | |
| 251 | return [CPUExpertsFp8W8A8] |
| 252 | |
| 253 | elif backend == Fp8MoeBackend.HPC: |
| 254 | from vllm.model_executor.layers.fused_moe.hpc_moe import ( |
| 255 | HPCExperts, |
| 256 | ) |
| 257 | |
| 258 | return [HPCExperts] |
| 259 | |
| 260 | else: |
| 261 | raise ValueError(f"Unknown FP8 MoE backend: {backend.value}") |
| 262 | |
| 263 | |
| 264 | def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend: |
| 265 | """Map user's MoEBackend to Fp8MoeBackend.""" |
| 266 | mapping = { |
| 267 | "triton": Fp8MoeBackend.TRITON, |
| 268 | "deep_gemm": Fp8MoeBackend.DEEPGEMM, |
| 269 | "cutlass": Fp8MoeBackend.VLLM_CUTLASS, |
| 270 | "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 271 | "flashinfer_cutlass": Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 272 | "marlin": Fp8MoeBackend.MARLIN, |
| 273 | "humming": Fp8MoeBackend.HUMMING, |
| 274 | "aiter": Fp8MoeBackend.AITER, |
| 275 | "hpc": Fp8MoeBackend.HPC, |
| 276 | } |
| 277 | if backend := mapping.get(runner_backend): |
| 278 | return backend |
| 279 | raise ValueError( |
| 280 | f"moe_backend='{runner_backend}' is not supported for FP8 MoE. " |
| 281 | f"Expected one of {list(mapping.keys())}." |
| 282 | ) |
| 283 | |
| 284 | |
| 285 | def refine_fp8_moe_block_shape( |
| 286 | config: FusedMoEConfig, |
| 287 | weight_block_size: list[int], |
| 288 | ) -> list[int] | None: |
| 289 | """Compute a refined block shape for block-quantized FP8 MoE weights whose |
| 290 | checkpoint blocks cannot be sharded exactly across TP ranks. |
| 291 | |
| 292 | TP shards the intermediate dim of the expert weights, so a per-shard size |
| 293 | that is not a multiple of the checkpoint's block size makes the |
| 294 | checkpoint's block scales impossible to shard exactly. When a finer block |
| 295 | size (>= 32) divides both the checkpoint blocks and all involved dims, |
| 296 | the weight scales can be refined to that granularity at load time (a |
| 297 | lossless upsampling, since the refined block divides the checkpoint |
| 298 | block). Only Triton-based kernels can consume the refined block shape: |
| 299 | they take it as a runtime argument, while the other backends require the |
| 300 | native 128x128 blocks. The refined shape is encoded in the QuantKey used |
| 301 | for backend selection, so backends that only support 128x128 blocks are |
| 302 | rejected by the oracle automatically. |
| 303 | |
| 304 | Returns the refined [block_n, block_k] shape, or None if no refinement |
| 305 | is needed or possible. |
| 306 | """ |
| 307 | block_n, block_k = weight_block_size |
| 308 | ispp = config.intermediate_size_per_partition |
| 309 | if ispp % block_n == 0 and (config.tp_size == 1 or ispp % block_k == 0): |
| 310 | return None |
| 311 | refine = math.gcd(block_n, block_k, ispp, config.hidden_dim) |
| 312 | if refine < 32: |
| 313 | return None |
| 314 | return [refine, refine] |
| 315 | |
| 316 | |
| 317 | def pad_tp_shard_to_weight_blocks( |
| 318 | config: FusedMoEConfig, |
| 319 | weight_block_size: list[int], |
| 320 | ) -> bool: |
| 321 | """Pad the TP shard to whole checkpoint blocks, keeping scales rank-local.""" |
| 322 | block_n, block_k = weight_block_size |
| 323 | if ( |
| 324 | block_n != block_k |
| 325 | or config.tp_size == 1 |
| 326 | or config.intermediate_size_per_partition % block_n == 0 |
| 327 | ): |
| 328 | return False |
| 329 | if ( |
| 330 | config.intermediate_size % block_n != 0 |
| 331 | or config.hidden_dim % block_n != 0 |
| 332 | or config.ep_size != 1 |
| 333 | or config.is_lora_enabled |
| 334 | or config.has_bias |
| 335 | ): |
| 336 | raise ValueError( |
| 337 | f"Block-aligned FP8 TP sharding requires {block_n}-aligned " |
| 338 | "global expert dimensions, pure TP, and no LoRA or " |
| 339 | "expert bias." |
| 340 | ) |
| 341 | config.tp_shard_with_padding = True |
| 342 | config.intermediate_size_per_partition = round_up( |
| 343 | config.intermediate_size_per_partition, block_n |
| 344 | ) |
| 345 | return True |
| 346 | |
| 347 | |
| 348 | def resolve_fp8_moe_weight_block_shape( |
| 349 | config: FusedMoEConfig, |
| 350 | weight_block_size: list[int], |
| 351 | activation_key: QuantKey, |
| 352 | is_checkpoint_fp8_serialized: bool, |
| 353 | ) -> tuple[list[int], tuple[int, int] | None]: |
| 354 | """Return the TP-adapted block shape and refine factor: |
| 355 | refine if kernels allow, else pad to the TP shard.""" |
| 356 | refined_shape = refine_fp8_moe_block_shape(config, weight_block_size) |
| 357 | if is_checkpoint_fp8_serialized and config.moe_backend != "auto": |
| 358 | kernel_classes = backend_to_kernel_cls(map_fp8_backend(config.moe_backend)) |
| 359 | can_refine = refined_shape is not None and any( |
| 360 | k_cls._supports_quant_scheme( |
| 361 | create_fp8_quant_key( |
| 362 | static=True, group_shape=GroupShape(*refined_shape) |
| 363 | ), |
| 364 | activation_key, |
| 365 | ) |
| 366 | for k_cls in kernel_classes |
| 367 | ) |
| 368 | if not can_refine and pad_tp_shard_to_weight_blocks(config, weight_block_size): |
| 369 | logger.info_once( |
| 370 | "FP8 %s TP loading uses complete checkpoint blocks: " |
| 371 | "local allocation %d, without weight requantization.", |
| 372 | config.moe_backend, |
| 373 | config.intermediate_size_per_partition, |
| 374 | ) |
| 375 | return weight_block_size, None |
| 376 | if refined_shape is None: |
| 377 | return weight_block_size, None |
| 378 | logger.info_once( |
| 379 | "FP8 MoE block scales refined from %s to %s to fit " |
| 380 | "the TP-sharded intermediate size %d.", |
| 381 | str(weight_block_size), |
| 382 | str(refined_shape), |
| 383 | config.intermediate_size_per_partition, |
| 384 | ) |
| 385 | return refined_shape, ( |
| 386 | weight_block_size[0] // refined_shape[0], |
| 387 | weight_block_size[1] // refined_shape[1], |
| 388 | ) |
| 389 | |
| 390 | |
| 391 | def select_fp8_moe_backend( |
| 392 | config: FusedMoEConfig, |
| 393 | weight_key: QuantKey | None, |
| 394 | activation_key: QuantKey | None, |
| 395 | allow_vllm_cutlass: bool = False, |
| 396 | ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts] | None]: |
| 397 | """Select the primary FP8 MoE backend |
| 398 | Note: Shape-specific fallbacks may still occur at runtime. |
| 399 | """ |
| 400 | # NOTE: the kernels are selected in the following order. |
| 401 | AVAILABLE_BACKENDS = _get_priority_backends(config, weight_key, activation_key) |
| 402 | |
| 403 | # NOTE(rob): We need to peak into the P/F selection to determine |
| 404 | # if we are using the batched or standard expert format, which |
| 405 | # if not ideal. Once we unify TP + DP/EP, we can select P/F first. |
| 406 | activation_format = ( |
| 407 | mk.FusedMoEActivationFormat.BatchedExperts |
| 408 | if config.moe_parallel_config.use_batched_activation_format |
| 409 | else mk.FusedMoEActivationFormat.Standard |
| 410 | ) |
| 411 | |
| 412 | def _make_log_backend(backend: Fp8MoeBackend): |
| 413 | available_backend_strs = [b.value for b in AVAILABLE_BACKENDS] |
| 414 | return ( |
| 415 | f"Using {backend.value} Fp8 MoE backend out " |
| 416 | f"of potential backends: {available_backend_strs}." |
| 417 | ) |
| 418 | |
| 419 | def _make_log_unsupported(backend: Fp8MoeBackend, reason: str | None) -> str: |
| 420 | if reason: |
| 421 | return ( |
| 422 | f"FP8 MoE backend {backend.value} does not support the " |
| 423 | f"deployment configuration since {reason}." |
| 424 | ) |
| 425 | else: |
| 426 | return ( |
| 427 | f"FP8 MoE backend '{backend.value}' does not support the " |
| 428 | "deployment configuration." |
| 429 | ) |
| 430 | |
| 431 | def _return_or_raise( |
| 432 | backend: Fp8MoeBackend, |
| 433 | config: FusedMoEConfig, |
| 434 | weight_key: QuantKey | None, |
| 435 | activation_key: QuantKey | None, |
| 436 | activation_format: mk.FusedMoEActivationFormat, |
| 437 | ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: |
| 438 | for k_cls in backend_to_kernel_cls(backend): |
| 439 | supported, reason = k_cls.is_supported_config( |
| 440 | k_cls, config, weight_key, activation_key, activation_format |
| 441 | ) |
| 442 | if supported: |
| 443 | logger.info_once(_make_log_backend(backend)) |
| 444 | return backend, k_cls |
| 445 | raise ValueError(_make_log_unsupported(backend, reason)) |
| 446 | |
| 447 | # Handle explicit moe_backend from user. |
| 448 | runner_backend = config.moe_backend |
| 449 | if runner_backend != "auto": |
| 450 | requested_backend = map_fp8_backend(runner_backend) |
| 451 | # For batched activation format, use batched variants if available. |
| 452 | if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 453 | if requested_backend == Fp8MoeBackend.DEEPGEMM: |
| 454 | requested_backend = Fp8MoeBackend.BATCHED_DEEPGEMM |
| 455 | elif requested_backend == Fp8MoeBackend.TRITON: |
| 456 | requested_backend = Fp8MoeBackend.BATCHED_TRITON |
| 457 | elif requested_backend == Fp8MoeBackend.VLLM_CUTLASS: |
| 458 | requested_backend = Fp8MoeBackend.BATCHED_VLLM_CUTLASS |
| 459 | |
| 460 | if ( |
| 461 | requested_backend |
| 462 | in [ |
| 463 | Fp8MoeBackend.VLLM_CUTLASS, |
| 464 | Fp8MoeBackend.BATCHED_VLLM_CUTLASS, |
| 465 | ] |
| 466 | and not allow_vllm_cutlass |
| 467 | ): |
| 468 | raise ValueError( |
| 469 | "vLLM CUTLASS FP8 MoE backend is disabled for this configuration." |
| 470 | ) |
| 471 | |
| 472 | return _return_or_raise( |
| 473 | requested_backend, config, weight_key, activation_key, activation_format |
| 474 | ) |
| 475 | |
| 476 | # Handle explicit DeepGEMM FP8 configuration. |
| 477 | if envs.is_set("VLLM_USE_DEEP_GEMM") or envs.is_set("VLLM_MOE_USE_DEEP_GEMM"): |
| 478 | if not envs.VLLM_USE_DEEP_GEMM or not envs.VLLM_MOE_USE_DEEP_GEMM: |
| 479 | AVAILABLE_BACKENDS.remove(Fp8MoeBackend.DEEPGEMM) |
| 480 | AVAILABLE_BACKENDS.remove(Fp8MoeBackend.BATCHED_DEEPGEMM) |
| 481 | else: |
| 482 | backend = ( |
| 483 | Fp8MoeBackend.DEEPGEMM |
| 484 | if activation_format == mk.FusedMoEActivationFormat.Standard |
| 485 | else Fp8MoeBackend.BATCHED_DEEPGEMM |
| 486 | ) |
| 487 | return _return_or_raise( |
| 488 | backend, config, weight_key, activation_key, activation_format |
| 489 | ) |
| 490 | |
| 491 | # Handle explicit AITER FP8 configuration. |
| 492 | if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"): |
| 493 | skip_aiter_moe = ( |
| 494 | not envs.VLLM_ROCM_USE_AITER |
| 495 | or not envs.VLLM_ROCM_USE_AITER_MOE |
| 496 | or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 497 | ) |
| 498 | if skip_aiter_moe: |
| 499 | if Fp8MoeBackend.AITER in AVAILABLE_BACKENDS: |
| 500 | AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER) |
| 501 | else: |
| 502 | backend = Fp8MoeBackend.AITER |
| 503 | return _return_or_raise( |
| 504 | backend, config, weight_key, activation_key, activation_format |
| 505 | ) |
| 506 | |
| 507 | if not allow_vllm_cutlass: |
| 508 | AVAILABLE_BACKENDS.remove(Fp8MoeBackend.VLLM_CUTLASS) |
| 509 | AVAILABLE_BACKENDS.remove(Fp8MoeBackend.BATCHED_VLLM_CUTLASS) |
| 510 | |
| 511 | # Select kernels in order of backend. |
| 512 | for backend in AVAILABLE_BACKENDS: |
| 513 | for k_cls in backend_to_kernel_cls(backend): |
| 514 | supported, reason = k_cls.is_supported_config( |
| 515 | k_cls, |
| 516 | config, |
| 517 | weight_key, |
| 518 | activation_key, |
| 519 | activation_format, |
| 520 | ) |
| 521 | if supported: |
| 522 | logger.info_once(_make_log_backend(backend)) |
| 523 | return backend, k_cls |
| 524 | else: |
| 525 | logger.debug_once(_make_log_unsupported(backend, reason)) |
| 526 | |
| 527 | # TODO(rob): per discussion with TPU team, we need a way to register |
| 528 | # MoE backends by OOT plugins, rather than having an explicit list |
| 529 | # of AVAILABLE_BACKENDS. Enabling returning `Fp8MoeBackend.NONE` is |
| 530 | # a temporary measure until these register APIs are complete. |
| 531 | if current_platform.is_cuda() or current_platform.is_rocm(): |
| 532 | raise NotImplementedError( |
| 533 | "No FP8 MoE backend supports the deployment configuration." |
| 534 | ) |
| 535 | |
| 536 | return Fp8MoeBackend.NONE, None |
| 537 | |
| 538 | |
| 539 | def _humming_fp8_weight_schema( |
| 540 | layer: RoutedExperts, weight: torch.Tensor, weight_scale: torch.Tensor |
| 541 | ) -> dict[str, Any]: |
| 542 | """Build the humming weight schema from the canonical on-device fp8/mxfp8 |
| 543 | tensors (scale dtype/shape, block size), not the producing quant method.""" |
| 544 | # mxfp8: e8m0 group-32 scales (stored as uint8 bytes or e8m0). humming has |
| 545 | # no compressed-tensors mxfp8 loader; its modelopt schema fits both sources. |
| 546 | if weight_scale.dtype in (torch.uint8, torch.float8_e8m0fnu): |
| 547 | return {"quant_method": "modelopt", "quant_algo": "mxfp8"} |
| 548 | |
| 549 | if hasattr(layer, "w13_weight_scale_inv"): |
| 550 | assert hasattr(layer, "weight_block_size") |
| 551 | return {"quant_method": "fp8", "weight_block_size": layer.weight_block_size} |
| 552 | |
| 553 | # fp8 (e4m3): recover the strategy from the scale layout (block from |
| 554 | # weight_block_size, else channel vs tensor by per-expert scale count). |
| 555 | config: dict[str, Any] = { |
| 556 | "quant_method": "compressed-tensors", |
| 557 | "format": "float-quantized", |
| 558 | "type": "float", |
| 559 | "num_bits": 8, |
| 560 | "symmetric": True, |
| 561 | } |
| 562 | weight_block_size = getattr(layer, "weight_block_size", None) |
| 563 | num_experts, num_output = weight.shape[0], weight.shape[-2] |
| 564 | if weight_block_size is not None: |
| 565 | config["strategy"] = "block" |
| 566 | config["block_structure"] = list(weight_block_size) |
| 567 | elif weight_scale.numel() >= num_experts * num_output: |
| 568 | config["strategy"] = "channel" |
| 569 | else: |
| 570 | config["strategy"] = "tensor" |
| 571 | return config |
| 572 | |
| 573 | |
| 574 | def convert_to_fp8_moe_kernel_format( |
| 575 | fp8_backend: Fp8MoeBackend, |
| 576 | # TODO(bnell): replace layer with weight_block_size |
| 577 | layer: RoutedExperts, |
| 578 | w13: torch.Tensor, |
| 579 | w2: torch.Tensor, |
| 580 | w13_scale: torch.Tensor, |
| 581 | w2_scale: torch.Tensor, |
| 582 | w13_input_scale: torch.Tensor | None, |
| 583 | w2_input_scale: torch.Tensor | None, |
| 584 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 585 | block_quant = hasattr(layer, "weight_block_size") |
| 586 | if fp8_backend in [Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.BATCHED_DEEPGEMM]: |
| 587 | assert block_quant |
| 588 | w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_deepgemm( |
| 589 | w13, |
| 590 | w2, |
| 591 | w13_scale, |
| 592 | w2_scale, |
| 593 | tuple(layer.weight_block_size), |
| 594 | ) |
| 595 | elif fp8_backend == Fp8MoeBackend.AITER: |
| 596 | w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2) |
| 597 | w13.is_shuffled = True |
| 598 | w2.is_shuffled = True |
| 599 | elif fp8_backend == Fp8MoeBackend.AITER_MXFP8: |
| 600 | w13, w2, w13_scale, w2_scale = rocm_aiter_ops.shuffle_mxfp8_moe_weights( |
| 601 | w13, w2, w13_scale, w2_scale |
| 602 | ) |
| 603 | w13.is_shuffled = True |
| 604 | w2.is_shuffled = True |
| 605 | elif fp8_backend == Fp8MoeBackend.HUMMING: |
| 606 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 607 | convert_to_humming_moe_kernel_format, |
| 608 | ) |
| 609 | |
| 610 | convert_to_humming_moe_kernel_format( |
| 611 | layer, quant_config=_humming_fp8_weight_schema(layer, w13, w13_scale) |
| 612 | ) |
| 613 | w13 = layer.w13_weight |
| 614 | w2 = layer.w2_weight |
| 615 | w13_scale = layer.w13_weight_scale |
| 616 | w2_scale = layer.w2_weight_scale |
| 617 | elif fp8_backend == Fp8MoeBackend.MARLIN: |
| 618 | weight_block_size = getattr(layer, "weight_block_size", None) |
| 619 | if weight_block_size == [1, 32]: |
| 620 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 621 | prepare_mxfp8_moe_layer_for_marlin, |
| 622 | ) |
| 623 | |
| 624 | w13, w2, w13_scale, w2_scale = prepare_mxfp8_moe_layer_for_marlin( |
| 625 | layer, |
| 626 | w13, |
| 627 | w2, |
| 628 | w13_scale, |
| 629 | w2_scale, |
| 630 | ) |
| 631 | else: |
| 632 | w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_marlin( |
| 633 | layer, |
| 634 | w13, |
| 635 | w2, |
| 636 | w13_scale, |
| 637 | w2_scale, |
| 638 | ) |
| 639 | elif fp8_backend in [ |
| 640 | Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 641 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 642 | ]: |
| 643 | w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_fi( |
| 644 | layer=layer, |
| 645 | w13=w13, |
| 646 | w2=w2, |
| 647 | w13_scale=w13_scale, |
| 648 | w13_input_scale=w13_input_scale, |
| 649 | w2_scale=w2_scale, |
| 650 | w2_input_scale=w2_input_scale, |
| 651 | is_trtllm=(fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM), |
| 652 | ) |
| 653 | elif fp8_backend == Fp8MoeBackend.XPU: |
| 654 | from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( |
| 655 | prepare_fp8_moe_layer_for_xpu, |
| 656 | ) |
| 657 | |
| 658 | w13, w13_scale, w2, w2_scale = prepare_fp8_moe_layer_for_xpu( |
| 659 | w13, w13_scale, w2, w2_scale |
| 660 | ) |
| 661 | elif fp8_backend == Fp8MoeBackend.CPU: |
| 662 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 663 | prepare_fp8_moe_layer_for_cpu, |
| 664 | ) |
| 665 | |
| 666 | w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2) |
| 667 | elif fp8_backend == Fp8MoeBackend.CPU_W8A8: |
| 668 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 669 | prepare_fp8_w8a8_moe_layer_for_cpu, |
| 670 | ) |
| 671 | |
| 672 | w13, w13_scale, w2, w2_scale = prepare_fp8_w8a8_moe_layer_for_cpu( |
| 673 | w13, w2, w13_scale, w2_scale |
| 674 | ) |
| 675 | else: |
| 676 | if fp8_backend not in [ |
| 677 | Fp8MoeBackend.TRITON, |
| 678 | Fp8MoeBackend.BATCHED_TRITON, |
| 679 | Fp8MoeBackend.VLLM_CUTLASS, |
| 680 | Fp8MoeBackend.BATCHED_VLLM_CUTLASS, |
| 681 | Fp8MoeBackend.XPU, |
| 682 | Fp8MoeBackend.HPC, |
| 683 | # EMULATION dequantizes weights at runtime; NATIVE_MXFP8 consumes |
| 684 | # the MXFP8 weights as-is — neither needs a load-time layout change. |
| 685 | Fp8MoeBackend.EMULATION, |
| 686 | Fp8MoeBackend.TRITON_MXFP8, |
| 687 | ]: |
| 688 | raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}") |
| 689 | |
| 690 | return w13, w2, w13_scale, w2_scale |
| 691 | |
| 692 | |
| 693 | def make_fp8_moe_quant_config( |
| 694 | fp8_backend: Fp8MoeBackend, |
| 695 | w1_scale: torch.Tensor, |
| 696 | w2_scale: torch.Tensor, |
| 697 | a1_scale: torch.Tensor | None, |
| 698 | a2_scale: torch.Tensor | None, |
| 699 | w1_bias: torch.Tensor | None = None, |
| 700 | w2_bias: torch.Tensor | None = None, |
| 701 | block_shape: list[int] | None = None, |
| 702 | per_act_token_quant: bool = False, |
| 703 | per_out_ch_quant: bool = False, |
| 704 | swiglu_limit: float | None = None, |
| 705 | gemm1_alpha: float | None = None, |
| 706 | gemm1_beta: float | None = None, |
| 707 | layer: torch.nn.Module | None = None, |
| 708 | ) -> FusedMoEQuantConfig: |
| 709 | """Create FusedMoEQuantConfig for the specified FP8 Backend. |
| 710 | The FusedMoEQuantConfig holds the scales that are used |
| 711 | at runtime by the Modular Kernel abstraction. |
| 712 | |
| 713 | Note that certain kernels (e.g. Flashinfer CUTLASS) need |
| 714 | special Quant configs to handle non-standard inputs to |
| 715 | their kernel interfaces. |
| 716 | |
| 717 | In a future PR, we will have this function should be |
| 718 | a method of the modular kernel itself. |
| 719 | """ |
| 720 | if fp8_backend == Fp8MoeBackend.CPU_W8A8: |
| 721 | return fp8_w8a8_moe_quant_config( |
| 722 | w1_scale=w1_scale, |
| 723 | w2_scale=w2_scale, |
| 724 | a1_scale=a1_scale, |
| 725 | block_shape=block_shape, |
| 726 | ) |
| 727 | |
| 728 | # MARLIN and CPU (W8A16) are mixed precision W8A16 configs. |
| 729 | if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU: |
| 730 | return fp8_w8a16_moe_quant_config( |
| 731 | w1_scale=w1_scale, |
| 732 | w2_scale=w2_scale, |
| 733 | w1_bias=w1_bias, |
| 734 | w2_bias=w2_bias, |
| 735 | block_shape=block_shape, |
| 736 | gemm1_alpha=gemm1_alpha, |
| 737 | gemm1_beta=gemm1_beta, |
| 738 | gemm1_clamp_limit=swiglu_limit, |
| 739 | ) |
| 740 | elif fp8_backend == Fp8MoeBackend.HUMMING: |
| 741 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 742 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 743 | get_humming_moe_quant_config, |
| 744 | ) |
| 745 | |
| 746 | assert isinstance(layer, RoutedExperts) |
| 747 | return get_humming_moe_quant_config( |
| 748 | layer, |
| 749 | gemm1_alpha=gemm1_alpha, |
| 750 | gemm1_beta=gemm1_beta, |
| 751 | gemm1_clamp_limit=swiglu_limit, |
| 752 | ) |
| 753 | |
| 754 | # Flashinfer CUTLASS or HPC per-tensor uses single dq scale |
| 755 | # (alpha = w_scale * a_scale) and inverse a2 scale. |
| 756 | if ( |
| 757 | fp8_backend in [Fp8MoeBackend.FLASHINFER_CUTLASS, Fp8MoeBackend.HPC] |
| 758 | and block_shape is None |
| 759 | ): |
| 760 | assert a1_scale is not None and a2_scale is not None |
| 761 | g1_alphas = w1_scale * a1_scale |
| 762 | g2_alphas = w2_scale * a2_scale |
| 763 | if layer is not None: |
| 764 | layer.register_parameter( |
| 765 | "g1_alphas", torch.nn.Parameter(g1_alphas, requires_grad=False) |
| 766 | ) |
| 767 | layer.register_parameter( |
| 768 | "g2_alphas", torch.nn.Parameter(g2_alphas, requires_grad=False) |
| 769 | ) |
| 770 | g1_alphas = layer.g1_alphas |
| 771 | g2_alphas = layer.g2_alphas |
| 772 | return fp8_w8a8_moe_quant_config( |
| 773 | w1_scale=w1_scale, |
| 774 | w2_scale=w2_scale, |
| 775 | w1_bias=w1_bias, |
| 776 | w2_bias=w2_bias, |
| 777 | a1_scale=a1_scale, |
| 778 | a2_scale=a2_scale, |
| 779 | a1_gscale=(1.0 / a1_scale), |
| 780 | a2_gscale=(1.0 / a2_scale), |
| 781 | g1_alphas=g1_alphas, |
| 782 | g2_alphas=g2_alphas, |
| 783 | gemm1_clamp_limit=swiglu_limit, |
| 784 | ) |
| 785 | # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are |
| 786 | # the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all |
| 787 | # backends; the DeepGEMM expert permute repacks them for the grouped GEMM. |
| 788 | if block_shape == [1, 32]: |
| 789 | return FusedMoEQuantConfig.make( |
| 790 | "mxfp8", |
| 791 | w1_scale=w1_scale, |
| 792 | w2_scale=w2_scale, |
| 793 | w1_bias=w1_bias, |
| 794 | w2_bias=w2_bias, |
| 795 | a1_scale=a1_scale, |
| 796 | a2_scale=a2_scale, |
| 797 | block_shape=block_shape, |
| 798 | is_scale_swizzled=False, |
| 799 | gemm1_alpha=gemm1_alpha, |
| 800 | gemm1_beta=gemm1_beta, |
| 801 | gemm1_clamp_limit=swiglu_limit, |
| 802 | ) |
| 803 | |
| 804 | # All other backends use normal config. |
| 805 | return fp8_w8a8_moe_quant_config( |
| 806 | w1_scale=w1_scale, |
| 807 | w2_scale=w2_scale, |
| 808 | w1_bias=w1_bias, |
| 809 | w2_bias=w2_bias, |
| 810 | a1_scale=a1_scale, |
| 811 | a2_scale=a2_scale, |
| 812 | block_shape=block_shape, |
| 813 | per_act_token_quant=per_act_token_quant, |
| 814 | per_out_ch_quant=per_out_ch_quant, |
| 815 | gemm1_alpha=gemm1_alpha, |
| 816 | gemm1_beta=gemm1_beta, |
| 817 | gemm1_clamp_limit=swiglu_limit, |
| 818 | ) |
| 819 | |
| 820 | |
| 821 | def make_fp8_moe_kernel( |
| 822 | moe_quant_config: FusedMoEQuantConfig, |
| 823 | moe_config: FusedMoEConfig, |
| 824 | experts_cls: type[mk.FusedMoEExperts], |
| 825 | fp8_backend: Fp8MoeBackend, |
| 826 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 827 | ) -> mk.FusedMoEKernel: |
| 828 | # Create Prepare/Finalize. |
| 829 | prepare_finalize = maybe_make_prepare_finalize( |
| 830 | moe=moe_config, |
| 831 | quant_config=moe_quant_config, |
| 832 | routing_tables=routing_tables, |
| 833 | allow_new_interface=True, |
| 834 | use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), |
| 835 | ) |
| 836 | assert prepare_finalize is not None |
| 837 | |
| 838 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 839 | |
| 840 | # Create Experts. |
| 841 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 842 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 843 | assert max_num_tokens is not None |
| 844 | experts = experts_cls( |
| 845 | moe_config=moe_config, |
| 846 | quant_config=moe_quant_config, |
| 847 | max_num_tokens=max_num_tokens, |
| 848 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 849 | ) |
| 850 | else: |
| 851 | experts = experts_cls( |
| 852 | moe_config=moe_config, |
| 853 | quant_config=moe_quant_config, |
| 854 | ) |
| 855 | |
| 856 | kernel = mk.FusedMoEKernel( |
| 857 | prepare_finalize, |
| 858 | experts, |
| 859 | ) |
| 860 | |
| 861 | return kernel |