| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | from torch.nn.parameter import Parameter |
| 5 | |
| 6 | import vllm._custom_ops as ops |
| 7 | from vllm.model_executor.custom_op import PluggableLayer |
| 8 | from vllm.model_executor.layers.linear import ( |
| 9 | ReplicatedLinear, |
| 10 | UnquantizedLinearMethod, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.torch_utils import direct_register_custom_op |
| 15 | |
| 16 | |
| 17 | @PluggableLayer.register("gate_linear") |
| 18 | class GateLinear(ReplicatedLinear): |
| 19 | """MoE gate linear layer with multi-tier GEMM dispatch: |
| 20 | |
| 21 | 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, |
| 22 | K divisible by 8) |
| 23 | 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, |
| 24 | M<=32, model-specific shapes) |
| 25 | 3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) |
| 26 | 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) |
| 27 | 5. F.linear via ReplicatedLinear (ultimate fallback) |
| 28 | |
| 29 | The ``out_dtype`` attribute is mutable and can be set after init |
| 30 | (e.g. when the required dtype depends on the expert quantization |
| 31 | method which is only known later). |
| 32 | |
| 33 | A ``quant_config`` that actually quantizes the gate disables every |
| 34 | specialized tier, leaving plain ``ReplicatedLinear`` behavior. |
| 35 | """ |
| 36 | |
| 37 | # (hidden_size, num_experts) pairs with an instantiated fp32 kernel: |
| 38 | # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 |
| 39 | FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} |
| 40 | FP32_MAX_TOKENS = 32 |
| 41 | # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses |
| 42 | # whose router shape stays ahead further up the range may raise it. |
| 43 | LL_BF16_MAX_TOKENS = 16 |
| 44 | |
| 45 | def __init__( |
| 46 | self, |
| 47 | input_size: int, |
| 48 | output_size: int, |
| 49 | bias: bool = False, |
| 50 | out_dtype: torch.dtype | None = None, |
| 51 | params_dtype: torch.dtype | None = None, |
| 52 | force_fp32_compute: bool = False, |
| 53 | skip_bias_add: bool = False, |
| 54 | quant_config: QuantizationConfig | None = None, |
| 55 | prefix: str = "", |
| 56 | return_bias: bool = True, |
| 57 | ): |
| 58 | is_hopper = current_platform.is_device_capability((9, 0)) |
| 59 | is_blackwell = current_platform.is_device_capability_family(100) |
| 60 | is_gfx950 = False |
| 61 | if current_platform.is_rocm(): |
| 62 | from vllm.platforms.rocm import on_gfx950 |
| 63 | |
| 64 | is_gfx950 = on_gfx950() |
| 65 | is_rocm_fp32_shape = False |
| 66 | if is_gfx950: |
| 67 | from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 |
| 68 | ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES, |
| 69 | ) |
| 70 | |
| 71 | is_rocm_fp32_shape = ( |
| 72 | input_size, |
| 73 | output_size, |
| 74 | ) in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES |
| 75 | can_use_specialized_kernels = ( |
| 76 | current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias |
| 77 | ) |
| 78 | |
| 79 | # If fp32 compute is required and no specialized kernel is available, |
| 80 | # store weights in fp32 so the fallback linear path computes in fp32. |
| 81 | if force_fp32_compute and not can_use_specialized_kernels: |
| 82 | params_dtype = torch.float32 |
| 83 | |
| 84 | super().__init__( |
| 85 | input_size, |
| 86 | output_size, |
| 87 | bias=bias, |
| 88 | skip_bias_add=skip_bias_add, |
| 89 | params_dtype=params_dtype, |
| 90 | quant_config=quant_config, |
| 91 | prefix=prefix, |
| 92 | return_bias=return_bias, |
| 93 | ) |
| 94 | self.out_dtype = out_dtype |
| 95 | |
| 96 | # A quantized gate exposes no plain ``weight``, so every specialized |
| 97 | # tier below is disabled and forward falls back to ReplicatedLinear. |
| 98 | self.is_unquantized = isinstance(self.quant_method, UnquantizedLinearMethod) |
| 99 | can_use_specialized_kernels &= self.is_unquantized |
| 100 | |
| 101 | self.allow_specialized_router_gemm = can_use_specialized_kernels |
| 102 | |
| 103 | # fp32 specialized kernel eligibility (exact dims, fp32 weight) |
| 104 | self.allow_fp32_router_gemm = ( |
| 105 | self.is_unquantized |
| 106 | and not bias |
| 107 | and self.weight.dtype == torch.float32 |
| 108 | and ( |
| 109 | ( |
| 110 | current_platform.is_cuda() |
| 111 | and (is_hopper or is_blackwell) |
| 112 | and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES |
| 113 | ) |
| 114 | or (is_gfx950 and is_rocm_fp32_shape) |
| 115 | ) |
| 116 | ) |
| 117 | self.allow_bf16x3_router_gemm = ( |
| 118 | self.is_unquantized |
| 119 | and not bias |
| 120 | and self.weight.dtype == torch.float32 |
| 121 | and current_platform.is_cuda() |
| 122 | and is_blackwell |
| 123 | and input_size % 8 == 0 |
| 124 | ) |
| 125 | |
| 126 | # Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype |
| 127 | # epilogue folds the fp32 cast into the GEMM, removing the standalone |
| 128 | # bf16->fp32 copy kernel that otherwise runs before grouped_topk. This is |
| 129 | # the plain cuBLAS (CUDA) / hipBLASLt (ROCm) out_dtype epilogue, so it |
| 130 | # applies on any CUDA-alike device (no bias, since torch.mm has no bias |
| 131 | # term). The specialized-kernel gate above excludes family-120 Blackwell |
| 132 | # (GB10 / DGX Spark), which this tier still covers. See #49921. |
| 133 | self._router_gemm_no_bias = self.is_unquantized and not bias |
| 134 | self._router_gemm_cublas_capable = ( |
| 135 | current_platform.is_cuda() or current_platform.is_rocm() |
| 136 | ) and self._router_gemm_no_bias |
| 137 | self.allow_cublas_router_gemm = ( |
| 138 | self._router_gemm_cublas_capable |
| 139 | and self.weight.dtype == torch.bfloat16 |
| 140 | and self.out_dtype == torch.float32 |
| 141 | ) |
| 142 | |
| 143 | # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc: |
| 144 | # 1. PDL support. Both dot-product and split-K kernels. |
| 145 | # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction. |
| 146 | self.allow_ll_bf16_gemm = False |
| 147 | if can_use_specialized_kernels: |
| 148 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( |
| 149 | is_available, |
| 150 | ) |
| 151 | |
| 152 | self.allow_ll_bf16_gemm = ( |
| 153 | self.weight.dtype == torch.bfloat16 |
| 154 | and self.out_dtype == torch.float32 |
| 155 | and is_available() |
| 156 | ) |
| 157 | |
| 158 | def set_out_dtype(self, out_dtype: torch.dtype) -> None: |
| 159 | """Set output dtype for the router logits after init. |
| 160 | |
| 161 | Useful when the required dtype depends on the expert quantization |
| 162 | method which is only known after the gate is constructed. |
| 163 | """ |
| 164 | if self.out_dtype is not None: |
| 165 | raise ValueError("out_dtype has already been set") |
| 166 | self.out_dtype = out_dtype |
| 167 | |
| 168 | if ( |
| 169 | not self.allow_cublas_router_gemm |
| 170 | and self._router_gemm_cublas_capable |
| 171 | and out_dtype == torch.float32 |
| 172 | ): |
| 173 | self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16 |
| 174 | |
| 175 | # out_dtype may start as None -> recompute eligibility here |
| 176 | if self.allow_specialized_router_gemm: |
| 177 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( |
| 178 | is_available, |
| 179 | ) |
| 180 | |
| 181 | self.allow_ll_bf16_gemm = ( |
| 182 | self.weight.dtype == torch.bfloat16 |
| 183 | and out_dtype == torch.float32 |
| 184 | and is_available() |
| 185 | ) |
| 186 | |
| 187 | def _return( |
| 188 | self, output: torch.Tensor |
| 189 | ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: |
| 190 | return output if not self.return_bias else (output, None) |
| 191 | |
| 192 | def forward( |
| 193 | self, x: torch.Tensor |
| 194 | ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: |
| 195 | # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims) |
| 196 | if ( |
| 197 | self.allow_ll_bf16_gemm |
| 198 | and x.shape[0] <= self.LL_BF16_MAX_TOKENS |
| 199 | and x.dtype == torch.bfloat16 |
| 200 | ): |
| 201 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( |
| 202 | ll_bf16_gemm, |
| 203 | ) |
| 204 | |
| 205 | output = ll_bf16_gemm(x, self.weight) |
| 206 | return self._return(output) |
| 207 | |
| 208 | # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32) |
| 209 | # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph |
| 210 | # capture does not freeze the runtime num_tokens branch. |
| 211 | if self.allow_fp32_router_gemm and x.dtype in ( |
| 212 | torch.float32, |
| 213 | torch.bfloat16, |
| 214 | ): |
| 215 | output = torch.ops.vllm.fp32_router_gemm_dispatch( |
| 216 | x, self.weight, self.allow_bf16x3_router_gemm |
| 217 | ) |
| 218 | return self._return(output) |
| 219 | |
| 220 | # Tier 3: bf16x3 CuteDSL kernel for fp32 router weights |
| 221 | if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: |
| 222 | from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 |
| 223 | bf16x3_router_gemm, |
| 224 | ) |
| 225 | |
| 226 | output = bf16x3_router_gemm(x, self.weight) |
| 227 | return self._return(output) |
| 228 | |
| 229 | # Tier 4: cuBLAS bf16→fp32 |
| 230 | if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: |
| 231 | output = torch.mm(x, self.weight.T, out_dtype=torch.float32) |
| 232 | return self._return(output) |
| 233 | |
| 234 | # Tier 5: F.linear (ReplicatedLinear) |
| 235 | if ( |
| 236 | self.out_dtype is not None |
| 237 | and self.is_unquantized |
| 238 | and x.dtype != self.weight.dtype |
| 239 | ): |
| 240 | x = x.to(self.weight.dtype) |
| 241 | base_output = super().forward(x) |
| 242 | output_bias: Parameter | None = None |
| 243 | if isinstance(base_output, tuple): |
| 244 | output, output_bias = base_output |
| 245 | else: |
| 246 | output = base_output |
| 247 | if self.out_dtype is not None and output.dtype != self.out_dtype: |
| 248 | output = output.to(self.out_dtype) |
| 249 | return output if not self.return_bias else (output, output_bias) |
| 250 | |
| 251 | |
| 252 | _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS |
| 253 | |
| 254 | |
| 255 | def fp32_router_gemm_dispatch_impl( |
| 256 | x: torch.Tensor, |
| 257 | weight: torch.Tensor, |
| 258 | allow_bf16x3_router_gemm: bool, |
| 259 | ) -> torch.Tensor: |
| 260 | """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, |
| 261 | otherwise optionally run the experimental BF16x3 kernel for medium/large |
| 262 | SM100 router batches, then fall back to F.linear. |
| 263 | This must be wrapped in a custom op because our torch.compile integration |
| 264 | does not support runtime dispatching on num_tokens. |
| 265 | """ |
| 266 | if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: |
| 267 | if current_platform.is_rocm(): |
| 268 | from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 |
| 269 | can_use_rocm_fp32_router_gemm, |
| 270 | rocm_fp32_router_gemm, |
| 271 | ) |
| 272 | |
| 273 | x = x.contiguous() |
| 274 | if can_use_rocm_fp32_router_gemm(x, weight): |
| 275 | return rocm_fp32_router_gemm(x, weight) |
| 276 | return torch.nn.functional.linear(x.float(), weight) |
| 277 | return ops.fp32_router_gemm(x, weight) |
| 278 | |
| 279 | if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: |
| 280 | from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 |
| 281 | _bf16x3_router_gemm, |
| 282 | ) |
| 283 | |
| 284 | return _bf16x3_router_gemm(x, weight) |
| 285 | |
| 286 | return torch.nn.functional.linear(x.float(), weight) |
| 287 | |
| 288 | |
| 289 | def fp32_router_gemm_dispatch_fake( |
| 290 | x: torch.Tensor, |
| 291 | weight: torch.Tensor, |
| 292 | allow_bf16x3_router_gemm: bool, |
| 293 | ) -> torch.Tensor: |
| 294 | return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32) |
| 295 | |
| 296 | |
| 297 | direct_register_custom_op( |
| 298 | op_name="fp32_router_gemm_dispatch", |
| 299 | op_func=fp32_router_gemm_dispatch_impl, |
| 300 | fake_impl=fp32_router_gemm_dispatch_fake, |
| 301 | ) |