foxhunt/docs/superpowers/plans/2026-03-27-gpu-hot-path-performance.md

# GPU Hot Path Performance — Zero-Sync Training Step

> **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:** Eliminate all CPU sync points and F32 compute waste from the per-step training hot path, restoring 3s/epoch performance on H100.

**Architecture:** Cache all CUDA device pointers at construction/graph-capture time into a `CachedPtrs` struct. Replace per-step `cuStreamSynchronize` + `EventTrackingGuard` + `raw_device_ptr` calls with cached u64 pointer lookups. Convert cuBLAS `cublasSgemm` (F32) to `cublasGemmEx` (BF16 compute with F32 accumulate) for 2-3× tensor core throughput. Batch 8 head spectral norm launches into a single multi-matrix kernel.

**Tech Stack:** Rust, CUDA, cudarc 0.19.3, cuBLAS (cublasGemmEx), BF16 tensor cores

**Verification requirement:** Every task MUST compile (`SQLX_OFFLINE=true cargo check -p ml`) and pass tests (`SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`). No exceptions. No skipping.

---

## Root Cause Analysis

**Measured:** Epoch time regressed from 3s → 30s. Per-step hot path has:

| Issue | Per-step cost | Source | Lines |
|---|---|---|---|
| `cuStreamSynchronize` in `apply_iqn_trunk_gradient` | ~50μs (stalls entire GPU pipeline) | Pre-existing | `gpu_dqn_trainer.rs:708` |
| `cuStreamSynchronize` in `apply_ensemble_trunk_gradient` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:911` |
| `cuStreamSynchronize` in `run_ensemble_step` | ~50μs | Pre-existing | `fused_training.rs:875` |
| `cuStreamSynchronize` in `replay_adam_and_readback` | ~50μs | Pre-existing | `gpu_dqn_trainer.rs:2369` |
| 99× `raw_device_ptr()` calls | ~1μs each but forces `device_ptr()` event machinery | Per-step | Throughout |
| 11× `EventTrackingGuard` create/drop | ~0.1μs each (atomic ops) | Per-step | Throughout |
| F32 `cublasSgemm` instead of BF16 `cublasGemmEx` | 2-3× slower SGEMM on H100 | Pre-existing | `batched_forward.rs:513` |
| 8× single-block spectral norm launches | ~5μs each | New | `gpu_dqn_trainer.rs:1439-1453` |
| 10× DtoD spectral norm sync-back | ~2μs each | New | `gpu_dqn_trainer.rs:1490-1503` |

**Total per-step overhead: ~200μs sync stalls + 2-3× slower SGEMM.** Over 10K steps/epoch: 2s sync + 2-3× slower forward/backward = easily 30s epochs.

---

## File Map

| File | Action |
|---|---|
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Major: add `CachedPtrs` struct, remove syncs from IQN/ensemble/CQL paths, cache all pointers at capture time |
| `crates/ml/src/cuda_pipeline/batched_forward.rs` | Major: replace `cublasSgemm` with `cublasGemmEx` (BF16 input, F32 accumulate) |
| `crates/ml/src/cuda_pipeline/batched_backward.rs` | Major: same `cublasSgemm` → `cublasGemmEx` conversion |
| `crates/ml/src/trainers/dqn/fused_training.rs` | Moderate: remove `cuStreamSynchronize` from `run_ensemble_step`, dynamic budget fix |
| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | Minor: add batched spectral norm kernel |
| `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs` | Minor: add performance regression test |

---

### Task 1: Cache all device pointers at graph-capture time

The root cause of most overhead: `raw_device_ptr()` calls `slice.device_ptr(stream)` which goes through cudarc's event tracking machinery. Even with events disabled, it creates a `SyncOnDrop` guard that's immediately leaked via `ManuallyDrop`. This happens 99 times per step.

**Fix:** Compute all u64 pointers ONCE at construction (or re-capture), store them in a `CachedPtrs` struct, and use the cached values in all per-step methods.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Define `CachedPtrs` struct**

After the `GpuDqnTrainConfig` struct (~line 219), add:

```rust
/// Pre-computed raw CUDA device pointers for all buffers used in the per-step hot path.
/// Eliminates 99 `raw_device_ptr()` → `device_ptr()` calls per training step.
/// Pointers are invalidated when buffers are reallocated (which never happens after init).
struct CachedPtrs {
    // Flat parameter buffers
    params: u64,
    target_params: u64,
    grad: u64,
    grad_norm: u64,
    m: u64,
    v: u64,
    t: u64,
    total_loss: u64,
    cql_grad_scratch: u64,

