`vllm.model_executor.layers.fused_moe.router.gate_linear` ¶

Classes:

GateLinear –

MoE gate linear layer with multi-tier GEMM dispatch:

Functions:

fp32_router_gemm_dispatch_impl –

Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,

`GateLinear` ¶

Bases: ReplicatedLinear

MoE gate linear layer with multi-tier GEMM dispatch:

cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, K divisible by 8)
DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
F.linear via ReplicatedLinear (ultimate fallback)

The out_dtype attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization method which is only known later).

Methods:

set_out_dtype –

Set output dtype for the router logits after init.

Source code in vllm/model_executor/layers/fused_moe/router/gate_linear.py

@PluggableLayer.register("gate_linear")
class GateLinear(ReplicatedLinear):
    """MoE gate linear layer with multi-tier GEMM dispatch:

    1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
       K divisible by 8)
    2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256)
    3. fp32 specialized kernel  (SM90+, bf16/fp32 in, fp32 out, M<=32,
       (H, E) in {(3072, 256), (6144, 128), (6144, 256)})
    4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
    5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
    6. F.linear via ReplicatedLinear (ultimate fallback)

    The ``out_dtype`` attribute is mutable and can be set after init
    (e.g. when the required dtype depends on the expert quantization
    method which is only known later).
    """

    # Dimensions supported by the DSV3 specialized kernel.
    # Valid (hidden_size, num_experts) combinations:
    #   (7168, 256) -> DeepSeek-V3,  (7168, 384) -> Kimi-K2,
    #   (6144, 256) -> GLM-5
    DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
    DSV3_SUPPORTED_HIDDEN_SIZES = [7168, 6144]
    # num_experts=384 is only instantiated for hidden_size=7168.
    DSV3_UNSUPPORTED_SHAPES = {(6144, 384)}

    # (hidden_size, num_experts) pairs with an instantiated fp32 kernel:
    #   (3072, 256) -> MiniMax-M2/M2.5,  (6144, 128) -> MiniMax-M3
    FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (6144, 256)}
    FP32_MAX_TOKENS = 32

    def __init__(
        self,
        input_size: int,
        output_size: int,
        bias: bool = False,
        out_dtype: torch.dtype | None = None,
        params_dtype: torch.dtype | None = None,
        force_fp32_compute: bool = False,
        prefix: str = "",
    ):
        is_hopper = current_platform.is_device_capability((9, 0))
        is_blackwell = current_platform.is_device_capability_family(100)
        can_use_specialized_kernels = (
            current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias
        )

        # If fp32 compute is required and no specialized kernel is available,
        # store weights in fp32 so the fallback linear path computes in fp32.
        if force_fp32_compute and not can_use_specialized_kernels:
            params_dtype = torch.float32

        super().__init__(
            input_size,
            output_size,
            bias=bias,
            params_dtype=params_dtype,
            quant_config=None,
            prefix=prefix,
        )
        self.out_dtype = out_dtype

        # DSV3 specialized kernel eligibility (SM90+, exact dims)
        self.allow_specialized_router_gemm = can_use_specialized_kernels
        self.allow_dsv3_router_gemm = (
            self.allow_specialized_router_gemm
            and self.weight.dtype == torch.bfloat16
            and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
            and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
            and (input_size, output_size) not in self.DSV3_UNSUPPORTED_SHAPES
        )
        # See https://github.com/vllm-project/vllm/pull/44217
        # for more details.
        self._dsv3_max_batch = 16 if is_hopper else 8

        # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
        vllm_config = get_current_vllm_config_or_none()
        enable_bf16x3_router_gemm = (
            vllm_config is not None
            and vllm_config.kernel_config.enable_bf16x3_router_gemm
        )
        self.allow_fp32_router_gemm = (
            not bias
            and self.weight.dtype == torch.float32
            and current_platform.is_cuda()
            and (is_hopper or is_blackwell)
            and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
        )
        self.allow_bf16x3_router_gemm = (
            not bias
            and self.weight.dtype == torch.float32
            and current_platform.is_cuda()
            and is_blackwell
            and input_size % 8 == 0
            and enable_bf16x3_router_gemm
        )
        if self.allow_bf16x3_router_gemm:
            logger.info_once("Enabled experimental SM100 BF16x3 router GEMM.")

        # cuBLAS bf16→fp32 eligibility
        self.allow_cublas_router_gemm = (
            self.allow_specialized_router_gemm
            and self.weight.dtype == torch.bfloat16
            and self.out_dtype == torch.float32
        )

        # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc:
        # 1. PDL support. Both dot-product and split-K kernels.
        # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction.
        self.allow_ll_bf16_gemm = False
        if can_use_specialized_kernels:
            from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
                is_available,
            )

            self.allow_ll_bf16_gemm = (
                self.weight.dtype == torch.bfloat16
                and self.out_dtype == torch.float32
                and is_available()
            )

    def set_out_dtype(self, out_dtype: torch.dtype) -> None:
        """Set output dtype for the router logits after init.

        Useful when the required dtype depends on the expert quantization
        method which is only known after the gate is constructed.
        """
        if self.out_dtype is not None:
            raise ValueError("out_dtype has already been set")
        self.out_dtype = out_dtype

        if (
            not self.allow_cublas_router_gemm
            and self.allow_specialized_router_gemm
            and out_dtype == torch.float32
        ):
            self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16

        # out_dtype may start as None -> recompute eligibility here
        if self.allow_specialized_router_gemm:
            from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
                is_available,
            )

            self.allow_ll_bf16_gemm = (
                self.weight.dtype == torch.bfloat16
                and out_dtype == torch.float32
                and is_available()
            )

    def forward(
        self, x: torch.Tensor
    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
        # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims)
        if self.allow_ll_bf16_gemm and x.shape[0] <= 16 and x.dtype == torch.bfloat16:
            from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
                ll_bf16_gemm,
            )

            output = ll_bf16_gemm(x, self.weight)
            return output, None

        # Tier 2: DSV3 specialized kernel (fallback for when cuteDSL unavailable)
        if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch:
            output = ops.dsv3_router_gemm(
                hidden_states=x,
                router_weight=self.weight,
                output_dtype=self.out_dtype,
            )
            return output, None

        # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32)
        # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
        # capture does not freeze the runtime num_tokens branch.
        if self.allow_fp32_router_gemm and x.dtype in (
            torch.float32,
            torch.bfloat16,
        ):
            output = torch.ops.vllm.fp32_router_gemm_dispatch(
                x, self.weight, self.allow_bf16x3_router_gemm
            )
            return output, None

        # Tier 4: experimental bf16x3 CuteDSL kernel for fp32 router weights
        if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
            from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import (  # noqa: E501
                bf16x3_router_gemm,
            )

            output = bf16x3_router_gemm(x, self.weight)
            return output, None

        # Tier 5: cuBLAS bf16→fp32
        if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
            output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
            return output, None

        # Tier 6: F.linear (ReplicatedLinear)
        if self.out_dtype is not None and x.dtype != self.weight.dtype:
            x = x.to(self.weight.dtype)
        output, output_bias = super().forward(x)
        if self.out_dtype is not None and output.dtype != self.out_dtype:
            output = output.to(self.out_dtype)
        return output, output_bias

