[Bugfix][MoE] Enable cuBLAS out_dtype router GEMM on all CUDA archs (fixes family-120/GB10) (#54048)
Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>
hcl committed
Aug 30, 2026 at 10:18 UTC
b5707bf994cb968adfc7a29fbb80b0582f53f38d
1 file changed
+10
-11
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+10
-11
@@ -116,15 +116,17 @@ class GateLinear(ReplicatedLinear):
116
117
# Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype
118
# epilogue folds the fp32 cast into the GEMM, removing the standalone
119
- # bf16->fp32 copy kernel that otherwise runs before grouped_topk.
120
- # cuBLAS on CUDA (SM90+, via allow_specialized_router_gemm); hipBLASLt on
121
- # ROCm, which supports the same out_dtype epilogue.
119
+ # bf16->fp32 copy kernel that otherwise runs before grouped_topk. This is
120
+ # the plain cuBLAS (CUDA) / hipBLASLt (ROCm) out_dtype epilogue, so it
121
+ # applies on any CUDA-alike device (no bias, since torch.mm has no bias
122
+ # term). The specialized-kernel gate above excludes family-120 Blackwell
123
+ # (GB10 / DGX Spark), which this tier still covers. See #49921.
124
self._router_gemm_no_bias = not bias
125
+ self._router_gemm_cublas_capable = (
126
+ current_platform.is_cuda() or current_platform.is_rocm()
127
+ ) and self._router_gemm_no_bias
128
self.allow_cublas_router_gemm = (
124
- (
125
- self.allow_specialized_router_gemm
126
- or (current_platform.is_rocm() and self._router_gemm_no_bias)
127
- )
129
+ self._router_gemm_cublas_capable
130
and self.weight.dtype == torch.bfloat16
131
and self.out_dtype == torch.float32
132
)
@@ -156,10 +158,7 @@ class GateLinear(ReplicatedLinear):
158
159
if (
160
not self.allow_cublas_router_gemm
159
- and (
160
- self.allow_specialized_router_gemm
161
- or (current_platform.is_rocm() and self._router_gemm_no_bias)
162
- )
161
+ and self._router_gemm_cublas_capable
162
and out_dtype == torch.float32
163
):
164
self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16