gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: TF32 tensor cores for forward GEMMs

Browse Source 
Switch all 3 forward-pass cublasGemmEx calls from CUBLAS_COMPUTE_32F to
CUBLAS_COMPUTE_32F_FAST_TF32, enabling H100 TF32 tensor cores for 2-3x
throughput. Backward pass remains at full F32 precision for gradient stability.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 11:39:10 +02:00
parent c30713c5a6
commit 4d78f90cc0
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -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 {