diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 58d3fcf00..2e1465da2 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -681,8 +681,9 @@ impl CublasForward { // SAFETY: w_bf16_ptr, a_bf16_ptr, c_bf16_ptr are valid CUDA device pointers // in BF16 (u16) format. The cuBLAS handle is bound to the correct stream. - // cublasGemmEx with CUDA_R_16BF + CUBLAS_COMPUTE_32F uses tensor cores - // with F32 accumulation for numerical stability. + // cublasGemmEx with CUDA_R_16BF + CUBLAS_COMPUTE_32F_FAST_TF32 uses + // H100 TF32 tensor cores (~10-bit mantissa, 2-3x throughput vs F32). + // Sufficient precision for forward-pass Q-value computation. unsafe { let status = cublas_sys::cublasGemmEx( self.handle.0, @@ -702,7 +703,7 @@ impl CublasForward { c_bf16_ptr as *mut std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_16BF, n as i32, // ldc: leading dim of C = N - cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32, cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, ); if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { @@ -831,7 +832,7 @@ impl CublasForward { c_bf16_ptr as *mut std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_16BF, n as i32, // ldc: leading dim of C = N - cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32, cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, ); if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { @@ -888,7 +889,7 @@ impl CublasForward { c_f32_ptr as *mut std::ffi::c_void, cublas_sys::cudaDataType_t::CUDA_R_32F, // C type: F32 output n as i32, // ldc - cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32, cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, ); if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {