vllm.model_executor.layers.fused_moe.utils

Functions:

• count_expert_num_tokens –

  Count the number to tokens assigned to each expert.

• fi_moe_largest_bucket –

  Estimate FlashInfer's MoE autotuning maximum token count.

• masked_amax_scale –

  Per-tensor activation quant scale that skips rows without a route.

• moe_use_td_hw_supported –

  Whether the current device can run the TD (gather) path of

• remap_topk_to_local –

  Map global expert ids to local ids; non-local or invalid routes become -1.

• resolve_layer_fused_shared_expert –

  Resolve whether AITER fused shared-expert execution is enabled.

• resolve_moe_use_td –

  Tri-state resolver for VLLM_TRITON_USE_TD.

• warn_if_moe_use_td_ineffective –

  One-shot warning when VLLM_TRITON_USE_TD is set but ignored.

_fp8_quantize(A, A_scale, per_act_token, block_shape=None, is_scale_transposed=False)

Perform fp8 quantization on the inputs. If a block_shape is provided, the output will be blocked.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def _fp8_quantize(
    A: torch.Tensor,
    A_scale: torch.Tensor | None,
    per_act_token: bool,
    block_shape: list[int] | None = None,
    is_scale_transposed: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Perform fp8 quantization on the inputs.  If a block_shape
    is provided, the output will be blocked.
    """
    if block_shape is None:
        # TODO(luka): use QuantFP8 custom op
        #  https://github.com/vllm-project/vllm/issues/20711
        A, A_scale = ops.scaled_fp8_quant(
            A, A_scale, use_per_token_if_dynamic=per_act_token
        )
    else:
        assert not per_act_token
        assert len(block_shape) == 2
        _, block_k = block_shape[0], block_shape[1]
        A, A_scale = per_token_group_quant_fp8(
            A, block_k, column_major_scales=is_scale_transposed
        )
        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

    return A, A_scale

_int8_quantize(A, A_scale, per_act_token, block_shape=None)

Perform int8 quantization on the inputs. If a block_shape is provided, the output will be blocked.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def _int8_quantize(
    A: torch.Tensor,
    A_scale: torch.Tensor | None,
    per_act_token: bool,
    block_shape: list[int] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Perform int8 quantization on the inputs.  If a block_shape
    is provided, the output will be blocked.
    """
    # If weights are per-channel (per_channel_quant=True), then
    # activations apply per-token quantization. Otherwise, assume
    # activation tensor-wise fp8/int8 quantization, dynamic or static
    if block_shape is None:
        if per_act_token:
            A, A_scale = per_token_quant_int8(A)
        elif A_scale is not None:
            # Static per-tensor: use the optimized CUDA kernel
            A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
        elif A_scale is None:
            # Dynamic per-tensor: compute scale then quantize via kernel
            A_scale = torch.clamp(A.abs().max() / 127.0, min=1e-10)
            A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
    else:
        assert not per_act_token
        assert len(block_shape) == 2
        _, block_k = block_shape[0], block_shape[1]
        A, A_scale = per_token_group_quant_int8(A, block_k)
        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

    return A, A_scale

_resize_cache(x, v)

Shrink the given tensor and apply the given view to it. This is used to resize the intermediate fused_moe caches.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
    """Shrink the given tensor and apply the given view to it.  This is
    used to resize the intermediate fused_moe caches.
    """
    assert prod(v) <= x.numel(), (
        f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})"
    )  # CUDAGRAPH unfriendly?
    return x.flatten()[: prod(v)].view(*v)

count_expert_num_tokens(topk_ids, num_local_experts, expert_map)

Count the number to tokens assigned to each expert.

Parameters:

• topk_ids

  (Tensor) –

  Tensor mapping each token to its list of experts.

• num_local_experts

  (int) –

  Number of experts in this rank.

• expert_map

  (Optional[Tensor]) –

  A tensor mapping expert indices from the global expert space to the local expert space of the expert parallel shard.

Returns: A tensor of size num_local_experts, where tensor[i] holds the number of tokens assigned to the ith expert.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def count_expert_num_tokens(
    topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None
) -> torch.Tensor:
    """Count the number to tokens assigned to each expert.

    Args:
        topk_ids (torch.Tensor): Tensor mapping each token to its list of experts.
        num_local_experts (int): Number of experts in this rank.
        expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
            from the global expert space to the local expert space of the expert
            parallel shard.

    Returns:
    A tensor of size num_local_experts, where tensor[i] holds the number
    of tokens assigned to the ith expert.

    """
    assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
    expert_num_tokens = torch.empty(
        (num_local_experts), device=topk_ids.device, dtype=torch.int32
    )

    grid = num_local_experts
    BLOCK_SIZE = min(topk_ids.numel(), 1024)
    BLOCK_SIZE = triton.next_power_of_2(BLOCK_SIZE)

    _count_expert_num_tokens[(grid,)](
        topk_ids,
        expert_num_tokens,
        num_local_experts,
        topk_ids.numel(),
        expert_map,
        HAS_EXPERT_MAP=expert_map is not None,
        BLOCK_SIZE=BLOCK_SIZE,
    )

    return expert_num_tokens

fi_moe_largest_bucket(moe_config)

Estimate FlashInfer's MoE autotuning maximum token count.

All DP ranks may contribute max_num_tokens to one invocation. Keep FlashInfer's default moe tune_max_num_tokens=8192 floor to avoid over-underestimation. DeepEP, SP, or PCP may make this underestimate, however overestimation may be dangerous, increasing tuning- cost and memory use.

NOTE: The DP factor applies even when EP is disabled:

Without --enable-expert-parallel, MoE layers would use tensor parallelism.

For a detailed explanation, see: docs/serving/data_parallel_deployment.md

Source code in vllm/model_executor/layers/fused_moe/utils.py

def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
    """Estimate FlashInfer's MoE autotuning maximum token count.

    All DP ranks may contribute `max_num_tokens` to one invocation.
    Keep FlashInfer's default moe `tune_max_num_tokens=8192`
    floor to avoid over-underestimation.
    DeepEP, SP, or PCP may make this underestimate, however overestimation
    may be dangerous, increasing tuning- cost and memory use.

    NOTE: The DP factor applies even when EP is disabled:
    > Without `--enable-expert-parallel`, MoE layers would use tensor parallelism.

    For a detailed explanation, see: `docs/serving/data_parallel_deployment.md`
    """
    dp_size = moe_config.elastic_ep_max_dp_size or moe_config.dp_size
    return max(moe_config.max_num_tokens * dp_size, 8192)

masked_amax_scale(x, quant_dtype, *, row_valid=None, num_valid=None, expert_num_tokens=None, rows_per_expert=None, eps=1e-10)

Per-tensor activation quant scale that skips rows without a route.

PaddedStandard buffers carry rows that hold no route (topk_ids == -1) and whose contents are uninitialized, so a tensor-wide absmax would be poisoned by them. Batched dispatch buffers have the same problem per expert: rows past expert_num_tokens[e] hold stale tokens. Those rows are dropped from the reduction instead; their quantized values are never consumed by an expert GEMM.

Exactly one of row_valid, num_valid and expert_num_tokens selects the valid rows.

Parameters:

• x

  (Tensor) –

  2D activations to reduce.

• quant_dtype

  (dtype) –

  Target quantization dtype (torch.int8 or a float8).

• row_valid

  (Tensor | None, default: None ) –

  Boolean flags, one per row of x. Use when the valid rows are not a prefix.

• num_valid

  (Tensor | None, default: None ) –

  1-element device tensor holding the number of leading valid rows. Cheaper than row_valid when padding trails the rows.

• expert_num_tokens

  (Tensor | None, default: None ) –

  Device tensor of received token counts, one per expert, for a buffer laid out as [expert, rows_per_expert, K].

• rows_per_expert

  (int | None, default: None ) –

  Row stride between experts, i.e. the worst-case per-expert capacity of the buffer.