    // Batch input buffers
    states: u64,
    next_states: u64,
    actions: u64,
    rewards: u64,
    dones: u64,
    is_weights: u64,

    // Activation save buffers
    save_h_s1: u64,
    save_h_s2: u64,
    save_h_v: u64,
    save_h_b0: u64,
    save_h_b1: u64,
    save_h_b2: u64,

    // Backward scratch
    bw_d_h_s1: u64,
    bw_d_h_s2: u64,
    bw_d_h_v: u64,
    bw_d_h_b0: u64,
    bw_d_h_b1: u64,
    bw_d_h_b2: u64,

    // IQN trunk scratch
    iqn_trunk_m: u64,
    iqn_trunk_grad_norm: u64,

    // Loss outputs
    td_errors: u64,

    // Per-weight pointers (20 online + 20 target, computed from params base)
    on_weight_ptrs: [u64; 20],
    tg_weight_ptrs: [u64; 20],
}
```

- [ ] **Step 2: Build `CachedPtrs` in constructor**

At the end of `new()`, before `Ok(Self { ... })`, compute all pointers. Since event tracking is enabled at this point and all buffers were just allocated on the same stream, `device_ptr()` is safe:

```rust
        let cached_ptrs = {
            let _evt_guard = EventTrackingGuard::new(stream.context());
            let param_sizes_arr = compute_param_sizes(&config);
            CachedPtrs {
                params: raw_device_ptr(&params_buf, &stream),
                target_params: raw_device_ptr(&target_params_buf, &stream),
                grad: raw_device_ptr(&grad_buf, &stream),
                grad_norm: raw_device_ptr(&grad_norm_buf, &stream),
                m: raw_device_ptr(&m_buf, &stream),
                v: raw_device_ptr(&v_buf, &stream),
                t: raw_device_ptr_i32(&t_buf, &stream),
                total_loss: raw_device_ptr(&total_loss_buf, &stream),
                cql_grad_scratch: raw_device_ptr(&cql_grad_scratch, &stream),
                states: raw_device_ptr(&states_buf, &stream),
                next_states: raw_device_ptr(&next_states_buf, &stream),
                actions: raw_device_ptr_i32(&actions_buf, &stream),
                rewards: raw_device_ptr(&rewards_buf, &stream),
                dones: raw_device_ptr(&dones_buf, &stream),
                is_weights: raw_device_ptr(&is_weights_buf, &stream),
                save_h_s1: raw_device_ptr(&save_h_s1, &stream),
                save_h_s2: raw_device_ptr(&save_h_s2, &stream),
                save_h_v: raw_device_ptr(&save_h_v, &stream),
                save_h_b0: raw_device_ptr(&save_h_b0, &stream),
                save_h_b1: raw_device_ptr(&save_h_b1, &stream),
                save_h_b2: raw_device_ptr(&save_h_b2, &stream),
                bw_d_h_s1: raw_device_ptr(&bw_d_h_s1, &stream),
                bw_d_h_s2: raw_device_ptr(&bw_d_h_s2, &stream),
                bw_d_h_v: raw_device_ptr(&bw_d_h_v, &stream),
                bw_d_h_b0: raw_device_ptr(&bw_d_h_b0, &stream),
                bw_d_h_b1: raw_device_ptr(&bw_d_h_b1, &stream),
                bw_d_h_b2: raw_device_ptr(&bw_d_h_b2, &stream),
                iqn_trunk_m: raw_device_ptr(&iqn_trunk_m, &stream),
                iqn_trunk_grad_norm: raw_device_ptr(&iqn_trunk_grad_norm, &stream),
                td_errors: raw_device_ptr(&td_errors_buf, &stream),
                on_weight_ptrs: f32_weight_ptrs(&params_buf, &param_sizes_arr, &stream),
                tg_weight_ptrs: f32_weight_ptrs(&target_params_buf, &param_sizes_arr, &stream),
            }
        };
