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

feat: rewrite IQL trainer to batched cuBLAS GEMMs

Browse Source 
Replace per-sample CUDA kernels (16,384 launches per GEMM) with batched
cublasLtMatmul — ONE call per GEMM layer. Forward pass uses 3 GEMMs +
SiLU + bias_add + expectile_loss. Backward pass uses 5 GEMMs + SiLU
backward + bias_grad_reduce. Eliminates grads_per_sample buffer (was
27MB), tiled backward loop, and weight_grad_reduce kernel.

Constructor now accepts Arc<SharedCublasHandle> instead of bare stream,
sharing the cuBLAS handle with the trunk forward/backward pipeline.
Eight IqlGemmDesc descriptors are cached at init with heuristic algo
selection for CUDA Graph compatibility.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 23:38:00 +02:00
parent 196c390946
commit d799bfafbe
2 changed files with 683 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

952
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

File diff suppressed because it is too large Load Diff

5
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -486,14 +486,15 @@ impl FusedTrainingCtx {
gamma: hyperparams.gamma as f32,
..GpuIqlConfig::default()
};
let gpu_iql = GpuIqlTrainer::new(stream.clone(), iql_config.clone())
let shared_cublas_for_iql = Arc::clone(trainer.shared_cublas());
let gpu_iql = GpuIqlTrainer::new(Arc::clone(&shared_cublas_for_iql), iql_config.clone())
.map_err(|e| anyhow::anyhow!("GPU IQL (high-tau) init: {e}"))?;
let iql_low_config = GpuIqlConfig {
expectile_tau: 1.0 - hyperparams.iql_expectile_tau,
..iql_config
};
let gpu_iql_low = GpuIqlTrainer::new(stream.clone(), iql_low_config)
let gpu_iql_low = GpuIqlTrainer::new(shared_cublas_for_iql, iql_low_config)
.map_err(|e| anyhow::anyhow!("GPU IQL (low-tau) init: {e}"))?;
// Initialize GPU IQN dual-head when iqn_lambda > 0 (CVaR risk sizing)