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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user