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

perf(phase-h): add cuBLAS-Lt BIAS epilogue infrastructure

Browse Source 
Adds gemm_cache_bias field, create_cached_fwd_gemm_desc_bias builder,
and sgemm_f32_fused_bias helper. Pre-creates descriptors for every
BIAS-only output shape used in the forward pass. No call sites
migrated yet — subsequent commits flip individual sites to use the
fused helper.

Mirrors the established RELU_BIAS infrastructure pattern: deterministic
algorithm selection, per-call bias-pointer wiring, Err-on-missing-cache
contract for graceful fall-back.
This commit is contained in:
jgrusewski
2026-04-28 12:27:48 +02:00
parent 756b1ef317
commit 50e6cdc906
2 changed files with 259 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -181,6 +181,13 @@ pub struct CublasGemmSet {
/// Cached descriptors with RELU_BIAS epilogue (for hidden layers).
/// Bias pointer is set dynamically per-call. Same key as gemm_cache.
gemm_cache_relu_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
/// Cached descriptors with BIAS epilogue (for linear output layers — value
/// logits, advantage logits, GRN Linear_a/Linear_b, VSN linear2, GLU value/gate).
/// Bias pointer is set dynamically per-call. Same key as `gemm_cache`.
/// Phase H Site 3: fuses GEMM + bias-add into one cublasLtMatmul call,
/// eliminating the standalone `add_bias_f32_kernel` launch (one less kernel
/// + memory round-trip per BIAS-only output layer).
gemm_cache_bias: HashMap<FwdGemmKey, CachedGemmDesc>,
/// Cached descriptors with RELU_AUX_BIAS epilogue (writes the dReLU bit-mask
/// for backward fusion via DRELU_BGRAD). Currently populated only for the
/// online value-FC layer's shape — the only site whose backward uses
@@ -483,6 +490,58 @@ impl CublasGemmSet {
}
}
// ── Phase H Site 3: BIAS-only epilogue variants for linear output layers ──
//
// Fuses GEMM + bias-add into a single cublasLtMatmul call (no ReLU,
// because these layers are either output logits or feed into a separate
// gating activation downstream — KAN/GLU/softmax). Each fused descriptor
// eliminates one `add_bias_f32_kernel` launch + one memory round-trip
// per call.
//
// Shapes covered:
// - VSN linear2 (group logit producer): (1, B, VSN_HIDDEN, VSN_HIDDEN)
// - GRN Linear_a / Linear_b (h_s1, h_s2 trunks)
// - Value-head v_logits: (num_atoms, B, value_h, value_h)
// - Branch adv_logits (×4 shapes — branch_size_d * num_atoms output dim)
// - Branch h_bd / GLU value & gate (no-ReLU when GLU-gated downstream):
// same K dims as RELU_BIAS but BIAS-only epilogue for the GLU path.
const VSN_HIDDEN: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM;
let mut bias_shapes: Vec<FwdGemmKey> = vec![
// VSN linear2 — projects [B, VSN_HIDDEN] → [B, 1] per group.
(1, batch_size, VSN_HIDDEN, VSN_HIDDEN),
// GRN h_s1 Linear_a / Linear_residual share K=s1_input_dim, ldb=s1_ldb.
// (Linear_residual has no bias — only Linear_a uses the BIAS epilogue.)
(shared_h1, batch_size, s1_input_dim, s1_ldb),
// GRN h_s1 Linear_b: M=2*shared_h1, K=shared_h1.
(2 * shared_h1, batch_size, shared_h1, shared_h1),
// GRN h_s2 Linear_a: M=shared_h2, K=shared_h1.
(shared_h2, batch_size, shared_h1, shared_h1),
// GRN h_s2 Linear_b: M=2*shared_h2, K=shared_h2.
(2 * shared_h2, batch_size, shared_h2, shared_h2),
// Value-head v_logits: M=num_atoms, K=value_h.
(num_atoms, batch_size, value_h, value_h),
// Branch h_bd / GLU value & gate — BIAS only (GLU/KAN provides gating):
(adv_h, batch_size, shared_h2, shared_h2),
(adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size),
(adv_h, batch_size, shared_h2 + 3, shared_h2 + 3),
];
// Branch adv_logits (×4 distinct branch sizes): M=bs*num_atoms, K=adv_h.
for &bs in &branch_sizes {
bias_shapes.push((bs * num_atoms, batch_size, adv_h, adv_h));
}
bias_shapes.sort();
bias_shapes.dedup();
let mut gemm_cache_bias = HashMap::new();
for &(n, b, k, ldb) in &bias_shapes {
match create_cached_fwd_gemm_desc_bias(lt_raw_handle, n, b, k, ldb, lt_ws_size) {
Ok(desc) => { gemm_cache_bias.insert((n, b, k, ldb), desc); }
Err(e) => {
tracing::warn!("BIAS epilogue not available for ({n},{b},{k},{ldb}): {e}");
// Fall back to separate kernels — no entry in cache
}
}
}
// ── Phase H Site 2: RELU_AUX_BIAS variant for the value-FC layer ──
//
// Writes the dReLU bit-mask aux buffer alongside the post-ReLU output.
@@ -645,6 +704,7 @@ impl CublasGemmSet {
branch_done_events,
gemm_cache,
gemm_cache_relu_bias,
gemm_cache_bias,
gemm_cache_relu_aux_bias,
_value_fc_relu_mask_buf: value_fc_relu_mask_buf,
value_fc_relu_mask_ptr_raw: value_fc_relu_mask_ptr,
@@ -2900,6 +2960,81 @@ impl CublasGemmSet {
}
}
/// Fused GEMM + bias via cublasLt epilogue (no ReLU).
/// Uses cached descriptor with BIAS epilogue. Sets bias pointer per-call.
/// Falls back to an `Err` if the shape is not pre-cached at init, signalling
/// the caller to use the slower separate-kernel path (`sgemm_f32` +
/// `launch_add_bias_f32_raw`).
///
/// Phase H Site 3: collapses GEMM + add_bias into one cublasLtMatmul call,
/// eliminating the standalone `add_bias_f32_kernel` launch and one memory
/// round-trip per BIAS-only output layer (value logits, advantage logits,
/// GRN Linear_a/b, VSN linear2, GLU value/gate).
#[allow(clippy::too_many_arguments)]
pub(crate) fn sgemm_f32_fused_bias(
&self,
stream: &CudaStream,
w_ptr: u64,
input_ptr: u64,
output_ptr: u64,
bias_ptr: u64,
n: usize, // out_dim
b: usize, // batch
k: usize, // in_dim
ldb: usize,
ws_ptr: u64,
ws_size: usize,
label: &str,
) -> Result<(), MLError> {
let key: FwdGemmKey = (n, b, k, ldb);
if let Some(cached) = self.gemm_cache_bias.get(&key) {
// Set bias pointer on the descriptor (lightweight CPU write).
unsafe {
cublaslt_result::set_matmul_desc_attribute(
cached.matmul_desc,
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_POINTER,
&bias_ptr as *const u64 as *const std::ffi::c_void,
std::mem::size_of::<u64>(),
).map_err(|e| MLError::ModelError(format!("set bias ptr {label}: {e:?}")))?;
}
// Launch fused GEMM+bias.
let alpha = 1.0_f32;
let beta = 0.0_f32;
let lt_handle = self.handle.lt_handle_for(stream)?;
unsafe {
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
let status = cublaslt_sys::cublasLtMatmul(
lt_handle,
cached.matmul_desc,
&alpha as *const f32 as *const std::ffi::c_void,
w_ptr as *const std::ffi::c_void,
cached.a_layout,
input_ptr as *const std::ffi::c_void,
cached.b_layout,
&beta as *const f32 as *const std::ffi::c_void,
output_ptr as *const std::ffi::c_void,
cached.c_layout,
output_ptr as *mut std::ffi::c_void,
cached.d_layout,
&cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
ws_ptr as *mut std::ffi::c_void,
ws_size,
cu_stream,
);
if status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
let e = cublaslt_result::CublasError(status);
tracing::error!(m=n, n_batch=b, k=k, ?e, "cublasLtMatmul FUSED BIAS FAILED for {label}");
return Err(MLError::ModelError(format!("cublasLtMatmul fused-bias {label}: {e:?}")));
}
}
Ok(())
} else {
// No cached epilogue descriptor — caller should use separate GEMM + bias.
Err(MLError::ModelError(format!("no BIAS epilogue for {label} (n={n},b={b},k={k},ldb={ldb})")))
}
}
/// Fused GEMM + bias + ReLU + dReLU bit-mask aux output via cublasLt epilogue.
///
/// Same as `sgemm_f32_fused_relu_bias` but uses `EPILOGUE_RELU_AUX_BIAS` —
@@ -3298,6 +3433,106 @@ fn create_cached_fwd_gemm_desc(
}
}
/// Create a cached GEMM descriptor with BIAS epilogue (no ReLU).
/// The bias pointer is set dynamically per-call via `set_matmul_desc_attribute`.
/// This fuses GEMM + bias-add into a single cublasLtMatmul kernel, eliminating
/// the separate `add_bias_f32_kernel` launch (Phase H Site 3).
fn create_cached_fwd_gemm_desc_bias(
lt_handle: cublaslt_sys::cublasLtHandle_t,
n: usize, b: usize, k: usize, ldb: usize, ws_size: usize,
) -> Result<CachedGemmDesc, MLError> {
let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;
unsafe {
let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
.map_err(|e| MLError::ModelError(format!("cached+bias MatmulDescCreate (n={n},b={b},k={k}): {e:?}")))?;
let transa: i32 = 1; // CUBLAS_OP_T
cublaslt_result::set_matmul_desc_attribute(
matmul_desc,
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA,
&transa as *const i32 as *const std::ffi::c_void,
std::mem::size_of::<i32>(),
).map_err(|e| MLError::ModelError(format!("cached+bias set TRANSA: {e:?}")))?;
let transb: i32 = 0; // CUBLAS_OP_N
cublaslt_result::set_matmul_desc_attribute(
matmul_desc,
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB,
&transb as *const i32 as *const std::ffi::c_void,
std::mem::size_of::<i32>(),
).map_err(|e| MLError::ModelError(format!("cached+bias set TRANSB: {e:?}")))?;
// Set BIAS epilogue — fuses bias add into the matmul kernel (no ReLU).
let epilogue: i32 = cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS as i32;
cublaslt_result::set_matmul_desc_attribute(
matmul_desc,
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE,
&epilogue as *const i32 as *const std::ffi::c_void,
std::mem::size_of::<i32>(),
).map_err(|e| MLError::ModelError(format!("cached+bias set EPILOGUE: {e:?}")))?;
// Set bias data type to F32.
let bias_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
cublaslt_result::set_matmul_desc_attribute(
matmul_desc,
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
&bias_type as *const _ as *const std::ffi::c_void,
std::mem::size_of::<cublaslt_sys::cudaDataType_t>(),
).map_err(|e| MLError::ModelError(format!("cached+bias set BIAS_DATA_TYPE: {e:?}")))?;
let a_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, k as i64)
.map_err(|e| MLError::ModelError(format!("cached+bias A layout: {e:?}")))?;
let b_layout = cublaslt_result::create_matrix_layout(f32_type, k as u64, b as u64, ldb as i64)
.map_err(|e| MLError::ModelError(format!("cached+bias B layout: {e:?}")))?;
let c_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
.map_err(|e| MLError::ModelError(format!("cached+bias C layout: {e:?}")))?;
let d_layout = cublaslt_result::create_matrix_layout(f32_type, n as u64, b as u64, n as i64)
.map_err(|e| MLError::ModelError(format!("cached+bias D layout: {e:?}")))?;
// Deterministic algorithm selection — keys on the BIAS epilogue value
// so the cache produces an algo selection valid for BIAS (independent
// from DEFAULT and RELU_BIAS entries for the same shape).
let shape = super::cublas_algo_deterministic::ShapeKey::with_epilogue(
transa, transb,
n as i32, b as i32, k as i32,
k as i32, ldb as i32, n as i32,
cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS,
ws_size,
);
let heuristic = super::cublas_algo_deterministic::get_matmul_algo_f32_tf32(
lt_handle, matmul_desc,
a_layout, b_layout, c_layout, d_layout,
shape,
);
let heuristic = match heuristic {
Ok(h) => h,
Err(e) => {
let _ = cublaslt_result::destroy_matrix_layout(d_layout);
let _ = cublaslt_result::destroy_matrix_layout(c_layout);
let _ = cublaslt_result::destroy_matrix_layout(b_layout);
let _ = cublaslt_result::destroy_matrix_layout(a_layout);
let _ = cublaslt_result::destroy_matmul_desc(matmul_desc);
return Err(MLError::ModelError(format!(
"cached+bias deterministic algo (n={n},b={b},k={k},ldb={ldb}): {e}"
)));
}
};
let algo = heuristic.algo;
tracing::info!(
n, b, k, ldb,
ws_needed = heuristic.workspaceSize,
"cached fwd GEMM+BIAS desc created (deterministic algorithm)"
);
Ok(CachedGemmDesc {
matmul_desc, a_layout, b_layout, c_layout, d_layout,
algo,
})
}
}
/// Create a cached GEMM descriptor with RELU_BIAS epilogue.
/// The bias pointer is set dynamically per-call via set_matmul_desc_attribute.
/// This fuses GEMM + bias-add + ReLU into a single cublasLtMatmul kernel,

