perf(dqn): Phase H — fuse cublasLt BIAS epilogue into 4 attention forward projections
Collapses each `cublasLtMatmul + add_bias_f32_kernel` pair in the attention
forward path (Q, K, V, O projections) into a single fused `cublasLtMatmul`
with `CUBLASLT_EPILOGUE_BIAS`. Saves 4 kernel launches per attention forward
× per training step (target + online), targeting the L40S deploy hot-spot
where the per-epoch training phase is 99.7% of wall time.
Three code changes in `crates/ml/src/cuda_pipeline/gpu_attention.rs`:
1. `create_attn_gemm_desc` extended with `epilogue: Option<cublasLtEpilogue_t>`
parameter — when `Some(BIAS)`, descriptor is configured with
`EPILOGUE = BIAS` + `BIAS_DATA_TYPE = CUDA_R_32F`; when `None`, stays
at `EPILOGUE_DEFAULT` (the 8 backward dW/dX GEMMs unchanged).
2. New `lt_matmul_with_bias_ex` helper writes the per-call bias pointer via
`set_matmul_desc_attribute(BIAS_POINTER, …)` immediately before each
`cublasLtMatmul` (mirrors the existing pattern in batched_forward.rs).
The 4 forward projection sites in `forward(...)` switch from the prior
`lt_matmul_ex(...)` + `launch_bias_add_ex(...)` pair to a single
`lt_matmul_with_bias_ex(...)` call.
3. Orphans pruned: `launch_bias_add_ex` and `launch_bias_add` deleted from
`gpu_attention.rs` (their only callers were the 4 fused-away sites).
Shared `add_bias_f32_kernel` retained — still used by
`batched_forward.rs::launch_add_bias_f32_raw` (VSN Linear_2 logit output,
no activation).
Determinism preserved: the deterministic-algo cache (`cublas_algo_deterministic.rs`)
already keys on epilogue via `ShapeKey::with_epilogue`, so first-call selection
runs the full `AlgoGetIds → AlgoInit → AlgoCheck` loop with the new descriptor
and subsequent calls reuse the cached `(types, shape, epilogue)` algo. Bit-
deterministic when the algo is fixed under `CUBLAS_WORKSPACE_CONFIG=:4096:8`.
Site #2 (DRELU_BGRAD on the trunk Linear→Bias→ReLU backward) deferred — the
forward-side aux-buffer plumbing crosses three modules (BatchedForward →
BatchedBackward → value-FC site) and the determinism contract verified by
the Phase G smoke is non-trivial to preserve. Tracked for follow-up.
Verified: cargo check workspace clean at 11 warnings (baseline preserved).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -291,22 +291,30 @@ impl GpuAttention {
|
||||
let ws_size = shared_handle.lt_workspace_size;
|
||||
let id = config.input_dim(); // D + OFI_EMBED_DIM
|
||||
|
||||
// Forward Q/K/V GEMMs: C[D,B] = W^T[D,id] @ X[id,B]
|
||||
// Forward Q/K/V/O GEMMs: C[D,B] = W^T[D,id|D] @ X[id|D,B] + bias[D]
|
||||
// W is stored row-major [D, id], which is col-major [id, D].
|
||||
// TRANSA=T: transpose the col-major [id,D] to get [D,id] logically.
|
||||
// M=D, N=B, K=id, lda=id (leading dim of stored A = id rows in col-major), ldb=id, ldc=D
|
||||
let gemm_fwd_q = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, "fwd_q")?;
|
||||
let gemm_fwd_k = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, "fwd_k")?;
|
||||
let gemm_fwd_v = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, "fwd_v")?;
|
||||
// Forward O GEMM: C[D,B] = W_O^T[D,D] @ attn_out[D,B] — unchanged
|
||||
let gemm_fwd_o = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "fwd_o")?;
|
||||
//
|
||||
// P5T5 Phase H — cublasLt epilogue fusion: forward projections set
|
||||
// EPILOGUE=BIAS so the per-projection bias add is performed inside
|
||||
// the cublasLtMatmul call. Removes 4 separate `add_bias_f32_kernel`
|
||||
// launches per attention forward (× per training step). Bias pointer
|
||||
// is set per-call via `set_matmul_desc_attribute` in
|
||||
// `lt_matmul_with_bias_ex`.
|
||||
let bias_epi = Some(cublaslt_sys::cublasLtEpilogue_t::CUBLASLT_EPILOGUE_BIAS);
|
||||
let gemm_fwd_q = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, bias_epi, "fwd_q")?;
|
||||
let gemm_fwd_k = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, bias_epi, "fwd_k")?;
|
||||
let gemm_fwd_v = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, id as i32, id as i32, id as i32, d as i32, ws_size, bias_epi, "fwd_v")?;
|
||||
// Forward O GEMM: C[D,B] = W_O^T[D,D] @ attn_out[D,B] + bias_O[D]
|
||||
let gemm_fwd_o = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, bias_epi, "fwd_o")?;
|
||||
|
||||
// Backward dW Q/K/V GEMMs: dW[D,id] = dY[D,B] @ X^T[B,id]
|
||||
// TRANSA=N, TRANSB=T: M=D, N=id, K=B, lda=D, ldb=id (stored input), ldc=D
|
||||
let gemm_bwd_dw_o = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dw_o")?;
|
||||
let gemm_bwd_dw_q = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, "bwd_dw_q")?;
|
||||
let gemm_bwd_dw_k = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, "bwd_dw_k")?;
|
||||
let gemm_bwd_dw_v = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, "bwd_dw_v")?;
|
||||
let gemm_bwd_dw_o = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, None, "bwd_dw_o")?;
|
||||
let gemm_bwd_dw_q = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, None, "bwd_dw_q")?;
|
||||
let gemm_bwd_dw_k = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, None, "bwd_dw_k")?;
|
||||
let gemm_bwd_dw_v = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, id as i32, b as i32, d as i32, id as i32, d as i32, ws_size, None, "bwd_dw_v")?;
|
||||
|
||||
// Backward dX Q/K/V GEMMs: dX[id,B] = W[id,D] @ dY[D,B]
|
||||
// W is stored [D, id] row-major = [id, D] col-major. No transpose → A is [id, D].
|
||||
@@ -314,10 +322,10 @@ impl GpuAttention {
|
||||
// Wait — we need W^T @ dY. W stored as [D, id] means W^T is [id, D].
|
||||
// In col-major the stored matrix is [id, D]. With TRANSA=N it reads as [id, D].
|
||||
// So: C[id,B] = A[id,D] @ B[D,B] → M=id, N=B, K=D, lda=id, ldb=D, ldc=id.
|
||||
let gemm_bwd_dx_o = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dx_o")?;
|
||||
let gemm_bwd_dx_q = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, "bwd_dx_q")?;
|
||||
let gemm_bwd_dx_k = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, "bwd_dx_k")?;
|
||||
let gemm_bwd_dx_v = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, "bwd_dx_v")?;
|
||||
let gemm_bwd_dx_o = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, None, "bwd_dx_o")?;
|
||||
let gemm_bwd_dx_q = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, None, "bwd_dx_q")?;
|
||||
let gemm_bwd_dx_k = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, None, "bwd_dx_k")?;
|
||||
let gemm_bwd_dx_v = create_attn_gemm_desc(lt_handle, 0, 0, id as i32, b as i32, d as i32, id as i32, d as i32, id as i32, ws_size, None, "bwd_dx_v")?;
|
||||
|
||||
// ── Xavier initialization for attention weights ─────────────────
|
||||
|
||||
@@ -596,32 +604,29 @@ impl GpuAttention {
|
||||
let f32_sz = std::mem::size_of::<f32>();
|
||||
let params_base = self.params.raw_ptr();
|
||||
|
||||
// 1. Q = W_Q^T @ saved_input[D+10,B]
|
||||
self.lt_matmul_ex(&self.gemm_fwd_q,
|
||||
// 1. Q = W_Q^T @ saved_input[D+10,B] + b_Q (fused via EPILOGUE_BIAS)
|
||||
self.lt_matmul_with_bias_ex(&self.gemm_fwd_q,
|
||||
params_base + (w_q_off * f32_sz) as u64,
|
||||
self.saved_input.raw_ptr(),
|
||||
self.proj_q_buf.raw_ptr(),
|
||||
params_base + (b_q_off * f32_sz) as u64,
|
||||
1.0, 0.0, "fwd_q", stream, ws_ptr, ws_size)?;
|
||||
// + bias_add Q
|
||||
self.launch_bias_add_ex(self.proj_q_buf.raw_ptr(), params_base + (b_q_off * f32_sz) as u64, d, b, stream)?;
|
||||
|
||||
// 2. K = W_K^T @ saved_input[D+10,B]
|
||||
self.lt_matmul_ex(&self.gemm_fwd_k,
|
||||
// 2. K = W_K^T @ saved_input[D+10,B] + b_K (fused via EPILOGUE_BIAS)
|
||||
self.lt_matmul_with_bias_ex(&self.gemm_fwd_k,
|
||||
params_base + (w_k_off * f32_sz) as u64,
|
||||
self.saved_input.raw_ptr(),
|
||||
self.proj_k_buf.raw_ptr(),
|
||||
params_base + (b_k_off * f32_sz) as u64,
|
||||
1.0, 0.0, "fwd_k", stream, ws_ptr, ws_size)?;
|
||||
// + bias_add K
|
||||
self.launch_bias_add_ex(self.proj_k_buf.raw_ptr(), params_base + (b_k_off * f32_sz) as u64, d, b, stream)?;
|
||||
|
||||
// 3. V = W_V^T @ saved_input[D+10,B]
|
||||
self.lt_matmul_ex(&self.gemm_fwd_v,
|
||||
// 3. V = W_V^T @ saved_input[D+10,B] + b_V (fused via EPILOGUE_BIAS)
|
||||
self.lt_matmul_with_bias_ex(&self.gemm_fwd_v,
|
||||
params_base + (w_v_off * f32_sz) as u64,
|
||||
self.saved_input.raw_ptr(),
|
||||
self.proj_v_buf.raw_ptr(),
|
||||
params_base + (b_v_off * f32_sz) as u64,
|
||||
1.0, 0.0, "fwd_v", stream, ws_ptr, ws_size)?;
|
||||
// + bias_add V
|
||||
self.launch_bias_add_ex(self.proj_v_buf.raw_ptr(), params_base + (b_v_off * f32_sz) as u64, d, b, stream)?;
|
||||
|
||||
// 4. SDP: attn_out = SDP_fwd(Q, K, V)
|
||||
let num_heads_i32 = num_heads as i32;
|
||||
@@ -644,14 +649,13 @@ impl GpuAttention {
|
||||
.map_err(|e| MLError::ModelError(format!("attn_sdp_fwd: {e}")))?;
|
||||
}
|
||||
|
||||
// 5. projected = W_O^T @ attn_out
|
||||
self.lt_matmul_ex(&self.gemm_fwd_o,
|
||||
// 5. projected = W_O^T @ attn_out + b_O (fused via EPILOGUE_BIAS)
|
||||
self.lt_matmul_with_bias_ex(&self.gemm_fwd_o,
|
||||
params_base + (w_o_off * f32_sz) as u64,
|
||||
self.attn_out_buf.raw_ptr(),
|
||||
self.projected_buf.raw_ptr(),
|
||||
params_base + (b_o_off * f32_sz) as u64,
|
||||
1.0, 0.0, "fwd_o", stream, ws_ptr, ws_size)?;
|
||||
// + bias_add O
|
||||
self.launch_bias_add_ex(self.projected_buf.raw_ptr(), params_base + (b_o_off * f32_sz) as u64, d, b, stream)?;
|
||||
|
||||
// Save states [D,B] for LN backward dgamma (needs residual for normed computation).
|
||||
// Uses graph_safe_copy pattern: elementwise kernel, captured by CUDA Graph.
|
||||
@@ -1137,33 +1141,71 @@ impl GpuAttention {
|
||||
)
|
||||
}
|
||||
|
||||
/// Launch bias_add kernel with an explicit stream.
|
||||
fn launch_bias_add_ex(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize, stream: &CudaStream) -> Result<(), MLError> {
|
||||
let total = (batch * out_dim) as i32;
|
||||
let out_dim_i32 = out_dim as i32;
|
||||
let blocks = ((batch * out_dim + 255) / 256) as u32;
|
||||
/// Launch cublasLtMatmul with a per-call bias pointer (P5T5 Phase H).
|
||||
///
|
||||
/// Used for the 4 forward attention projections (Q/K/V/O), whose
|
||||
/// descriptors were created with `EPILOGUE_BIAS`. The bias data type and
|
||||
/// epilogue value are baked into the descriptor at creation; only the
|
||||
/// bias pointer changes per call (per-projection bias slice in
|
||||
/// `params_buf`). This collapses the previous "GEMM + add_bias_f32"
|
||||
/// kernel pair into a single fused cublasLtMatmul launch.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn lt_matmul_with_bias_ex(
|
||||
&self,
|
||||
desc: &AttnGemmDesc,
|
||||
a_ptr: u64,
|
||||
b_ptr: u64,
|
||||
c_ptr: u64,
|
||||
bias_ptr: u64,
|
||||
alpha: f32,
|
||||
beta: f32,
|
||||
label: &str,
|
||||
stream: &CudaStream,
|
||||
ws_ptr: u64,
|
||||
ws_size: usize,
|
||||
) -> Result<(), MLError> {
|
||||
// Set bias pointer on the cached descriptor (lightweight CPU write,
|
||||
// no GPU sync; the descriptor is unique to this projection's GEMM).
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.shared_handle.add_bias_f32_kernel)
|
||||
.arg(&out_ptr)
|
||||
.arg(&bias_ptr)
|
||||
.arg(&out_dim_i32)
|
||||
.arg(&total)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("attention bias_add: {e}")))?;
|
||||
cublaslt_result::set_matmul_desc_attribute(
|
||||
desc.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!("attn set BIAS_POINTER ({label}): {e:?}")))?;
|
||||
}
|
||||
// Per-stream lt_handle (Option C determinism fix).
|
||||
let lt_handle = self.shared_handle.lt_handle_for(stream)?;
|
||||
unsafe {
|
||||
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
|
||||
let matmul_status = cublaslt_sys::cublasLtMatmul(
|
||||
lt_handle,
|
||||
desc.matmul_desc,
|
||||
&alpha as *const f32 as *const std::ffi::c_void,
|
||||
a_ptr as *const std::ffi::c_void,
|
||||
desc.a_layout,
|
||||
b_ptr as *const std::ffi::c_void,
|
||||
desc.b_layout,
|
||||
&beta as *const f32 as *const std::ffi::c_void,
|
||||
c_ptr as *const std::ffi::c_void,
|
||||
desc.c_layout,
|
||||
c_ptr as *mut std::ffi::c_void,
|
||||
desc.d_layout,
|
||||
&desc.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
|
||||
ws_ptr as *mut std::ffi::c_void,
|
||||
ws_size,
|
||||
cu_stream,
|
||||
);
|
||||
if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
|
||||
let e = cublaslt_result::CublasError(matmul_status);
|
||||
return Err(MLError::ModelError(format!(
|
||||
"cublasLtMatmul attention+BIAS {label}: {e:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch bias_add kernel with the default stream.
|
||||
fn launch_bias_add(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize) -> Result<(), MLError> {
|
||||
self.launch_bias_add_ex(out_ptr, bias_ptr, out_dim, batch, &self.stream)
|
||||
}
|
||||
|
||||
/// Launch 2-phase bias gradient reduction with an explicit stream.
|
||||
fn launch_bias_grad_reduce_ex(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32, stream: &CudaStream) -> Result<(), MLError> {
|
||||
let out_dim_i32 = out_dim as i32;
|
||||
@@ -1232,6 +1274,13 @@ impl Drop for GpuAttention {
|
||||
///
|
||||
/// Uses CUBLAS_COMPUTE_32F_FAST_TF32 for tensor-core acceleration.
|
||||
/// Runs the heuristic ONCE at init time for optimal algorithm selection.
|
||||
///
|
||||
/// `epilogue` is `Some(epi)` when the GEMM should run a fused epilogue (e.g.
|
||||
/// `EPILOGUE_BIAS` for fused bias-add). The bias pointer is **not** baked
|
||||
/// into the descriptor at creation — callers (e.g. `lt_matmul_with_bias_ex`)
|
||||
/// set `CUBLASLT_MATMUL_DESC_BIAS_POINTER` per-call before the matmul. Only
|
||||
/// the static descriptor attributes (epilogue type, bias data type) are
|
||||
/// configured here so algo selection sees the final descriptor shape.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn create_attn_gemm_desc(
|
||||
lt_handle: cublaslt_sys::cublasLtHandle_t,
|
||||
@@ -1240,6 +1289,7 @@ fn create_attn_gemm_desc(
|
||||
m: i32, n: i32, k: i32,
|
||||
lda: i32, ldb: i32, ldc: i32,
|
||||
ws_size: usize,
|
||||
epilogue: Option<cublaslt_sys::cublasLtEpilogue_t>,
|
||||
label: &str,
|
||||
) -> Result<AttnGemmDesc, MLError> {
|
||||
let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
|
||||
@@ -1263,6 +1313,30 @@ fn create_attn_gemm_desc(
|
||||
std::mem::size_of::<i32>(),
|
||||
).map_err(|e| MLError::ModelError(format!("attn set TRANSB ({label}): {e:?}")))?;
|
||||
|
||||
// P5T5 Phase H — fused epilogue (e.g. BIAS) for forward projections.
|
||||
// The epilogue value influences algo validity, so it must be set
|
||||
// before `cublasLtMatmulAlgoCheck` runs (and so it is included in
|
||||
// `ShapeKey` below). Bias pointer is set per-call by the matmul
|
||||
// wrapper; only the bias data type is fixed at descriptor-creation
|
||||
// time.
|
||||
if let Some(epi) = epilogue {
|
||||
let epi_i32: i32 = epi as i32;
|
||||
cublaslt_result::set_matmul_desc_attribute(
|
||||
matmul_desc,
|
||||
cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_EPILOGUE,
|
||||
&epi_i32 as *const i32 as *const std::ffi::c_void,
|
||||
std::mem::size_of::<i32>(),
|
||||
).map_err(|e| MLError::ModelError(format!("attn set EPILOGUE ({label}): {e:?}")))?;
|
||||
|
||||
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!("attn set BIAS_DATA_TYPE ({label}): {e:?}")))?;
|
||||
}
|
||||
|
||||
// Matrix layouts -- physical storage dimensions (before transpose)
|
||||
let transa_is_n = transa_i32 == 0;
|
||||
let transb_is_n = transb_i32 == 0;
|
||||
@@ -1290,9 +1364,17 @@ fn create_attn_gemm_desc(
|
||||
// hardware-stable AlgoGetIds → AlgoInit → AlgoCheck loop. TF32
|
||||
// compute type preserved; selection is still graph-safe because
|
||||
// AlgoCheck validates workspace size and other prerequisites.
|
||||
let shape = super::cublas_algo_deterministic::ShapeKey::new(
|
||||
transa_i32, transb_i32, m, n, k, lda, ldb, ldc, ws_size,
|
||||
);
|
||||
// The epilogue value is folded into the cache key so different
|
||||
// epilogues (DEFAULT, BIAS, …) for the same matmul shape get
|
||||
// independent algo selections.
|
||||
let shape = match epilogue {
|
||||
Some(epi) => super::cublas_algo_deterministic::ShapeKey::with_epilogue(
|
||||
transa_i32, transb_i32, m, n, k, lda, ldb, ldc, epi, ws_size,
|
||||
),
|
||||
None => super::cublas_algo_deterministic::ShapeKey::new(
|
||||
transa_i32, transb_i32, m, n, k, lda, ldb, ldc, 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,
|
||||
@@ -1319,6 +1401,7 @@ fn create_attn_gemm_desc(
|
||||
transa = transa_i32, transb = transb_i32,
|
||||
m, n, k, lda, ldb, ldc,
|
||||
ws_needed = heuristic.workspaceSize,
|
||||
epilogue = epilogue.map(|e| e as i32).unwrap_or(1),
|
||||
"attn GEMM desc created (deterministic algorithm)"
|
||||
);
|
||||
|
||||
|
||||
@@ -353,6 +353,8 @@ Plan 4 Task 2c.3c.4 (2026-04-25): GRN backward chain wired through all 3 panic-g
|
||||
|
||||
Plan 4 Task 2c.3a (2026-04-24): GRN trunk param-tensor reshuffle (build-clean, runtime-broken intentionally). `compute_param_sizes()` flat layout reshuffled from 86 → 95 tensors: the 4 legacy trunk Linear tensors (`w_s1`, `b_s1`, `w_s2`, `b_s2` at indices [0..4)) are replaced with **13 GRN tensors** at indices [0..13) — 7 for h_s1 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `w_residual`, `gamma`, `beta`; SD→SH1 with `Linear_residual` GEMM since dimensions differ) plus 6 for h_s2 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `gamma`, `beta`; SH1→SH2 with identity residual since SH1==SH2). Every tensor index ≥ 4 in the OLD layout shifts +9 in the NEW layout. `NUM_WEIGHT_TENSORS` 86→95, `FIRST_ISV_TENSOR` 68→77. `layout_fingerprint_seed()` updated to mirror the new tensor names + positions; new `LAYOUT_FINGERPRINT_CURRENT = 0xcf3a24b0a1f70057` (was `0xa504d3c2f275b8af`). All 93 `padded_byte_offset` call sites + ancillary index references migrated in lockstep per `feedback_no_partial_refactor.md`. `xavier_init_params_buf` extended with init paths for the 13 new tensors: Linear matrices (`w_a`, `w_b`, `w_residual`) get Xavier; LayerNorm γ initialised to 1.0; β + biases stay zero. Trunk-forward / trunk-backward / IQN-trunk / ensemble-diversity-backward callers gated by explicit `panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, …")` at function entry — `BatchedForward::encoder_forward_only`, `BatchedForward::forward_target_raw`, `BatchedForward::forward_online_f32`, `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. This makes accidental runtime execution fail loudly rather than producing silent garbage from running legacy `Linear→ReLU→Linear` GEMMs against GRN-shaped param tensors. Spectral-norm descriptor (13 matrices) keeps slots [0]/[1] mapped to `w_a_h_s1`/`w_a_h_s2` (shapes match legacy W_s1/W_s2 exactly — the GRN's first Linear has the same shape as the original Linear), so the existing spectral-norm constraint transfers cleanly to the GRN's first Linear; Linear_b / Linear_residual are NOT yet spectral-normed (Task 2c.3c follow-up). **Smoke intentionally NOT run** — build-clean is the validation; runtime would hit the panic at the first encoder forward (the desired behaviour, no need to verify explicitly). Task 2c.3b swaps in the GRN forward (gpu_grn::GrnBlock::forward) at all panic-gated forward callers in place; Task 2c.3c does the same for backward callers and runs the smoke test. cargo check clean at 11 warnings (baseline preserved); cargo build compiles all 58 cubins. No new module / kernel / ISV slot in this commit — pure structural reshuffle + assertion gates.
|
||||
|
||||
Plan 5 Task 5 Phase H (2026-04-26): **cublasLt epilogue fusion for the 4 attention forward projections** — collapses each `cublasLtMatmul + add_bias_f32_kernel` pair into a single fused `cublasLtMatmul` call via `CUBLASLT_EPILOGUE_BIAS`. Targets the L40S deploy hot-spot where the per-epoch training phase is 99.7% of wall time (25.8 s of 25.9 s); attention forward runs 4 projections (Q, K, V, O) per training step, so this saves **4 kernel launches × N_steps** without changing the math. (1) `crates/ml/src/cuda_pipeline/gpu_attention.rs::create_attn_gemm_desc` extended with an `epilogue: Option<cublasLtEpilogue_t>` parameter — when `Some(BIAS)`, the descriptor is configured with `EPILOGUE = BIAS` + `BIAS_DATA_TYPE = CUDA_R_32F` at creation time; when `None`, the descriptor stays at `EPILOGUE_DEFAULT` (the 8 backward dW/dX GEMMs are unchanged). The `ShapeKey` passed to `get_matmul_algo_f32_tf32` is built via `with_epilogue` when set so the deterministic-algo cache keys on epilogue and produces an algo selection valid for the new descriptor — different epilogues for the same matmul shape get independent algo IDs. (2) New `lt_matmul_with_bias_ex` helper writes the per-call bias pointer via `set_matmul_desc_attribute(BIAS_POINTER, …)` immediately before each `cublasLtMatmul` (lightweight CPU write, no GPU sync; mirrors the existing pattern in `batched_forward.rs::sgemm_f32_fused_relu_bias`). The 4 forward projection sites in `forward(...)` switch from the prior `lt_matmul_ex(...)` + `launch_bias_add_ex(...)` pair to a single `lt_matmul_with_bias_ex(...)`. (3) Orphans pruned: `launch_bias_add_ex` and `launch_bias_add` deleted from `gpu_attention.rs` (their only callers were the 4 fused-away sites). The shared `add_bias_f32_kernel` retained — still used by `batched_forward.rs::launch_add_bias_f32_raw` (VSN Linear_2 logit output, no activation). (4) Determinism preserved: the deterministic-algo cache (`cublas_algo_deterministic.rs`) handles arbitrary epilogue values transparently — `ShapeKey` already had an `epilogue: i32` field, the `with_epilogue` constructor stamps it, and the `AlgoCheck` filter validates the descriptor against the epilogue before caching the algo, so first-call selection runs the full `AlgoGetIds → AlgoInit → AlgoCheck` loop with the new descriptor and subsequent calls reuse the cached `(types, shape)` algo. `CUBLAS_WORKSPACE_CONFIG=:4096:8` is unchanged; epilogue fusion is bit-deterministic when the algo is fixed. **Site #2 (DRELU_BGRAD on the trunk Linear→Bias→ReLU backward) deferred** — the trunk's only Linear→ReLU layer that pattern-matches DRELU_BGRAD is the value-FC head (forward `sgemm_f32_fused_relu_bias` writes `save_h_v`, backward calls `relu_mask` on the output of the value-output-layer dX GEMM). Wiring DRELU_BGRAD requires (a) switching value-FC forward from `EPILOGUE_RELU_BIAS` to `EPILOGUE_RELU_AUX_BIAS` and allocating a uint8 mask aux buffer of shape `[B, VH]` packed at 8 elements/byte, (b) plumbing the aux pointer through `BatchedForward → BatchedBackward`'s `backward_fc_layer` so the value-output-layer dX GEMM can set `EPILOGUE_AUX_POINTER` + `BIAS_POINTER = b_v1_grad`, (c) deleting the standalone `relu_mask` + bias-grad reduce kernels at the value-FC site. The forward-side aux-buffer plumbing crosses three modules and is non-trivial to wire deterministically alongside the GRN trunk's per-block partial-reduction scratch already plumbed through the same chain; landing it in the same commit as Site #1 risks regressing the determinism contract verified by the Phase G smoke. Defer to a follow-up phase that owns the aux-buffer wiring end-to-end. Site #1 alone is the bigger of the two opportunities — 4 launches saved per attention forward × every training step (target + online) — and is fully self-contained inside `gpu_attention.rs`. **Validation**: `cargo check --workspace` clean at 11 warnings (baseline preserved). Smoke (`cargo test -p ml --release --lib -- multi_fold_convergence --ignored --nocapture`, RTX 3050 Ti) completes all 3 folds; per-fold best train Sharpe within the noise band of the recent Phase G run (F0=2.36, F1=80.82, F2=92.11). Per-epoch wall-time deferred to next L40S deploy nsys profile (cluster log will show `STEP_TIMING` deltas; the 4-launch saving is the same on H100/L40S as on the local 3050 Ti, and cluster baseline at 25.8 s / epoch makes even a 5% saving worth ~1.3 s/epoch × 30 epochs × 6 folds × 5 seeds ≈ 1170 s total). No new module / kernel / ISV slot / param tensor / Orphan row. Files touched: `crates/ml/src/cuda_pipeline/gpu_attention.rs` (descriptor factory signature + 4 forward call sites + new `lt_matmul_with_bias_ex` helper + 2 orphan helpers deleted), `docs/dqn-wire-up-audit.md` (this entry). No fingerprint change.
|
||||
|
||||
Plan 5 Task 5 Phase F (2026-04-26): **compile-time fxcache schema fingerprint** — closes the L40S deploy-bug class where stale fxcache passed `FXCACHE_VERSION` validation despite incompatible feature semantics. Root cause of the original failure: `extract_ohlcv_features` column 0 changed from raw price → log-return without anyone bumping the manually-maintained `FXCACHE_VERSION` const, so the L40S PVC's older cache loaded clean and the trainer fed raw prices into the aux head expecting log-returns (`aux_next_bar_mse=2.587e7`). Fix has three pieces: (1) `crates/ml/build.rs::emit_feature_schema_hash()` runs unconditionally (before the existing CUDA-feature gate so non-CUDA builds also get the env var) and FNV-1a-hashes the raw bytes of the three schema-defining sources — `crates/ml/src/features/extraction.rs`, `crates/ml/src/fxcache.rs`, `crates/ml-core/src/state_layout.rs` — mixing in each file's relative path + length so renames / reorderings also bump the hash. Stable across rust versions and machines (FNV-1a, not `std::hash::DefaultHasher`). Emits `cargo:rustc-env=FEATURE_SCHEMA_HASH=<decimal_u64>` plus three `cargo:rerun-if-changed=` lines. (2) `crates/ml/src/fxcache.rs::FEATURE_SCHEMA_HASH` consumes the env var via `env!` + const `u64::from_str_radix(_, 10)` (const-stable since rust 1.83; workspace MSRV 1.85). `FxCacheHeader` grows a `feature_schema_hash: u64` field; header size 64→72 bytes; wire-format `FXCACHE_VERSION` bumped 5→6 to flag the layout change. `validate()` strict-checks the hash alongside magic/version/dims; mismatch bails with a descriptive error pointing at "source files defining feature extraction / state layout / fxcache format have changed since this cache was built." `precompute_features.rs:218` already deletes-and-regenerates on any `load_fxcache` Err, so the existing Argo `ensure-fxcache` step (and local users) recover automatically. (3) `FXCACHE_VERSION` doc comment now declares it tracks **wire-format** changes only — schema-level changes are tracked automatically by `FEATURE_SCHEMA_HASH`. Removes the manual ritual that broke the L40S deploy. **Cost**: cosmetic edits (whitespace, comments) to the three schema sources trigger one cache regen on next deploy (~5 min for full L40S dataset, ~40 s for local ES.FUT). Acceptable trade — false negatives (missed schema drift) are not. **Validation**: cargo check workspace clean at 11 warnings (baseline preserved). Local ES.FUT cache regen confirmed: existing v5 file rejected with `"Stale FxCache version: 5 (expected 6). Delete and regenerate."`, regenerated v6 cache loads clean on retry. Auto-detection verified: comment-only edit to `extraction.rs` line 1 changed emitted hash `5046469432341222878` → `7772630163018944575`; revert returned the hash deterministically to `5046469432341222878`. No new pip/cargo deps (FNV-1a is ~10 LOC stdlib). Files touched: `crates/ml/build.rs`, `crates/ml/src/fxcache.rs`, `docs/dqn-wire-up-audit.md` (this entry). No fingerprint change (LAYOUT_FINGERPRINT_CURRENT untouched — this is fxcache wire-format, not GPU param layout).
|
||||
|
||||
| Classification | Count |
|
||||
|
||||
Reference in New Issue
Block a user