vllm.model_executor.layers.fused_moe.oracle.unquantized

Classes:

• UnquantizedMoEKernelOracle –

  Class-based view of the unquantized MoE kernel oracle.

Functions:

• aiter_moe_intermediate_alignment –

  Intermediate-size alignment required by AITER's CK 2stages MoE kernel.

• map_unquantized_backend –

  Map user's MoEBackend to UnquantizedMoeBackend.

• select_unquantized_moe_backend –

  Select the primary Unquantized MoE backend.

• unquantized_round_up_hidden_size_and_intermediate_size –

  Round up dimensions before allocation to satisfy the selected kernel.

UnquantizedMoEKernelOracle

Bases: MoEKernelOracle[UnquantizedMoeBackend]

Class-based view of the unquantized MoE kernel oracle.

Each method delegates to its module-level counterpart so that instantiating and calling this class is bit-identical to calling the standalone functions. Follow-up PRs may move logic from the module-level functions into these methods.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]):
    """Class-based view of the unquantized MoE kernel oracle.

    Each method delegates to its module-level counterpart so that
    instantiating and calling this class is bit-identical to calling
    the standalone functions. Follow-up PRs may move logic from the
    module-level functions into these methods.
    """

    name = "Unquantized"

    def backend_enum_cls(self) -> type[UnquantizedMoeBackend]:
        return UnquantizedMoeBackend

    def get_priority_backends(
        self, moe_config: FusedMoEConfig
    ) -> list[UnquantizedMoeBackend]:
        return _get_priority_backends(moe_config)

    def backend_to_kernel_cls(
        self, backend: UnquantizedMoeBackend
    ) -> list[type[mk.FusedMoEExperts]]:
        return backend_to_kernel_cls(backend)

    def map_backend(self, runner_backend: MoEBackend) -> UnquantizedMoeBackend:
        return map_unquantized_backend(runner_backend)

    def select_backend(
        self,
        moe_config: FusedMoEConfig,
        weight_key: "QuantKey | None" = None,
        activation_key: "QuantKey | None" = None,
    ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
        assert weight_key is None and activation_key is None, (
            "Weights and activations will never be quantized for "
            "UnquantizedMoEKernelOracle"
        )
        return select_unquantized_moe_backend(moe_config)

    def convert_to_kernel_format(
        self,
        backend: UnquantizedMoeBackend,
        moe_config: FusedMoEConfig,
        w13_weight: torch.Tensor,
        w2_weight: torch.Tensor,
        layer: torch.nn.Module | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        return convert_to_unquantized_kernel_format(
            backend, moe_config, w13_weight, w2_weight, layer=layer
        )

    def make_kernel(
        self,
        quant_config: FusedMoEQuantConfig,
        moe_config: FusedMoEConfig,
        backend: UnquantizedMoeBackend,
        experts_cls: type[mk.FusedMoEExperts],
        routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
    ) -> mk.FusedMoEKernel:
        return make_unquantized_moe_kernel(
            quant_config, moe_config, backend, experts_cls, routing_tables
        )

_get_priority_backends(moe_config)

Get available backends in priority order based on platform and config.

This function can be extended to become more complex as needed.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]:
    """Get available backends in priority order based on platform and config.

    This function can be extended to become more complex as needed.
    """

    def _move_to_back(
        backends: list[UnquantizedMoeBackend],
        backend: UnquantizedMoeBackend,
    ) -> None:
        backends.append(backends.pop(backends.index(backend)))

    if current_platform.is_rocm():
        _AVAILABLE_BACKENDS = [
            UnquantizedMoeBackend.AITER,
            UnquantizedMoeBackend.TRITON,
            UnquantizedMoeBackend.BATCHED_TRITON,
        ]
    elif current_platform.is_cuda():
        _AVAILABLE_BACKENDS = [
            UnquantizedMoeBackend.FLASHINFER_TRTLLM,
            UnquantizedMoeBackend.FLASHINFER_CUTLASS,
            UnquantizedMoeBackend.TRITON,
            UnquantizedMoeBackend.BATCHED_TRITON,
        ]

        # On Hopper (SM90), the FlashInfer unquantized MoE kernels are slower
        # than Triton, so prefer Triton by default.
        if current_platform.is_device_capability_family(90):
            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM)
            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)

        # HACK: Qwen3.5 has crash with FLASHINFER_CUTLASS BF16 if DEP.
        # Updating the oracle querying logic is out of the scope of this
        # PR. Need to fix the kernel or update structure in follow up.
        if moe_config.moe_parallel_config.dp_size > 1:
            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)

        # HACK: unquantized FlashInfer aliases SWIGLUOAI to plain Swiglu
        # (swiglu_alpha/limit only set on the MXFP4 branch). Route to
        # Triton's swigluoai_and_mul until that's plumbed through. Same
        # demotion pattern as the Qwen3.5/dp_size hack above.
        if moe_config.activation == MoEActivation.SWIGLUOAI:
            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM)
            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)

    elif current_platform.is_xpu():
        # XPUExperts does not support batch invariance; TRITON is kept as a
        # fallback so VLLM_BATCH_INVARIANT=1 still resolves to a valid
        # (batch-invariant) backend on XPU.
        _AVAILABLE_BACKENDS = [
            UnquantizedMoeBackend.XPU,
            UnquantizedMoeBackend.TRITON,
        ]
    elif current_platform.is_cpu():
        _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU]
    return _AVAILABLE_BACKENDS