```

Add `cached_ptrs: CachedPtrs` to the struct and `Ok(Self { ... })`.

- [ ] **Step 3: Replace `raw_device_ptr` calls in `apply_iqn_trunk_gradient`**

Replace all `raw_device_ptr(&self.xxx, &self.stream)` with `self.cached_ptrs.xxx`. Remove the `cuStreamSynchronize` (line 708) and `EventTrackingGuard` (line 712) — they were only needed because `device_ptr()` could trigger event validation. With cached pointers, no events are involved.

**CRITICAL:** The `cuStreamSynchronize` at line 708 has the comment "ensure graph_forward replay completed before we touch buffers." With cached pointers, we don't "touch" buffers via device_ptr — we just use pre-computed addresses. All kernel launches go to the same stream, so ordering is guaranteed by CUDA stream semantics. The sync is unnecessary.

- [ ] **Step 4: Replace `raw_device_ptr` calls in `apply_ensemble_trunk_gradient`**

Same pattern. Remove the `cuStreamSynchronize` (line 911) and `EventTrackingGuard` (line 914).

- [ ] **Step 5: Replace `raw_device_ptr` calls in `apply_cql_gradient` and `apply_cql_clipped_saxpy`**

Same pattern for both methods.

- [ ] **Step 6: Replace `raw_device_ptr` calls in `clip_grad_buf_inplace` and `read_grad_norm_sync`**

For `clip_grad_buf_inplace`: use `self.cached_ptrs.grad` and `self.cached_ptrs.grad_norm`. Remove `EventTrackingGuard`.

For `read_grad_norm_sync`: keep the `cuStreamSynchronize` (this is a diagnostic method that explicitly syncs for readback) but use cached pointers. This method only runs every 1000 steps.

- [ ] **Step 7: Replace `raw_device_ptr` calls in `apply_spectral_norm`**

Use `self.cached_ptrs.params` for `params_base` in the sync-back section. The spectral norm kernel launches on individual weight set buffers (DuelingWeightSet/BranchingWeightSet) — these are passed as references, so their `device_ptr()` calls remain. BUT: spectral norm runs ONCE per step before graph_forward. The sync at the method start is the graph_forward guard — remove it and rely on stream ordering.

- [ ] **Step 8: Replace `raw_device_ptr` in kernel launch methods**

Update `launch_grad_norm`, `launch_adam_update`, and the SAXPY kernel launches (IQN clipped, ensemble clipped) to use cached pointers. Remove the `EventTrackingGuard` in each.

- [ ] **Step 9: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`
Expected: Compiles clean, all 129+ tests pass.

- [ ] **Step 10: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "perf: cache all CUDA device pointers — eliminate 99 raw_device_ptr + 4 cuStreamSync per step"
```

---

### Task 2: Remove `cuStreamSynchronize` from `run_ensemble_step`

The ensemble step in `fused_training.rs:875` syncs before touching `save_h_s2`. With cached pointers and single-stream execution, this sync is unnecessary — CUDA guarantees ordering on the same stream.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Remove the sync + EventTrackingGuard in `run_ensemble_step`**

At ~line 874-876, remove:
```rust
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let _ = self.stream.context().check_err();
```

And the local `EvtGuard` struct + its creation (~lines 881-893). The ensemble heads use `device_ptr()` on their own weight buffers — since these are created after event tracking is disabled by the top-level guard in `apply_spectral_norm`, they have no events. Use raw u64 pointers from the ensemble head structs instead.

Note: the `device_ptr()` calls in the ensemble loop (on `ensemble_extra_heads[k]` weights) cannot be cached in `CachedPtrs` because ensemble heads are optional and variable-count. For these, use a SINGLE `EventTrackingGuard` scope at the top of `run_ensemble_step` and call `device_ptr()` within it — but do NOT sync the stream.

- [ ] **Step 2: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`
Expected: All pass.

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "perf: remove cuStreamSynchronize from run_ensemble_step"
```

---

### Task 3: Convert cuBLAS forward from `cublasSgemm` (F32) to `cublasGemmEx` (BF16 tensors, F32 accumulate)

The forward pass uses `cublasSgemm` (F32×F32→F32). On H100, `cublasGemmEx` with BF16 inputs and F32 accumulate uses tensor cores at 2-3× the throughput. The BF16 weight buffers (`bf16_params_buf`, `bf16_target_params_buf`) already exist but are never used.

**Key constraints:**
- Weights: BF16 (pre-converted from F32 via `f32_to_bf16_kernel`)
- Activations/states: MUST stay F32 (written by upstream kernels in F32)
- Accumulator: F32 (no precision loss in output)
- cuBLAS GemmEx can do BF16×F32→F32 using `CUBLAS_COMPUTE_32F`

Actually — `cublasGemmEx` requires both A and B to have the same type when using tensor cores. We need BF16×BF16→F32. This means **states and activations must also be BF16** for the forward SGEMM. But:
- States come from the experience collector in F32
- Activations (h_s1, h_s2) are written by SGEMM and consumed by the next layer

**Practical approach:** Use **mixed-precision GemmEx** with:
- A = weights (BF16, from bf16_params_buf)
- B = input (F32, states/activations)
- C = output (F32, activations/logits)
- Compute: `CUBLAS_COMPUTE_32F` (uses tensor cores on H100 for BF16×F32 with TF32 accumulate)

On H100, `CUBLAS_COMPUTE_32F` with mixed types uses TF32 tensor cores, which is ~1.5× faster than pure F32. For full BF16 tensor core speed (3×), both inputs must be BF16. That requires converting activations to BF16 too — a larger change.

**This task implements the TF32 mixed-precision path (1.5× speedup, minimal changes).**

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs`