• eps

  (float, default: 1e-10 ) –

  Lower bound on the returned scale, matching the dynamic kernels.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def masked_amax_scale(
    x: torch.Tensor,
    quant_dtype: torch.dtype,
    *,
    row_valid: torch.Tensor | None = None,
    num_valid: torch.Tensor | None = None,
    expert_num_tokens: torch.Tensor | None = None,
    rows_per_expert: int | None = None,
    eps: float = 1e-10,
) -> torch.Tensor:
    """Per-tensor activation quant scale that skips rows without a route.

    ``PaddedStandard`` buffers carry rows that hold no route (``topk_ids == -1``)
    and whose contents are uninitialized, so a tensor-wide absmax would be
    poisoned by them. Batched dispatch buffers have the same problem per
    expert: rows past ``expert_num_tokens[e]`` hold stale tokens. Those rows
    are dropped from the reduction instead; their quantized values are never
    consumed by an expert GEMM.

    Exactly one of ``row_valid``, ``num_valid`` and ``expert_num_tokens``
    selects the valid rows.

    Args:
        x: 2D activations to reduce.
        quant_dtype: Target quantization dtype (``torch.int8`` or a float8).
        row_valid: Boolean flags, one per row of ``x``. Use when the valid rows
            are not a prefix.
        num_valid: 1-element device tensor holding the number of leading valid
            rows. Cheaper than ``row_valid`` when padding trails the rows.
        expert_num_tokens: Device tensor of received token counts, one per
            expert, for a buffer laid out as ``[expert, rows_per_expert, K]``.
        rows_per_expert: Row stride between experts, i.e. the worst-case
            per-expert capacity of the buffer.
        eps: Lower bound on the returned scale, matching the dynamic kernels.

    """
    assert x.dim() == 2, f"expected 2D activations, got {x.shape}"
    assert sum(m is not None for m in (row_valid, num_valid, expert_num_tokens)) == 1, (
        "pass exactly one of row_valid / num_valid / expert_num_tokens"
    )
    assert (expert_num_tokens is None) == (rows_per_expert is None), (
        "rows_per_expert is only valid together with expert_num_tokens"
    )
    out = torch.zeros((1,), dtype=torch.float32, device=x.device)
    if row_valid is not None:
        assert row_valid.shape == (x.size(0),), (
            f"expected {x.size(0)} row flags, got {tuple(row_valid.shape)}"
        )
        row_valid = row_valid.contiguous()
        bound = out
    elif num_valid is not None:
        assert num_valid.numel() == 1, f"expected 1 row count, got {num_valid}"
        bound = num_valid
    else:
        assert expert_num_tokens is not None and rows_per_expert is not None
        assert rows_per_expert > 0, "expected a positive rows_per_expert"
        assert x.size(0) % rows_per_expert == 0, (
            f"{x.size(0)} rows is not a multiple of {rows_per_expert}"
        )
        assert expert_num_tokens.numel() == x.size(0) // rows_per_expert, (
            f"expected {x.size(0) // rows_per_expert} expert counts, "
            f"got {expert_num_tokens.numel()}"
        )
        bound = expert_num_tokens
    scale_max = (
        127.0 if quant_dtype == torch.int8 else float(torch.finfo(quant_dtype).max)
    )
    block_m, block_k = 32, 128
    _masked_amax_scale_kernel[(triton.cdiv(x.size(0), block_m),)](
        x,
        # The pointers of the inactive constexpr branches are never
        # dereferenced, so `out` doubles as a dummy.
        out if row_valid is None else row_valid,
        bound,
        out,
        x.size(0),
        x.size(1),
        x.stride(0),
        x.stride(1),
        rows_per_expert or 0,
        scale_max,
        eps,
        HAS_MASK=row_valid is not None,
        HAS_EXPERT_BOUND=expert_num_tokens is not None,
        BLOCK_M=block_m,
        BLOCK_K=block_k,
    )
    return out

moe_use_td_hw_supported()

Whether the current device can run the TD (gather) path of fused_moe_kernel (ignores the VLLM_TRITON_USE_TD override).

The A-load uses tensor_descriptor.gather, which lowers to the PTX tile::gather4 instruction. That instruction is part of the tcgen05/Tensor Memory (TMEM) family introduced with Blackwell and has no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike scatter4, gather4 is supported across the whole sm100+ range including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498, which enables gather4 on sm120/sm121 while leaving scatter4 unsupported there. So this gates on a blanket has_device_capability(100) rather than the sm100 family check used for the scatter store path.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def moe_use_td_hw_supported() -> bool:
    """Whether the current device can run the TD (gather) path of
    ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override).

    The A-load uses ``tensor_descriptor.gather``, which lowers to the PTX
    ``tile::gather4`` instruction. That instruction is part of the
    ``tcgen05``/Tensor Memory (TMEM) family introduced with Blackwell and has
    no Hopper (sm90) equivalent -- ptxas rejects it there ("Feature
    '.tile::gather4 ...' requires .target sm_100 or higher"). Unlike
    ``scatter4``, ``gather4`` is supported across the whole sm100+ range
    including consumer Blackwell (sm120/sm121): see triton-lang/triton#8498,
    which enables ``gather4`` on sm120/sm121 while leaving ``scatter4``
    unsupported there. So this gates on a blanket ``has_device_capability(100)``
    rather than the sm100 *family* check used for the scatter store path.
    """
    if current_platform.is_xpu():
        return True
    if current_platform.is_cuda():
        return current_platform.has_device_capability(100)
    return False

remap_topk_to_local(topk_ids, expert_map, *, out_dtype=torch.int64)

