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

perf: TF32 tensor cores for backward GEMMs (2-3x throughput on H100)

Browse Source 
Changed 2 backward cuBLAS GEMMs from CUBLAS_COMPUTE_32F to
CUBLAS_COMPUTE_32F_FAST_TF32. Both activation gradient (BF16→BF16)
and weight gradient (BF16→F32) paths now use H100 tensor cores.

TF32's 10-bit mantissa is sufficient for gradients — Adam's momentum
exponential moving average smooths any additional noise. Forward
GEMMs already use TF32 (unchanged).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 22:49:13 +02:00
parent c5b4046169
commit 3d0824eb9b
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -237,7 +237,7 @@ impl CublasBackward {
c as *mut std::ffi::c_void,
cublas_sys::cudaDataType_t::CUDA_R_16BF,
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 {
@@ -281,7 +281,7 @@ impl CublasBackward {
c as *mut std::ffi::c_void,
cublas_sys::cudaDataType_t::CUDA_R_32F,
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 {