_trtllm_bf16_lora_supported(moe_config)

Gate for routing LoRA-enabled BF16 MoE to the FlashInfer TRT-LLM gemm1_lora_delta path (PR #3153).

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def _trtllm_bf16_lora_supported(moe_config: FusedMoEConfig) -> bool:
    """Gate for routing LoRA-enabled BF16 MoE to the FlashInfer TRT-LLM
    gemm1_lora_delta path (PR #3153).
    """
    from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
        TrtLlmBf16LoRAExperts,
    )

    # LoRA path returns before the oracle loop; reuse is_supported_config here.
    supported, _ = TrtLlmBf16LoRAExperts.is_supported_config(
        TrtLlmBf16LoRAExperts,
        moe_config,
        None,
        None,
        mk.FusedMoEActivationFormat.Standard,
    )
    if not supported:
        return False
    # The flashinfer trtllm fused-MoE kernel requires the per-partition
    # intermediate size to be a multiple of 128. Plain TP shards the MoE
    # intermediate dim (e.g. 768 -> 192 at tp=4), which would crash the kernel
    # at runtime; fall back to Triton in that case.
    return moe_config.intermediate_size_per_partition % 128 == 0

_zero_intermediate_padding(moe_config, w13_weight, w2_weight)

Zero intermediate padding; reloads rewrite only the checkpoint slices.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def _zero_intermediate_padding(
    moe_config: FusedMoEConfig,
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
) -> None:
    """Zero intermediate padding; reloads rewrite only the checkpoint slices."""
    unpadded = moe_config.intermediate_size_per_partition_unpadded
    assert unpadded is not None
    padded = moe_config.intermediate_size_per_partition
    if padded <= unpadded:
        return
    w13_weight[:, unpadded:padded].zero_()
    if moe_config.is_act_and_mul:
        w13_weight[:, padded + unpadded :].zero_()
    w2_weight[:, :, unpadded:].zero_()

aiter_moe_intermediate_alignment(intermediate)

Intermediate-size alignment required by AITER's CK 2stages MoE kernel.

AITER dispatches on inter_dim <= 192: below the threshold both stages use 64-wide tiles, above it at least one stage uses a 128-wide tile, and CK's IsSupportedArgument rejects a size not divisible by that width.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def aiter_moe_intermediate_alignment(intermediate: int) -> int:
    """Intermediate-size alignment required by AITER's CK 2stages MoE kernel.

    AITER dispatches on ``inter_dim <= 192``: below the threshold both stages
    use 64-wide tiles, above it at least one stage uses a 128-wide tile, and
    CK's ``IsSupportedArgument`` rejects a size not divisible by that width.
    """
    return 64 if intermediate <= 192 else 128

map_unquantized_backend(runner_backend)

Map user's MoEBackend to UnquantizedMoeBackend.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend:
    """Map user's MoEBackend to UnquantizedMoeBackend."""
    mapping = {
        "triton": UnquantizedMoeBackend.TRITON,
        "batched_triton": UnquantizedMoeBackend.BATCHED_TRITON,
        "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM,
        "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS,
        "aiter": UnquantizedMoeBackend.AITER,
    }
    if backend := mapping.get(runner_backend):
        return backend
    raise ValueError(
        f"moe_backend='{runner_backend}' is not supported for unquantized MoE. "
        f"Expected one of {list(mapping.keys())}."
    )

select_unquantized_moe_backend(moe_config)