- [ ] **Step 1: Replace `cublasSgemm` with `cublasGemmEx` using TF32 compute**

In `sgemm_layer` (~line 495), change from:

```rust
cublas_result::sgemm(
    self.handle.0,
    cublas_sys::cublasOperation_t::CUBLAS_OP_T,
    cublas_sys::cublasOperation_t::CUBLAS_OP_N,
    n as i32, b as i32, k as i32,
    &alpha,
    w_ptr as *const f32, k as i32,
    a_ptr as *const f32, k as i32,
    &beta,
    c_ptr as *mut f32, n as i32,
)
```

To:

```rust
cublas_sys::cublasGemmEx(
    self.handle.0,
    cublas_sys::cublasOperation_t::CUBLAS_OP_T,
    cublas_sys::cublasOperation_t::CUBLAS_OP_N,
    n as i32, b as i32, k as i32,
    &alpha as *const f32 as *const std::ffi::c_void,
    w_ptr as *const std::ffi::c_void,
    cublas_sys::cudaDataType_t::CUDA_R_32F,  // A type
    k as i32,
    a_ptr as *const std::ffi::c_void,
    cublas_sys::cudaDataType_t::CUDA_R_32F,  // B type
    k as i32,
    &beta as *const f32 as *const std::ffi::c_void,
    c_ptr as *mut std::ffi::c_void,
    cublas_sys::cudaDataType_t::CUDA_R_32F,  // C type
    n as i32,
    cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32,  // TF32 tensor cores
    cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP,
)
```

`CUBLAS_COMPUTE_32F_FAST_TF32` enables TF32 tensor cores on H100/A100. With F32 inputs, cuBLAS internally truncates to TF19 mantissa for the multiply, accumulates in F32. ~1.5× speedup over pure F32 SGEMM.

Do the same for `sgemm_layer_raw`.

- [ ] **Step 2: Set cuBLAS math mode for tensor ops**

In the `CublasForward::new()` constructor, set the math mode:

```rust
unsafe {
    cublas_sys::cublasSetMathMode(
        handle.0,
        cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
    );
}
```

This allows `cublasSgemm` to use TF32 even without switching to `cublasGemmEx`.

- [ ] **Step 3: Do the same for `batched_backward.rs`**

The backward pass also uses `cublasSgemm`. Apply the same TF32 math mode or `cublasGemmEx` conversion.

- [ ] **Step 4: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`
Expected: All pass. TF32 has ≤0.1% numerical difference from F32 — tests should still pass.

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/batched_forward.rs crates/ml/src/cuda_pipeline/batched_backward.rs
git commit -m "perf: enable TF32 tensor cores for cuBLAS SGEMM — 1.5x forward+backward throughput"
```

---

### Task 4: Batch spectral norm into fewer kernel launches

Currently 8 separate kernel launches for head spectral norm, each with `grid=(1,1,1)` — massively underutilized GPU. Batch them into a single kernel that processes all 8 matrices sequentially (same single block, but loop over matrices using a matrix descriptor array).

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

