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

[Bugfix][MoE] Enable cuBLAS out_dtype router GEMM on all CUDA archs (fixes family-120/GB10) (#54048)

Browse Source 
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
This commit is contained in:
hcl
2026-08-29 19:18:42 -07:00
committed by GitHub
parent 680e2177e4
commit b5707bf994
1 changed files with 10 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+10 -11
View File
@@ -116,15 +116,17 @@ class GateLinear(ReplicatedLinear):
# Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype
# epilogue folds the fp32 cast into the GEMM, removing the standalone
# bf16->fp32 copy kernel that otherwise runs before grouped_topk.
# cuBLAS on CUDA (SM90+, via allow_specialized_router_gemm); hipBLASLt on
# ROCm, which supports the same out_dtype epilogue.
# bf16->fp32 copy kernel that otherwise runs before grouped_topk. This is
# the plain cuBLAS (CUDA) / hipBLASLt (ROCm) out_dtype epilogue, so it
# applies on any CUDA-alike device (no bias, since torch.mm has no bias
# term). The specialized-kernel gate above excludes family-120 Blackwell
# (GB10 / DGX Spark), which this tier still covers. See #49921.
self._router_gemm_no_bias = not bias
self._router_gemm_cublas_capable = (
current_platform.is_cuda() or current_platform.is_rocm()
) and self._router_gemm_no_bias
self.allow_cublas_router_gemm = (
(
self.allow_specialized_router_gemm
or (current_platform.is_rocm() and self._router_gemm_no_bias)
)
self._router_gemm_cublas_capable
and self.weight.dtype == torch.bfloat16
and self.out_dtype == torch.float32
)
@@ -156,10 +158,7 @@ class GateLinear(ReplicatedLinear):
if (
not self.allow_cublas_router_gemm
and (
self.allow_specialized_router_gemm
or (current_platform.is_rocm() and self._router_gemm_no_bias)
)
and self._router_gemm_cublas_capable
and out_dtype == torch.float32
):
self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16