Map global expert ids to local ids; non-local or invalid routes become -1.

Replaces torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1) with one kernel. Returns a NEW tensor -- the caller keeps the original topk_ids as global_topk_ids, so this must not write in place.

Parameters:

• topk_ids

  (Tensor) –

  Global expert ids; -1 marks invalid slots. Must be contiguous; the Triton kernel indexes by flat numel() offsets.

• expert_map

  (Tensor) –

  Global-to-local mapping; -1 marks non-local experts.

• out_dtype

  (dtype, default: int64 ) –

  Output dtype, torch.int32 or torch.int64.

Returns:

• Tensor –

  Remapped local expert ids with the same shape as topk_ids.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def remap_topk_to_local(
    topk_ids: torch.Tensor,
    expert_map: torch.Tensor,
    *,
    out_dtype: torch.dtype = torch.int64,
) -> torch.Tensor:
    """Map global expert ids to local ids; non-local or invalid routes become -1.

    Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)``
    with one kernel. Returns a NEW tensor -- the caller keeps the original
    ``topk_ids`` as ``global_topk_ids``, so this must not write in place.

    Args:
        topk_ids: Global expert ids; -1 marks invalid slots. Must be contiguous;
            the Triton kernel indexes by flat ``numel()`` offsets.
        expert_map: Global-to-local mapping; -1 marks non-local experts.
        out_dtype: Output dtype, ``torch.int32`` or ``torch.int64``.

    Returns:
        Remapped local expert ids with the same shape as ``topk_ids``.

    """
    if out_dtype not in (torch.int32, torch.int64):
        raise ValueError(f"out_dtype must be int32 or int64, got {out_dtype}")
    assert topk_ids.is_contiguous()
    out = torch.empty(topk_ids.shape, dtype=out_dtype, device=topk_ids.device)
    n = topk_ids.numel()
    BLOCK = 1024
    grid = (triton.cdiv(n, BLOCK),)
    _remap_topk_to_local_kernel[grid](
        topk_ids,
        expert_map,
        out,
        n,
        BLOCK=BLOCK,
        OUT_INT32=out_dtype == torch.int32,
    )
    return out

resolve_layer_fused_shared_expert(quant_config, prefix, shared_expert_name='shared_experts', moe_backend=None, disable_under_ep=False)

Resolve whether AITER fused shared-expert execution is enabled.

Parameters:

• quant_config

  (QuantizationConfig | None) –

  Model quantization configuration.

• prefix

  (str) –

  MoE module prefix.

• shared_expert_name

  (str, default: 'shared_experts' ) –

  Shared-expert module name under prefix.

• moe_backend

  (str | None, default: None ) –

  MoE backend that runs the fused shared expert.

• disable_under_ep

  (bool, default: False ) –

  Whether this model family does not support FSE under expert parallelism. Kept as an explicit call-site policy instead of branching on model type here.

Returns:

• bool –

  Whether AITER fused shared experts are enabled.

Raises:

• ValueError –

  If requested shared-expert fusion is quantization-incompatible.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def resolve_layer_fused_shared_expert(
    quant_config: "QuantizationConfig | None",
    prefix: str,
    shared_expert_name: str = "shared_experts",
    moe_backend: str | None = None,
    disable_under_ep: bool = False,
) -> bool:
    """Resolve whether AITER fused shared-expert execution is enabled.

    Args:
        quant_config: Model quantization configuration.
        prefix: MoE module prefix.
        shared_expert_name: Shared-expert module name under ``prefix``.
        moe_backend: MoE backend that runs the fused shared expert.
        disable_under_ep: Whether this model family does not support FSE under
            expert parallelism. Kept as an explicit call-site policy instead of
            branching on model type here.

    Returns:
        Whether AITER fused shared experts are enabled.

    Raises:
        ValueError: If requested shared-expert fusion is quantization-incompatible.

    """
    # NOTE: is_fusion_moe_shared_experts_enabled is decorated with @if_aiter_supported
    # that returns None if AITER is not available.
    fse_requested = bool(rocm_aiter_ops.is_fusion_moe_shared_experts_enabled())

    # TODO: should it be disabled for all models?
    skip_reason: str | None = None
    if fse_requested and disable_under_ep:
        from vllm.config import get_current_vllm_config_or_none

        vllm_config = get_current_vllm_config_or_none()
        parallel_config = (
            vllm_config.parallel_config if vllm_config is not None else None
        )
        if parallel_config is not None and parallel_config.enable_expert_parallel:
            fse_requested = False
            skip_reason = "expert parallelism is enabled"

    if fse_requested:
        fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
            quant_config,
            f"{prefix}.experts",
            f"{prefix}.{shared_expert_name}",
            moe_backend=moe_backend,
        )
        if not fse_compatible:
            skip_reason = fse_reason or "unknown reason"
    else:
        fse_compatible = True
    is_fused_shared_expert_enabled = fse_requested and fse_compatible

    # In online shared_expert quantization is used, register it to
    # `online_quantization_config.quantized_layers` ahead of weight loading.
    if is_fused_shared_expert_enabled and quant_config is not None:
        online_quantization_config = quant_config.online_quantization_config
        if online_quantization_config is not None:
            from vllm.model_executor.layers.linear import (
                LinearBase,
                UnquantizedLinearMethod,
            )

            for projection_name in ("gate_up_proj", "down_proj"):
                projection_prefix = f"{prefix}.{shared_expert_name}.{projection_name}"
                quant_method_metadata = (
                    online_quantization_config.resolve_quant_method_cls(
                        LinearBase, projection_prefix
                    )
                )

                if quant_method_metadata is None:
                    continue

                source, quant_key_str, target_pattern, _, quant_method_cls = (
                    quant_method_metadata
                )

                if quant_method_cls in (None, UnquantizedLinearMethod):
                    continue

                online_quantization_config.quantized_layers[projection_prefix] = (
                    source.value,
                    quant_key_str,
                    target_pattern,
                )
    if skip_reason is not None:
        # Keep the full per-layer reason at debug level; the warning below is
        # deduplicated by cause so a long model logs it once.
        logger.debug("Skipping shared-expert fusion at %s: %s", prefix, skip_reason)
        logger.warning_once(
            "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
            "cannot be enabled - skipping shared-expert fusion: %s.",
            _normalize_skip_reason(skip_reason),
        )
    return is_fused_shared_expert_enabled

resolve_moe_use_td()

Tri-state resolver for VLLM_TRITON_USE_TD.

Unset auto-selects the TD path on XPU only, mirroring the attention dispatcher in triton_attn.py. 1/0 force it on/off regardless of hardware; forcing 1 where it cannot compile (see moe_use_td_hw_supported) fails at ptxas. Blackwell CUDA (sm100+) can compile it but is opt-in only, pending validation.

Source code in vllm/model_executor/layers/fused_moe/utils.py

def resolve_moe_use_td() -> bool:
    """Tri-state resolver for ``VLLM_TRITON_USE_TD``.

    Unset auto-selects the TD path on XPU only, mirroring the attention
    dispatcher in ``triton_attn.py``. ``1``/``0`` force it on/off regardless
    of hardware; forcing ``1`` where it cannot compile (see
    ``moe_use_td_hw_supported``) fails at ptxas. Blackwell CUDA (sm100+) can
    compile it but is opt-in only, pending validation.
    """
    override = envs.VLLM_TRITON_USE_TD
    if override is None:
        return current_platform.is_xpu()
    return override

warn_if_moe_use_td_ineffective(active_backend, is_quantized=False)

One-shot warning when VLLM_TRITON_USE_TD is set but ignored.

Fires when the user set the env explicitly and either (a) the active MoE backend is not the fused Triton kernel, or (b) the model is quantized (the TD path falls back to the pointer path under any quantization).

Source code in vllm/model_executor/layers/fused_moe/utils.py

def warn_if_moe_use_td_ineffective(
    active_backend: str, is_quantized: bool = False
) -> None:
    """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored.

    Fires when the user set the env explicitly and either (a) the active
    MoE backend is not the fused Triton kernel, or (b) the model is
    quantized (the TD path falls back to the pointer path under any
    quantization).
    """
    global _warned_moe_use_td_ineffective
    if _warned_moe_use_td_ineffective:
        return
    if envs.VLLM_TRITON_USE_TD is None:
        return
    is_triton = active_backend.upper() == "TRITON"
    if is_triton and not is_quantized:
        return
    if not is_triton:
        reason = (
            f"the active MoE backend is {active_backend!r}; pass "
            "`--moe-backend triton` to enable the tensor-descriptor path"
        )
    else:
        reason = (
            "the model uses quantized MoE weights; the TD path is "
            "currently restricted to non-quantized weights and falls "
            "back to the pointer path"
        )
    logger.warning(
        "VLLM_TRITON_USE_TD is set to %s but %s.",
        envs.VLLM_TRITON_USE_TD,
        reason,
    )
    _warned_moe_use_td_ineffective = True