- [ ] **Step 1: Add batched spectral norm kernel**

In `dqn_utility_kernels.cu`, add:

```cuda
/* ══════════════════════════════════════════════════════════════════════
 * BATCHED SPECTRAL NORM KERNEL
 *
 * Processes N weight matrices in a single kernel launch.
 * Grid = (N, 1, 1), Block = (256, 1, 1).
 * Each block handles one matrix via the same power iteration as
 * spectral_norm_kernel but indexed by blockIdx.x into descriptor arrays.
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void batched_spectral_norm_kernel(
    /* Arrays of N pointers, one per matrix */
    float** __restrict__ W_arr,       /* [N] pointers to weight matrices */
    float** __restrict__ u_arr,       /* [N] pointers to left singular vectors */
    float** __restrict__ v_arr,       /* [N] pointers to right singular vectors */
    const int* __restrict__ out_dims, /* [N] out_dim per matrix */
    const int* __restrict__ in_dims,  /* [N] in_dim per matrix */
    float sigma_max,
    int n_matrices
) {
    int mat_idx = blockIdx.x;
    if (mat_idx >= n_matrices) return;

    float* W = W_arr[mat_idx];
    float* u = u_arr[mat_idx];
    float* v = v_arr[mat_idx];
    int out_dim = out_dims[mat_idx];
    int in_dim = in_dims[mat_idx];
    int n_total = out_dim * in_dim;

    /* Same power iteration logic as spectral_norm_kernel... */
    /* (copy the body of spectral_norm_kernel here, replacing
       the fixed arguments with the per-matrix values) */
    __shared__ float shmem[256];
    int tid = threadIdx.x;
    int bd = blockDim.x;

    // Step 1: v_new = W^T u
    for (int col = tid; col < in_dim; col += bd) {
        float val = 0.0f;
        for (int row = 0; row < out_dim; row++)
            val += W[row * in_dim + col] * u[row];
        v[col] = val;
    }
    __syncthreads();

    // Normalize v
    float local_v2 = 0.0f;
    for (int col = tid; col < in_dim; col += bd)
        local_v2 += v[col] * v[col];
    shmem[tid] = local_v2;
    __syncthreads();
    for (int s = bd / 2; s > 0; s >>= 1) {
        if (tid < s) shmem[tid] += shmem[tid + s];
        __syncthreads();
    }
    float v_norm = sqrtf(shmem[0] + 1e-12f);
    for (int col = tid; col < in_dim; col += bd)
        v[col] /= v_norm;
    __syncthreads();

    // Step 2: u_new = W v_new
    for (int row = tid; row < out_dim; row += bd) {
        float val = 0.0f;
        for (int col = 0; col < in_dim; col++)
            val += W[row * in_dim + col] * v[col];
        u[row] = val;
    }
    __syncthreads();

    // sigma = ||u_new||
    float local_u2 = 0.0f;
    for (int row = tid; row < out_dim; row += bd)
        local_u2 += u[row] * u[row];
    shmem[tid] = local_u2;
    __syncthreads();
    for (int s = bd / 2; s > 0; s >>= 1) {
        if (tid < s) shmem[tid] += shmem[tid + s];
        __syncthreads();
    }
    float sigma = sqrtf(shmem[0] + 1e-12f);

    // Normalize u
    for (int row = tid; row < out_dim; row += bd)
        u[row] /= (sigma + 1e-12f);
    __syncthreads();

    // Scale W if sigma > sigma_max
    if (sigma < 1e-6f) return;
    float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f;
    if (scale < 1.0f) {
        for (int i = tid; i < n_total; i += bd)
            W[i] *= scale;
    }
}
```

- [ ] **Step 2: Load the kernel in Rust**

Add `batched_spectral_norm_kernel` loading in `compile_training_kernels`. Add the function field to the struct.

- [ ] **Step 3: Allocate descriptor arrays on GPU**

Allocate `CudaSlice<u64>` arrays for the 10 W/u/v pointer arrays and `CudaSlice<i32>` for out_dims/in_dims. Upload once at construction.

- [ ] **Step 4: Replace 10 individual spec_norm! calls with single batched launch**

In `apply_spectral_norm`, replace the 2 trunk + 8 head launches with:

