iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[DeepSeek] Use torch.mm for bf16xbf16->fp32 gemm (#41300)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This commit is contained in:
Woosuk Kwon
2026-04-30 16:28:57 -07:00
committed by GitHub
parent 71725f6730
commit 9c61864bf8
8 changed files with 9 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
-1
View File
@@ -1053,7 +1053,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_MOE_EXT_SRC
"csrc/moe/moe_wna16.cu"
"csrc/moe/grouped_topk_kernels.cu"
"csrc/moe/router_gemm.cu"
"csrc/moe/topk_softplus_sqrt_kernels.cu")
endif()
csrc/moe/moe_ops.h
-4
View File
@@ -67,10 +67,6 @@ void shuffle_rows(const torch::Tensor& input_tensor,
torch::Tensor& output_tensor);
#ifndef USE_ROCM
// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
torch::Tensor const& weight);
// DeepSeek V3 optimized router GEMM kernel for SM90+
// Computes output = mat_a @ mat_b.T where:
// mat_a: [num_tokens, hidden_dim] in bf16
csrc/moe/router_gemm.cu
-52
View File
@@ -1,52 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// bf16 x bf16 -> fp32 router GEMM via cuBLAS.
// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32,
// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp.
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
// cuBLAS column-major math for row-major PyTorch tensors:
// weight[N,K]_row lda=K -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T ->
// (N,K) input[M,K]_row ldb=K -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N
// -> (K,M) out[M,N]_row ldc=N -> cuBLAS sees (N,M) col-major (written as
// output^T)
// cuBLAS: C(N,M) = weight(N,K) @ input(K,M) => C^T = output[M,N]
// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output)
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
torch::Tensor const& weight) {
TORCH_CHECK(input.dtype() == torch::kBFloat16,
"router_gemm_bf16_fp32: input must be bfloat16");
TORCH_CHECK(weight.dtype() == torch::kBFloat16,
"router_gemm_bf16_fp32: weight must be bfloat16");
TORCH_CHECK(input.dim() == 2 && weight.dim() == 2,
"router_gemm_bf16_fp32: input and weight must be 2-D");
TORCH_CHECK(input.size(1) == weight.size(1),
"router_gemm_bf16_fp32: inner dimensions must match");
int64_t const M = input.size(0);
int64_t const N = weight.size(0);
int64_t const K = input.size(1);
auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32));
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
TORCH_CUDABLAS_CHECK(
cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
float const alpha = 1.0f;
float const beta = 0.0f;
TORCH_CUDABLAS_CHECK(cublasGemmEx(
handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast<int>(N),
static_cast<int>(M), static_cast<int>(K), &alpha, weight.data_ptr(),
CUDA_R_16BF, static_cast<int>(K), input.data_ptr(), CUDA_R_16BF,
static_cast<int>(K), &beta, out.data_ptr(), CUDA_R_32F,
static_cast<int>(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
return out;
}
csrc/moe/torch_bindings.cpp
-4
View File
@@ -133,10 +133,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
"Tensor)");
m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor");
m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32);
// DeepSeek V3 optimized router GEMM for SM90+
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
// conditionally compiled so impl registration is in source file
vllm/_custom_ops.py
-17
View File
@@ -2449,23 +2449,6 @@ def moe_wna16_gemm(
)
def router_gemm_bf16_fp32(input: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
"""bf16 x bf16 -> fp32 GEMM via cuBLAS. weight shape: (N, K)."""
return torch.ops._moe_C.router_gemm_bf16_fp32(input, weight)
if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "router_gemm_bf16_fp32"):
@register_fake("_moe_C::router_gemm_bf16_fp32")
def router_gemm_bf16_fp32_fake(
input: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
return torch.empty(
input.shape[0], weight.shape[0], dtype=torch.float32, device=input.device
)
def dsv3_router_gemm(
hidden_states: torch.Tensor,
router_weight: torch.Tensor,
vllm/model_executor/layers/deepseek_v4_attention.py
+8 -5
View File
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.linear import (
ReplicatedLinear,
)
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
from vllm.model_executor.layers.utils import cublas_gemm_bf16_bf16_fp32
from vllm.utils.deep_gemm import fp8_einsum
from vllm.utils.torch_utils import direct_register_custom_op
from vllm.v1.attention.ops.deepseek_v4_ops import (
@@ -349,8 +348,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
compressor = self.compressor
def compressor_kv_score() -> torch.Tensor:
return cublas_gemm_bf16_bf16_fp32(
hidden_states, compressor.fused_wkv_wgate.weight
return torch.mm(
hidden_states,
compressor.fused_wkv_wgate.weight.T,
out_dtype=torch.float32,
)
aux_fns[0] = compressor_kv_score
@@ -364,8 +365,10 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
return weights
def indexer_compressor_kv_score() -> torch.Tensor:
return cublas_gemm_bf16_bf16_fp32(
hidden_states, indexer.compressor.fused_wkv_wgate.weight
return torch.mm(
hidden_states,
indexer.compressor.fused_wkv_wgate.weight.T,
out_dtype=torch.float32,
)
aux_fns[1] = indexer_weights_proj
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+1 -1
View File
@@ -105,7 +105,7 @@ class GateLinear(ReplicatedLinear):
# Tier 2: cuBLAS bf16→fp32
if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16:
output = ops.router_gemm_bf16_fp32(x, self.weight)
output = torch.mm(x, self.weight.T, out_dtype=torch.float32)
return output, None
# Tier 3: F.linear (ReplicatedLinear)
vllm/model_executor/layers/utils.py
-7
View File
@@ -304,13 +304,6 @@ def cpu_unquantized_gemm(
return layer.cpu_linear(x, weight, bias)
def cublas_gemm_bf16_bf16_fp32(
x: torch.Tensor,
weight: torch.Tensor,
):
return ops.router_gemm_bf16_fp32(x, weight)
def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
if current_platform.is_rocm():
return rocm_unquantized_gemm