From 9c61864bf8a911a8369f35d79d538c7f11cf3dc2 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 30 Apr 2026 16:28:57 -0700 Subject: [PATCH] [DeepSeek] Use torch.mm for bf16xbf16->fp32 gemm (#41300) Signed-off-by: Woosuk Kwon --- CMakeLists.txt | 1 - csrc/moe/moe_ops.h | 4 -- csrc/moe/router_gemm.cu | 52 ------------------- csrc/moe/torch_bindings.cpp | 4 -- vllm/_custom_ops.py | 17 ------ .../layers/deepseek_v4_attention.py | 13 +++-- .../layers/fused_moe/router/gate_linear.py | 2 +- vllm/model_executor/layers/utils.py | 7 --- 8 files changed, 9 insertions(+), 91 deletions(-) delete mode 100644 csrc/moe/router_gemm.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index fb8a1d7e1e..bf4ac05e4f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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() diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index 973190935d..ac0e8d59f6 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -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 diff --git a/csrc/moe/router_gemm.cu b/csrc/moe/router_gemm.cu deleted file mode 100644 index a939f8846f..0000000000 --- a/csrc/moe/router_gemm.cu +++ /dev/null @@ -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 -#include -#include - -// 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(N), - static_cast(M), static_cast(K), &alpha, weight.data_ptr(), - CUDA_R_16BF, static_cast(K), input.data_ptr(), CUDA_R_16BF, - static_cast(K), &beta, out.data_ptr(), CUDA_R_32F, - static_cast(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); - - return out; -} diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 7bf56ba7a2..b737cb5435 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -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 diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 428d9b5d2e..ffea35de88 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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, diff --git a/vllm/model_executor/layers/deepseek_v4_attention.py b/vllm/model_executor/layers/deepseek_v4_attention.py index 175074cd87..74b494dc48 100644 --- a/vllm/model_executor/layers/deepseek_v4_attention.py +++ b/vllm/model_executor/layers/deepseek_v4_attention.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 77d8e75602..a868c9c848 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -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) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index f7fb2ad047..63a79f668e 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -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