From 45d8f4661cae30d8f273c311d5cdfd49ba3f447c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 17 Apr 2026 23:58:32 +0200 Subject: [PATCH] plan: rewrite Task 9 for child graph architecture Co-Authored-By: Claude Opus 4.6 (1M context) --- ...026-04-17-unified-cublas-training-graph.md | 2198 +++++++++++++++++ 1 file changed, 2198 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md diff --git a/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md b/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md new file mode 100644 index 000000000..e82120005 --- /dev/null +++ b/docs/superpowers/plans/2026-04-17-unified-cublas-training-graph.md @@ -0,0 +1,2198 @@ +# Unified cuBLAS Training Graph Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace fragmented graph capture with single unified cudarc graph. Rewrite IQL/IQN/Attention from per-sample CUDA kernels to batched cuBLAS GEMMs. Wire temporal ops inside unified graph. Target: <80s epochs from 685s. + +**Architecture:** Single cudarc begin_capture/end_capture captures spectral + trunk forward + temporal + loss + backward + aux (IQL/IQN/attention cuBLAS) + adam. IQL/IQN/Attention forward/backward use cublasLtMatmul with cached descriptors via SharedCublasHandle. Element-wise ops (SiLU, expectile loss) remain as precompiled cubin kernels. + +**Tech Stack:** Rust, cudarc, cublasLt (cublaslt-sys), CUDA cubins via build.rs + +--- + +## Task 1: IQL cuBLAS Rewrite + +**File:** `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` + +**Summary:** Replace per-sample `forward_loss_kernel` and `backward_per_sample_kernel` + `weight_grad_reduce_kernel` with batched cuBLAS GEMMs. Keep the same public interface (`train_value_step`, `gather_q_taken`, `compute_advantage_weights`, etc.). Replace internals only. + +### Step 1.1: Add cuBLAS GEMM descriptor cache to GpuIqlTrainer + +- [ ] Add new imports and types to the top of the file. +- [ ] Add GEMM descriptor fields to `GpuIqlTrainer` struct. +- [ ] Add new activation save buffers (h1_pre, h2_pre are the existing `save_pre1`/`save_pre2`; keep them). +- [ ] Add forward intermediate output buffers. + +**Add to imports at top of `gpu_iql_trainer.rs`:** + +```rust +use std::collections::HashMap; +use std::sync::Arc; + +use cudarc::cublaslt::result as cublaslt_result; +use cudarc::cublaslt::sys as cublaslt_sys; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use tracing::info; + +use crate::MLError; +use super::shared_cublas_handle::SharedCublasHandle; +``` + +**Add GEMM descriptor type (same pattern as `batched_forward.rs`):** + +```rust +/// Cached cuBLAS GEMM descriptor for IQL forward/backward. +/// Same structure as CachedGemmDesc in batched_forward.rs. +struct IqlGemmDesc { + matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t, + a_layout: cublaslt_sys::cublasLtMatrixLayout_t, + b_layout: cublaslt_sys::cublasLtMatrixLayout_t, + c_layout: cublaslt_sys::cublasLtMatrixLayout_t, + d_layout: cublaslt_sys::cublasLtMatrixLayout_t, + algo: cublaslt_sys::cublasLtMatmulAlgo_t, +} + +unsafe impl Send for IqlGemmDesc {} +unsafe impl Sync for IqlGemmDesc {} + +impl Drop for IqlGemmDesc { + fn drop(&mut self) { + unsafe { + let _ = cublaslt_result::destroy_matrix_layout(self.d_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.c_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.b_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.a_layout); + let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc); + } + } +} +``` + +**Add new fields to `GpuIqlTrainer` struct (replace old per-sample kernel fields):** + +```rust +pub struct GpuIqlTrainer { + config: GpuIqlConfig, + stream: Arc, + + // ── cuBLAS handle (shared with trunk) ────────────────────────── + shared_handle: Arc, + + // ── cuBLAS GEMM descriptors (cached at init) ────────────────── + /// h1_pre = W1^T @ states: M=H, N=B, K=state_dim, ldb=state_dim_padded + gemm_fwd_h1: IqlGemmDesc, + /// h2_pre = W2^T @ h1: M=H, N=B, K=H, ldb=H + gemm_fwd_h2: IqlGemmDesc, + /// v = W3^T @ h2: M=1, N=B, K=H, ldb=H + gemm_fwd_v: IqlGemmDesc, + /// dW3 = dv @ h2^T: M=1, N=H, K=B (TRANSA=N, TRANSB=T) + gemm_bwd_dw3: IqlGemmDesc, + /// dh2 = W3 @ dv: M=H, N=B, K=1 (TRANSA=N, TRANSB=N) + gemm_bwd_dh2: IqlGemmDesc, + /// dW2 = dh2_pre @ h1^T: M=H, N=H, K=B (TRANSA=N, TRANSB=T) + gemm_bwd_dw2: IqlGemmDesc, + /// dh1 = W2 @ dh2_pre: M=H, N=B, K=H (TRANSA=N, TRANSB=N) + gemm_bwd_dh1: IqlGemmDesc, + /// dW1 = dh1_pre @ states^T: M=H, N=state_dim, K=B (TRANSA=N, TRANSB=T) + gemm_bwd_dw1: IqlGemmDesc, + + // ── Element-wise cubin kernels ───────────────────────────────── + silu_fwd_kernel: CudaFunction, + silu_bwd_kernel: CudaFunction, + expectile_loss_kernel: CudaFunction, + bias_add_kernel: CudaFunction, + loss_reduce_kernel: CudaFunction, + grad_norm_phase1_kernel: CudaFunction, + grad_norm_phase2_kernel: CudaFunction, + adam_kernel: CudaFunction, + // Keep all non-training kernels unchanged: + forward_kernel: CudaFunction, + gather_q_taken_kernel: CudaFunction, + advantage_weight_kernel: CudaFunction, + modulate_td_kernel: CudaFunction, + adv_variance_kernel: CudaFunction, + per_sample_support_kernel: CudaFunction, + branch_advantage_kernel: CudaFunction, + expectile_gap_kernel: CudaFunction, + gap_mean_kernel: CudaFunction, + per_sample_epsilon_kernel: CudaFunction, + support_floor_kernel: CudaFunction, + adv_sigma_ema_kernel: CudaFunction, + + // ── V network parameters (flat f32 on GPU) ───────────────────── + // Layout: W1[H*SD] + b1[H] + W2[H*H] + b2[H] + W3[H] + b3[1] + params_buf: CudaSlice, + + // ── Adam optimizer state ──────────────────────────────────────── + m_buf: CudaSlice, + v_buf: CudaSlice, + grad_buf: CudaSlice, // [total_params] reduced gradients + grad_norm_buf: CudaSlice, + grad_norm_partials: CudaSlice, + + // ── cuBLAS forward intermediate buffers ───────────────────────── + h1_pre_buf: CudaSlice, // [H, B] pre-activation layer 1 (GEMM output) + h1_buf: CudaSlice, // [H, B] post-SiLU layer 1 + h2_pre_buf: CudaSlice, // [H, B] pre-activation layer 2 + h2_buf: CudaSlice, // [H, B] post-SiLU layer 2 + v_pre_buf: CudaSlice, // [1, B] output layer pre-bias + + // ── cuBLAS backward intermediate buffers ──────────────────────── + dv_buf: CudaSlice, // [1, B] d_expectile_loss output + dh2_buf: CudaSlice, // [H, B] dh2 (pre silu_bwd) + dh2_pre_buf: CudaSlice, // [H, B] dh2_pre (post silu_bwd) + dh1_buf: CudaSlice, // [H, B] dh1 + dh1_pre_buf: CudaSlice, // [H, B] dh1_pre + + // ── Output buffers (same as before) ──────────────────────────── + v_out_buf: CudaSlice, + loss_buf: CudaSlice, + total_loss_buf: CudaSlice, + q_taken_buf: CudaSlice, + advantage_weights_buf: CudaSlice, + + // ── Integration buffers (unchanged) ──────────────────────────── + adv_stats_buf: CudaSlice, + adv_sigma_ema_buf: CudaSlice, + readiness_buf: CudaSlice, + p5_state_buf: CudaSlice, + per_sample_support_buf: CudaSlice, + branch_scales_buf: CudaSlice, + expectile_gap_buf: CudaSlice, + gap_mean_buf: CudaSlice, + per_sample_epsilon_buf: CudaSlice, + + // ── Training state ───────────────────────────────────────────── + adam_step: i32, + t_buf: CudaSlice, + total_params: usize, + grad_norm_blocks: usize, + state_dim_padded: usize, +} +``` + +**Removed fields** (compared to current struct): +- `forward_loss_kernel` -- replaced by cuBLAS GEMMs + `silu_fwd_kernel` + `expectile_loss_kernel` +- `backward_per_sample_kernel` -- replaced by cuBLAS backward GEMMs + `silu_bwd_kernel` +- `weight_grad_reduce_kernel` -- no longer needed (cuBLAS backward produces [H, K] directly) +- `grads_per_sample: CudaSlice` -- was [tile * total_params], no longer needed +- `grad_tile_size: usize` -- no longer needed +- `save_pre1`, `save_pre2`, `save_h1`, `save_h2` -- replaced by `h1_pre_buf`, `h1_buf`, `h2_pre_buf`, `h2_buf` (same data, different naming, col-major [H, B] instead of [B, H]) + +### Step 1.2: Create GEMM descriptors in constructor + +- [ ] Modify `GpuIqlTrainer::new()` to accept `shared_handle: Arc` instead of bare `stream: Arc`. +- [ ] Create 8 GEMM descriptors using the same `create_cached_fwd_gemm_desc` pattern from `batched_forward.rs`. +- [ ] Load new element-wise cubin kernels (silu_fwd, silu_bwd, expectile_loss, bias_add). +- [ ] Allocate new intermediate buffers. + +**Constructor signature change:** + +```rust +pub fn new( + shared_handle: Arc, + config: GpuIqlConfig, +) -> Result { + let stream = Arc::clone(&shared_handle.stream); + let total_params = config.total_params(); + let b = config.batch_size; + let h = config.value_hidden_dim; + let sd = config.state_dim; + let state_dim_padded = (sd + 127) & !127; + let lt_handle = shared_handle.lt_handle.0; + let lt_ws_size = shared_handle.lt_workspace_size; + + // Create 3 forward GEMM descriptors + let gemm_fwd_h1 = create_iql_gemm_desc(lt_handle, h, b, sd, state_dim_padded, lt_ws_size, + 1, 0, "iql_fwd_h1")?; // TRANSA=T, TRANSB=N + let gemm_fwd_h2 = create_iql_gemm_desc(lt_handle, h, b, h, h, lt_ws_size, + 1, 0, "iql_fwd_h2")?; + let gemm_fwd_v = create_iql_gemm_desc(lt_handle, 1, b, h, h, lt_ws_size, + 1, 0, "iql_fwd_v")?; + + // Create 5 backward GEMM descriptors + // dW3 = dv[1,B] @ h2^T[B,H] -> [1,H] TRANSA=N, TRANSB=T + let gemm_bwd_dw3 = create_iql_gemm_desc_nt(lt_handle, 1, h, b, 1, h, lt_ws_size, "iql_bwd_dw3")?; + // dh2 = W3[H,1]^T inverted: dh2[H,B] = W3_col[H,1] @ dv[1,B] TRANSA=N, TRANSB=N + let gemm_bwd_dh2 = create_iql_gemm_desc(lt_handle, h, b, 1, 1, lt_ws_size, + 0, 0, "iql_bwd_dh2")?; + // dW2 = dh2_pre[H,B] @ h1^T[B,H] -> [H,H] + let gemm_bwd_dw2 = create_iql_gemm_desc_nt(lt_handle, h, h, b, h, h, lt_ws_size, "iql_bwd_dw2")?; + // dh1 = W2[H,H] @ dh2_pre[H,B] -> [H,B] + let gemm_bwd_dh1 = create_iql_gemm_desc(lt_handle, h, b, h, h, lt_ws_size, + 0, 0, "iql_bwd_dh1")?; + // dW1 = dh1_pre[H,B] @ states^T[B,SD] -> [H,SD] + let gemm_bwd_dw1 = create_iql_gemm_desc_nt(lt_handle, h, sd, b, h, state_dim_padded, lt_ws_size, "iql_bwd_dw1")?; + + // Load element-wise cubin kernels + let iql_cubin = include_bytes!(concat!(env!("OUT_DIR"), "/iql_value_kernel.cubin")); + let context = stream.context(); + let module = context.load_cubin(iql_cubin.to_vec()) + .map_err(|e| MLError::ModelError(format!("iql cubin: {e}")))?; + let silu_fwd_kernel = module.load_function("iql_silu_fwd") + .map_err(|e| MLError::ModelError(format!("iql_silu_fwd: {e}")))?; + let silu_bwd_kernel = module.load_function("iql_silu_bwd") + .map_err(|e| MLError::ModelError(format!("iql_silu_bwd: {e}")))?; + let expectile_loss_kernel = module.load_function("iql_expectile_loss") + .map_err(|e| MLError::ModelError(format!("iql_expectile_loss: {e}")))?; + let bias_add_kernel = module.load_function("iql_bias_add") + .map_err(|e| MLError::ModelError(format!("iql_bias_add: {e}")))?; + // ... keep loading existing kernels (loss_reduce, grad_norm, adam, etc.) ... + + // Allocate cuBLAS intermediate buffers (col-major [out_dim, B]) + let h1_pre_buf = alloc_f32(&stream, h * b, "iql_h1_pre")?; + let h1_buf = alloc_f32(&stream, h * b, "iql_h1")?; + let h2_pre_buf = alloc_f32(&stream, h * b, "iql_h2_pre")?; + let h2_buf = alloc_f32(&stream, h * b, "iql_h2")?; + let v_pre_buf = alloc_f32(&stream, b, "iql_v_pre")?; + let dv_buf = alloc_f32(&stream, b, "iql_dv")?; + let dh2_buf = alloc_f32(&stream, h * b, "iql_dh2")?; + let dh2_pre_buf = alloc_f32(&stream, h * b, "iql_dh2_pre")?; + let dh1_buf = alloc_f32(&stream, h * b, "iql_dh1")?; + let dh1_pre_buf = alloc_f32(&stream, h * b, "iql_dh1_pre")?; + + // ... rest of constructor (params init, adam state, integration buffers unchanged) ... +} +``` + +**GEMM descriptor creation functions (add at bottom of file):** + +```rust +/// Create a GEMM descriptor: C = alpha * op(A) @ op(B) + beta * C +/// transa/transb: 0=N, 1=T +fn create_iql_gemm_desc( + lt_handle: cublaslt_sys::cublasLtHandle_t, + m: usize, n: usize, k: usize, ldb: usize, ws_size: usize, + transa: i32, transb: i32, label: &str, +) -> 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!("{label} MatmulDescCreate: {e:?}")))?; + + 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!("{label} set TRANSA: {e:?}")))?; + + 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!("{label} set TRANSB: {e:?}")))?; + + // Physical A layout: [K, M] col-major when TRANSA=T, [M, K] when TRANSA=N + let (a_rows, a_cols, a_ld) = if transa == 1 { (k, m, k) } else { (m, k, m) }; + let a_layout = cublaslt_result::create_matrix_layout( + f32_type, a_rows as u64, a_cols as u64, a_ld as i64, + ).map_err(|e| MLError::ModelError(format!("{label} A layout: {e:?}")))?; + + // Physical B layout: [K, N] col-major when TRANSB=N, [N, K] when TRANSB=T + let (b_rows, b_cols, b_ld) = if transb == 1 { (n, k, ldb) } else { (k, n, ldb) }; + let b_layout = cublaslt_result::create_matrix_layout( + f32_type, b_rows as u64, b_cols as u64, b_ld as i64, + ).map_err(|e| MLError::ModelError(format!("{label} B layout: {e:?}")))?; + + let c_layout = cublaslt_result::create_matrix_layout( + f32_type, m as u64, n as u64, m as i64, + ).map_err(|e| MLError::ModelError(format!("{label} C layout: {e:?}")))?; + + let d_layout = cublaslt_result::create_matrix_layout( + f32_type, m as u64, n as u64, m as i64, + ).map_err(|e| MLError::ModelError(format!("{label} D layout: {e:?}")))?; + + let matmul_pref = cublaslt_result::create_matmul_pref() + .map_err(|e| MLError::ModelError(format!("{label} MatmulPrefCreate: {e:?}")))?; + cublaslt_result::set_matmul_pref_attribute( + matmul_pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| { + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + MLError::ModelError(format!("{label} set pref ws: {e:?}")) + })?; + + let heuristic = cublaslt_result::get_matmul_algo_heuristic( + lt_handle, matmul_desc, a_layout, b_layout, c_layout, d_layout, matmul_pref, + ); + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + let algo = heuristic.map_err(|e| MLError::ModelError(format!("{label} heuristic: {e:?}")))?.algo; + + tracing::info!(%label, m, n, k, ldb, "IQL GEMM desc created"); + Ok(IqlGemmDesc { matmul_desc, a_layout, b_layout, c_layout, d_layout, algo }) + } +} + +/// Create GEMM desc for dW = A[M,K] @ B^T[K,N] where B has ldb leading dim. +fn create_iql_gemm_desc_nt( + lt_handle: cublaslt_sys::cublasLtHandle_t, + m: usize, n: usize, k: usize, lda: usize, ldb: usize, ws_size: usize, label: &str, +) -> Result { + // dW = A @ B^T: TRANSA=N, TRANSB=T + create_iql_gemm_desc(lt_handle, m, n, k, ldb, ws_size, 0, 1, label) +} +``` + +### Step 1.3: Rewrite `train_value_step()` to use cuBLAS + +- [ ] Replace the per-sample forward+loss kernel with 3 cuBLAS GEMMs + element-wise kernels. +- [ ] Replace tiled backward with 5 cuBLAS backward GEMMs + activation backward kernels. +- [ ] Keep loss_reduce, grad_norm, Adam unchanged. + +**New `train_value_step()`:** + +```rust +pub fn train_value_step( + &mut self, + states_f32: &CudaSlice, +) -> Result<(), MLError> { + let b = self.config.batch_size; + let h = self.config.value_hidden_dim; + let sd = self.config.state_dim; + let total_params_i32 = self.total_params as i32; + let batch_size_i32 = b as i32; + let expectile_tau = self.config.expectile_tau; + + let lt_handle = self.shared_handle.lt_handle.0; + let lt_ws_ptr = self.shared_handle.lt_workspace_ptr; + let lt_ws_size = self.shared_handle.lt_workspace_size; + let cu_stream = self.stream.cu_stream() as *mut cublaslt_sys::CUstream_st; + + let alpha: f32 = 1.0; + let beta: f32 = 0.0; + + // Compute param offsets (same layout as current flat buffer) + let w1_off = 0; + let b1_off = h * sd; + let w2_off = b1_off + h; + let b2_off = w2_off + h * h; + let w3_off = b2_off + h; + let b3_off = w3_off + h; + let f32_sz = std::mem::size_of::(); + + // Raw pointers into params_buf + let w1_ptr = self.params_buf.raw_ptr() + (w1_off * f32_sz) as u64; + let b1_ptr = self.params_buf.raw_ptr() + (b1_off * f32_sz) as u64; + let w2_ptr = self.params_buf.raw_ptr() + (w2_off * f32_sz) as u64; + let b2_ptr = self.params_buf.raw_ptr() + (b2_off * f32_sz) as u64; + let w3_ptr = self.params_buf.raw_ptr() + (w3_off * f32_sz) as u64; + let b3_ptr = self.params_buf.raw_ptr() + (b3_off * f32_sz) as u64; + + // ── Forward pass (3 GEMMs + 2 SiLU + 1 bias_add) ── + + // 1. h1_pre[H,B] = W1^T @ states: cublasLtMatmul + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_fwd_h1.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w1_ptr as *const std::ffi::c_void, // A = W1 + self.gemm_fwd_h1.a_layout, + states_f32.raw_ptr() as *const std::ffi::c_void, // B = states + self.gemm_fwd_h1.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + self.h1_pre_buf.raw_ptr() as *mut std::ffi::c_void, // C + self.gemm_fwd_h1.c_layout, + self.h1_pre_buf.raw_ptr() as *mut std::ffi::c_void, // D = C + self.gemm_fwd_h1.d_layout, + &self.gemm_fwd_h1.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + // bias_add: h1_pre += b1 (broadcast along batch dimension) + let n_h1 = (h * b) as i32; + let h_i32 = h as i32; + let blocks_h1 = ((h * b + 255) / 256) as u32; + unsafe { + self.stream.launch_builder(&self.bias_add_kernel) + .arg(&self.h1_pre_buf) // in-place + .arg(&b1_ptr) + .arg(&h_i32) // out_dim (repeating pattern) + .arg(&n_h1) // total elements + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_add h1: {e}")))?; + } + // SiLU: h1 = silu(h1_pre), saves h1_pre for backward + unsafe { + self.stream.launch_builder(&self.silu_fwd_kernel) + .arg(&self.h1_pre_buf) + .arg(&mut self.h1_buf) + .arg(&n_h1) + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL silu_fwd h1: {e}")))?; + } + + // 2. h2_pre[H,B] = W2^T @ h1 + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_fwd_h2.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w2_ptr as *const std::ffi::c_void, + self.gemm_fwd_h2.a_layout, + self.h1_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_fwd_h2.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + self.h2_pre_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_fwd_h2.c_layout, + self.h2_pre_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_fwd_h2.d_layout, + &self.gemm_fwd_h2.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + // bias_add + SiLU for layer 2 (same pattern as layer 1) + unsafe { + self.stream.launch_builder(&self.bias_add_kernel) + .arg(&self.h2_pre_buf) + .arg(&b2_ptr) + .arg(&h_i32) + .arg(&n_h1) + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_add h2: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.silu_fwd_kernel) + .arg(&self.h2_pre_buf) + .arg(&mut self.h2_buf) + .arg(&n_h1) + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL silu_fwd h2: {e}")))?; + } + + // 3. v_pre[1,B] = W3^T @ h2 + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_fwd_v.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w3_ptr as *const std::ffi::c_void, + self.gemm_fwd_v.a_layout, + self.h2_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_fwd_v.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + self.v_out_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_fwd_v.c_layout, + self.v_out_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_fwd_v.d_layout, + &self.gemm_fwd_v.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + // Add output bias (b3, scalar) and compute v_out + let n_v = b as i32; + let one_i32 = 1_i32; + let v_blocks = ((b + 255) / 256) as u32; + unsafe { + self.stream.launch_builder(&self.bias_add_kernel) + .arg(&self.v_out_buf) + .arg(&b3_ptr) + .arg(&one_i32) + .arg(&n_v) + .launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_add v: {e}")))?; + } + + // 4. Expectile loss + dv: combined kernel + // loss[B] = |tau - 1(u<0)| * u^2 where u = q_taken - v_out + // dv[B] = d(loss)/d(v_out) + unsafe { + self.stream.launch_builder(&self.expectile_loss_kernel) + .arg(&self.v_out_buf) + .arg(&self.q_taken_buf) + .arg(&expectile_tau) + .arg(&mut self.loss_buf) + .arg(&mut self.dv_buf) + .arg(&batch_size_i32) + .launch(LaunchConfig { grid_dim: (v_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL expectile_loss: {e}")))?; + } + + // ── Backward pass (5 GEMMs + 2 SiLU backward) ── + + // Zero grad_buf before accumulating weight gradients + self.stream.memset_zeros(&mut self.grad_buf) + .map_err(|e| MLError::ModelError(format!("IQL zero grad_buf: {e}")))?; + + // inv_batch for 1/N mean reduction + let inv_batch: f32 = 1.0 / b as f32; + let beta_zero: f32 = 0.0; + + // dW3[1,H] = dv[1,B] @ h2^T[B,H] (alpha=1/B for mean reduction) + let dw3_ptr = self.grad_buf.raw_ptr() + (w3_off * f32_sz) as u64; + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_bwd_dw3.matmul_desc, + &inv_batch as *const f32 as *const std::ffi::c_void, + self.dv_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw3.a_layout, + self.h2_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw3.b_layout, + &beta_zero as *const f32 as *const std::ffi::c_void, + dw3_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw3.c_layout, + dw3_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw3.d_layout, + &self.gemm_bwd_dw3.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // db3: sum(dv) / B -- scalar bias gradient (launch single-block reduce kernel) + // (use existing loss_reduce pattern to sum dv into grad_buf[b3_off]) + let db3_ptr = self.grad_buf.raw_ptr() + (b3_off * f32_sz) as u64; + // ... bias gradient reduce kernel ... + + // dh2[H,B] = W3[H,1] @ dv[1,B] + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_bwd_dh2.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w3_ptr as *const std::ffi::c_void, + self.gemm_bwd_dh2.a_layout, + self.dv_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dh2.b_layout, + &beta_zero as *const f32 as *const std::ffi::c_void, + self.dh2_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_bwd_dh2.c_layout, + self.dh2_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_bwd_dh2.d_layout, + &self.gemm_bwd_dh2.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // dh2_pre = dh2 * d_silu(h2_pre) -- element-wise + unsafe { + self.stream.launch_builder(&self.silu_bwd_kernel) + .arg(&self.h2_pre_buf) // pre-activation (saved from forward) + .arg(&self.dh2_buf) // upstream gradient + .arg(&mut self.dh2_pre_buf) // output + .arg(&n_h1) + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL silu_bwd h2: {e}")))?; + } + + // dW2[H,H] = dh2_pre[H,B] @ h1^T[B,H] (alpha=1/B) + let dw2_ptr = self.grad_buf.raw_ptr() + (w2_off * f32_sz) as u64; + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_bwd_dw2.matmul_desc, + &inv_batch as *const f32 as *const std::ffi::c_void, + self.dh2_pre_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw2.a_layout, + self.h1_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw2.b_layout, + &beta_zero as *const f32 as *const std::ffi::c_void, + dw2_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw2.c_layout, + dw2_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw2.d_layout, + &self.gemm_bwd_dw2.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // dh1[H,B] = W2[H,H] @ dh2_pre[H,B] + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_bwd_dh1.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w2_ptr as *const std::ffi::c_void, + self.gemm_bwd_dh1.a_layout, + self.dh2_pre_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dh1.b_layout, + &beta_zero as *const f32 as *const std::ffi::c_void, + self.dh1_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_bwd_dh1.c_layout, + self.dh1_buf.raw_ptr() as *mut std::ffi::c_void, + self.gemm_bwd_dh1.d_layout, + &self.gemm_bwd_dh1.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // dh1_pre = dh1 * d_silu(h1_pre) + unsafe { + self.stream.launch_builder(&self.silu_bwd_kernel) + .arg(&self.h1_pre_buf) + .arg(&self.dh1_buf) + .arg(&mut self.dh1_pre_buf) + .arg(&n_h1) + .launch(LaunchConfig { grid_dim: (blocks_h1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL silu_bwd h1: {e}")))?; + } + + // dW1[H,SD] = dh1_pre[H,B] @ states^T[B,SD] (alpha=1/B) + let dw1_ptr = self.grad_buf.raw_ptr() + (w1_off * f32_sz) as u64; + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.gemm_bwd_dw1.matmul_desc, + &inv_batch as *const f32 as *const std::ffi::c_void, + self.dh1_pre_buf.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw1.a_layout, + states_f32.raw_ptr() as *const std::ffi::c_void, + self.gemm_bwd_dw1.b_layout, + &beta_zero as *const f32 as *const std::ffi::c_void, + dw1_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw1.c_layout, + dw1_ptr as *mut std::ffi::c_void, + self.gemm_bwd_dw1.d_layout, + &self.gemm_bwd_dw1.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // Bias gradients: db1, db2 use column-sum reduce kernels + // (sum along batch dimension of dh1_pre, dh2_pre respectively) + // ... (launch iql_bias_grad_reduce kernel for each bias) ... + + // ── Loss reduce + grad norm + Adam (unchanged from current code) ── + // 4. Loss reduce (deterministic sequential sum) + unsafe { + self.stream.launch_builder(&self.loss_reduce_kernel) + .arg(&self.loss_buf) + .arg(&mut self.total_loss_buf) + .arg(&batch_size_i32) + .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("IQL loss_reduce: {e}")))?; + } + + // 5-6. Grad norm phase 1+2 (same as current) + let norm_blocks = self.grad_norm_blocks; + // ... (identical to current code) ... + + // 7. Adam (same as current) + self.adam_step += 1; + // ... (identical to current code) ... + + Ok(()) +} +``` + +### Step 1.4: Update call sites + +- [ ] Update `FusedTrainingCtx::new()` in `fused_training.rs` to pass `shared_handle` to `GpuIqlTrainer::new()`. +- [ ] Update both `gpu_iql` and `gpu_iql_low` construction. + +**File:** `crates/ml/src/trainers/dqn/fused_training.rs` + +Find the lines where `GpuIqlTrainer::new()` is called (in `FusedTrainingCtx::new()`) and add the `shared_handle` parameter: + +```rust +// Before: +let gpu_iql = GpuIqlTrainer::new(Arc::clone(&stream), iql_config)?; +// After: +let gpu_iql = GpuIqlTrainer::new(Arc::clone(&shared_handle), iql_config)?; +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 2: IQL Element-wise cubin Kernels + +**File:** `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` + +**Summary:** Add 4 new element-wise CUDA kernel functions to the existing .cu file. These replace the per-sample matrix multiplication logic that is now handled by cuBLAS. Keep ALL existing kernels (forward_only, gather_q_taken, advantage_weight, etc.) that are NOT related to training forward/backward. + +### Step 2.1: Add SiLU forward kernel + +- [ ] Add `iql_silu_fwd` kernel to `iql_value_kernel.cu`. + +```cuda +/* ------------------------------------------------------------------ */ +/* Element-wise SiLU forward (used after cuBLAS GEMM) */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iql_silu_fwd( + const float* __restrict__ x, // [N] input (pre-activation) + float* __restrict__ out, // [N] output (post-activation) + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + float val = x[idx]; + out[idx] = val / (1.0f + expf(-val)); + } +} +``` + +### Step 2.2: Add SiLU backward kernel + +- [ ] Add `iql_silu_bwd` kernel. + +```cuda +/* ------------------------------------------------------------------ */ +/* Element-wise SiLU backward */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iql_silu_bwd( + const float* __restrict__ x, // [N] saved pre-activation + const float* __restrict__ d_out, // [N] upstream gradient + float* __restrict__ d_x, // [N] output gradient + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + float val = x[idx]; + float s = 1.0f / (1.0f + expf(-val)); + float grad = s * (1.0f + val * (1.0f - s)); + d_x[idx] = d_out[idx] * grad; + } +} +``` + +### Step 2.3: Add expectile loss + gradient kernel + +- [ ] Add `iql_expectile_loss` kernel that computes both loss and dv in one pass. + +```cuda +/* ------------------------------------------------------------------ */ +/* Expectile loss + gradient (fused) */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iql_expectile_loss( + const float* __restrict__ v_out, // [B] V(s) predictions + const float* __restrict__ q_taken, // [B] Q(s,a) targets + float tau, // expectile parameter + float* __restrict__ loss_out, // [B] per-sample loss + float* __restrict__ dv_out, // [B] dL/dv per sample + int B +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B) { + float v = v_out[idx]; + float q = q_taken[idx]; + float u = q - v; + float weight = (u >= 0.0f) ? tau : (1.0f - tau); + loss_out[idx] = weight * u * u; + // dL/dv = -2 * weight * u (chain rule: d/dv of weight*(q-v)^2) + dv_out[idx] = -2.0f * weight * u; + } +} +``` + +### Step 2.4: Add bias-add kernel + +- [ ] Add `iql_bias_add` kernel for broadcasting bias along batch dimension. + +```cuda +/* ------------------------------------------------------------------ */ +/* Bias add: in-place x[i] += bias[i % out_dim] */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iql_bias_add( + float* __restrict__ x, // [out_dim, B] col-major, modified in-place + const float* __restrict__ bias, // [out_dim] + int out_dim, // repeating pattern length + int N // total elements = out_dim * B +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + x[idx] += bias[idx % out_dim]; + } +} +``` + +### Step 2.5: Add bias gradient reduce kernel + +- [ ] Add `iql_bias_grad_reduce` kernel for computing bias gradients from activation gradients. + +```cuda +/* ------------------------------------------------------------------ */ +/* Bias gradient: sum columns of d_pre[out_dim, B] -> db[out_dim] */ +/* Deterministic: each thread handles one output neuron. */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iql_bias_grad_reduce( + const float* __restrict__ d_pre, // [out_dim, B] col-major + float* __restrict__ d_bias, // [out_dim] output + int out_dim, + int B, + float inv_batch // 1.0 / B for mean reduction +) { + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j < out_dim) { + float sum = 0.0f; + for (int b = 0; b < B; b++) { + sum += d_pre[j + b * out_dim]; // col-major stride + } + d_bias[j] = sum * inv_batch; + } +} +``` + +### Step 2.6: Verify build.rs compiles new kernels + +- [ ] Verify the existing `build.rs` compiles `iql_value_kernel.cu` and the new kernel names are exported in the cubin. + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +``` + +--- + +## Task 3: IQN cuBLAS Rewrite + +**File:** `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` + +**Summary:** Replace per-sample `forward_loss_kernel`, `backward_kernel`, and `weight_grad_reduce_kernel` with batched cuBLAS GEMMs. Keep the same public interface (`train_iqn_step_gpu`, `execute_training_pipeline`). + +### Step 3.1: Add GEMM descriptors to GpuIqnHead + +- [ ] Add `shared_handle: Arc` field. +- [ ] Add `IqnGemmDesc` type (same pattern as IQL). +- [ ] Create GEMM descriptors for IQN-specific shapes: + - Embedding forward: `M=H, N=B*Q, K=embed_dim, ldb=embed_dim` (TRANSA=T, TRANSB=N) + - Per-branch forward (4x): `M=branch_size, N=B*Q, K=H, ldb=H` + - Embedding backward dW: `M=H, N=embed_dim, K=B*Q` (TRANSA=N, TRANSB=T) + - Per-branch backward dW (4x): `M=branch_size, N=H, K=B*Q` + - Per-branch backward dX (4x): `M=H, N=B*Q, K=branch_size` + +```rust +// GpuIqnConfig dimensions: +// H = hidden_dim = 256 +// D = embed_dim = 64 +// Q = num_quantiles = 32 +// B = batch_size +// b0..b3 = branch_0..3_size = 3 each + +// Forward GEMMs: +// embed: [H, B*Q] = W_embed^T[H,D] @ cos_features[D, B*Q] +gemm_fwd_embed: IqnGemmDesc, // M=H, N=B*Q, K=D +// branch_k: [bk, B*Q] = W_bk^T[bk,H] @ combined[H, B*Q] +gemm_fwd_branch: [IqnGemmDesc; 4], // M=bk, N=B*Q, K=H + +// Backward GEMMs: +// dW_embed: [H, D] = d_embed[H, B*Q] @ cos^T[B*Q, D] +gemm_bwd_dw_embed: IqnGemmDesc, +// dW_bk: [bk, H] = dq_bk[bk, B*Q] @ combined^T[B*Q, H] +gemm_bwd_dw_branch: [IqnGemmDesc; 4], +// d_combined_bk: [H, B*Q] = W_bk[H,bk] @ dq_bk[bk, B*Q] +gemm_bwd_dx_branch: [IqnGemmDesc; 4], +``` + +### Step 3.2: Add intermediate buffers + +- [ ] Add cuBLAS output buffers for each GEMM output. + +```rust +// Forward intermediates +embed_pre_buf: CudaSlice, // [H, B*Q] embedding GEMM output (before ReLU) +embed_buf: CudaSlice, // [H, B*Q] post-ReLU +combined_buf: CudaSlice, // [H, B*Q] h_s2 * sigmoid(embed) +branch_logits_bufs: [CudaSlice; 4], // [bk, B*Q] per-branch logits + +// Backward intermediates +d_combined_buf: CudaSlice, // [H, B*Q] accumulated dX from all branches +d_embed_pre_buf: CudaSlice, // [H, B*Q] gradient before ReLU backward +``` + +### Step 3.3: Rewrite `execute_training_pipeline()` forward path + +- [ ] Replace `forward_loss_kernel` launch with cuBLAS GEMMs + element-wise kernels. + +The forward path becomes: +1. `embed_pre = cublasLtMatmul(W_embed^T, cos_features)` -- one GEMM +2. `embed = relu(embed_pre) + bias_embed` -- cubin kernel +3. `combined = h_s2_tiled * sigmoid(embed)` -- `iqn_hadamard_sigmoid` cubin kernel +4. For each branch k: `logits_bk = cublasLtMatmul(W_bk^T, combined) + bias_bk` -- 4 GEMMs +5. `loss = iqn_quantile_huber_loss(logits, target_logits, taus, actions)` -- cubin kernel + +### Step 3.4: Rewrite backward path + +- [ ] Replace `backward_kernel` + `weight_grad_reduce_kernel` with cuBLAS backward GEMMs. + +The backward path becomes: +1. For each branch k: `dW_bk = dq_bk @ combined^T` (alpha=1/B) +2. For each branch k: `d_combined += W_bk @ dq_bk` +3. `d_embed = d_combined * h_s2_tiled * d_sigmoid(embed) * d_relu(embed_pre)` -- cubin +4. `dW_embed = d_embed @ cos^T` (alpha=1/B) +5. Bias gradients via column-sum reduce kernels + +### Step 3.5: Update constructor signature + +- [ ] Change `GpuIqnHead::new()` to accept `shared_handle: Arc`. + +```rust +pub fn new( + shared_handle: Arc, + config: GpuIqnConfig, +) -> Result { + let stream = Arc::clone(&shared_handle.stream); + // ... create GEMM descriptors, load cubin kernels, allocate buffers ... +} +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 4: IQN Element-wise cubin Kernels + +**File:** `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` + +**Summary:** Add element-wise CUDA kernels for IQN operations that are not matrix multiplications. + +### Step 4.1: Add hadamard-sigmoid kernel + +- [ ] Add `iqn_hadamard_sigmoid` kernel. + +```cuda +/* ------------------------------------------------------------------ */ +/* Hadamard product with sigmoid: out = h_s2_tiled * sigmoid(embed) */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iqn_hadamard_sigmoid( + const float* __restrict__ h_s2_tiled, // [H, B*Q] h_s2 replicated Q times + const float* __restrict__ embed, // [H, B*Q] post-ReLU embedding + float* __restrict__ out, // [H, B*Q] combined + int N // H * B * Q +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + float sig = 1.0f / (1.0f + expf(-embed[idx])); + out[idx] = h_s2_tiled[idx] * sig; + } +} +``` + +### Step 4.2: Add ReLU forward/backward kernels + +- [ ] Add `iqn_relu_fwd` and `iqn_relu_bwd` kernels. + +```cuda +extern "C" __global__ +void iqn_relu_fwd( + const float* __restrict__ x, + float* __restrict__ out, + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + out[idx] = fmaxf(x[idx], 0.0f); + } +} + +extern "C" __global__ +void iqn_relu_bwd( + const float* __restrict__ x, // saved pre-activation + const float* __restrict__ d_out, // upstream gradient + float* __restrict__ d_x, // output gradient + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + d_x[idx] = (x[idx] > 0.0f) ? d_out[idx] : 0.0f; + } +} +``` + +### Step 4.3: Add quantile Huber loss kernel + +- [ ] Add `iqn_quantile_huber_loss` kernel that computes loss + per-branch dq gradients. + +```cuda +/* ------------------------------------------------------------------ */ +/* Quantile Huber loss: per-sample, per-quantile, per-branch */ +/* Replaces the fused forward+loss portion of iqn_forward_loss_kernel */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void iqn_quantile_huber_loss( + const float* __restrict__ q_online, // [total_branch_actions, B*Q] online Qs + const float* __restrict__ q_target, // [total_branch_actions, B*Q] target Qs + const float* __restrict__ taus, // [B*Q] quantile values + const int* __restrict__ actions, // [B, 4] per-branch taken actions + float* __restrict__ per_sample_loss, // [B] output loss + float* __restrict__ d_q_online, // [total_branch_actions, B*Q] output gradient + int B, int Q, int b0, int b1, int b2, int b3, + float kappa, // Huber threshold + float gamma // discount factor +) { + int sample = blockIdx.x; + if (sample >= B) return; + int tid = threadIdx.x; + + // Each thread handles a subset of quantiles/actions for this sample + // ... deterministic accumulation (no atomicAdd) ... + // Writes per_sample_loss[sample] and d_q_online columns for this sample +} +``` + +### Step 4.4: Add hadamard backward kernel + +- [ ] Add `iqn_hadamard_sigmoid_bwd` kernel for backward through `combined = h_s2 * sigmoid(embed)`. + +```cuda +extern "C" __global__ +void iqn_hadamard_sigmoid_bwd( + const float* __restrict__ h_s2_tiled, // [H, B*Q] + const float* __restrict__ embed, // [H, B*Q] saved embedding + const float* __restrict__ d_combined, // [H, B*Q] upstream gradient + float* __restrict__ d_embed, // [H, B*Q] gradient w.r.t. embed + float* __restrict__ d_h_s2_tiled, // [H, B*Q] gradient w.r.t. h_s2 + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + float sig = 1.0f / (1.0f + expf(-embed[idx])); + float dsig = sig * (1.0f - sig); + d_embed[idx] = d_combined[idx] * h_s2_tiled[idx] * dsig; + d_h_s2_tiled[idx] = d_combined[idx] * sig; + } +} +``` + +### Step 4.5: Add bias-add and bias-grad-reduce kernels (IQN variants) + +- [ ] Add `iqn_bias_add` and `iqn_bias_grad_reduce` kernels (same pattern as IQL variants). + +```cuda +extern "C" __global__ +void iqn_bias_add( + float* __restrict__ x, + const float* __restrict__ bias, + int out_dim, + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + x[idx] += bias[idx % out_dim]; + } +} + +extern "C" __global__ +void iqn_bias_grad_reduce( + const float* __restrict__ d_pre, + float* __restrict__ d_bias, + int out_dim, + int BQ, + float inv_bq +) { + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j < out_dim) { + float sum = 0.0f; + for (int i = 0; i < BQ; i++) { + sum += d_pre[j + i * out_dim]; + } + d_bias[j] = sum * inv_bq; + } +} +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +``` + +--- + +## Task 5: Attention cuBLAS Rewrite + +**File:** `crates/ml/src/cuda_pipeline/gpu_attention.rs` + +**Summary:** Replace per-sample attention forward and backward kernels with cuBLAS GEMMs for Q/K/V/O projections. Keep softmax + scaled dot-product as element-wise cubin kernels (not matmul-shaped). + +### Step 5.1: Add GEMM descriptors to GpuAttention + +- [ ] Add `shared_handle: Arc` field. +- [ ] Create GEMM descriptors for 4 forward projections (Q, K, V, O): + - `M=D, N=B, K=D, ldb=D` (TRANSA=T, TRANSB=N) for W_Q, W_K, W_V, W_O + +```rust +// Forward GEMMs (4 projections): +gemm_fwd_q: AttnGemmDesc, // Q[D,B] = W_Q^T[D,D] @ states[D,B] +gemm_fwd_k: AttnGemmDesc, // K[D,B] = W_K^T[D,D] @ states[D,B] +gemm_fwd_v: AttnGemmDesc, // V[D,B] = W_V^T[D,D] @ states[D,B] +gemm_fwd_o: AttnGemmDesc, // O[D,B] = W_O^T[D,D] @ attn_out[D,B] + +// Backward GEMMs (4 dW + 4 dX = 8 projections): +gemm_bwd_dw_o: AttnGemmDesc, // dW_O[D,D] = d_output[D,B] @ attn_out^T[B,D] +gemm_bwd_dx_o: AttnGemmDesc, // d_attn_out[D,B] = W_O[D,D] @ d_output[D,B] +gemm_bwd_dw_v: AttnGemmDesc, // dW_V[D,D] = d_v[D,B] @ states^T[B,D] +gemm_bwd_dx_v: AttnGemmDesc, // d_states_v[D,B] = W_V[D,D] @ d_v[D,B] +gemm_bwd_dw_k: AttnGemmDesc, // ... same pattern +gemm_bwd_dx_k: AttnGemmDesc, +gemm_bwd_dw_q: AttnGemmDesc, +gemm_bwd_dx_q: AttnGemmDesc, +``` + +### Step 5.2: Add intermediate buffers + +- [ ] Add projection output buffers and backward gradient buffers. + +```rust +proj_q_buf: CudaSlice, // [D, B] +proj_k_buf: CudaSlice, // [D, B] +proj_v_buf: CudaSlice, // [D, B] +attn_scores_buf: CudaSlice, // [num_heads, B, B] or [B, num_heads, head_dim] +attn_out_buf: CudaSlice, // [D, B] post-attention, pre-output projection +d_proj_q_buf: CudaSlice, // [D, B] backward +d_proj_k_buf: CudaSlice, +d_proj_v_buf: CudaSlice, +d_attn_out_buf: CudaSlice, +``` + +### Step 5.3: Rewrite `forward()` method + +- [ ] Replace per-sample `multihead_feature_attention` kernel with 4 cuBLAS GEMMs + softmax cubin kernel. + +```rust +pub fn forward(&mut self, states: &CudaSlice, batch_size: usize) -> Result<&CudaSlice, MLError> { + let d = self.config.state_dim; + let b = batch_size; + // ... same DtoD save_input copy ... + + // 1. Q = W_Q^T @ states (cuBLAS) + // 2. K = W_K^T @ states (cuBLAS) + // 3. V = W_V^T @ states (cuBLAS) + // 4. + bias_add for Q, K, V + // 5. attn_scores + softmax + V aggregation (cubin kernel: attn_sdp_fwd) + // 6. O = W_O^T @ attn_out (cuBLAS) + // 7. + bias_add for O + // 8. LayerNorm + residual (cubin kernel: attn_layer_norm) + + Ok(&self.output_buf) +} +``` + +### Step 5.4: Rewrite `backward()` method + +- [ ] Replace per-sample + tiled backward with cuBLAS backward GEMMs. +- [ ] Remove `d_params_per_sample` buffer and tiling loop. + +```rust +pub fn backward(&mut self, d_output: &CudaSlice, batch_size: usize) -> Result<(), MLError> { + let d = self.config.state_dim; + let b = batch_size; + + // Zero grad_buf + self.stream.memset_zeros(&mut self.d_params)?; + + // 1. LayerNorm backward (cubin) + // 2. dW_O = d_output @ attn_out^T (cuBLAS, alpha=1/B) + // 3. d_attn_out = W_O @ d_output (cuBLAS) + // 4. Attention backward through softmax (cubin: attn_sdp_bwd) + // 5. dW_V = d_v @ states^T (cuBLAS, alpha=1/B) + // 6. dW_K = d_k @ states^T (cuBLAS, alpha=1/B) + // 7. dW_Q = d_q @ states^T (cuBLAS, alpha=1/B) + // 8. Bias gradient reduces (cubin: attn_bias_grad_reduce) + // 9. LayerNorm parameter gradients (cubin) + + Ok(()) +} +``` + +### Step 5.5: Update constructor + +- [ ] Change `GpuAttention::new()` to accept `shared_handle: Arc`. + +```rust +pub fn new( + shared_handle: Arc, + config: GpuAttentionConfig, +) -> Result { + let stream = Arc::clone(&shared_handle.stream); + // ... create GEMM descriptors, load cubin kernels ... +} +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 6: Attention Element-wise cubin Kernels + +**File:** `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` (add to existing) + +**Summary:** Add element-wise CUDA kernels for attention ops that are not matmul-shaped. + +### Step 6.1: Add scaled dot-product attention forward kernel + +- [ ] Add `attn_sdp_fwd` kernel that computes multi-head attention scores, softmax, and V aggregation. + +```cuda +/* ------------------------------------------------------------------ */ +/* Scaled dot-product attention forward (multi-head, feature-level) */ +/* Q[D,B], K[D,B], V[D,B] -> attn_out[D,B] */ +/* Per-head: score = Q_h^T @ K_h / sqrt(head_dim), softmax, @ V_h */ +/* ------------------------------------------------------------------ */ +extern "C" __global__ +void attn_sdp_fwd( + const float* __restrict__ Q, // [D, B] + const float* __restrict__ K, // [D, B] + const float* __restrict__ V, // [D, B] + float* __restrict__ attn_out, // [D, B] + float* __restrict__ save_scores, // [num_heads, B, B] saved for backward + int B, int D, int num_heads +) { + // Each block handles one (head, sample) pair + int head_dim = D / num_heads; + float scale = 1.0f / sqrtf((float)head_dim); + // ... compute attention scores, softmax, V aggregation ... + // No atomicAdd -- each output element written by exactly one thread +} +``` + +### Step 6.2: Add scaled dot-product attention backward kernel + +- [ ] Add `attn_sdp_bwd` kernel. + +```cuda +extern "C" __global__ +void attn_sdp_bwd( + const float* __restrict__ d_attn_out, // [D, B] + const float* __restrict__ Q, // [D, B] saved + const float* __restrict__ K, // [D, B] saved + const float* __restrict__ V, // [D, B] saved + const float* __restrict__ scores, // [num_heads, B, B] saved + float* __restrict__ d_Q, // [D, B] output + float* __restrict__ d_K, // [D, B] output + float* __restrict__ d_V, // [D, B] output + int B, int D, int num_heads +) { + // Reverse of sdp_fwd -- deterministic, no atomicAdd +} +``` + +### Step 6.3: Add LayerNorm forward+backward kernel + +- [ ] Add `attn_layer_norm_fwd` and `attn_layer_norm_bwd` kernels. + +```cuda +extern "C" __global__ +void attn_layer_norm_fwd( + const float* __restrict__ x, // [D, B] input (pre-residual) + const float* __restrict__ residual,// [D, B] skip connection input + const float* __restrict__ gamma, // [D] scale + const float* __restrict__ beta, // [D] shift + float* __restrict__ out, // [D, B] output + float* __restrict__ save_mean, // [B] saved mean + float* __restrict__ save_rstd, // [B] saved 1/std + int D, int B +) { + // Per-sample (one block per sample): compute mean, variance, normalize +} + +extern "C" __global__ +void attn_layer_norm_bwd( + const float* __restrict__ d_out, // [D, B] upstream gradient + const float* __restrict__ x_norm, // [D, B] normalized values (saved) + const float* __restrict__ gamma, // [D] + const float* __restrict__ save_rstd, // [B] + float* __restrict__ d_x, // [D, B] input gradient + float* __restrict__ d_gamma, // [D] gamma gradient + float* __restrict__ d_beta, // [D] beta gradient + int D, int B +) { + // Deterministic -- each d_gamma[j] sums along B dimension sequentially +} +``` + +### Step 6.4: Add attention bias-add and bias-grad kernels + +- [ ] Add `attn_bias_add` and `attn_bias_grad_reduce` kernels (same pattern as IQL). + +```cuda +extern "C" __global__ +void attn_bias_add( + float* __restrict__ x, + const float* __restrict__ bias, + int out_dim, int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { x[idx] += bias[idx % out_dim]; } +} + +extern "C" __global__ +void attn_bias_grad_reduce( + const float* __restrict__ d_pre, + float* __restrict__ d_bias, + int out_dim, int B, float inv_batch +) { + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j < out_dim) { + float sum = 0.0f; + for (int b = 0; b < B; b++) { sum += d_pre[j + b * out_dim]; } + d_bias[j] = sum * inv_batch; + } +} +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +``` + +--- + +## Task 7: Temporal Ops Integration + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +**Summary:** Expose temporal ops (Mamba2, ISV temporal route, regime dropout, predictive coding, risk budget) as a callable method that can be inserted into the unified graph capture sequence. Currently these are individual methods on GpuDqnTrainer; wrap them in a single `submit_temporal_ops()` method. + +### Step 7.1: Create `submit_temporal_ops()` method + +- [ ] Add a new public method to `GpuDqnTrainer` that calls the temporal ops in order. + +```rust +/// Submit all temporal enrichment ops on the training stream. +/// Called between ISV/plan forward and loss computation. +/// All ops operate on h_s2 in-place (enrich with temporal context). +pub(crate) fn submit_temporal_ops(&mut self) -> Result<(), MLError> { + let batch_size = self.config.batch_size; + + // 1. Mamba2 temporal scan (enriches h_s2 with K=8 history) + self.launch_mamba2_scan(batch_size)?; + + // 2. ISV temporal route (per-feature temporal weights) + self.launch_isv_temporal_route(batch_size)?; + + // 3. Regime dropout (regime-conditioned) + self.launch_regime_dropout(batch_size)?; + + // 4. Predictive coding loss (temporal smoothness) + self.launch_predictive_coding_forward(batch_size)?; + + // 5. Risk budget forward + self.launch_risk_budget_forward(batch_size)?; + + Ok(()) +} +``` + +### Step 7.2: Wire into `submit_forward_ops_main()` + +- [ ] Insert `submit_temporal_ops()` call between ISV/plan heads (step 1b) and loss computation (step 2+3) in `submit_forward_ops_main()`. + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`, in `submit_forward_ops_main()` (around line 7884): + +```rust +// After step 1b (ISV + plan heads): +// ── 1c. Temporal enrichment ops ── +self.submit_temporal_ops()?; + +// ── 2+3. Loss + gradient (blended MSE + C51) ── +self.launch_curiosity_inference()?; +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 8: ISV Signal Update Wiring + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +**Summary:** Move ISV signal update into `submit_forward_ops_main()` so it can be captured in the unified graph. Currently it runs after Adam readback (`run_full_step` step 5a-isv), which is outside any graph. In the unified graph, ISV reads from pinned device-mapped scalars (loss, grad_norm, Q-mean) written by earlier nodes. + +### Step 8.1: Create `submit_isv_signal_update()` method + +- [ ] Add method to GpuDqnTrainer that launches ISV signal update kernels. +- [ ] Must run AFTER Adam (which writes the scalars ISV reads). + +```rust +/// Submit ISV signal update kernels to the training stream. +/// Reads pinned device-mapped scalars (loss, grad_norm, Q-mean) written +/// by earlier ops in the same graph. Stream ordering guarantees visibility. +pub(crate) fn submit_isv_signal_update_graphed(&mut self) -> Result<(), MLError> { + self.update_isv_signals() +} +``` + +### Step 8.2: Wire into unified graph sequence + +- [ ] This will be called from FusedTrainingCtx during unified graph capture (Task 9), after the Adam ops. + +**No standalone test for this task -- tested as part of Task 9.** + +--- + +## Task 9: Child Graph Architecture + +**Files:** +- Modify: `crates/ml/src/trainers/dqn/fused_training.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +**Summary:** Replace all 6 graph fields with composable child sub-graphs assembled into a single parent graph via `cudaGraphAddChildGraphNode`. Each child is captured independently with cudarc, tested in isolation, then composed. Single `parent_graph.launch()` per step. + +### Step 9.1: Replace graph fields with child graph architecture + +- [ ] Remove `graph_aux`, `graph_spectral`, `graph_mega`, `aux_frequency` fields. +- [ ] Remove `graph_forward`, `graph_forward_ddqn`, `graph_adam` from GpuDqnTrainer. +- [ ] Add child sub-graph fields and parent graph to FusedTrainingCtx. + +**In `FusedTrainingCtx` struct definition:** + +Remove ALL old graph fields. Add: +```rust + // ── Child graph architecture ────────────────────────────────────── + // Each child is captured independently via cudarc, then composed + // into parent_graph via cudaGraphAddChildGraphNode. Single launch. + spectral_child: Option, + forward_child: Option, + ddqn_child: Option, + aux_child: Option, + adam_child: Option, + /// Composed parent graph — single launch replays all children. + parent_graph: Option, +``` + +### Step 9.2: Create child graph capture functions + +- [ ] Add `capture_child()` helper that captures a single sub-graph via cudarc. +- [ ] Add `capture_all_children()` that captures each child independently. +- [ ] Add `compose_parent_graph()` that builds the parent from children via `cudaGraphAddChildGraphNode`. + +```rust +/// Capture a single child sub-graph. Returns the CudaGraph (not instantiated). +fn capture_child( + &mut self, + label: &str, + submit_fn: F, +) -> Result +where + F: FnOnce(&mut Self) -> Result<()>, +{ + self.trainer.sync_all_streams() + .map_err(|e| anyhow::anyhow!("{label} pre-sync: {e}"))?; + let _ = self.stream.context().check_err(); + unsafe { self.stream.context().disable_event_tracking(); } + + self.stream.begin_capture( + cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL, + ).map_err(|e| { + unsafe { self.stream.context().enable_event_tracking(); } + anyhow::anyhow!("{label} begin_capture: {e}") + })?; + + let submit_result = submit_fn(self); + + let graph_result = self.stream.end_capture( + cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + unsafe { self.stream.context().enable_event_tracking(); } + let _ = self.stream.context().check_err(); + + submit_result?; + let graph = graph_result + .map_err(|e| anyhow::anyhow!("{label} end_capture: {e}"))? + .ok_or_else(|| anyhow::anyhow!("{label} returned None"))?; + + info!("{label} child graph captured"); + Ok(graph) +} + +/// Capture all 5 children independently, then compose parent. +fn capture_training_graph( + &mut self, + agent: &mut DQNAgentType, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, +) -> Result<()> { + // 1. Capture each child independently + self.spectral_child = Some(self.capture_child("spectral", |s| { + s.trainer.apply_spectral_norm(&s.online_dueling, &s.online_branching) + .map_err(|e| anyhow::anyhow!("{e}")) + })?); + + self.forward_child = Some(self.capture_child("forward", |s| { + s.trainer.submit_forward_ops_main() + .map_err(|e| anyhow::anyhow!("{e}")) + })?); + + self.ddqn_child = Some(self.capture_child("ddqn", |s| { + s.trainer.submit_forward_ops_ddqn() + .map_err(|e| anyhow::anyhow!("{e}")) + })?); + + self.aux_child = Some(self.capture_child("aux", |s| { + s.submit_aux_ops(agent, gpu_batch) + })?); + + self.adam_child = Some(self.capture_child("adam", |s| { + s.mamba2_backward(s.batch_size)?; + s.step_mamba2_adam()?; + s.trainer.apply_pruning_mask() + .map_err(|e| anyhow::anyhow!("{e}"))?; + s.trainer.submit_adam_ops(&s.online_dueling, &s.online_branching) + .map_err(|e| anyhow::anyhow!("{e}"))?; + s.trainer.update_isv_signals() + .map_err(|e| anyhow::anyhow!("{e}")) + })?); + + // 2. Compose parent graph from children + self.compose_parent_graph()?; + + info!("parent_graph composed: 5 children (spectral + forward + ddqn + aux + adam)"); + Ok(()) +} + +/// Build parent graph by adding each child as a child node. +fn compose_parent_graph(&mut self) -> Result<()> { + use cudarc::driver::sys as cuda_sys; + + unsafe { + let mut parent: cuda_sys::CUgraph = std::ptr::null_mut(); + let result = cuda_sys::cuGraphCreate(&mut parent, 0); + if result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + anyhow::bail!("cuGraphCreate failed: {result:?}"); + } + + // Add children in dependency order (each depends on previous) + let mut prev_node: cuda_sys::CUgraphNode = std::ptr::null_mut(); + let children = [ + &self.spectral_child, + &self.forward_child, + &self.ddqn_child, + &self.aux_child, + &self.adam_child, + ]; + + for (i, child_opt) in children.iter().enumerate() { + let child = child_opt.as_ref() + .ok_or_else(|| anyhow::anyhow!("child {i} not captured"))?; + let mut node: cuda_sys::CUgraphNode = std::ptr::null_mut(); + let deps = if i == 0 { std::ptr::null() } else { &prev_node }; + let num_deps = if i == 0 { 0 } else { 1 }; + let result = cuda_sys::cuGraphAddChildGraphNode( + &mut node, parent, deps, num_deps, child.graph, + ); + if result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + cuda_sys::cuGraphDestroy(parent); + anyhow::bail!("cuGraphAddChildGraphNode {i} failed: {result:?}"); + } + prev_node = node; + } + + // Instantiate parent + let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut(); + let result = cuda_sys::cuGraphInstantiateWithFlags(&mut exec, parent, 0); + if result != cuda_sys::cudaError_enum::CUDA_SUCCESS { + cuda_sys::cuGraphDestroy(parent); + anyhow::bail!("parent cuGraphInstantiateWithFlags failed: {result:?}"); + } + + self.parent_graph = Some(/* wrap exec + parent */); + } + Ok(()) +} +``` +} +``` + +### Step 9.3: Simplify `run_full_step()` + +- [ ] Replace the entire graph_mega/graph_forward/graph_aux/graph_adam dispatch with a single unified graph launch. + +**New `run_full_step()` core logic:** + +```rust +pub(crate) fn run_full_step( + &mut self, + gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, + agent: &mut DQNAgentType, + _device: &MlDevice, +) -> Result { + let step = self.steps_since_varmap_sync; + let cu_stream = self.stream.cu_stream(); + + // Upload batch data (outside graph -- buffer contents change each step) + if let Some(ref pe) = self.phase_events { + PhaseEvents::record(pe.upload_start, cu_stream); + } + self.trainer.upload_batch_gpu(gpu_batch) + .map_err(|e| anyhow::anyhow!("batch upload: {e}"))?; + self.trainer.update_stochastic_depth_mask() + .map_err(|e| anyhow::anyhow!("stochastic depth mask: {e}"))?; + self.trainer.adam_step_async(); + + // HER donor computation (outside graph -- indices vary per step) + if let Some(ref mut her) = self.gpu_her { + use crate::cuda_pipeline::gpu_her::HerGpuStrategy; + let batch_size = self.batch_size; + match her.config.strategy { + HerGpuStrategy::Random => { + her.generate_random_donors_gpu(batch_size) + .map_err(|e| anyhow::anyhow!("HER random donors GPU: {e}"))?; + } + HerGpuStrategy::Future | HerGpuStrategy::Final => { + let source_indices = self.her_source_indices_gpu.as_ref() + .ok_or_else(|| anyhow::anyhow!("HER source indices not pre-computed"))?; + let her_batch_size = her.config.her_batch_size(); + her.relabel_batch_with_strategy( + gpu_batch.episode_ids_ptr, source_indices, + 1, source_indices.len(), her_batch_size, + ).map_err(|e| anyhow::anyhow!("HER donor selection: {e}"))?; + } + } + } + + // Pre-replay state updates (host-side counters for graphed HtoD) + self.pre_replay_state_update(agent); + + if let Some(ref pe) = self.phase_events { + PhaseEvents::record(pe.upload_end, cu_stream); + PhaseEvents::record(pe.fwd_bwd_start, cu_stream); + } + + // ── Main training: single unified graph ── + if let Some(ref graph) = self.graph_unified { + graph.0.launch() + .map_err(|e| anyhow::anyhow!("graph_unified replay: {e}"))?; + } else { + // Ungraphed path (step 0): run everything, then capture + self.trainer.apply_spectral_norm(&self.online_dueling, &self.online_branching) + .map_err(|e| anyhow::anyhow!("spectral norm: {e}"))?; + + self.trainer.run_nan_checks_post_forward(self.batch_size)?; + let _fused = self.trainer.train_step_gpu( + gpu_batch, + &self.online_dueling, &self.online_branching, + &self.target_dueling, &self.target_branching, + ).map_err(|e| anyhow::anyhow!("train_step_gpu: {e}"))?; + + self.submit_aux_ops(agent, gpu_batch)?; + + self.trainer.apply_pruning_mask() + .map_err(|e| anyhow::anyhow!("pruning: {e}"))?; + let fused_result = self.trainer.replay_adam_and_readback() + .map_err(|e| anyhow::anyhow!("adam: {e}"))?; + + self.trainer.update_isv_signals() + .map_err(|e| anyhow::anyhow!("ISV signal update: {e}"))?; + + // Capture unified graph after first successful step + if step == 0 { + if let Err(e) = self.capture_graph_unified(agent, gpu_batch) { + tracing::warn!("graph_unified capture failed: {e}"); + } + } + } + + if let Some(ref pe) = self.phase_events { + PhaseEvents::record(pe.fwd_bwd_end, cu_stream); + } + + // ── Conditional ops (outside graph) ── + // Ensemble diversity, causal intervention, gradient vaccine + if !self.ensemble_extra_heads.is_empty() { + if let Err(e) = self.run_ensemble_step() { + tracing::warn!("Ensemble diversity step failed (non-fatal): {e}"); + } + } + if self.steps_since_varmap_sync % 20 == 0 { + if let Some(ref vaccine_batch) = self.pending_vaccine_batch { + if let Err(e) = self.trainer.apply_gradient_vaccine(vaccine_batch) { + tracing::debug!("Gradient vaccine failed (non-fatal): {e}"); + } + } + } + self.pending_vaccine_batch = None; + + // Readback scalars (loss, grad_norm) from pinned device-mapped memory + let fused_result = self.trainer.read_scalars_from_pinned()?; + + // Mamba2 BPTT + SGD (separate from main graph -- variable topology) + self.mamba2_backward(self.batch_size)?; + self.step_mamba2_adam()?; + + // PER priority update + agent.update_priorities_from_td( + self.trainer.td_errors_buf(), + self.batch_size, + &self.stream, + ).map_err(|e| anyhow::anyhow!("PER update: {e}"))?; + + // Bookkeeping + agent.fused_post_step_gpu_bookkeeping() + .map_err(|e| anyhow::anyhow!("GPU bookkeeping: {e}"))?; + self.steps_since_varmap_sync += 1; + self.last_combined_norm = fused_result.grad_norm; + self.phase_batch_count += 1; + + Ok(fused_result) +} +``` + +### Step 9.4: Expose `submit_adam_ops()` from GpuDqnTrainer + +- [ ] Extract the Adam portion of `replay_adam_and_readback()` into a separate `submit_adam_ops()` method that can be called during graph capture (without the readback). + +**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +```rust +/// Submit grad_norm + Adam + unflatten ops to the training stream. +/// Used during unified graph capture (no readback -- scalars read from pinned memory). +pub(crate) fn submit_adam_ops(&mut self) -> Result<(), MLError> { + // grad_norm phase 1+2 + self.launch_grad_norm()?; + // Adam update + self.launch_adam_update()?; + // Unflatten (42 DtoD copies from flat grad_buf back to weight tensors) + self.launch_unflatten()?; + Ok(()) +} +``` + +### Step 9.5: Add `read_scalars_from_pinned()` method + +- [ ] Add method to read loss/grad_norm from pinned device-mapped memory without GPU sync. + +```rust +/// Read training scalars from pinned device-mapped memory. +/// No GPU sync needed -- the values are written by the graph and visible via +/// coherent device-mapped memory. +pub(crate) fn read_scalars_from_pinned(&self) -> Result { + // Read from pinned host memory (total_loss_host_ptr, grad_norm_host_ptr) + let total_loss = unsafe { *self.total_loss_host_ptr }; + let grad_norm = unsafe { *self.grad_norm_host_ptr }; + Ok(FusedStepResult { total_loss, grad_norm, td_errors: vec![] }) +} +``` + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 10: Dead Code Cleanup + +**Files:** +- `crates/ml/src/trainers/dqn/fused_training.rs` +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs` +- `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` +- `crates/ml/src/cuda_pipeline/gpu_attention.rs` +- `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` +- `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` +- `crates/ml/src/cuda_pipeline/attention_kernel.cu` +- `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` + +**Summary:** Remove all dead code that was replaced by the cuBLAS rewrite and unified graph. + +### Step 10.1: Remove old graph capture functions from FusedTrainingCtx + +- [ ] Delete `capture_graph_spectral()` method (lines ~1414-1451). +- [ ] Delete `capture_graph_aux()` method (lines ~1457-1498). +- [ ] Delete `capture_graph_mega()` method (lines ~1505-1569). + +### Step 10.2: Remove old graph fields from FusedTrainingCtx + +- [ ] Remove `graph_aux`, `graph_spectral`, `graph_mega` field declarations. +- [ ] Remove initialization of these fields in the constructor (`graph_aux: None`, etc.). +- [ ] Remove `aux_frequency` field from struct and constructor. + +### Step 10.3: Remove old graph fields from GpuDqnTrainer + +- [ ] Remove `graph_forward`, `graph_forward_ddqn`, `graph_adam` fields. +- [ ] Remove `capture_training_graphs()` method. +- [ ] Remove `replay_forward()` and `replay_adam()` if they exist as separate methods. + +### Step 10.4: Remove per-sample kernels from IQL + +- [ ] Remove `forward_loss_kernel` field from `GpuIqlTrainer`. +- [ ] Remove `backward_per_sample_kernel` field. +- [ ] Remove `weight_grad_reduce_kernel` field. +- [ ] Remove `grads_per_sample: CudaSlice` field. +- [ ] Remove `grad_tile_size: usize` field. +- [ ] Remove `save_pre1`, `save_pre2`, `save_h1`, `save_h2` fields (replaced by cuBLAS intermediates). + +### Step 10.5: Remove per-sample kernels from IQN + +- [ ] Remove `forward_loss_kernel` field from `GpuIqnHead`. +- [ ] Remove `backward_kernel` field. +- [ ] Remove `weight_grad_reduce_kernel` field. + +### Step 10.6: Remove per-sample kernels from GpuAttention + +- [ ] Remove `forward_kernel` field from `GpuAttention`. +- [ ] Remove `backward_kernel` field. +- [ ] Remove `weight_grad_reduce_kernel` field. +- [ ] Remove `d_params_per_sample: CudaSlice` field. + +### Step 10.7: Remove per-sample CUDA kernels from .cu files + +**File:** `crates/ml/src/cuda_pipeline/iql_value_kernel.cu` +- [ ] Remove `iql_forward_loss_kernel` function. +- [ ] Remove `iql_backward_per_sample` function. +- [ ] Remove `iql_weight_grad_reduce_kernel` function. +- [ ] Keep: `iql_forward_kernel` (inference-only), `iql_loss_reduce`, `iql_grad_norm_phase1`, `iql_grad_norm_phase2`, `iql_adam_kernel`, `iql_gather_q_taken`, `iql_advantage_weights`, all integration kernels. +- [ ] Keep: new cuBLAS element-wise kernels added in Task 2. + +**File:** `crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu` +- [ ] Remove `iqn_forward_loss_kernel` function. +- [ ] Remove `iqn_backward_kernel` function (the tiled per-sample backward). +- [ ] Remove `iqn_weight_grad_reduce_kernel` function. +- [ ] Keep: `iqn_trunk_forward_kernel`, `iqn_loss_reduce`, `iqn_grad_norm_phase1`, `iqn_grad_norm_phase2`, `iqn_adam_kernel`, `iqn_ema_kernel`, `iqn_decode_actions`, `iqn_sample_taus`, `iqn_forward_kernel` (inference). +- [ ] Keep: new cuBLAS element-wise kernels added in Task 4. + +**File:** `crates/ml/src/cuda_pipeline/attention_kernel.cu` +- [ ] Remove `multihead_feature_attention` kernel (the per-sample forward). +- [ ] Keep: any utility kernels that are still used. + +**File:** `crates/ml/src/cuda_pipeline/attention_backward_kernel.cu` +- [ ] Remove `attention_backward_kernel` function (the tiled per-sample backward). +- [ ] Remove `attn_weight_grad_reduce` function. +- [ ] Keep: `attn_grad_norm_phase1`, `attn_grad_norm_phase2`, `attn_adam_kernel`. +- [ ] Keep: new cuBLAS element-wise kernels added in Task 6. + +### Step 10.8: Remove `pre_replay_state_update()` body simplification + +- [ ] The method stays but is simplified: it only increments adam_step counters and updates tau values. No graph_aux-specific logic needed. + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 11: Config Cleanup + +**File:** `crates/ml/src/trainers/dqn/config.rs` + +**Summary:** Remove `aux_frequency` from `DQNHyperparameters` and all references. + +### Step 11.1: Remove `aux_frequency` field + +- [ ] Remove `pub aux_frequency: usize` field from `DQNHyperparameters` struct (line ~571). +- [ ] Remove `aux_frequency: 4` from `Default` impl (line ~1352). + +### Step 11.2: Remove all references + +- [ ] Search for `aux_frequency` across the workspace and remove/update all references. + +**File:** `crates/ml/src/trainers/dqn/fused_training.rs` +- [ ] Remove `aux_frequency: hyperparams.aux_frequency.max(1)` from constructor (line ~736). +- [ ] Remove `let aux_freq = self.aux_frequency;` and `step % aux_freq == 0` check (line ~980-981). + +**File:** Any config serialization/deserialization files that reference `aux_frequency`. + +### Step 11.3: Search for any remaining references + +- [ ] Run: `grep -rn "aux_frequency" crates/` to find any remaining references. +- [ ] Update or remove each one. + +**Test command:** +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +--- + +## Task 12: Smoke Test + H100 Validation + +**Summary:** End-to-end validation on both RTX 3050 (local) and H100 (production). + +### Step 12.1: Local smoke test (RTX 3050) + +- [ ] Run the full smoke test suite. + +```bash +FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture +``` + +- [ ] Verify no panics, no CUDA errors, no NaN values. +- [ ] Verify training completes at least 1 epoch. + +### Step 12.2: Workspace build check + +- [ ] Full workspace build with all features. + +```bash +SQLX_OFFLINE=true cargo check --workspace +SQLX_OFFLINE=true cargo test -p ml --lib +``` + +### Step 12.3: H100 validation (3-epoch baseline) + +- [ ] Deploy to H100 via Argo workflow. + +```bash +./scripts/argo-train.sh --model dqn +``` + +- [ ] Monitor epoch timing: target <80s (from 685s). +- [ ] Compare val_Sharpe, WinRate, Trades against pre-rewrite baseline metrics. +- [ ] Verify determinism: two consecutive runs produce identical `val_Sharpe_raw`. + +### Step 12.4: Performance regression check + +- [ ] Measure per-step timing breakdown: + - Aux ops total time (target: ~5ms, from ~3.5s) + - Graph launches per step (target: 1, from 4) + - Total kernel launches per step (target: ~60, from ~71 per-sample) + +### Step 12.5: Memory usage check + +- [ ] Verify VRAM usage is lower than before (removed per-sample gradient buffers: ~1.7GB for IQL + similar for IQN + ~26MB for attention). +- [ ] Check no OOM on H100 (80GB) with B=16384. + +--- + +## Dependency Graph + +``` +Task 1 (IQL cuBLAS) ──┐ +Task 2 (IQL cubins) ──┤ + ├── Task 7 (Temporal) ──┐ +Task 3 (IQN cuBLAS) ──┤ │ +Task 4 (IQN cubins) ──┤ ├── Task 9 (Unified graph) ──┐ + │ │ │ +Task 5 (Attention cuBLAS)──┤ Task 8 (ISV wire) ──┘ ├── Task 12 (Validation) +Task 6 (Attention cubins)──┘ │ + Task 10 (Dead code) ─────┤ + Task 11 (Config) ─────┘ +``` + +**Parallelizable tasks:** Tasks 1+2, 3+4, 5+6 can be developed in parallel. Tasks 7, 8 can be developed in parallel with each other. Task 9 depends on all of 1-8. Tasks 10, 11 can run in parallel after Task 9. Task 12 runs last. + +--- + +## AMENDMENTS (from plan review) + +### Amendment A: train_step_gpu must not capture its own graphs + +**Affects:** Task 9 (Step 9.3) and Task 10. + +The ungraphed path at step 0 calls `self.trainer.train_step_gpu()` which internally calls `capture_training_graphs()` and creates `graph_forward`, `graph_forward_ddqn`, `graph_adam`. These MUST NOT be created during the unified path — the unified graph captures everything. + +**Fix:** In Task 10, when removing `capture_training_graphs()` from `gpu_dqn_trainer.rs`, also modify `train_step_gpu()` to ONLY submit ops (call `submit_forward_ops_main()` + `submit_forward_ops_ddqn()`) WITHOUT capturing any graphs. Rename to `submit_training_ops_ungraphed()` for clarity. The graph capture happens ONLY in `capture_graph_unified()`. + +```rust +/// Submit all training ops ungraphed (step 0 only, before unified graph capture). +pub(crate) fn submit_training_ops_ungraphed( + &mut self, + gpu_batch: &GpuBatch, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, + target_d: &DuelingWeightSet, + target_b: &BranchingWeightSet, +) -> Result<(), MLError> { + self.upload_batch_gpu(gpu_batch)?; + self.update_stochastic_depth_mask()?; + self.adam_step += 1; + unsafe { *self.t_pinned = self.adam_step; } + self.submit_forward_ops_main()?; + self.submit_forward_ops_ddqn()?; + Ok(()) +} +``` + +### Amendment B: Mamba2 backward must be inside unified graph + +**Affects:** Task 7 and Task 9. + +Mamba2 forward is inside the unified graph (via `submit_forward_ops_main` from Task 7). Mamba2 backward is currently outside the graph (line 1693 in Step 9.3). This means during replay, the backward operates on stale activations from the PREVIOUS capture. + +**Fix:** Move `mamba2_backward()` + `step_mamba2_adam()` inside the unified graph capture in Task 9 (after Phase 6 Adam, before Phase 7 ISV). Update `capture_graph_unified()`: + +```rust + // ── Phase 6: Pruning + grad norm + Adam + unflatten ── + let adam_result = if conf_result.is_ok() { + self.trainer.apply_pruning_mask() + .and_then(|_| self.trainer.submit_adam_ops()) + .map_err(|e| anyhow::anyhow!("unified adam: {e}")) + } else { Ok(()) }; + + // ── Phase 6b: Mamba2 BPTT + SGD (must be inside graph with forward) ── + let mamba_result = if adam_result.is_ok() { + self.mamba2_backward(self.batch_size) + .and_then(|_| self.step_mamba2_adam()) + .map_err(|e| anyhow::anyhow!("unified mamba2 bwd: {e}")) + } else { Ok(()) }; + + // ── Phase 7: ISV signal update ── + let isv_result = if mamba_result.is_ok() { + // ... +``` + +Remove `self.mamba2_backward()` and `self.step_mamba2_adam()` from the post-graph section of `run_full_step()`. + +### Amendment C: submit_adam_ops must exist as public method + +**Affects:** Task 9 (Step 9.2). + +The capture references `self.trainer.submit_adam_ops()` but the current code's public method is `replay_adam_and_readback()` which replays a pre-captured graph_adam. Task 10 removes graph_adam. The underlying op submission function `submit_adam_ops()` already exists at `gpu_dqn_trainer.rs:8125` — it just needs to be `pub(crate)` (currently it is). + +**Verify:** `submit_adam_ops()` signature at gpu_dqn_trainer.rs:8125: +```rust +pub(crate) fn submit_adam_ops( + &mut self, + online_d: &DuelingWeightSet, + online_b: &BranchingWeightSet, +) -> Result<(), MLError> +``` + +Note: it takes weight set refs. In the unified capture, these are available as `self.online_dueling` / `self.online_branching`. Update the capture call: +```rust +self.trainer.submit_adam_ops(&self.online_dueling, &self.online_branching) +``` + +### Amendment D: cuBLAS workspace sharing verification + +**Affects:** Tasks 1, 3, 5. + +All IQL (8 GEMMs), IQN (~10 GEMMs), and Attention (~12 GEMMs) share the single `lt_workspace_ptr` (32 MB) from `SharedCublasHandle`. During unified graph capture, all GEMMs are stream-ordered (one at a time) so workspace sharing is safe — no concurrent access. + +**Verification step:** Add to each GEMM descriptor creation: log the `heuristic.workspaceSize` returned by the heuristic. If ANY GEMM needs >32MB, increase `SharedCublasHandle::lt_workspace_size`. + +```rust +let heuristic = cublaslt_result::get_matmul_algo_heuristic(...)?; +tracing::info!(%label, ws_needed = heuristic.workspaceSize, "GEMM desc created"); +if heuristic.workspaceSize > lt_ws_size { + return Err(MLError::ModelError(format!( + "{label}: needs {}MB workspace but only {}MB available", + heuristic.workspaceSize / (1024*1024), lt_ws_size / (1024*1024) + ))); +} +``` + +Add this to Task 1 Step 1.2, Task 3, and Task 5 constructors. + +--- + +## Task 13: Hardcoded Dimension Cleanup + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/common_device_functions.cuh` +- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` +- Modify: `crates/ml/src/trainers/dqn/trainer/metrics.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs` +- Modify: `crates/ml/src/trainers/dqn/config.rs` + +- [ ] **Step 1: Delete unused `#define` dimensions from common_device_functions.cuh** + +Remove `STATE_DIM`, `PORTFOLIO_DIM`, `OFI_DIM` defines — all kernels receive these as runtime parameters. Verify no kernel uses the defines by grepping for each. + +```bash +grep -rn "STATE_DIM\b" crates/ml/src/cuda_pipeline/*.cu crates/ml/src/cuda_pipeline/*.cuh | grep -v "state_dim\|//\|#define" +grep -rn "PORTFOLIO_DIM\b" crates/ml/src/cuda_pipeline/*.cu crates/ml/src/cuda_pipeline/*.cuh | grep -v "portfolio_dim\|//\|#define" +grep -rn "OFI_DIM\b" crates/ml/src/cuda_pipeline/*.cu crates/ml/src/cuda_pipeline/*.cuh | grep -v "ofi_dim\|//\|#define" +``` + +If any kernel uses the define (not the runtime parameter), change it to use the parameter. Then delete the defines. + +- [ ] **Step 2: Delete dead `_portfolio_dim` in gpu_experience_collector.rs** + +Remove line 758: `let _portfolio_dim: usize = 12;` — dead code, wrong value, confusing. + +- [ ] **Step 3: Fix hardcoded feature_dim in validation backtest (metrics.rs)** + +Replace hardcoded `feature_dim: usize = 62` at line 446 with computed value: +```rust +let market_dim: usize = 42; +let ofi_dim: usize = crate::fxcache::OFI_DIM; +let feature_dim: usize = market_dim + ofi_dim; // 42 + 20 = 62 +``` + +- [ ] **Step 4: Fix portfolio_dim in gpu_backtest_evaluator.rs** + +The `PORTFOLIO_AND_MTF_DIM` constant should match the training experience collector's portfolio features. Verify it includes plan features (14 portfolio) and MTF features (16). + +- [ ] **Step 5: Remove misleading default state_dim in attention config** + +In `gpu_attention.rs`, remove the `Default` impl for `GpuAttentionConfig` or set `state_dim: 0` to force explicit configuration: +```rust +impl Default for GpuAttentionConfig { + fn default() -> Self { + Self { + state_dim: 0, // MUST be set explicitly — no safe default + num_heads: 4, + batch_size: 0, // MUST be set explicitly + } + } +} +``` + +- [ ] **Step 6: Change bars_per_day default to 0.0 (force explicit set)** + +In `config.rs`, change `bars_per_day: 390.0` default to `bars_per_day: 0.0`. Add validation in constructor: +```rust +if hyperparams.bars_per_day < 1.0 { + return Err(anyhow::anyhow!( + "bars_per_day must be set from data (got {:.0}). Use fxcache timestamps.", + hyperparams.bars_per_day + )); +} +``` + +- [ ] **Step 7: Fix stale comments** + +Search for comments with wrong dimensions and fix: +```bash +grep -rn "42 market.*8 portfolio\|8 portfolio\|OFI_DIM.*8\b" crates/ml/src/ | grep -v "base\|OFI_DIM: usize = 20" +``` + +Update all to reflect actual dimensions: 42 market + 14 portfolio + 16 MTF + 20 OFI = 92, aligned 96. + +- [ ] **Step 8: Compile and test** + +Run: `SQLX_OFFLINE=true cargo check --workspace` +Expected: Clean compilation with no warnings about unused variables. + +- [ ] **Step 9: Commit** + +```bash +git add -A +git commit -m "cleanup: remove hardcoded dimensions — all dims from config or computed + +Deleted unused #define STATE_DIM/PORTFOLIO_DIM/OFI_DIM from .cuh files. +Removed dead _portfolio_dim in experience collector. +feature_dim in validation computed from market_dim + OFI_DIM. +bars_per_day default 0.0 (must be set from data). +Fixed stale dimension comments across codebase." +```