| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from enum import Enum |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 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.activation import MoEActivation |
| 15 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 16 | maybe_make_prepare_finalize, |
| 17 | ) |
| 18 | from vllm.model_executor.layers.fused_moe.config import ( |
| 19 | FusedMoEConfig, |
| 20 | FusedMoEQuantConfig, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle |
| 23 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 24 | align_moe_weights_for_fi, |
| 25 | convert_moe_weights_to_flashinfer_trtllm_block_layout, |
| 26 | swap_w13_to_w31, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.math_utils import round_up |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | class UnquantizedMoeBackend(Enum): |
| 38 | FLASHINFER_TRTLLM = "FlashInfer TRTLLM" |
| 39 | FLASHINFER_CUTLASS = "FlashInfer CUTLASS" |
| 40 | AITER = "ROCm AITER" |
| 41 | TRITON = "TRITON" |
| 42 | BATCHED_TRITON = "BATCHED_TRITON" |
| 43 | MOONEP = "MOONEP" |
| 44 | CPU = "CPU" |
| 45 | XPU = "XPU" |
| 46 | TPU = "TPU" |
| 47 | OOT = "OOT" |
| 48 | |
| 49 | |
| 50 | def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]: |
| 51 | """Get available backends in priority order based on platform and config. |
| 52 | |
| 53 | This function can be extended to become more complex as needed. |
| 54 | """ |
| 55 | |
| 56 | def _move_to_back( |
| 57 | backends: list[UnquantizedMoeBackend], |
| 58 | backend: UnquantizedMoeBackend, |
| 59 | ) -> None: |
| 60 | backends.append(backends.pop(backends.index(backend))) |
| 61 | |
| 62 | if current_platform.is_rocm(): |
| 63 | _AVAILABLE_BACKENDS = [ |
| 64 | UnquantizedMoeBackend.AITER, |
| 65 | UnquantizedMoeBackend.TRITON, |
| 66 | UnquantizedMoeBackend.BATCHED_TRITON, |
| 67 | ] |
| 68 | elif current_platform.is_cuda(): |
| 69 | _AVAILABLE_BACKENDS = [ |
| 70 | UnquantizedMoeBackend.FLASHINFER_TRTLLM, |
| 71 | UnquantizedMoeBackend.FLASHINFER_CUTLASS, |
| 72 | UnquantizedMoeBackend.TRITON, |
| 73 | UnquantizedMoeBackend.BATCHED_TRITON, |
| 74 | ] |
| 75 | |
| 76 | # On Hopper (SM90), the FlashInfer unquantized MoE kernels are slower |
| 77 | # than Triton, so prefer Triton by default. |
| 78 | if current_platform.is_device_capability_family(90): |
| 79 | _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM) |
| 80 | _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) |
| 81 | |
| 82 | # HACK: Qwen3.5 has crash with FLASHINFER_CUTLASS BF16 if DEP. |
| 83 | # Updating the oracle querying logic is out of the scope of this |
| 84 | # PR. Need to fix the kernel or update structure in follow up. |
| 85 | if moe_config.moe_parallel_config.dp_size > 1: |
| 86 | _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) |
| 87 | |
| 88 | # HACK: unquantized FlashInfer aliases SWIGLUOAI to plain Swiglu |
| 89 | # (swiglu_alpha/limit only set on the MXFP4 branch). Route to |
| 90 | # Triton's swigluoai_and_mul until that's plumbed through. Same |
| 91 | # demotion pattern as the Qwen3.5/dp_size hack above. |
| 92 | if moe_config.activation == MoEActivation.SWIGLUOAI: |
| 93 | _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM) |
| 94 | _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS) |
| 95 | |
| 96 | elif current_platform.is_xpu(): |
| 97 | # XPUExperts does not support batch invariance; TRITON is kept as a |
| 98 | # fallback so VLLM_BATCH_INVARIANT=1 still resolves to a valid |
| 99 | # (batch-invariant) backend on XPU. |
| 100 | _AVAILABLE_BACKENDS = [ |
| 101 | UnquantizedMoeBackend.XPU, |
| 102 | UnquantizedMoeBackend.TRITON, |
| 103 | ] |
| 104 | elif current_platform.is_cpu(): |
| 105 | _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU] |
| 106 | return _AVAILABLE_BACKENDS |
| 107 | |
| 108 | |
| 109 | def backend_to_kernel_cls( |
| 110 | backend: UnquantizedMoeBackend, |
| 111 | ) -> list[type[mk.FusedMoEExperts]]: |
| 112 | if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: |
| 113 | from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( |
| 114 | TrtLlmBf16ExpertsModular, |
| 115 | TrtLlmBf16ExpertsMonolithic, |
| 116 | ) |
| 117 | |
| 118 | return [TrtLlmBf16ExpertsMonolithic, TrtLlmBf16ExpertsModular] |
| 119 | |
| 120 | elif backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: |
| 121 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501 |
| 122 | FlashInferExperts, |
| 123 | ) |
| 124 | |
| 125 | return [FlashInferExperts] |
| 126 | |
| 127 | elif backend == UnquantizedMoeBackend.AITER: |
| 128 | from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( |
| 129 | AiterExperts, |
| 130 | ) |
| 131 | |
| 132 | return [AiterExperts] |
| 133 | |
| 134 | elif backend == UnquantizedMoeBackend.TRITON: |
| 135 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( |
| 136 | TritonExperts, |
| 137 | ) |
| 138 | |
| 139 | return [TritonExperts] |
| 140 | |
| 141 | elif backend == UnquantizedMoeBackend.BATCHED_TRITON: |
| 142 | from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( |
| 143 | BatchedTritonExperts, |
| 144 | ) |
| 145 | |
| 146 | return [BatchedTritonExperts] |
| 147 | |
| 148 | elif backend == UnquantizedMoeBackend.MOONEP: |
| 149 | from vllm.model_executor.layers.fused_moe.experts.moonep_experts import ( |
| 150 | MoonEPExperts, |
| 151 | ) |
| 152 | |
| 153 | return [MoonEPExperts] |
| 154 | |
| 155 | elif backend == UnquantizedMoeBackend.XPU: |
| 156 | from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExperts |
| 157 | |
| 158 | return [XPUExperts] |
| 159 | |
| 160 | elif backend == UnquantizedMoeBackend.CPU: |
| 161 | from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( |
| 162 | ArmCPUUnquantizedExperts, |
| 163 | CPUUnquantizedExperts, |
| 164 | PowerCPUUnquantizedExperts, |
| 165 | X86CPUUnquantizedExperts, |
| 166 | ) |
| 167 | |
| 168 | # Prefer architecture-specific kernels before the portable vector path. |
| 169 | return [ |
| 170 | X86CPUUnquantizedExperts, |
| 171 | ArmCPUUnquantizedExperts, |
| 172 | PowerCPUUnquantizedExperts, |
| 173 | CPUUnquantizedExperts, |
| 174 | ] |
| 175 | |
| 176 | else: |
| 177 | raise ValueError(f"Unknown unquantized MoE backend: {backend.value}") |
| 178 | |
| 179 | |
| 180 | def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend: |
| 181 | """Map user's MoEBackend to UnquantizedMoeBackend.""" |
| 182 | mapping = { |
| 183 | "triton": UnquantizedMoeBackend.TRITON, |
| 184 | "batched_triton": UnquantizedMoeBackend.BATCHED_TRITON, |
| 185 | "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM, |
| 186 | "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS, |
| 187 | "aiter": UnquantizedMoeBackend.AITER, |
| 188 | } |
| 189 | if backend := mapping.get(runner_backend): |
| 190 | return backend |
| 191 | raise ValueError( |
| 192 | f"moe_backend='{runner_backend}' is not supported for unquantized MoE. " |
| 193 | f"Expected one of {list(mapping.keys())}." |
| 194 | ) |
| 195 | |
| 196 | |
| 197 | def _trtllm_bf16_lora_supported(moe_config: FusedMoEConfig) -> bool: |
| 198 | """Gate for routing LoRA-enabled BF16 MoE to the FlashInfer TRT-LLM |
| 199 | gemm1_lora_delta path (PR #3153). |
| 200 | """ |
| 201 | from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import ( |
| 202 | TrtLlmBf16LoRAExperts, |
| 203 | ) |
| 204 | |
| 205 | # LoRA path returns before the oracle loop; reuse is_supported_config here. |
| 206 | supported, _ = TrtLlmBf16LoRAExperts.is_supported_config( |
| 207 | TrtLlmBf16LoRAExperts, |
| 208 | moe_config, |
| 209 | None, |
| 210 | None, |
| 211 | mk.FusedMoEActivationFormat.Standard, |
| 212 | ) |
| 213 | if not supported: |
| 214 | return False |
| 215 | # The flashinfer trtllm fused-MoE kernel requires the per-partition |
| 216 | # intermediate size to be a multiple of 128. Plain TP shards the MoE |
| 217 | # intermediate dim (e.g. 768 -> 192 at tp=4), which would crash the kernel |
| 218 | # at runtime; fall back to Triton in that case. |
| 219 | return moe_config.intermediate_size_per_partition % 128 == 0 |
| 220 | |
| 221 | |
| 222 | def select_unquantized_moe_backend( |
| 223 | moe_config: FusedMoEConfig, |
| 224 | ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]: |
| 225 | """Select the primary Unquantized MoE backend. |
| 226 | Note: Shape-specific fallbacks may still occur at runtime. |
| 227 | """ |
| 228 | if current_platform.is_tpu(): |
| 229 | return UnquantizedMoeBackend.TPU, None |
| 230 | |
| 231 | if current_platform.is_out_of_tree(): |
| 232 | return UnquantizedMoeBackend.OOT, None |
| 233 | |
| 234 | if moe_config.is_lora_enabled: |
| 235 | if _trtllm_bf16_lora_supported(moe_config): |
| 236 | from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import ( |
| 237 | TrtLlmBf16LoRAExperts, |
| 238 | ) |
| 239 | |
| 240 | logger.info_once( |
| 241 | "Using TrtLlmBf16LoRAExperts Unquantized MoE LoRA backend " |
| 242 | "(TrtLlmBf16LoRAExperts)." |
| 243 | ) |
| 244 | return UnquantizedMoeBackend.FLASHINFER_TRTLLM, TrtLlmBf16LoRAExperts |
| 245 | logger.info_once("Using TRITON Unquantized MoE LoRA backend") |
| 246 | return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls( |
| 247 | UnquantizedMoeBackend.TRITON |
| 248 | )[0] |
| 249 | |
| 250 | # NOTE: the kernels are selected in the following order. |
| 251 | AVAILABLE_BACKENDS = _get_priority_backends(moe_config) |
| 252 | |
| 253 | # NOTE(rob): We need to peak into the P/F selection to determine |
| 254 | # if we are using the batched or standard expert format, which |
| 255 | # if not ideal. Once we unify TP + DP/EP, we can select P/F first. |
| 256 | activation_format = ( |
| 257 | mk.FusedMoEActivationFormat.BatchedExperts |
| 258 | if moe_config.moe_parallel_config.use_batched_activation_format |
| 259 | or moe_config.moe_backend == "batched_triton" |
| 260 | else mk.FusedMoEActivationFormat.Standard |
| 261 | ) |
| 262 | |
| 263 | def _make_log_backend(backend: UnquantizedMoeBackend) -> str: |
| 264 | available_strs = [b.value for b in AVAILABLE_BACKENDS] |
| 265 | return ( |
| 266 | f"Using {backend.value} Unquantized MoE backend out " |
| 267 | f"of potential backends: {available_strs}." |
| 268 | ) |
| 269 | |
| 270 | def _make_log_unsupported( |
| 271 | backend: UnquantizedMoeBackend, reason: str | None |
| 272 | ) -> str: |
| 273 | if reason: |
| 274 | return ( |
| 275 | f"Unquantized MoE backend {backend.value} does not support the " |
| 276 | f"deployment configuration since {reason}." |
| 277 | ) |
| 278 | return ( |
| 279 | f"Unquantized MoE backend '{backend.value}' does not support the " |
| 280 | "deployment configuration." |
| 281 | ) |
| 282 | |
| 283 | def _return_or_raise( |
| 284 | backend: UnquantizedMoeBackend, |
| 285 | config: FusedMoEConfig, |
| 286 | activation_format: mk.FusedMoEActivationFormat, |
| 287 | ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]: |
| 288 | reason = None |
| 289 | for k_cls in backend_to_kernel_cls(backend): |
| 290 | supported, reason = k_cls.is_supported_config( |
| 291 | k_cls, config, None, None, activation_format |
| 292 | ) |
| 293 | if supported: |
| 294 | logger.info_once(_make_log_backend(backend)) |
| 295 | return backend, k_cls |
| 296 | raise ValueError(_make_log_unsupported(backend, reason)) |
| 297 | |
| 298 | # MoonEP owns the expert layout (expert-grouped [NvS, H] + cu_seqlens), |
| 299 | # so it is not interchangeable with the token-major experts backends. |
| 300 | if moe_config.moe_parallel_config.use_moonep_kernels: |
| 301 | return _return_or_raise( |
| 302 | UnquantizedMoeBackend.MOONEP, moe_config, activation_format |
| 303 | ) |
| 304 | |
| 305 | runner_backend = moe_config.moe_backend |
| 306 | # 'humming' is quantization-only; an unquantized layer (e.g. excluded via |
| 307 | # modules_to_not_convert) falls through to auto instead of erroring. |
| 308 | if runner_backend not in ["auto", "humming"]: |
| 309 | requested_backend = map_unquantized_backend(runner_backend) |
| 310 | if ( |
| 311 | activation_format == mk.FusedMoEActivationFormat.BatchedExperts |
| 312 | and requested_backend == UnquantizedMoeBackend.TRITON |
| 313 | ): |
| 314 | requested_backend = UnquantizedMoeBackend.BATCHED_TRITON |
| 315 | |
| 316 | return _return_or_raise(requested_backend, moe_config, activation_format) |
| 317 | |
| 318 | # Handle explicit AITER FP8 configuration. |
| 319 | if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"): |
| 320 | skip_aiter_moe = ( |
| 321 | not envs.VLLM_ROCM_USE_AITER |
| 322 | or not envs.VLLM_ROCM_USE_AITER_MOE |
| 323 | or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 324 | ) |
| 325 | if skip_aiter_moe: |
| 326 | if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS: |
| 327 | AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER) |
| 328 | else: |
| 329 | backend = UnquantizedMoeBackend.AITER |
| 330 | reason = None |
| 331 | for k_cls in backend_to_kernel_cls(backend): |
| 332 | supported, reason = k_cls.is_supported_config( |
| 333 | k_cls, moe_config, None, None, activation_format |
| 334 | ) |
| 335 | if supported: |
| 336 | logger.info_once(_make_log_backend(backend)) |
| 337 | return backend, k_cls |
| 338 | # AITER was explicitly requested but does not support this |
| 339 | # deployment configuration (e.g. a non-gated MoE activation, |
| 340 | # which AiterExperts._supports_no_act_and_mul() rejects |
| 341 | # unconditionally). Rather than hard-crashing at engine init, |
| 342 | # warn — explicitly stating that this is falling back, and |
| 343 | # which backends will be tried next — and fall back to the |
| 344 | # remaining backends in priority order below. |
| 345 | logger.warning_once( |
| 346 | "VLLM_ROCM_USE_AITER_MOE=1 was requested, but %s " |
| 347 | "Falling back to try the remaining available MoE backends: %s.", |
| 348 | _make_log_unsupported(backend, reason), |
| 349 | ", ".join( |
| 350 | b.value |
| 351 | for b in AVAILABLE_BACKENDS |
| 352 | if b != UnquantizedMoeBackend.AITER |
| 353 | ), |
| 354 | ) |
| 355 | if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS: |
| 356 | AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER) |
| 357 | |
| 358 | for backend in AVAILABLE_BACKENDS: |
| 359 | for k_cls in backend_to_kernel_cls(backend): |
| 360 | supported, reason = k_cls.is_supported_config( |
| 361 | k_cls, moe_config, None, None, activation_format |
| 362 | ) |
| 363 | if supported: |
| 364 | logger.info_once(_make_log_backend(backend)) |
| 365 | return backend, k_cls |
| 366 | |
| 367 | logger.debug_once(_make_log_unsupported(backend, reason)) |
| 368 | |
| 369 | raise NotImplementedError( |
| 370 | "No Unquantized MoE backend supports the deployment configuration." |
| 371 | ) |
| 372 | |
| 373 | |
| 374 | def aiter_moe_intermediate_alignment(intermediate: int) -> int: |
| 375 | """Intermediate-size alignment required by AITER's CK 2stages MoE kernel. |
| 376 | |
| 377 | AITER dispatches on ``inter_dim <= 192``: below the threshold both stages |
| 378 | use 64-wide tiles, above it at least one stage uses a 128-wide tile, and |
| 379 | CK's ``IsSupportedArgument`` rejects a size not divisible by that width. |
| 380 | """ |
| 381 | return 64 if intermediate <= 192 else 128 |
| 382 | |
| 383 | |
| 384 | def unquantized_round_up_hidden_size_and_intermediate_size( |
| 385 | backend: UnquantizedMoeBackend, |
| 386 | hidden_size: int, |
| 387 | intermediate_size: int, |
| 388 | ) -> tuple[int, int]: |
| 389 | """Round up dimensions before allocation to satisfy the selected kernel.""" |
| 390 | if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: |
| 391 | intermediate_size = round_up(intermediate_size, 128) |
| 392 | elif backend == UnquantizedMoeBackend.AITER: |
| 393 | intermediate_size = round_up( |
| 394 | intermediate_size, aiter_moe_intermediate_alignment(intermediate_size) |
| 395 | ) |
| 396 | return hidden_size, intermediate_size |
| 397 | |
| 398 | |
| 399 | def _zero_intermediate_padding( |
| 400 | moe_config: FusedMoEConfig, |
| 401 | w13_weight: torch.Tensor, |
| 402 | w2_weight: torch.Tensor, |
| 403 | ) -> None: |
| 404 | """Zero intermediate padding; reloads rewrite only the checkpoint slices.""" |
| 405 | unpadded = moe_config.intermediate_size_per_partition_unpadded |
| 406 | assert unpadded is not None |
| 407 | padded = moe_config.intermediate_size_per_partition |
| 408 | if padded <= unpadded: |
| 409 | return |
| 410 | w13_weight[:, unpadded:padded].zero_() |
| 411 | if moe_config.is_act_and_mul: |
| 412 | w13_weight[:, padded + unpadded :].zero_() |
| 413 | w2_weight[:, :, unpadded:].zero_() |
| 414 | |
| 415 | |
| 416 | def convert_to_unquantized_kernel_format( |
| 417 | unquantized_backend: UnquantizedMoeBackend, |
| 418 | moe_config: FusedMoEConfig, |
| 419 | w13_weight: torch.Tensor, |
| 420 | w2_weight: torch.Tensor, |
| 421 | layer: torch.nn.Module | None = None, |
| 422 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 423 | if unquantized_backend == UnquantizedMoeBackend.MOONEP: |
| 424 | if layer is None: |
| 425 | raise ValueError( |
| 426 | "MoonEP weight conversion requires the layer module to stash " |
| 427 | "the [E+B] weight layout on" |
| 428 | ) |
| 429 | # MoonEP addresses expert weights by global [E+B] row and needs one |
| 430 | # contiguous tensor per projection (gate / up / down). The layer's |
| 431 | # w13_weight becomes the gate tensor and w2_weight the down tensor; |
| 432 | # the up tensor and the full layout are stashed on the layer for |
| 433 | # the kernel wiring in _setup_kernel. |
| 434 | from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import ( |
| 435 | MOONEP_DEFAULT_NUM_PREFETCH_SLOTS, |
| 436 | gather_moonep_weight_layout, |
| 437 | ) |
| 438 | |
| 439 | if getattr(layer, "_moonep_weight_layout", None) is not None: |
| 440 | # On a reload the layer parameters were already replaced by the |
| 441 | # converted [E+B] gate/down tensors, so the source-format weights |
| 442 | # this conversion needs are gone (and the unregistered up |
| 443 | # projection cannot be recovered from them at all). Tracked in |
| 444 | # RFC #52095. |
| 445 | raise NotImplementedError( |
| 446 | "MoonEP does not support in-place weight reloads yet" |
| 447 | ) |
| 448 | layout = gather_moonep_weight_layout( |
| 449 | w13_weight, |
| 450 | w2_weight, |
| 451 | num_global_experts=moe_config.num_experts, |
| 452 | num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS, |
| 453 | ) |
| 454 | layer._moonep_weight_layout = layout |
| 455 | return layout.full_gate_weight, layout.full_down_weight |
| 456 | |
| 457 | if unquantized_backend == UnquantizedMoeBackend.AITER: |
| 458 | _zero_intermediate_padding(moe_config, w13_weight, w2_weight) |
| 459 | w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight) |
| 460 | w13_weight.is_shuffled = True |
| 461 | w2_weight.is_shuffled = True |
| 462 | return w13_weight, w2_weight |
| 463 | |
| 464 | elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: |
| 465 | if moe_config.is_act_and_mul: |
| 466 | # Swap halves to arrange as [w3; w1] (kernel expectation) |
| 467 | # Non-gated MoE: w13 is a single projection, no need to swap. |
| 468 | w13_weight = swap_w13_to_w31(w13_weight) |
| 469 | |
| 470 | elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: |
| 471 | is_act_and_mul = moe_config.is_act_and_mul |
| 472 | # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK |
| 473 | # weight layout uses block_k=128). Pad along the intermediate dim when |
| 474 | # the model + TP split don't satisfy the constraint. |
| 475 | w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi( |
| 476 | w13_weight, w2_weight, is_act_and_mul, min_alignment=128 |
| 477 | ) |
| 478 | moe_config.intermediate_size_per_partition = padded_intermediate |
| 479 | |
| 480 | _zero_intermediate_padding(moe_config, w13_weight, w2_weight) |
| 481 | |
| 482 | _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} |
| 483 | convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 484 | _cache_permute_indices, |
| 485 | w13_weight, |
| 486 | w2_weight, |
| 487 | is_gated_act_gemm=is_act_and_mul, |
| 488 | ) |
| 489 | # Keep checkpoint-shaped parameters for reload/IPC. The experts create |
| 490 | # BlockMajorK views at dispatch without changing the underlying storage. |
| 491 | return w13_weight, w2_weight |
| 492 | |
| 493 | if ( |
| 494 | unquantized_backend == UnquantizedMoeBackend.TRITON |
| 495 | and current_platform.is_rocm() |
| 496 | and envs.VLLM_ROCM_MOE_PADDING |
| 497 | ): |
| 498 | # Skip .contiguous(): it would undo the ROCm MoE weight padding. |
| 499 | return w13_weight, w2_weight |
| 500 | return w13_weight.contiguous(), w2_weight.contiguous() |
| 501 | |
| 502 | |
| 503 | def make_unquantized_moe_kernel( |
| 504 | quant_config: FusedMoEQuantConfig, |
| 505 | moe_config: FusedMoEConfig, |
| 506 | backend: UnquantizedMoeBackend, |
| 507 | experts_cls: type[mk.FusedMoEExperts], |
| 508 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 509 | ) -> mk.FusedMoEKernel: |
| 510 | from vllm.model_executor.layers.fused_moe.utils import ( |
| 511 | warn_if_moe_use_td_ineffective, |
| 512 | ) |
| 513 | |
| 514 | # Warn against the selected backend, not each probed candidate. |
| 515 | warn_if_moe_use_td_ineffective(backend.value, is_quantized=False) |
| 516 | |
| 517 | # Create Prepare/Finalize |
| 518 | is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) |
| 519 | prepare_finalize = maybe_make_prepare_finalize( |
| 520 | moe=moe_config, |
| 521 | quant_config=quant_config, |
| 522 | routing_tables=routing_tables, |
| 523 | allow_new_interface=True, |
| 524 | use_monolithic=is_monolithic, |
| 525 | ) |
| 526 | assert prepare_finalize is not None |
| 527 | |
| 528 | logger.info_once("Using %s", prepare_finalize.__class__.__name__) |
| 529 | logger.info_once("Using %s MoE backend", experts_cls.__name__) |
| 530 | |
| 531 | # Create Experts |
| 532 | if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: |
| 533 | max_num_tokens = prepare_finalize.max_num_tokens_per_rank() |
| 534 | assert max_num_tokens is not None |
| 535 | experts = experts_cls( |
| 536 | moe_config=moe_config, |
| 537 | quant_config=quant_config, |
| 538 | max_num_tokens=max_num_tokens, |
| 539 | num_dispatchers=prepare_finalize.num_dispatchers(), |
| 540 | ) |
| 541 | else: |
| 542 | experts = experts_cls( |
| 543 | moe_config=moe_config, |
| 544 | quant_config=quant_config, |
| 545 | ) |
| 546 | |
| 547 | kernel = mk.FusedMoEKernel( |
| 548 | prepare_finalize, |
| 549 | experts, |
| 550 | ) |
| 551 | |
| 552 | return kernel |
| 553 | |
| 554 | |
| 555 | # --------------------------------------------------------------------------- |
| 556 | # Class-based view (first PR of the #37753 series; see oracle/base.py). |
| 557 | # Methods delegate to the module-level functions above so behaviour is |
| 558 | # bit-identical with pre-class code. |
| 559 | # --------------------------------------------------------------------------- |
| 560 | |
| 561 | |
| 562 | class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]): |
| 563 | """Class-based view of the unquantized MoE kernel oracle. |
| 564 | |
| 565 | Each method delegates to its module-level counterpart so that |
| 566 | instantiating and calling this class is bit-identical to calling |
| 567 | the standalone functions. Follow-up PRs may move logic from the |
| 568 | module-level functions into these methods. |
| 569 | """ |
| 570 | |
| 571 | def backend_enum_cls(self) -> type[UnquantizedMoeBackend]: |
| 572 | return UnquantizedMoeBackend |
| 573 | |
| 574 | def get_priority_backends( |
| 575 | self, moe_config: FusedMoEConfig |
| 576 | ) -> list[UnquantizedMoeBackend]: |
| 577 | return _get_priority_backends(moe_config) |
| 578 | |
| 579 | def backend_to_kernel_cls( |
| 580 | self, backend: UnquantizedMoeBackend |
| 581 | ) -> list[type[mk.FusedMoEExperts]]: |
| 582 | return backend_to_kernel_cls(backend) |
| 583 | |
| 584 | def map_backend(self, runner_backend: MoEBackend) -> UnquantizedMoeBackend: |
| 585 | return map_unquantized_backend(runner_backend) |
| 586 | |
| 587 | def select_backend( |
| 588 | self, |
| 589 | moe_config: FusedMoEConfig, |
| 590 | weight_key: "QuantKey | None" = None, |
| 591 | activation_key: "QuantKey | None" = None, |
| 592 | ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]: |
| 593 | assert weight_key is None and activation_key is None, ( |
| 594 | "Weights and activations will never be quantized for " |
| 595 | "UnquantizedMoEKernelOracle" |
| 596 | ) |
| 597 | return select_unquantized_moe_backend(moe_config) |
| 598 | |
| 599 | def convert_to_kernel_format( |
| 600 | self, |
| 601 | backend: UnquantizedMoeBackend, |
| 602 | moe_config: FusedMoEConfig, |
| 603 | w13_weight: torch.Tensor, |
| 604 | w2_weight: torch.Tensor, |
| 605 | layer: torch.nn.Module | None = None, |
| 606 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 607 | return convert_to_unquantized_kernel_format( |
| 608 | backend, moe_config, w13_weight, w2_weight, layer=layer |
| 609 | ) |
| 610 | |
| 611 | def make_kernel( |
| 612 | self, |
| 613 | quant_config: FusedMoEQuantConfig, |
| 614 | moe_config: FusedMoEConfig, |
| 615 | backend: UnquantizedMoeBackend, |
| 616 | experts_cls: type[mk.FusedMoEExperts], |
| 617 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 618 | ) -> mk.FusedMoEKernel: |
| 619 | return make_unquantized_moe_kernel( |
| 620 | quant_config, moe_config, backend, experts_cls, routing_tables |
| 621 | ) |