fix: GEMV NaN bug — skip custom kernel for small N (<256)

The custom launch_gemv_bf16 kernel produces NaN when output dimension N
is small (e.g. N=32 for the MoE router). Fall back to cuBLAS GemmEx for
N < 256. Also removes the padding workaround in gpt_oss MoE forward.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-30 15:20:04 +08:00
parent 9ad91a4a92
commit d29c39d74e
2 changed files with 3 additions and 20 deletions

View File

@@ -173,7 +173,7 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
}
}
GemmBackend::CuBlas => {
if m == 1 && dtype == DType::BF16 {
if m == 1 && dtype == DType::BF16 && n >= 256 {
let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
unsafe {
launch_gemv_bf16(