24
docs/dqn-wire-up-audit.md
View File

@@ -2,6 +2,30 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
P5T5 Phase H Site 3 — BIAS epilogue infrastructure (2026-04-28):
adds the `gemm_cache_bias` field on `CublasGemmSet`, a
`create_cached_fwd_gemm_desc_bias` builder mirroring the existing
`_relu_bias` variant, and a `sgemm_f32_fused_bias` helper that fuses
GEMM + bias-add into one `cublasLtMatmul` call (`CUBLASLT_EPILOGUE_BIAS`).
Pre-creates descriptors at `CublasGemmSet::new` for every BIAS-only
output shape used in the forward pass: VSN linear2 `(1,B,VSN_HIDDEN,
VSN_HIDDEN)`; GRN h_s1/h_s2 Linear_a/Linear_b across the trunk; the
value-head `v_logits` `(num_atoms,B,value_h,value_h)`; the four
`adv_logits` shapes `(branch_size_d * num_atoms, B, adv_h, adv_h)`;
and the GLU/KAN value+gate shapes `(adv_h, B, K, K)` for K ∈ {sh2,
sh2+branch_0, sh2+3}. Mirrors the established `gemm_cache_relu_bias`
infrastructure — same descriptor creation pattern (deterministic
algo selection via `cublas_algo_deterministic`), same per-call bias-
pointer wiring via `set_matmul_desc_attribute`, same fall-back
contract on missing-cache (`Err` → caller routes to slow path).
Sites are not yet migrated in this commit (this is just the
infrastructure); subsequent commits flip individual call sites from
`sgemm_f32 + launch_add_bias_f32_raw` pairs to
`sgemm_f32_fused_bias`. Touched: `crates/ml/src/cuda_pipeline/
batched_forward.rs` (+136 LOC: field, init block, helper fn, builder
fn). cargo check clean at 13 warnings (workspace baseline);
`cargo test --no-run` clean at 24 warnings.
NaN diagnostic wire-up + label fix (2026-04-28): fold 1 of
`train-multi-seed-72fl6` hit NaN at step 5 with `flagged=[]` because the
8 NaN-check kernels in `run_nan_checks_pre_forward` /