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:

/// 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:

        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

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:

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

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:

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:

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:

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

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:

/* ══════════════════════════════════════════════════════════════════════
 * 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:

// 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

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:

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:

#[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

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:

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:

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:

// 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

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

/// 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

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