A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 15:15 UTC
[Performance] Cublas Bf16 Gate with Fp32 Output (#35121)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Feb 27, 2026 at 02:51 UTC
38c498b8e3aaec95049f384edfc56ca12cbe1839
9 files changed
+206
-80
CMakeLists.txt
+2
-1
@@ -971,7 +971,8 @@ set(VLLM_MOE_EXT_SRC
971
if(VLLM_GPU_LANG STREQUAL "CUDA")
972
list(APPEND VLLM_MOE_EXT_SRC
973
"csrc/moe/moe_wna16.cu"
974
- "csrc/moe/grouped_topk_kernels.cu")
974
+ "csrc/moe/grouped_topk_kernels.cu"
975
+ "csrc/moe/router_gemm.cu")
976
endif()
977
978
if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/moe/moe_ops.h
+4
@@ -58,6 +58,10 @@ void shuffle_rows(const torch::Tensor& input_tensor,
58
torch::Tensor& output_tensor);
59
60
#ifndef USE_ROCM
61
+// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
62
+torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
63
+ torch::Tensor const& weight);
64
+
65
// DeepSeek V3 optimized router GEMM kernel for SM90+
66
// Computes output = mat_a @ mat_b.T where:
67
// mat_a: [num_tokens, hidden_dim] in bf16
csrc/moe/router_gemm.cu
new
+52
@@ -0,0 +1,52 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+// bf16 x bf16 -> fp32 router GEMM via cuBLAS.
5
+// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32,
6
+// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp.
7
+
8
+#include <torch/all.h>
9
+#include <ATen/cuda/CUDAContext.h>
10
+#include <cublas_v2.h>
11
+
12
+// cuBLAS column-major math for row-major PyTorch tensors:
13
+// weight[N,K]_row lda=K -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T ->
14
+// (N,K) input[M,K]_row ldb=K -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N
15
+// -> (K,M) out[M,N]_row ldc=N -> cuBLAS sees (N,M) col-major (written as
16
+// output^T)
17
+// cuBLAS: C(N,M) = weight(N,K) @ input(K,M) => C^T = output[M,N]
18
+// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output)
19
+
20
+torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
21
+ torch::Tensor const& weight) {
22
+ TORCH_CHECK(input.dtype() == torch::kBFloat16,
23
+ "router_gemm_bf16_fp32: input must be bfloat16");
24
+ TORCH_CHECK(weight.dtype() == torch::kBFloat16,
25
+ "router_gemm_bf16_fp32: weight must be bfloat16");
26
+ TORCH_CHECK(input.dim() == 2 && weight.dim() == 2,
27
+ "router_gemm_bf16_fp32: input and weight must be 2-D");
28
+ TORCH_CHECK(input.size(1) == weight.size(1),
29
+ "router_gemm_bf16_fp32: inner dimensions must match");
30
+
31
+ int64_t const M = input.size(0);
32
+ int64_t const N = weight.size(0);
33
+ int64_t const K = input.size(1);
34
+
35
+ auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32));
36
+
37
+ cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
38
+ TORCH_CUDABLAS_CHECK(
39
+ cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
40
+
41
+ float const alpha = 1.0f;
42
+ float const beta = 0.0f;
43
+
44
+ TORCH_CUDABLAS_CHECK(cublasGemmEx(
45
+ handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast<int>(N),
46
+ static_cast<int>(M), static_cast<int>(K), &alpha, weight.data_ptr(),
47
+ CUDA_R_16BF, static_cast<int>(K), input.data_ptr(), CUDA_R_16BF,
48
+ static_cast<int>(K), &beta, out.data_ptr(), CUDA_R_32F,
49
+ static_cast<int>(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
50
+
51
+ return out;
52
+}
csrc/moe/torch_bindings.cpp
+4
@@ -125,6 +125,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
125
"Tensor)");
126
m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
127
128
+ // cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
129
+ m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor");
130
+ m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32);
131
+
132
// DeepSeek V3 optimized router GEMM for SM90+
133
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
134
// conditionally compiled so impl registration is in source file
vllm/_custom_ops.py
+17
@@ -2190,6 +2190,23 @@ def moe_wna16_gemm(
2190
)
2191
2192
2193
+def router_gemm_bf16_fp32(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
2194
+ """bf16 x bf16 -> fp32 GEMM via cuBLAS. weight shape: (N, K)."""
2195
+ return torch.ops._moe_C.router_gemm_bf16_fp32(input, weight)
2196
+
2197
+
2198
+if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "router_gemm_bf16_fp32"):
2199
+
2200
+ @register_fake("_moe_C::router_gemm_bf16_fp32")
2201
+ def router_gemm_bf16_fp32_fake(
2202
+ input: torch.Tensor,
2203
+ weight: torch.Tensor,
2204
+ ) -> torch.Tensor:
2205
+ return torch.empty(
2206
+ input.shape[0], weight.shape[0], dtype=torch.float32, device=input.device
2207
+ )
2208
+
2209
+
2210
def dsv3_router_gemm(
2211
hidden_states: torch.Tensor,
2212
router_weight: torch.Tensor,
vllm/model_executor/layers/fused_moe/__init__.py
+2
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
28
from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
29
FusedMoERouter,
30
)
31
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
32
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
33
from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
34
UnquantizedFusedMoEMethod,
@@ -64,6 +65,7 @@ __all__ = [
65
"FusedMoEPermuteExpertsUnpermute",
66
"FusedMoEActivationFormat",
67
"FusedMoEPrepareAndFinalize",
68
+ "GateLinear",
69
"RoutingMethodType",
70
"SharedFusedMoE",
71
"ZeroExpertFusedMoE",
vllm/model_executor/layers/fused_moe/router/gate_linear.py
new
+117
@@ -0,0 +1,117 @@
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
+from vllm.model_executor.custom_op import PluggableLayer
7
+from vllm.model_executor.layers.linear import ReplicatedLinear
8
+from vllm.platforms import current_platform
9
+
10
+
11
+@PluggableLayer.register("gate_linear")
12
+class GateLinear(ReplicatedLinear):
13
+ """MoE gate linear layer with three-tier GEMM dispatch:
14
+
15
+ 1. DSV3 specialized kernel (SM90+, batch<=16, supported dims)
16
+ 2. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 + fp32 out_dtype)
17
+ 3. F.linear via ReplicatedLinear (ultimate fallback)
18
+
19
+ The ``out_dtype`` attribute is mutable and can be set after init
20
+ (e.g. when the required dtype depends on the expert quantization
21
+ method which is only known later).
22
+ """
23
+
24
+ # Dimensions supported by the DSV3 specialized kernel
25
+ DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
26
+ DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
27
+
28
+ def __init__(
29
+ self,
30
+ input_size: int,
31
+ output_size: int,
32
+ bias: bool = False,
33
+ out_dtype: torch.dtype | None = None,
34
+ params_dtype: torch.dtype | None = None,
35
+ force_fp32_compute: bool = False,
36
+ prefix: str = "",
37
+ ):
38
+ is_hopper_or_blackwell = current_platform.is_device_capability(
39
+ (9, 0)
40
+ ) or current_platform.is_device_capability_family(100)
41
+ can_use_specialized_kernels = (
42
+ current_platform.is_cuda() and is_hopper_or_blackwell and not bias
43
+ )
44
+
45
+ # If fp32 compute is required and no specialized kernel is available,
46
+ # store weights in fp32 so Tier 3 computes in fp32 natively.
47
+ if force_fp32_compute and not can_use_specialized_kernels:
48
+ params_dtype = torch.float32
49
+
50
+ super().__init__(
51
+ input_size,
52
+ output_size,
53
+ bias=bias,
54
+ params_dtype=params_dtype,
55
+ quant_config=None,
56
+ prefix=prefix,
57
+ )
58
+ self.out_dtype = out_dtype
59
+
60
+ # DSV3 specialized kernel eligibility (SM90+, exact dims)
61
+ self.allow_specialized_router_gemm = can_use_specialized_kernels
62
+ self.allow_dsv3_router_gemm = (
63
+ self.allow_specialized_router_gemm
64
+ and output_size in self.DSV3_SUPPORTED_NUM_EXPERTS
65
+ and input_size in self.DSV3_SUPPORTED_HIDDEN_SIZES
66
+ )
67
+
68
+ # cuBLAS bf16→fp32 eligibility
69
+ self.allow_cublas_router_gemm = (
70
+ self.allow_specialized_router_gemm
71
+ and self.weight.dtype == torch.bfloat16
72
+ and self.out_dtype == torch.float32
73
+ )
74
+
75
+ def set_out_dtype(self, out_dtype: torch.dtype) -> None:
76
+ """Set output dtype for the router logits after init.
77
+
78
+ Useful when the required dtype depends on the expert quantization
79
+ method which is only known after the gate is constructed.
80
+ """
81
+ if self.out_dtype is not None:
82
+ raise ValueError("out_dtype has already been set")
83
+ self.out_dtype = out_dtype
84
+
85
+ if (
86
+ not self.allow_cublas_router_gemm
87
+ and self.allow_specialized_router_gemm
88
+ and out_dtype == torch.float32
89
+ ):
90
+ self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16
91
+
92
+ def forward(
93
+ self, x: torch.Tensor
94
+ ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
95
+ import vllm._custom_ops as ops
96
+
97
+ # Tier 1: DSV3 specialized kernel
98
+ if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
99
+ output = ops.dsv3_router_gemm(
100
+ hidden_states=x,
101
+ router_weight=self.weight,
102
+ output_dtype=self.out_dtype,
103
+ )
104
+ return output, None
105
+
106
+ # Tier 2: cuBLAS bf16→fp32
107
+ if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
108
+ output = ops.router_gemm_bf16_fp32(x, self.weight)
109
+ return output, None
110
+
111
+ # Tier 3: F.linear (ReplicatedLinear)
112
+ if self.out_dtype is not None and x.dtype != self.weight.dtype:
113
+ x = x.to(self.weight.dtype)
114
+ output, output_bias = super().forward(x)
115
+ if self.out_dtype is not None and output.dtype != self.out_dtype:
116
+ output = output.to(self.out_dtype)
117
+ return output, output_bias
vllm/model_executor/models/deepseek_v2.py
+2
-70
@@ -47,7 +47,7 @@ from vllm.logger import init_logger
47
from vllm.model_executor.layers.activation import SiluAndMul
48
from vllm.model_executor.layers.attention import Attention
49
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
50
-from vllm.model_executor.layers.fused_moe import SharedFusedMoE
50
+from vllm.model_executor.layers.fused_moe import GateLinear, SharedFusedMoE
51
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
52
from vllm.model_executor.layers.linear import (
53
ColumnParallelLinear,
@@ -221,73 +221,6 @@ class DeepseekV2MLP(nn.Module):
221
return x
222
223
224
-class DeepSeekV2Gate(ReplicatedLinear):
225
- def __init__(
226
- self,
227
- hidden_size: int,
228
- n_experts: int,
229
- quant_config: QuantizationConfig | None = None,
230
- prefix: str = "",
231
- ):
232
- assert quant_config is None
233
- super().__init__(
234
- hidden_size,
235
- n_experts,
236
- bias=False,
237
- quant_config=quant_config,
238
- prefix=f"{prefix}.gate",
239
- )
240
-
241
- # Unquantized only, will be called "weight".
242
- assert hasattr(self, "weight")
243
- is_hopper_or_blackwell = current_platform.is_device_capability(
244
- (9, 0)
245
- ) or current_platform.is_device_capability_family(100)
246
- SUPPORTED_NUM_EXPERTS = [256, 384]
247
- SUPPORTED_HIDDEN_SIZES = [7168]
248
-
249
- self.allow_dsv3_router_gemm = (
250
- current_platform.is_cuda()
251
- and is_hopper_or_blackwell
252
- and n_experts in SUPPORTED_NUM_EXPERTS
253
- and hidden_size in SUPPORTED_HIDDEN_SIZES
254
- )
255
-
256
- self._out_dtype: torch.dtype | None = None
257
-
258
- def set_out_dtype(self, out_dtype: torch.dtype) -> None:
259
- """
260
- Set out dtype for the router logits. This is needed after
261
- __init__, b/c we need to check if the trtllm kernel is
262
- selected before we decide between bf16 and fp32.
263
- """
264
-
265
- if self._out_dtype is not None:
266
- raise ValueError("out_dtype has already been set")
267
- else:
268
- self._out_dtype = out_dtype
269
-
270
- @property
271
- def out_dtype(self) -> torch.dtype:
272
- if self._out_dtype is None:
273
- raise ValueError("out_dtype has not been set yet")
274
- return self._out_dtype
275
-
276
- def forward(
277
- self,
278
- x: torch.Tensor,
279
- ) -> tuple[torch.Tensor, None]:
280
- """
281
- Use specialized GEMM for low batch size for DSV3 and KIMI.
282
- """
283
- if self.allow_dsv3_router_gemm and x.shape[0] <= 16:
284
- return ops.dsv3_router_gemm(
285
- hidden_states=x, router_weight=self.weight, output_dtype=self.out_dtype
286
- ), None
287
- else:
288
- return super().forward(x)
289
-
290
-
224
class DeepseekV2MoE(nn.Module):
225
def __init__(
226
self,
@@ -316,10 +249,9 @@ class DeepseekV2MoE(nn.Module):
249
"Only silu is supported for now."
250
)
251
319
- self.gate = DeepSeekV2Gate(
252
+ self.gate = GateLinear(
253
config.hidden_size,
254
config.n_routed_experts,
322
- quant_config=None,
255
prefix=f"{prefix}.gate",
256
)
257
if getattr(config, "topk_method", None) == "noaux_tc":
vllm/model_executor/models/nemotron_h.py
+6
-9
@@ -34,7 +34,7 @@ from vllm.distributed.parallel_state import get_pp_group
34
from vllm.model_executor.layers.activation import ReLUSquaredActivation
35
from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.fused_moe import (
37
- FusedMoE,
37
+ GateLinear,
38
SharedFusedMoE,
39
activation_without_mul,
40
)
@@ -148,13 +148,11 @@ class NemotronHMoE(nn.Module):
148
149
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
150
151
- router_logits_dtype = torch.float32
152
- self.gate = ReplicatedLinear(
151
+ self.gate = GateLinear(
152
config.hidden_size,
153
config.n_routed_experts,
155
- bias=False,
156
- params_dtype=router_logits_dtype,
157
- quant_config=None,
154
+ out_dtype=torch.float32,
155
+ force_fp32_compute=True,
156
prefix=f"{prefix}.gate",
157
)
158
@@ -232,7 +230,6 @@ class NemotronHMoE(nn.Module):
230
enable_eplb=self.enable_eplb,
231
num_redundant_experts=self.n_redundant_experts,
232
is_sequence_parallel=self.is_sequence_parallel,
235
- router_logits_dtype=router_logits_dtype,
233
routed_input_transform=self.fc1_latent_proj,
234
)
235
@@ -244,7 +241,7 @@ class NemotronHMoE(nn.Module):
241
hidden_states = sequence_parallel_chunk(hidden_states)
242
243
# router_logits: (num_tokens, n_experts)
247
- router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
244
+ router_logits, _ = self.gate(hidden_states)
245
246
# SharedFusedMoE handles:
247
# - shared experts (with original hidden_states)
@@ -675,7 +672,7 @@ class NemotronHModel(nn.Module):
672
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
673
if self.has_moe:
674
# (param_name, weight_name, expert_id, shard_id)
678
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
675
+ expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
676
# - FusedMoe.w1 (aka gate_proj) should be up_proj since that's
677
# what the activation is applied to
678
# - FusedMoe.w3 (aka up_proj) should be ignored since we're