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

perf(cuda): enable TF32 Tensor Core math on all cuBLAS handles

Browse Source 
DQN head (dqn.rs) and Mamba2 encoder (mamba2_block.rs) were forcing
CUBLAS_COMPUTE_32F — pure FP32 scalar, bypassing Tensor Cores entirely.
TF32 (19-bit mantissa) gives 2-3× SGEMM throughput on L40S/H100 with
negligible precision loss (well within RL gradient noise).

8+ SGEMMs per step now use Tensor Cores: DQN forward/backward (4) +
Mamba2 W_in/W_a/W_b/W_out projections (4+).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-27 20:19:45 +02:00
parent e730b5cbd1
commit 4fa9da9fc1
2 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml-alpha/src/mamba2_block.rs
View File

@@ -547,6 +547,12 @@ impl Mamba2Block {
) )
.result() .result()
.map_err(|e| anyhow!("Mamba2Block: cublasSetWorkspace_v2: {e:?}"))?; .map_err(|e| anyhow!("Mamba2Block: cublasSetWorkspace_v2: {e:?}"))?;
cudarc::cublas::sys::cublasSetMathMode(
*cublas.handle(),
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
)
.result()
.map_err(|e| anyhow!("Mamba2Block: cublasSetMathMode TF32: {e:?}"))?;
} }
// ── Parameter allocation + initialisation. ────────────────────── // ── Parameter allocation + initialisation. ──────────────────────

6
crates/ml-alpha/src/rl/dqn.rs
View File

@@ -311,6 +311,12 @@ impl DqnHead {
) )
.result() .result()
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetWorkspace_v2: {e:?}"))?; .map_err(|e| anyhow::anyhow!("DqnHead: cublasSetWorkspace_v2: {e:?}"))?;
cudarc::cublas::sys::cublasSetMathMode(
*cublas.handle(),
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
)
.result()
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetMathMode TF32: {e:?}"))?;
} }
// Bias-add and reduce-sum-axis0 kernel handles from ml-core. // Bias-add and reduce-sum-axis0 kernel handles from ml-core.