`set_out_dtype(out_dtype)` ¶

Set output dtype for the router logits after init.

Useful when the required dtype depends on the expert quantization method which is only known after the gate is constructed.

Source code in vllm/model_executor/layers/fused_moe/router/gate_linear.py

def set_out_dtype(self, out_dtype: torch.dtype) -> None:
    """Set output dtype for the router logits after init.

    Useful when the required dtype depends on the expert quantization
    method which is only known after the gate is constructed.
    """
    if self.out_dtype is not None:
        raise ValueError("out_dtype has already been set")
    self.out_dtype = out_dtype

    if (
        not self.allow_cublas_router_gemm
        and self.allow_specialized_router_gemm
        and out_dtype == torch.float32
    ):
        self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16

    # out_dtype may start as None -> recompute eligibility here
    if self.allow_specialized_router_gemm:
        from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import (
            is_available,
        )

        self.allow_ll_bf16_gemm = (
            self.weight.dtype == torch.bfloat16
            and out_dtype == torch.float32
            and is_available()
        )

`fp32_router_gemm_dispatch_impl(x, weight, allow_bf16x3_router_gemm)` ¶

Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, otherwise optionally run the experimental BF16x3 kernel for medium/large SM100 router batches, then fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens.

Source code in vllm/model_executor/layers/fused_moe/router/gate_linear.py

def fp32_router_gemm_dispatch_impl(
    x: torch.Tensor,
    weight: torch.Tensor,
    allow_bf16x3_router_gemm: bool,
) -> torch.Tensor:
    """
    Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS,
    otherwise optionally run the experimental BF16x3 kernel for medium/large
    SM100 router batches, then fall back to F.linear.
    This must be wrapped in a custom op because our torch.compile integration
    does not support runtime dispatching on num_tokens.
    """
    if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
        return ops.fp32_router_gemm(x, weight)

    if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
        from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import (  # noqa: E501
            bf16x3_router_gemm,
        )

        return bf16x3_router_gemm(x, weight)

    return torch.nn.functional.linear(x.float(), weight)

vllm.model_executor.layers.fused_moe.router.gate_linear ¶

GateLinear ¶

set_out_dtype(out_dtype) ¶

fp32_router_gemm_dispatch_impl(x, weight, allow_bf16x3_router_gemm) ¶

`vllm.model_executor.layers.fused_moe.router.gate_linear` ¶

`GateLinear` ¶

`set_out_dtype(out_dtype)` ¶

`fp32_router_gemm_dispatch_impl(x, weight, allow_bf16x3_router_gemm)` ¶