```rust
// Single batched launch: 10 blocks × 256 threads
unsafe {
    self.stream
        .launch_builder(&self.batched_spectral_norm_kernel)
        .arg(&self.spec_w_ptrs_ptr)
        .arg(&self.spec_u_ptrs_ptr)
        .arg(&self.spec_v_ptrs_ptr)
        .arg(&self.spec_out_dims_ptr)
        .arg(&self.spec_in_dims_ptr)
        .arg(&sigma_max)
        .arg(&10_i32)
        .launch(LaunchConfig {
            grid_dim: (10, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 256 * 4,
        })?;
}
```

This replaces 10 kernel launches with 1.

- [ ] **Step 5: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`
Expected: All pass.

- [ ] **Step 6: Commit**

```bash
git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
git commit -m "perf: batch 10 spectral norm launches into single kernel"
```

---

### Task 5: Dynamic gradient budget allocation

When IQN/ensemble are disabled, their budget is wasted — C51 only gets 70% even as the sole gradient source.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

**NOTE:** This is already partially implemented (dynamic `c51_frac` computation at Step 2b). Verify it's correct and add a test.

- [ ] **Step 1: Verify the dynamic budget code**

Read the current Step 2b in `fused_training.rs` and confirm:
```rust
let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 };
let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 };
let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 };
let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac;
```

This is correct. C51 gets 100% when nothing else is active.

- [ ] **Step 2: Add test for dynamic budget**

In `gradient_budget.rs`, add a test that constructs a trainer with CQL disabled (cql_alpha=0) and verifies C51 can use the full max_grad_norm:

```rust
#[test]
fn test_c51_gets_full_budget_when_no_auxiliaries() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    let mut cfg = test_config();
    cfg.cql_alpha = 0.0; // Disable CQL
    cfg.iqn_lambda = 0.0; // Disable IQN
    let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg)?;

    // Fill grad_buf with large gradient
    let total_params = trainer.total_params();
    stream.memcpy_htod(&vec![1.0_f32; total_params], trainer.grad_buf_mut())
        .map_err(|e| anyhow::anyhow!("{e}"))?;

    // Clip to full max_grad_norm (what C51 should get when solo)
    trainer.clip_grad_buf_inplace(10.0)?; // max_grad_norm * 1.0

    let norm = trainer.read_grad_norm_sync()?;
    assert!(
        norm > 9.9,
        "C51 solo should use full max_grad_norm, got {norm}"
    );

    Ok(())
}
```

- [ ] **Step 3: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- gradient_budget`
Expected: All pass including new test.

- [ ] **Step 4: Commit**

```bash
git add crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs
git commit -m "fix: dynamic gradient budget — C51 gets full max_grad_norm when auxiliaries disabled"
```

---

### Task 6: Remove `cuStreamSynchronize` from `replay_adam_and_readback`

The Adam readback (`replay_adam_and_readback`) syncs the stream and reads back loss + grad_norm (8 bytes DtoH). This is the ONLY legitimate sync in the per-step path — we need the loss for logging. But we can defer it to the NEXT step's start, overlapping GPU work with CPU logging.

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Convert to async readback with deferred sync**

In `replay_adam_and_readback`, replace:
```rust
cuStreamSynchronize(self.stream.cu_stream());
cuMemcpyDtoH_v2(loss_host..., loss_ptr, 4);
cuMemcpyDtoH_v2(norm_host..., norm_ptr, 4);
```

With async DtoH into a pinned host buffer, returning the PREVIOUS step's values:

```rust
pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    // Return PREVIOUS step's readback (already synced by this step's graph_forward)
    let prev_loss = self.scalars_readback_host[0];
    let prev_norm = self.scalars_readback_host[1];

    // Replay Adam (async)
    self.replay_adam()?;

    // Queue async DtoH for THIS step's results (will be ready by next step)
    unsafe {
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            self.scalars_readback_host.as_mut_ptr().cast(),
            self.cached_ptrs.total_loss, 4,
            self.stream.cu_stream(),
        );
        cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
            self.scalars_readback_host.as_mut_ptr().add(1).cast(),
            self.cached_ptrs.grad_norm, 4,
            self.stream.cu_stream(),
        );
    }

    Ok(FusedTrainScalars {
        total_loss: prev_loss,
        grad_norm: prev_norm.sqrt(),
    })
}
```

