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 4, 11:15 UTC
[DeepSeek] Use torch.mm for bf16xbf16->fp32 gemm (#41300)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Apr 30, 2026 at 16:28 UTC
9c61864bf8a911a8369f35d79d538c7f11cf3dc2
8 files changed
+9
-91
CMakeLists.txt
-1
@@ -1053,7 +1053,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1053
list(APPEND VLLM_MOE_EXT_SRC
1054
"csrc/moe/moe_wna16.cu"
1055
"csrc/moe/grouped_topk_kernels.cu"
1056
- "csrc/moe/router_gemm.cu"
1056
"csrc/moe/topk_softplus_sqrt_kernels.cu")
1057
endif()
1058
csrc/moe/moe_ops.h
-4
@@ -67,10 +67,6 @@ void shuffle_rows(const torch::Tensor& input_tensor,
67
torch::Tensor& output_tensor);
68
69
#ifndef USE_ROCM
70
-// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
71
-torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
72
- torch::Tensor const& weight);
73
-
70
// DeepSeek V3 optimized router GEMM kernel for SM90+
71
// Computes output = mat_a @ mat_b.T where:
72
// mat_a: [num_tokens, hidden_dim] in bf16
csrc/moe/router_gemm.cu
deleted
-52
@@ -1,52 +0,0 @@
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
@@ -133,10 +133,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
133
"Tensor)");
134
m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
135
136
- // cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
137
- m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor");
138
- m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32);
139
-
136
// DeepSeek V3 optimized router GEMM for SM90+
137
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
138
// conditionally compiled so impl registration is in source file
vllm/_custom_ops.py
-17
@@ -2449,23 +2449,6 @@ def moe_wna16_gemm(
2449
)
2450
2451
2452
-def router_gemm_bf16_fp32(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
2453
- """bf16 x bf16 -> fp32 GEMM via cuBLAS. weight shape: (N, K)."""
2454
- return torch.ops._moe_C.router_gemm_bf16_fp32(input, weight)
2455
-
2456
-
2457
-if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "router_gemm_bf16_fp32"):
2458
-
2459
- @register_fake("_moe_C::router_gemm_bf16_fp32")
2460
- def router_gemm_bf16_fp32_fake(
2461
- input: torch.Tensor,
2462
- weight: torch.Tensor,
2463
- ) -> torch.Tensor:
2464
- return torch.empty(
2465
- input.shape[0], weight.shape[0], dtype=torch.float32, device=input.device
2466
- )
2467
-
2468
-
2452
def dsv3_router_gemm(
2453
hidden_states: torch.Tensor,
2454
router_weight: torch.Tensor,
vllm/model_executor/layers/deepseek_v4_attention.py
+8
-5
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.linear import (
17
ReplicatedLinear,
18
)
19
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
20
-from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
20
from vllm.utils.deep_gemm import fp8_einsum
21
from vllm.utils.torch_utils import direct_register_custom_op
22
from vllm.v1.attention.ops.deepseek_v4_ops import (
@@ -349,8 +348,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
348
compressor = self.compressor
349
350
def compressor_kv_score() -> torch.Tensor:
352
- return cublas_gemm_bf16_bf16_fp32(
353
- hidden_states, compressor.fused_wkv_wgate.weight
351
+ return torch.mm(
352
+ hidden_states,
353
+ compressor.fused_wkv_wgate.weight.T,
354
+ out_dtype=torch.float32,
355
)
356
357
aux_fns[0] = compressor_kv_score
@@ -364,8 +365,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
365
return weights
366
367
def indexer_compressor_kv_score() -> torch.Tensor:
367
- return cublas_gemm_bf16_bf16_fp32(
368
- hidden_states, indexer.compressor.fused_wkv_wgate.weight
368
+ return torch.mm(
369
+ hidden_states,
370
+ indexer.compressor.fused_wkv_wgate.weight.T,
371
+ out_dtype=torch.float32,
372
)
373
374
aux_fns[1] = indexer_weights_proj
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+1
-1
@@ -105,7 +105,7 @@ class GateLinear(ReplicatedLinear):
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)
108
+ output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
109
return output, None
110
111
# Tier 3: F.linear (ReplicatedLinear)
vllm/model_executor/layers/utils.py
-7
@@ -304,13 +304,6 @@ def cpu_unquantized_gemm(
304
return layer.cpu_linear(x, weight, bias)
305
306
307
-def cublas_gemm_bf16_bf16_fp32(
308
- x: torch.Tensor,
309
- weight: torch.Tensor,
310
-):
311
- return ops.router_gemm_bf16_fp32(x, weight)
312
-
313
-
307
def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
308
if current_platform.is_rocm():
309
return rocm_unquantized_gemm