From d29c39d74ea816fa1baac5358b77947da4a26f5c Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sat, 30 May 2026 15:20:04 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20GEMV=20NaN=20bug=20=E2=80=94=20skip=20cu?= =?UTF-8?q?stom=20kernel=20for=20small=20N=20(<256)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/xserv-kernels/src/gemm.rs | 2 +- crates/xserv-model/src/gpt_oss.rs | 21 ++------------------- 2 files changed, 3 insertions(+), 20 deletions(-) diff --git a/crates/xserv-kernels/src/gemm.rs b/crates/xserv-kernels/src/gemm.rs index df74bc2..5a9ad1a 100644 --- a/crates/xserv-kernels/src/gemm.rs +++ b/crates/xserv-kernels/src/gemm.rs @@ -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( diff --git a/crates/xserv-model/src/gpt_oss.rs b/crates/xserv-model/src/gpt_oss.rs index 634ee10..d50540f 100644 --- a/crates/xserv-model/src/gpt_oss.rs +++ b/crates/xserv-model/src/gpt_oss.rs @@ -411,30 +411,13 @@ impl GptOss { // Router: [tokens, hidden] @ [hidden, num_experts] + bias → [tokens, num_experts] unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); } - // Pad to 2 rows to avoid GEMV path (workaround for GEMV NaN bug with small N) - let x_padded = if num_tokens == 1 { - let x_cpu_tmp = x.to_device(Device::Cpu); - let xd = x_cpu_tmp.as_slice::(); - let mut padded = xd.to_vec(); - padded.extend(vec![bf16::ZERO; hidden]); - Tensor::from_slice(&padded, &[2, hidden]).to_device(x.device()) - } else { - x.clone() - }; - let router_logits_full = add_bias( - &matmul_2d(&x_padded, &layer.router_wt), + let router_logits = add_bias( + &matmul_2d(x, &layer.router_wt), &layer.router_bias, ); - let router_logits = if num_tokens == 1 { - router_logits_full.narrow(0, 0, 1).contiguous() - } else { - router_logits_full - }; unsafe { xserv_cuda::ffi::cudaDeviceSynchronize(); } let router_cpu = router_logits.to_device(Device::Cpu); let router_data = router_cpu.as_slice::(); - - // Copy x to CPU after all GPU ops are synced let x_cpu = x.to_device(Device::Cpu); let x_data = x_cpu.as_slice::();