This eliminates the per-step sync. The readback has 1-step lag (reporting previous step's metrics) which is standard and acceptable for monitoring.

- [ ] **Step 2: Allocate pinned host buffer**

Replace `scalars_readback_host: Vec<f32>` with pinned memory for async DtoH:

```rust
// In constructor:
let scalars_readback_host = vec![0.0_f32; 2]; // [loss, grad_norm_sq]
// cudarc doesn't expose cuMemAllocHost directly — keep Vec but ensure
// the async DtoH is followed by a stream sync at epoch boundary.
```

Actually — `cuMemcpyDtoHAsync_v2` requires the destination to be page-locked (pinned). With a regular Vec, this is technically UB. The safe approach: keep the sync but move it to the epoch boundary, NOT per-step.

**Revised approach:** Batch the sync. Instead of syncing every step, sync once per N steps (e.g. every 100 steps or at epoch boundary). Use the existing `readback_host` Vec with synchronous DtoH but only when actually needed.

- [ ] **Step 3: Move sync to epoch boundary**

In `fused_training.rs`, remove the per-step sync from `replay_adam_and_readback`. Instead, sync only when `steps_since_varmap_sync % 100 == 0` or at epoch boundary. For intermediate steps, return the last known values.

- [ ] **Step 4: Verify compilation + tests**

Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn`
Expected: All pass.

- [ ] **Step 5: Commit**

```bash
git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "perf: defer Adam readback sync to every 100 steps — eliminate per-step cuStreamSync"
```

---

### Task 7: Performance regression smoke test

Add a test that measures training step throughput and fails if it drops below a threshold.

**Files:**
- Modify: `crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs`

- [ ] **Step 1: Add throughput test**

```rust
/// Training step throughput must not regress.
/// At 3s/epoch with 10K steps: ~3333 steps/sec minimum on RTX 3050.
/// On H100: ~10K+ steps/sec. Test uses a conservative threshold.
#[test]
fn test_training_step_throughput() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    let cfg = test_config();
    let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?;
    let mut dueling = alloc_dueling(&stream, &cfg);
    let mut branching = alloc_branching(&stream, &cfg);

    // Warm up
    trainer.apply_spectral_norm(&mut dueling, &mut branching)?;
    trainer.clip_grad_buf_inplace(7.0)?;

    // Measure 1000 iterations of clip + spectral norm (the hot path additions)
    let n_iters = 1000;
    let start = std::time::Instant::now();
    for _ in 0..n_iters {
        trainer.apply_spectral_norm(&mut dueling, &mut branching)?;
        trainer.clip_grad_buf_inplace(7.0)?;
    }
    // Sync to ensure all GPU work completes
    unsafe { cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream()); }
    let elapsed = start.elapsed();

    let per_iter_us = elapsed.as_micros() as f64 / n_iters as f64;
    assert!(
        per_iter_us < 500.0, // 500μs budget for spectral norm + clip per step
        "Spectral norm + clip took {per_iter_us:.0}μs/iter — budget is 500μs"
    );

    Ok(())
}
```

- [ ] **Step 2: Verify compilation + run**

Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_training_step_throughput --nocapture`
Expected: PASS with per-iter time well under 500μs.

- [ ] **Step 3: Commit**

```bash
git add crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs
git commit -m "test: training step throughput regression test"
```

---

## Execution Dependencies

```
Task 1 (cache pointers) → Task 2 (remove ensemble sync) → Task 6 (defer Adam readback)
                         → Task 3 (TF32 tensor cores) — independent
                         → Task 4 (batch spectral norm) — independent
Task 5 (dynamic budget) — independent
Task 7 (throughput test) — after Tasks 1-4
```

Task 1 is the foundation — all other tasks benefit from cached pointers. Tasks 3 and 4 are independent GPU optimizations. Task 6 requires Task 1's cached pointers. Task 7 validates everything.

## Expected Impact

| Optimization | Estimated speedup |
|---|---|
| Eliminate 4× cuStreamSync (200μs/step) | 2s/epoch saved |
| TF32 tensor cores (1.5× SGEMM) | 30-40% forward+backward time |
| Batch spectral norm (10→1 launch) | ~50μs/step saved |
| Eliminate 99 raw_device_ptr calls | ~100μs/step saved |
| Defer Adam readback sync | ~50μs/step saved |
| **Combined** | **Target: ≤5s/epoch on H100** |