From 50e6cdc9062ba815e87d14ebf6962daee574cdcd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 12:27:48 +0200 Subject: [PATCH] perf(phase-h): add cuBLAS-Lt BIAS epilogue infrastructure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../ml/src/cuda_pipeline/batched_forward.rs | 235 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 24 ++ 2 files changed, 259 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index 0ce26908b..15c4afcb0 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -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, + /// 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, /// 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 = 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::(), + ).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 { + 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::(), + ).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::(), + ).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::(), + ).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::(), + ).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, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 67161ed4e..9fc53f33d 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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` /