@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. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out,
M<=32, model-specific shapes)
3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight)
4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
5. 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).
A ``quant_config`` that actually quantizes the gate disables every
specialized tier, leaving plain ``ReplicatedLinear`` behavior.
"""
# (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)}
FP32_MAX_TOKENS = 32
# Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses
# whose router shape stays ahead further up the range may raise it.
LL_BF16_MAX_TOKENS = 16
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,
skip_bias_add: bool = False,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
return_bias: bool = True,
):
is_hopper = current_platform.is_device_capability((9, 0))
is_blackwell = current_platform.is_device_capability_family(100)
is_gfx950 = False
if current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx950
is_gfx950 = on_gfx950()
is_rocm_fp32_shape = False
if is_gfx950:
from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501
ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES,
)
is_rocm_fp32_shape = (
input_size,
output_size,
) in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES
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,
skip_bias_add=skip_bias_add,
params_dtype=params_dtype,
quant_config=quant_config,
prefix=prefix,
return_bias=return_bias,
)
self.out_dtype = out_dtype
# A quantized gate exposes no plain ``weight``, so every specialized
# tier below is disabled and forward falls back to ReplicatedLinear.
self.is_unquantized = isinstance(self.quant_method, UnquantizedLinearMethod)
can_use_specialized_kernels &= self.is_unquantized
self.allow_specialized_router_gemm = can_use_specialized_kernels
# fp32 specialized kernel eligibility (exact dims, fp32 weight)
self.allow_fp32_router_gemm = (
self.is_unquantized
and 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
)
or (is_gfx950 and is_rocm_fp32_shape)
)
)
self.allow_bf16x3_router_gemm = (
self.is_unquantized
and not bias
and self.weight.dtype == torch.float32
and current_platform.is_cuda()
and is_blackwell
and input_size % 8 == 0
)
# Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype
# epilogue folds the fp32 cast into the GEMM, removing the standalone
# bf16->fp32 copy kernel that otherwise runs before grouped_topk. This is
# the plain cuBLAS (CUDA) / hipBLASLt (ROCm) out_dtype epilogue, so it
# applies on any CUDA-alike device (no bias, since torch.mm has no bias
# term). The specialized-kernel gate above excludes family-120 Blackwell
# (GB10 / DGX Spark), which this tier still covers. See #49921.
self._router_gemm_no_bias = self.is_unquantized and not bias
self._router_gemm_cublas_capable = (
current_platform.is_cuda() or current_platform.is_rocm()
) and self._router_gemm_no_bias
self.allow_cublas_router_gemm = (
self._router_gemm_cublas_capable
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._router_gemm_cublas_capable
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 _return(
self, output: torch.Tensor
) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
return output if not self.return_bias else (output, None)
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] <= self.LL_BF16_MAX_TOKENS
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 self._return(output)
# Tier 2: fp32 specialized kernel (model-specific shapes, 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 self._return(output)
# Tier 3: 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 self._return(output)
# Tier 4: 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 self._return(output)
# Tier 5: F.linear (ReplicatedLinear)
if (
self.out_dtype is not None
and self.is_unquantized
and x.dtype != self.weight.dtype
):
x = x.to(self.weight.dtype)
base_output = super().forward(x)
output_bias: Parameter | None = None
if isinstance(base_output, tuple):
output, output_bias = base_output
else:
output = base_output
if self.out_dtype is not None and output.dtype != self.out_dtype:
output = output.to(self.out_dtype)
return output if not self.return_bias else (output, output_bias)