vllm.model_executor.layers.fused_moe.utils ¶

def _fp8_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 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)
        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)

    return A, A_scale

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

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)

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.

    Parameters:
    - 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

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`
    """
    return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)

def is_model_fused_shared_expert_compatible(
    layers: nn.ModuleList | Iterable[nn.Module],
    moe_cls: type[nn.Module],
    moe_name: str,
) -> bool:
    """Resolve one fused-shared-expert state for a model's MoE layers."""

    def get_moe_layer(layer: nn.Module) -> nn.Module | None:
        for name in moe_name.split("."):
            layer = getattr(layer, name, None)
            if layer is None:
                return None
        return layer

    moe_layers = (
        moe_layer
        for layer in layers
        if not isinstance(layer, PPMissingLayer)
        and (moe_layer := get_moe_layer(layer)) is not None
        and isinstance(moe_layer, moe_cls)
    )

    enabled = [
        getattr(layer, "is_fused_shared_expert_enabled", False) for layer in moe_layers
    ]
    enabled_count = sum(enabled)
    disabled_count = len(enabled) - enabled_count
    if enabled_count > 0 and disabled_count > 0:
        raise NotImplementedError(
            "Fused shared experts must be enabled for all MoE layers; found "
            f"{enabled_count} enabled and {disabled_count} disabled layers. "
            "Per-layer fused shared experts is not yet supported. Please open "
            "an issue."
        )
    return enabled_count > 0 and disabled_count == 0

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

def resolve_layer_fused_shared_expert(
    quant_config: "QuantizationConfig | None",
    prefix: str,
    shared_expert_name: str = "shared_experts",
) -> 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``.

    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())
    fse_compatible, fse_reason = (
        is_shared_expert_quant_fse_compatible(
            quant_config,
            f"{prefix}.experts",
            f"{prefix}.{shared_expert_name}",
        )
        if fse_requested
        else (True, None)
    )
    is_fused_shared_expert_enabled = fse_requested and fse_compatible
    if fse_requested and not is_fused_shared_expert_enabled:
        logger.warning(
            "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but "
            "cannot be enabled: %s.",
            fse_reason,
        )
    return is_fused_shared_expert_enabled

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

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