Select the primary Unquantized MoE backend. Note: Shape-specific fallbacks may still occur at runtime.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def select_unquantized_moe_backend(
    moe_config: FusedMoEConfig,
) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
    """Select the primary Unquantized MoE backend.
    Note: Shape-specific fallbacks may still occur at runtime.
    """
    if current_platform.is_tpu():
        return UnquantizedMoeBackend.TPU, None

    if current_platform.is_out_of_tree():
        return UnquantizedMoeBackend.OOT, None

    if moe_config.is_lora_enabled:
        # The Punica path is token-major: it re-aligns `topk_ids` per adapter
        # and needs the token->LoRA mapping for the tokens this rank owns,
        # which only the Standard-format DP/EP prepare/finalize forwards after
        # dispatch (NaiveDPEPPrepareAndFinalize.set_lora_context). No
        # BatchedExperts implementation supports LoRA at all.
        if moe_config.activation_format != mk.FusedMoEActivationFormat.Standard:
            raise ValueError(
                "MoE LoRA requires a Standard activation layout, but this "
                "deployment produces "
                f"'{moe_config.activation_format.value[0]}' "
                "(all2all_backend="
                f"'{moe_config.moe_parallel_config.all2all_backend}', "
                f"moe_backend='{moe_config.moe_backend}'). Disable LoRA or "
                "select a backend that produces the Standard layout."
            )
        if _trtllm_bf16_lora_supported(moe_config):
            from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
                TrtLlmBf16LoRAExperts,
            )

            logger.info_once(
                "Using TrtLlmBf16LoRAExperts Unquantized MoE LoRA backend "
                "(TrtLlmBf16LoRAExperts)."
            )
            return UnquantizedMoeBackend.FLASHINFER_TRTLLM, TrtLlmBf16LoRAExperts

        return UnquantizedMoEKernelOracle(moe_config, None, None).return_or_raise(
            UnquantizedMoeBackend.TRITON,
            [UnquantizedMoeBackend.TRITON],
            is_lora=True,
        )

    # NOTE: the kernels are selected in the following order.
    AVAILABLE_BACKENDS = _get_priority_backends(moe_config)

    # NOTE(rob): We need to peak into the P/F selection to determine
    # if we are using the batched or standard expert format, which
    # if not ideal. Once we unify TP + DP/EP, we can select P/F first.
    oracle = UnquantizedMoEKernelOracle(moe_config, None, None)

    # MoonEP owns the expert layout (expert-grouped [NvS, H] + cu_seqlens),
    # so it is not interchangeable with the token-major experts backends.
    if moe_config.moe_parallel_config.use_moonep_kernels:
        return oracle.return_or_raise(UnquantizedMoeBackend.MOONEP, AVAILABLE_BACKENDS)

    runner_backend = moe_config.moe_backend
    # 'humming' and 'marlin' are quantization-only; an unquantized layer (e.g.
    # excluded via modules_to_not_convert, or an unquantized MTP draft of a
    # quantized model) falls through to auto instead of erroring.
    if runner_backend not in ["auto", "humming", "marlin"]:
        requested_backend = map_unquantized_backend(runner_backend)
        if (
            moe_config.activation_format.is_batched
            and requested_backend == UnquantizedMoeBackend.TRITON
        ):
            requested_backend = UnquantizedMoeBackend.BATCHED_TRITON

        return oracle.return_or_raise(requested_backend, AVAILABLE_BACKENDS)

    # Handle explicit AITER FP8 configuration.
    if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):
        skip_aiter_moe = (
            not envs.VLLM_ROCM_USE_AITER
            or not envs.VLLM_ROCM_USE_AITER_MOE
            or rocm_aiter_ops.is_rdna_aiter_enabled()
        )
        if skip_aiter_moe:
            if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:
                AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
        else:
            backend = UnquantizedMoeBackend.AITER
            kernel_cls, rejected = oracle.try_backend(backend)
            if kernel_cls is not None:
                logger.info_once(oracle.make_log_backend(backend, AVAILABLE_BACKENDS))
                return backend, kernel_cls
            # AITER was explicitly requested but does not support this
            # deployment configuration (e.g. a non-gated MoE activation,
            # which AiterExperts._supports_no_act_and_mul() rejects
            # unconditionally). Rather than hard-crashing at engine init,
            # warn — explicitly stating that this is falling back, and
            # which backends will be tried next — and fall back to the
            # remaining backends in priority order below.
            logger.warning_once(
                "VLLM_ROCM_USE_AITER_MOE=1 was requested, but %s "
                "Falling back to try the remaining available MoE backends: %s.",
                oracle.make_log_unsupported(
                    backend, rejected[-1][1] if rejected else None
                ),
                ", ".join(
                    b.value
                    for b in AVAILABLE_BACKENDS
                    if b != UnquantizedMoeBackend.AITER
                ),
            )
            if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:
                AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)

    for backend in AVAILABLE_BACKENDS:
        kernel_cls, _ = oracle.try_backend(backend)
        if kernel_cls is not None:
            logger.info_once(oracle.make_log_backend(backend, AVAILABLE_BACKENDS))
            return backend, kernel_cls

    raise NotImplementedError(
        "No Unquantized MoE backend supports the deployment configuration."
    )

unquantized_round_up_hidden_size_and_intermediate_size(backend, hidden_size, intermediate_size)

Round up dimensions before allocation to satisfy the selected kernel.

Source code in vllm/model_executor/layers/fused_moe/oracle/unquantized.py

def unquantized_round_up_hidden_size_and_intermediate_size(
    backend: UnquantizedMoeBackend,
    hidden_size: int,
    intermediate_size: int,
) -> tuple[int, int]:
    """Round up dimensions before allocation to satisfy the selected kernel."""
    if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
        intermediate_size = round_up(intermediate_size, 128)
    elif backend == UnquantizedMoeBackend.AITER:
        intermediate_size = round_up(
            intermediate_size, aiter_moe_intermediate_alignment(intermediate_size)
        )
    return hidden_size, intermediate_size