| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from enum import Enum |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 8 | from vllm import envs |
| 9 | from vllm.config.kernel import MoEBackend |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 12 | maybe_make_prepare_finalize, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.fused_moe.config import ( |
| 15 | FusedMoEConfig, |
| 16 | FusedMoEQuantConfig, |
| 17 | nvfp4_moe_quant_config, |
| 18 | nvfp4_w4a16_moe_quant_config, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 21 | from vllm.model_executor.layers.quantization.utils.b12x_moe import ( |
| 22 | prepare_nvfp4_moe_layer_for_b12x, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( |
| 25 | nvfp4_swizzled_scale_to_cutedsl_mma_view, |
| 26 | prepare_nvfp4_moe_layer_for_fi_or_cutlass, |
| 27 | prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( |
| 30 | prepare_nvfp4_moe_layer_for_marlin, |
| 31 | ) |
| 32 | from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( |
| 33 | kE2M1ToFloat_handle, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 36 | QuantKey, |
| 37 | ) |
| 38 | |
| 39 | logger = init_logger(__name__) |
| 40 | |
| 41 | |
| 42 | class NvFp4MoeBackend(Enum): |
| 43 | B12X = "B12X" |
| 44 | FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" |
| 45 | FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" |
| 46 | FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" |
| 47 | FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED" |
| 48 | FLASHINFER_B12X = "FLASHINFER_B12X" |
| 49 | VLLM_CUTLASS = "VLLM_CUTLASS" |
| 50 | MARLIN = "MARLIN" |
| 51 | HUMMING = "HUMMING" |
| 52 | EMULATION = "EMULATION" |
| 53 | |
| 54 | |
| 55 | FLASHINFER_NVFP4_MOE_BACKENDS = [ |
| 56 | NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 57 | NvFp4MoeBackend.FLASHINFER_CUTLASS, |
| 58 | NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 59 | NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, |
| 60 | NvFp4MoeBackend.FLASHINFER_B12X, |
| 61 | ] |
| 62 | |
| 63 | |
| 64 | def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool: |
| 65 | # Checks whether `backend` supports quantizing with scaling factors |
| 66 | # of all experts in Expert Parallel Mode when all experts are not |
| 67 | # on the same rank. |
| 68 | |
| 69 | return backend in FLASHINFER_NVFP4_MOE_BACKENDS |
| 70 | |
| 71 | |
| 72 | def backend_to_kernel_cls( |
| 73 | backend: NvFp4MoeBackend, |
| 74 | ) -> list[type[mk.FusedMoEExperts]]: |
| 75 | if backend == NvFp4MoeBackend.B12X: |
| 76 | from vllm.model_executor.layers.fused_moe.b12x import B12xExperts |
| 77 | |
| 78 | return [B12xExperts] |
| 79 | |
| 80 | elif backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: |
| 81 | from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( |
| 82 | TrtLlmNvFp4ExpertsModular, |
| 83 | TrtLlmNvFp4ExpertsMonolithic, |
| 84 | ) |
| 85 | |
| 86 | # NOTE: prefer Monolthic > Modular, so return Monolithic first. |
| 87 | return [ |
| 88 | TrtLlmNvFp4ExpertsMonolithic, |
| 89 | TrtLlmNvFp4ExpertsModular, |
| 90 | ] |
| 91 | |
| 92 | elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS: |
| 93 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501 |
| 94 | FlashInferExperts, |
| 95 | ) |
| 96 | |
| 97 | return [FlashInferExperts] |
| 98 | |
| 99 | elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: |
| 100 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( # noqa: E501 |
| 101 | FlashInferCuteDSLExperts, |
| 102 | ) |
| 103 | |
| 104 | return [FlashInferCuteDSLExperts] |
| 105 | |
| 106 | elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED: |
| 107 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import ( # noqa: E501 |
| 108 | FlashInferCuteDSLBatchedExperts, |
| 109 | ) |
| 110 | |
| 111 | return [FlashInferCuteDSLBatchedExperts] |
| 112 | |
| 113 | elif backend == NvFp4MoeBackend.FLASHINFER_B12X: |
| 114 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( # noqa: E501 |
| 115 | FlashInferB12xExperts, |
| 116 | ) |
| 117 | |
| 118 | return [FlashInferB12xExperts] |
| 119 | |
| 120 | elif backend == NvFp4MoeBackend.VLLM_CUTLASS: |
| 121 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 122 | CutlassExpertsFp4, |
| 123 | ) |
| 124 | |
| 125 | return [CutlassExpertsFp4] |
| 126 | |
| 127 | elif backend == NvFp4MoeBackend.MARLIN: |
| 128 | from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( |
| 129 | MarlinExperts, |
| 130 | ) |
| 131 | |
| 132 | return [MarlinExperts] |
| 133 | elif backend == NvFp4MoeBackend.HUMMING: |
| 134 | from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import ( |
| 135 | BatchedHummingGroupedExperts, |
| 136 | HummingGroupedExperts, |
| 137 | HummingIndexedExperts, |
| 138 | ) |
| 139 | |
| 140 | return [ |
| 141 | BatchedHummingGroupedExperts, |
| 142 | HummingGroupedExperts, |
| 143 | HummingIndexedExperts, |
| 144 | ] |
| 145 | elif backend == NvFp4MoeBackend.EMULATION: |
| 146 | from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import ( |
| 147 | Nvfp4QuantizationEmulationTritonExperts, |
| 148 | ) |
| 149 | |
| 150 | return [Nvfp4QuantizationEmulationTritonExperts] |
| 151 | else: |
| 152 | raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}") |
| 153 | |
| 154 | |
| 155 | def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: |
| 156 | """Map user's MoEBackend to NvFp4MoeBackend.""" |
| 157 | mapping = { |
| 158 | "b12x": NvFp4MoeBackend.B12X, |
| 159 | "cutlass": NvFp4MoeBackend.VLLM_CUTLASS, |
| 160 | "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 161 | "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, |
| 162 | "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 163 | "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, |
| 164 | "marlin": NvFp4MoeBackend.MARLIN, |
| 165 | "humming": NvFp4MoeBackend.HUMMING, |
| 166 | "emulation": NvFp4MoeBackend.EMULATION, |
| 167 | } |
| 168 | if backend := mapping.get(runner_backend): |
| 169 | return backend |
| 170 | raise ValueError( |
| 171 | f"moe_backend='{runner_backend}' is not supported for NvFP4 MoE. " |
| 172 | f"Expected one of {list(mapping.keys())}." |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | def _use_a16(backend: NvFp4MoeBackend, checkpoint_uses_a16: bool) -> bool: |
| 177 | return checkpoint_uses_a16 or ( |
| 178 | backend == NvFp4MoeBackend.B12X and envs.VLLM_B12X_MOE_FP4_FORCE_A16 |
| 179 | ) |
| 180 | |
| 181 | |
| 182 | def select_nvfp4_moe_backend( |
| 183 | config: FusedMoEConfig, |
| 184 | weight_key: QuantKey | None, |
| 185 | activation_key: QuantKey | None, |
| 186 | ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]: |
| 187 | """Select the primary NvFP4 MoE backend |
| 188 | Note: Shape-specific fallbacks may still occur at runtime. |
| 189 | """ |
| 190 | # NOTE: the kernels are selected in the following order. |
| 191 | # FLASHINFER_B12X is intentionally excluded from auto-selection until |
| 192 | # the upstream CUTLASS SM121 MMA op guard is resolved; use |
| 193 | # moe_backend="flashinfer_b12x" to opt in explicitly. |
| 194 | AVAILABLE_BACKENDS = [ |
| 195 | NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 196 | NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 197 | NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, |
| 198 | NvFp4MoeBackend.FLASHINFER_CUTLASS, |
| 199 | NvFp4MoeBackend.VLLM_CUTLASS, |
| 200 | NvFp4MoeBackend.MARLIN, |
| 201 | NvFp4MoeBackend.HUMMING, |
| 202 | NvFp4MoeBackend.EMULATION, |
| 203 | ] |
| 204 | |
| 205 | NVFP4_BACKENDS_WITH_CLAMP = { |
| 206 | NvFp4MoeBackend.B12X, |
| 207 | NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 208 | NvFp4MoeBackend.FLASHINFER_CUTLASS, |
| 209 | NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 210 | NvFp4MoeBackend.VLLM_CUTLASS, |
| 211 | NvFp4MoeBackend.MARLIN, |
| 212 | NvFp4MoeBackend.EMULATION, |
| 213 | NvFp4MoeBackend.HUMMING, |
| 214 | } |
| 215 | |
| 216 | if config.swiglu_limit is not None: |
| 217 | AVAILABLE_BACKENDS = [ |
| 218 | b for b in AVAILABLE_BACKENDS if b in NVFP4_BACKENDS_WITH_CLAMP |
| 219 | ] |
| 220 | |
| 221 | use_batched = config.moe_parallel_config.use_batched_activation_format |
| 222 | activation_format = ( |
| 223 | mk.FusedMoEActivationFormat.BatchedExperts |
| 224 | if use_batched |
| 225 | else mk.FusedMoEActivationFormat.Standard |
| 226 | ) |
| 227 | |
| 228 | def _make_log_backend(backend: NvFp4MoeBackend): |
| 229 | available_backend_strs = [b.value for b in AVAILABLE_BACKENDS] |
| 230 | return ( |
| 231 | f"Using '{backend.value}' NvFp4 MoE backend out " |
| 232 | f"of potential backends: {available_backend_strs}." |
| 233 | ) |
| 234 | |
| 235 | def _make_log_unsupported(backend: NvFp4MoeBackend, reason: str | None) -> str: |
| 236 | if reason: |
| 237 | return ( |
| 238 | f"NvFp4 MoE backend '{backend.value}' does not support the " |
| 239 | f"deployment configuration since {reason}." |
| 240 | ) |
| 241 | else: |
| 242 | return ( |
| 243 | f"NvFp4 MoE backend '{backend.value}' does not support the " |
| 244 | "deployment configuration." |
| 245 | ) |
| 246 | |
| 247 | def _return_or_raise( |
| 248 | backend: NvFp4MoeBackend, |
| 249 | config: FusedMoEConfig, |
| 250 | weight_key: QuantKey | None, |
| 251 | activation_key: QuantKey | None, |
| 252 | activation_format: mk.FusedMoEActivationFormat, |
| 253 | ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]: |
| 254 | for k_cls in backend_to_kernel_cls(backend): |
| 255 | supported, reason = k_cls.is_supported_config( |
| 256 | k_cls, config, weight_key, activation_key, activation_format |
| 257 | ) |
| 258 | if supported: |
| 259 | logger.info_once(_make_log_backend(backend)) |
| 260 | return backend, k_cls |
| 261 | |
| 262 | raise ValueError(_make_log_unsupported(backend, reason)) |
| 263 | |
| 264 | # Handle explicit moe_backend from user. |
| 265 | runner_backend = config.moe_backend |
| 266 | if runner_backend != "auto": |
| 267 | requested_backend = map_nvfp4_backend(runner_backend) |
| 268 | if _use_a16(requested_backend, False): |
| 269 | activation_key = None |
| 270 | # For batched activation format, use batched variant if available. |
| 271 | if ( |
| 272 | activation_format == mk.FusedMoEActivationFormat.BatchedExperts |
| 273 | and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL |
| 274 | ): |
| 275 | requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED |
| 276 | if ( |
| 277 | config.swiglu_limit is not None |
| 278 | and requested_backend not in NVFP4_BACKENDS_WITH_CLAMP |
| 279 | ): |
| 280 | raise ValueError( |
| 281 | f"Model sets swiglu_limit={config.swiglu_limit}, but the " |
| 282 | f"explicitly requested moe_backend={runner_backend!r} does " |
| 283 | f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', " |
| 284 | f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', " |
| 285 | f"'b12x', 'marlin', or 'humming' instead." |
| 286 | ) |
| 287 | return _return_or_raise( |
| 288 | requested_backend, config, weight_key, activation_key, activation_format |
| 289 | ) |
| 290 | |
| 291 | # Select kernels in order of backend. |
| 292 | for backend in AVAILABLE_BACKENDS: |
| 293 | for k_cls in backend_to_kernel_cls(backend): |
| 294 | supported, reason = k_cls.is_supported_config( |
| 295 | k_cls, |
| 296 | config, |
| 297 | weight_key, |
| 298 | activation_key, |
| 299 | activation_format, |
| 300 | ) |
| 301 | if supported: |
| 302 | logger.info_once(_make_log_backend(backend)) |
| 303 | return backend, k_cls |
| 304 | else: |
| 305 | logger.debug_once(_make_log_unsupported(backend, reason)) |
| 306 | |
| 307 | raise NotImplementedError( |
| 308 | "No NvFp4 MoE backend supports the deployment configuration." |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | def convert_to_nvfp4_moe_kernel_format( |
| 313 | nvfp4_backend: NvFp4MoeBackend, |
| 314 | layer: RoutedExperts, |
| 315 | w13: torch.Tensor, |
| 316 | w13_scale: torch.Tensor, |
| 317 | w13_scale_2: torch.Tensor, |
| 318 | a13_scale: torch.Tensor | None, |
| 319 | w2: torch.Tensor, |
| 320 | w2_scale: torch.Tensor, |
| 321 | w2_scale_2: torch.Tensor, |
| 322 | a2_scale: torch.Tensor | None, |
| 323 | is_act_and_mul: bool, |
| 324 | use_a16: bool = False, |
| 325 | ) -> tuple[ |
| 326 | torch.Tensor, |
| 327 | torch.Tensor, |
| 328 | torch.Tensor, |
| 329 | torch.Tensor, |
| 330 | torch.Tensor, |
| 331 | torch.Tensor, |
| 332 | torch.Tensor, |
| 333 | torch.Tensor, |
| 334 | ]: |
| 335 | use_a16 = _use_a16(nvfp4_backend, use_a16) |
| 336 | if nvfp4_backend == NvFp4MoeBackend.B12X: |
| 337 | if a13_scale is None or a2_scale is None: |
| 338 | if not use_a16: |
| 339 | raise ValueError("b12x NVFP4 MoE requires activation scales") |
| 340 | num_experts = w13.shape[0] |
| 341 | a13_scale = torch.ones(num_experts, dtype=torch.float32, device=w13.device) |
| 342 | a2_scale = torch.ones(num_experts, dtype=torch.float32, device=w2.device) |
| 343 | ( |
| 344 | w13, |
| 345 | w13_scale, |
| 346 | w13_scale_2, |
| 347 | a13_scale, |
| 348 | w2, |
| 349 | w2_scale, |
| 350 | w2_scale_2, |
| 351 | a2_scale, |
| 352 | ) = prepare_nvfp4_moe_layer_for_b12x( |
| 353 | w13=w13, |
| 354 | w13_scale=w13_scale, |
| 355 | w13_scale_2=w13_scale_2, |
| 356 | a13_scale=a13_scale, |
| 357 | w2=w2, |
| 358 | w2_scale=w2_scale, |
| 359 | w2_scale_2=w2_scale_2, |
| 360 | a2_scale=a2_scale, |
| 361 | is_act_and_mul=is_act_and_mul, |
| 362 | reorder_w13=use_a16, |
| 363 | ) |
| 364 | elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: |
| 365 | ( |
| 366 | w13, |
| 367 | w13_scale, |
| 368 | w13_scale_2, |
| 369 | a13_scale, |
| 370 | w2, |
| 371 | w2_scale, |
| 372 | w2_scale_2, |
| 373 | a2_scale, |
| 374 | ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( |
| 375 | layer=layer, |
| 376 | w13=w13, |
| 377 | w13_scale=w13_scale, |
| 378 | w13_scale_2=w13_scale_2, |
| 379 | a13_scale=a13_scale, |
| 380 | w2=w2, |
| 381 | w2_scale=w2_scale, |
| 382 | w2_scale_2=w2_scale_2, |
| 383 | a2_scale=a2_scale, |
| 384 | ) |
| 385 | elif ( |
| 386 | nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS |
| 387 | or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS |
| 388 | ): |
| 389 | ( |
| 390 | w13, |
| 391 | w13_scale, |
| 392 | w13_scale_2, |
| 393 | a13_scale, |
| 394 | w2, |
| 395 | w2_scale, |
| 396 | w2_scale_2, |
| 397 | a2_scale, |
| 398 | ) = prepare_nvfp4_moe_layer_for_fi_or_cutlass( |
| 399 | backend=nvfp4_backend, |
| 400 | layer=layer, |
| 401 | w13=w13, |
| 402 | w13_scale=w13_scale, |
| 403 | w13_scale_2=w13_scale_2, |
| 404 | a13_scale=a13_scale, |
| 405 | w2=w2, |
| 406 | w2_scale=w2_scale, |
| 407 | w2_scale_2=w2_scale_2, |
| 408 | a2_scale=a2_scale, |
| 409 | is_act_and_mul=is_act_and_mul, |
| 410 | ) |
| 411 | elif nvfp4_backend == NvFp4MoeBackend.HUMMING: |
| 412 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 413 | convert_to_humming_moe_kernel_format, |
| 414 | ) |
| 415 | |
| 416 | # Discriminate the source checkpoint layout by its global-scale param: |
| 417 | # compressed-tensors uses *_weight_global_scale, modelopt *_weight_scale_2. |
| 418 | # The logical schema is identical (nvfp4 group-16); only the on-layer |
| 419 | # param names differ. TODO: normalize both methods to a single canonical |
| 420 | # layout upstream so the oracle needs neither the probe nor the re-alias. |
| 421 | if hasattr(layer, "w13_weight_global_scale"): |
| 422 | quant_config = { |
| 423 | "quant_method": "compressed-tensors", |
| 424 | "format": "nvfp4-pack-quantized", |
| 425 | "type": "float", |
| 426 | "num_bits": 4, |
| 427 | "strategy": "group", |
| 428 | "group_size": 16, |
| 429 | } |
| 430 | # CT pack-quantized reads `weight_packed`; the method renamed it to |
| 431 | # `weight`. Re-alias (convert replaces all params anyway). |
| 432 | layer.w13_weight_packed = layer.w13_weight |
| 433 | layer.w2_weight_packed = layer.w2_weight |
| 434 | else: |
| 435 | quant_config = {"quant_method": "modelopt", "quant_algo": "nvfp4"} |
| 436 | |
| 437 | convert_to_humming_moe_kernel_format(layer, quant_config=quant_config) |
| 438 | a13_scale = None |
| 439 | a2_scale = None |
| 440 | w13 = layer.w13_weight |
| 441 | w13_scale = layer.w13_weight_scale |
| 442 | w13_scale_2 = getattr(layer, "w13_weight_scale_2", None) |
| 443 | w2 = layer.w2_weight |
| 444 | w2_scale = layer.w2_weight_scale |
| 445 | w2_scale_2 = getattr(layer, "w2_weight_scale_2", None) |
| 446 | elif nvfp4_backend == NvFp4MoeBackend.MARLIN: |
| 447 | a13_scale = None |
| 448 | a2_scale = None |
| 449 | ( |
| 450 | w13, |
| 451 | w13_scale, |
| 452 | w13_scale_2, |
| 453 | w2, |
| 454 | w2_scale, |
| 455 | w2_scale_2, |
| 456 | ) = prepare_nvfp4_moe_layer_for_marlin( |
| 457 | layer=layer, |
| 458 | w13=w13, |
| 459 | w13_scale=w13_scale, |
| 460 | w13_scale_2=w13_scale_2, |
| 461 | w2=w2, |
| 462 | w2_scale=w2_scale, |
| 463 | w2_scale_2=w2_scale_2, |
| 464 | is_act_and_mul=is_act_and_mul, |
| 465 | ) |
| 466 | elif nvfp4_backend == NvFp4MoeBackend.EMULATION: |
| 467 | # Move the E2M1 lookup table to the device now, because |
| 468 | # `.to(device)` is not allowed during CUDA graph capture. |
| 469 | kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(w13.device) |
| 470 | |
| 471 | if a13_scale is None or a2_scale is None: |
| 472 | raise ValueError( |
| 473 | "Activation global scales should not be None, got" |
| 474 | f" a13_scale={a13_scale}, a2_scale={a2_scale}" |
| 475 | ) |
| 476 | |
| 477 | if torch.unique(a13_scale).numel() != 1 or torch.unique(a2_scale).numel() != 1: |
| 478 | logger.warning_once( |
| 479 | "In NVFP4 linear, the activation global scale for inputs are different" |
| 480 | " for MOE w13 (gate_up_proj) layer or MOE w2 (down_proj). Using" |
| 481 | " a13_scale = a13_scale.max() and a2_scale = a2_scale.max()." |
| 482 | ) |
| 483 | |
| 484 | # 1. We take the max following e.g. quantization/utils/flashinfer_fp4_moe.py. |
| 485 | # 2. moe_kernel_quantize_input -> ref_nvfp4_quant_dequant |
| 486 | # use the inverse scale directly (large global scale). |
| 487 | # NOTE: Before this point, `a13_scale` and `a2_scale` are such that: |
| 488 | # `FP8_MAX = activation[expert_id].abs().max() * global_scale[expert_id]`, |
| 489 | # and `global_scale[expert_id]` are small (~1e-4). |
| 490 | # Taking the largest global scale likely results in overflowing the FP8 range |
| 491 | # for other experts - other selection strategies may be used. |
| 492 | a13_scale = 1.0 / a13_scale.max().to(torch.float32) |
| 493 | a2_scale = 1.0 / a2_scale.max().to(torch.float32) |
| 494 | else: |
| 495 | raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}") |
| 496 | |
| 497 | return ( |
| 498 | w13, |
| 499 | w13_scale, |
| 500 | w13_scale_2, |
| 501 | a13_scale, |
| 502 | w2, |
| 503 | w2_scale, |
| 504 | w2_scale_2, |
| 505 | a2_scale, |
| 506 | ) |
| 507 | |
| 508 | |
| 509 | def make_nvfp4_moe_quant_config( |
| 510 | backend: NvFp4MoeBackend, |
| 511 | w13_scale: torch.Tensor, |
| 512 | w2_scale: torch.Tensor, |
| 513 | w13_scale_2: torch.Tensor, |
| 514 | w2_scale_2: torch.Tensor, |
| 515 | a13_scale: torch.Tensor, |
| 516 | a2_scale: torch.Tensor, |
| 517 | swiglu_limit: float | None = None, |
| 518 | swiglu_alpha: float | None = None, |
| 519 | swiglu_beta: float | None = None, |
| 520 | layer: torch.nn.Module | None = None, |
| 521 | use_a16: bool = False, |
| 522 | ) -> FusedMoEQuantConfig: |
| 523 | use_a16 = _use_a16(backend, use_a16) |
| 524 | if backend == NvFp4MoeBackend.HUMMING: |
| 525 | from vllm.model_executor.layers.fused_moe import RoutedExperts |
| 526 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 527 | get_humming_moe_quant_config, |
| 528 | ) |
| 529 | |
| 530 | assert isinstance(layer, RoutedExperts) |
| 531 | return get_humming_moe_quant_config( |
| 532 | layer, |
| 533 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 534 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 535 | gemm1_clamp_limit=swiglu_limit, |
| 536 | ) |
| 537 | elif backend == NvFp4MoeBackend.MARLIN or ( |
| 538 | backend == NvFp4MoeBackend.B12X and use_a16 |
| 539 | ): |
| 540 | return nvfp4_w4a16_moe_quant_config( |
| 541 | g1_alphas=w13_scale_2, |
| 542 | g2_alphas=w2_scale_2, |
| 543 | w1_scale=w13_scale, |
| 544 | w2_scale=w2_scale, |
| 545 | gemm1_alpha=swiglu_alpha, |
| 546 | gemm1_beta=swiglu_beta, |
| 547 | gemm1_clamp_limit=swiglu_limit, |
| 548 | ) |
| 549 | elif backend == NvFp4MoeBackend.EMULATION: |
| 550 | return nvfp4_moe_quant_config( |
| 551 | g1_alphas=w13_scale_2, |
| 552 | g2_alphas=w2_scale_2, |
| 553 | a1_gscale=a13_scale, |
| 554 | a2_gscale=a2_scale, |
| 555 | w1_scale=w13_scale, |
| 556 | w2_scale=w2_scale, |
| 557 | gemm1_alpha=swiglu_alpha, |
| 558 | gemm1_beta=swiglu_beta, |
| 559 | gemm1_clamp_limit=swiglu_limit, |
| 560 | ) |
| 561 | |
| 562 | if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: |
| 563 | w13_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w13_scale) |
| 564 | w2_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w2_scale) |
| 565 | |
| 566 | # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas. |
| 567 | # The expert's process_weights_after_loading will fuse activation |
| 568 | # scales in-place. Since the quant config references the same tensor |
| 569 | # as the registered parameter, EPLB rearrangement stays in sync. |
| 570 | return nvfp4_moe_quant_config( |
| 571 | g1_alphas=w13_scale_2, |
| 572 | g2_alphas=w2_scale_2, |
| 573 | a1_gscale=(1.0 / a13_scale), |
| 574 | a2_gscale=(1.0 / a2_scale), |
| 575 | w1_scale=w13_scale, |
| 576 | w2_scale=w2_scale, |
| 577 | # NOTE(rob): this is a hack until the MoE kernels |
| 578 | # create their own quant configs. TRTLLM kernel |
| 579 | # does not accept swizzled input quant scales. |
| 580 | is_scale_swizzled=( |
| 581 | backend |
| 582 | not in ( |
| 583 | NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 584 | NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 585 | ) |
| 586 | ), |
| 587 | gemm1_alpha=swiglu_alpha, |
| 588 | gemm1_beta=swiglu_beta, |
| 589 | gemm1_clamp_limit=swiglu_limit, |
| 590 | ) |
| 591 | |
| 592 | |
| 593 | def make_nvfp4_moe_kernel( |
| 594 | moe_quant_config: FusedMoEQuantConfig, |
| 595 | moe_config: FusedMoEConfig, |
| 596 | experts_cls: type[mk.FusedMoEExperts], |
| 597 | backend: NvFp4MoeBackend, |
| 598 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 599 | per_token_activation: bool = False, |
| 600 | ) -> mk.FusedMoEKernel: |
| 601 | # Create Prepare/Finalize. |
| 602 | prepare_finalize = maybe_make_prepare_finalize( |
| 603 | moe=moe_config, |
| 604 | quant_config=moe_quant_config, |
| 605 | routing_tables=routing_tables, |
| 606 | allow_new_interface=True, |
| 607 | use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic), |
| 608 | input_dtype=moe_config.in_dtype if per_token_activation else None, |
| 609 | ) |
| 610 | assert prepare_finalize is not None |
| 611 | |
| 612 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 613 | |
| 614 | extra_kwargs = {} |
| 615 | if ( |
| 616 | backend |
| 617 | in ( |
| 618 | NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 619 | NvFp4MoeBackend.FLASHINFER_CUTEDSL, |
| 620 | ) |
| 621 | and per_token_activation |
| 622 | ): |
| 623 | extra_kwargs["per_token_activation"] = True |
| 624 | |
| 625 | # Create Experts. |
| 626 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 627 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 628 | assert max_num_tokens is not None |
| 629 | experts = experts_cls( |
| 630 | moe_config=moe_config, |
| 631 | quant_config=moe_quant_config, |
| 632 | max_num_tokens=max_num_tokens, |
| 633 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 634 | **extra_kwargs, |
| 635 | ) |
| 636 | else: |
| 637 | experts = experts_cls( |
| 638 | moe_config=moe_config, |
| 639 | quant_config=moe_quant_config, |
| 640 | **extra_kwargs, |
| 641 | ) |
| 642 | |
| 643 | kernel = mk.FusedMoEKernel( |
| 644 | prepare_finalize, |
| 645 | experts, |
| 646 | ) |
| 647 | |
| 648 | return kernel |