# 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<CudaStream>,

    // ── cuBLAS handle (shared with trunk) ──────────────────────────
    shared_handle: Arc<SharedCublasHandle>,

    // ── 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<f32>,

    // ── Adam optimizer state ────────────────────────────────────────
    m_buf: CudaSlice<f32>,
    v_buf: CudaSlice<f32>,
    grad_buf: CudaSlice<f32>,       // [total_params] reduced gradients
    grad_norm_buf: CudaSlice<f32>,
    grad_norm_partials: CudaSlice<f32>,

    // ── cuBLAS forward intermediate buffers ─────────────────────────
    h1_pre_buf: CudaSlice<f32>,     // [H, B] pre-activation layer 1 (GEMM output)
    h1_buf: CudaSlice<f32>,         // [H, B] post-SiLU layer 1
    h2_pre_buf: CudaSlice<f32>,     // [H, B] pre-activation layer 2
    h2_buf: CudaSlice<f32>,         // [H, B] post-SiLU layer 2
    v_pre_buf: CudaSlice<f32>,      // [1, B] output layer pre-bias

    // ── cuBLAS backward intermediate buffers ────────────────────────
    dv_buf: CudaSlice<f32>,         // [1, B] d_expectile_loss output
    dh2_buf: CudaSlice<f32>,        // [H, B] dh2 (pre silu_bwd)
    dh2_pre_buf: CudaSlice<f32>,    // [H, B] dh2_pre (post silu_bwd)
    dh1_buf: CudaSlice<f32>,        // [H, B] dh1
    dh1_pre_buf: CudaSlice<f32>,    // [H, B] dh1_pre

    // ── Output buffers (same as before) ────────────────────────────
    v_out_buf: CudaSlice<f32>,
    loss_buf: CudaSlice<f32>,
    total_loss_buf: CudaSlice<f32>,
    q_taken_buf: CudaSlice<f32>,
    advantage_weights_buf: CudaSlice<f32>,

    // ── Integration buffers (unchanged) ────────────────────────────
    adv_stats_buf: CudaSlice<f32>,
    adv_sigma_ema_buf: CudaSlice<f32>,
    readiness_buf: CudaSlice<f32>,
    p5_state_buf: CudaSlice<f32>,
    per_sample_support_buf: CudaSlice<f32>,
    branch_scales_buf: CudaSlice<f32>,
    expectile_gap_buf: CudaSlice<f32>,
    gap_mean_buf: CudaSlice<f32>,
    per_sample_epsilon_buf: CudaSlice<f32>,

    // ── Training state ─────────────────────────────────────────────
    adam_step: i32,
    t_buf: CudaSlice<i32>,
    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<f32>` -- 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<SharedCublasHandle>` instead of bare `stream: Arc<CudaStream>`.
- [ ] 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<SharedCublasHandle>,
    config: GpuIqlConfig,
) -> Result<Self, MLError> {
    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<IqlGemmDesc, MLError> {
    let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
    let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32;

    unsafe {
        let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type)
            .map_err(|e| MLError::ModelError(format!("{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::<i32>(),
        ).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::<i32>(),
        ).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::<usize>(),
        ).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<IqlGemmDesc, MLError> {
    // 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<f32>,
) -> 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::<f32>();

    // 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<SharedCublasHandle>` 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<f32>,     // [H, B*Q] embedding GEMM output (before ReLU)
embed_buf: CudaSlice<f32>,         // [H, B*Q] post-ReLU
combined_buf: CudaSlice<f32>,      // [H, B*Q] h_s2 * sigmoid(embed)
branch_logits_bufs: [CudaSlice<f32>; 4],  // [bk, B*Q] per-branch logits

// Backward intermediates
d_combined_buf: CudaSlice<f32>,    // [H, B*Q] accumulated dX from all branches
d_embed_pre_buf: CudaSlice<f32>,   // [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<SharedCublasHandle>`.

```rust
pub fn new(
    shared_handle: Arc<SharedCublasHandle>,
    config: GpuIqnConfig,
) -> Result<Self, MLError> {
    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<SharedCublasHandle>` 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<f32>,       // [D, B]
proj_k_buf: CudaSlice<f32>,       // [D, B]
proj_v_buf: CudaSlice<f32>,       // [D, B]
attn_scores_buf: CudaSlice<f32>,  // [num_heads, B, B] or [B, num_heads, head_dim]
attn_out_buf: CudaSlice<f32>,     // [D, B] post-attention, pre-output projection
d_proj_q_buf: CudaSlice<f32>,     // [D, B] backward
d_proj_k_buf: CudaSlice<f32>,
d_proj_v_buf: CudaSlice<f32>,
d_attn_out_buf: CudaSlice<f32>,
```

### 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<f32>, batch_size: usize) -> Result<&CudaSlice<f32>, 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<f32>, 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<SharedCublasHandle>`.

```rust
pub fn new(
    shared_handle: Arc<SharedCublasHandle>,
    config: GpuAttentionConfig,
) -> Result<Self, MLError> {
    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<cudarc::driver::CudaGraph>,
    forward_child: Option<cudarc::driver::CudaGraph>,
    ddqn_child: Option<cudarc::driver::CudaGraph>,
    aux_child: Option<cudarc::driver::CudaGraph>,
    adam_child: Option<cudarc::driver::CudaGraph>,
    /// Composed parent graph — single launch replays all children.
    parent_graph: Option<cudarc::driver::CudaGraph>,
```

### 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<F>(
    &mut self,
    label: &str,
    submit_fn: F,
) -> Result<cudarc::driver::CudaGraph>
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<FusedStepResult> {
    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<FusedStepResult, MLError> {
    // 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<f32>` 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<f32>` 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."
```
