foxhunt/docs/superpowers/plans/2026-04-02-mega-graph-sub10s-epochs.md

Mega-Graph Refactor: <10s DQN Epochs on H100

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: Achieve <10s epochs (2000 steps x 5ms/step) on H100 80GB for a 323K parameter DQN by eliminating all per-step kernel launch overhead through mega-graph capture and dual-stream parallelism.

Architecture: Currently two CUDA graphs (graph_forward + graph_adam) with 10+ ungraphed auxiliary operations between them. Refactor into a single mega-graph for the main C51+CQL+Adam path, with auxiliary heads (IQN, attention, IQL, EMA) running on a parallel CUDA stream. Fuse remaining kernels to minimize launch count.

Tech Stack: Rust 1.85, cudarc (CUDA driver bindings), CUDA 13.0 (H100 PCIe), precompiled cubins, CudaStream/CudaGraph/CudaEvent APIs.

---

Current Per-Step Breakdown (~25ms estimated after sync removal)

upload_batch_gpu (6 DtoD copies)                    ~0.1ms  [main stream]
spectral_norm (per weight matrix)                   ~0.2ms  [main stream]
graph_forward replay                                ~2.0ms  [main stream]
C51 gradient clip                                   ~0.1ms  [main stream]
target_ema_update (EMA kernel)                      ~0.2ms  [main stream]
attention fwd + bwd + Adam (cuBLAS)                 ~3.0ms  [main stream]
IQL value step (2 kernels)                          ~1.0ms  [main stream]
IQN train + trunk grad (cuBLAS + custom)            ~5.0ms  [main stream]
IQN loss DtoD for PER                               ~0.05ms [main stream]
ensemble diversity (if active)                      ~0.5ms  [main stream]
CQL gradient + clipped SAXPY                        ~0.5ms  [main stream]
causal intervention (1/100)                         ~0.05ms [main stream, amortized]
vaccine (1/10, ungraphed fwd+bwd)                   ~1.5ms  [main stream, amortized]
pruning mask                                        ~0.1ms  [main stream]
graph_adam replay                                   ~1.0ms  [main stream]
                                                   --------
TOTAL                                              ~15-25ms [sequential]

Target: 5ms/step = max(main_stream, aux_stream)

File Structure

Modified Files

File	Responsibility	Phase
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs	Graph capture, submit_ops, kernel launches	1,2,4
crates/ml/src/trainers/dqn/fused_training.rs	Per-step orchestration, stream dispatch	1,2
crates/ml/src/trainers/dqn/trainer/constructor.rs	Stream allocation, double-buffer init	2,4
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs	IQN fused warp kernel	3
crates/ml/src/cuda_pipeline/gpu_attention.rs	Attention fused kernel	3
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu	Fused EMA+spectral cubin	3

New Files

File	Responsibility	Phase
crates/ml/src/cuda_pipeline/ema_spectral_fused_kernel.cu	Fused EMA + spectral norm CUDA kernel	3

---

Phase 1: Merge graph_forward + graph_adam into graph_step (Week 1)

Rationale

Currently graph_forward and graph_adam are separate CUDA graphs with 10+ ungraphed operations between them. Several of those operations (CQL gradient, C51 clip, pruning mask, spectral norm) are pure element-wise kernels with NO data-dependent branching and NO external state mutation. They can be captured into the graph.

After Phase 1, the per-step flow becomes:

upload → graph_step(fwd+cql+clip+prune+adam) → [ungraphed: EMA, attn, IQL, IQN]

Expected: ~11ms/step (graph ~3ms + ungraphed aux ~8ms)

Task 1.1: Move CQL gradient into submit_forward_ops

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4001-4106 (submit_forward_ops)
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4108-4160 (submit_adam_ops)
• Modify: crates/ml/src/trainers/dqn/fused_training.rs:798-813 (CQL step)

CQL gradient is element-wise: reads logits → computes logsumexp penalty → writes to cql_grad_scratch. Then a clipped SAXPY adds it to grad_buf. Both operations have fixed control flow (no Option checks — has_cql() is always true in production).

• Step 1: Read current CQL flow in fused_training.rs:798-813

The CQL block calls two methods on self.trainer:

// fused_training.rs:802-807
self.trainer.apply_cql_gradient()     // → cql_grad_scratch
self.trainer.apply_cql_clipped_saxpy(cql_budget)  // cql_grad_scratch → grad_buf

Verify both methods are capturable: no host readbacks, no allocs, no branching on GPU data.

Run: grep -n "synchronize\|memcpy_dtoh\|alloc" crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | grep -i "cql" Expected: No matches — CQL is pure async kernels.

• Step 2: Move CQL kernel launches into submit_forward_ops

In gpu_dqn_trainer.rs, add CQL gradient + clipped SAXPY after the cuBLAS backward call at the end of submit_forward_ops (line ~4103). The CQL kernels read the same grad_buf and d_value_logits_buf that the C51 backward just wrote — stream ordering is sufficient.

// gpu_dqn_trainer.rs, end of submit_forward_ops(), after launch_cublas_backward:

// CQL conservative penalty gradient (captured in graph — fixed control flow)
self.launch_cql_logit_grad()?;
// CQL clipped SAXPY: cql_grad_scratch → grad_buf (budget-limited)
{
    let cql_budget = self.config.max_grad_norm * 0.25; // CQL_GRAD_BUDGET
    self.clip_grad_buf_inplace_impl(/* buffer= */ &self.cql_grad_scratch, cql_budget)?;
    self.clipped_saxpy_impl(/* dst= */ &self.grad_buf, /* src= */ &self.cql_grad_scratch, cql_budget)?;
}

Note: This requires extracting clip_grad_buf_inplace and apply_cql_clipped_saxpy internals into capturable _impl methods that take buffer references (no &mut self borrows that conflict with graph capture).

• Step 3: Remove CQL from run_full_step in fused_training.rs

Delete lines 798-813 (the CQL block in run_full_step). The CQL gradient is now inside the graph.

// DELETE from fused_training.rs:
// ── Step 5c: CQL conservative penalty (isolated gradient) ─────────
// ... entire block ...

• Step 4: Compile and test

Run: SQLX_OFFLINE=true cargo check -p ml Expected: Clean compilation.

Run: SQLX_OFFLINE=true cargo test -p ml --lib -- trainers::dqn --nocapture 2>&1 | tail -20 Expected: All tests 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: capture CQL gradient in CUDA graph — eliminate per-step kernel launches"

Task 1.2: Move C51 gradient clip + pruning mask into graph

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4001-4106 (submit_forward_ops)
• Modify: crates/ml/src/trainers/dqn/fused_training.rs:627-639 (C51 clip)
• Modify: crates/ml/src/trainers/dqn/fused_training.rs:857-859 (pruning mask)

C51 clip (clip_grad_buf_inplace) and pruning mask (apply_pruning_mask) are both element-wise on grad_buf. They currently run between graph_forward and graph_adam. Move them into the graph capture.

• Step 1: Move C51 clip into submit_forward_ops, after CQL

After the CQL SAXPY (added in Task 1.1), add the C51 budget clip:

// After CQL SAXPY in submit_forward_ops:
// C51 gradient budget clip (budget fraction baked into graph as literal)
{
    let c51_budget = self.config.max_grad_norm * 0.60; // 1.0 - CQL(0.25) - IQN(0.10) - ENS(0.05)
    self.clip_grad_buf_inplace_impl(&self.grad_buf, c51_budget)?;
}

Note: The budget fractions are compile-time constants (all auxiliaries are always active). Baking them as literals is correct.

• Step 2: Move pruning mask into submit_adam_ops, before Adam

// gpu_dqn_trainer.rs, start of submit_adam_ops(), before launch_adam_update:
self.apply_pruning_mask_impl()?; // element-wise: grad_buf *= mask

• Step 3: Remove both from run_full_step

Delete C51 clip block (lines 627-639) and pruning mask call (line 858-859) from fused_training.rs.

• Step 4: Merge graph_forward + graph_adam into single graph_step

In gpu_dqn_trainer.rs, modify capture_and_instantiate_graphs (line ~3880):

// BEFORE: two captures
// graph_forward = capture { submit_forward_ops() }
// graph_adam = capture { submit_adam_ops() }

// AFTER: one capture
// graph_step = capture { submit_forward_ops() + submit_adam_ops() }
let begin_result = self.stream.begin_capture(...);
self.submit_forward_ops()?;
self.submit_adam_ops(online_d, online_b)?;
let graph_step = self.stream.end_capture(...)?;

Update all references from graph_forward/graph_adam to graph_step:

• self.graph_forward: Option<CudaGraphExec> → self.graph_step: Option<CudaGraphExec>

• Delete self.graph_adam field

• replay_forward() → replay_step()

• replay_adam() → delete (merged into replay_step)

• replay_adam_and_readback() → replay_step_and_readback()

• Step 5: Update run_full_step to use single graph

// fused_training.rs, new flow:
// Step 1: HER relabeling (unchanged)
// Step 1b: Spectral norm (unchanged — runs before graph, modifies weights)
// Step 2: Upload batch + replay graph_step (fwd + CQL + clip + prune + adam)
self.trainer.train_step_gpu(...)  // uploads batch
let result = self.trainer.replay_step_and_readback()?;
// Step 3: EMA (unchanged)
// Step 3b: Attention (unchanged)
// Step 4: IQL (unchanged)
// Step 5: IQN (unchanged)
// DELETED: CQL, C51 clip, pruning — now inside graph

• Step 6: Compile and test

Run: SQLX_OFFLINE=true cargo check -p ml Run: SQLX_OFFLINE=true cargo test -p ml --lib -- trainers::dqn --nocapture 2>&1 | tail -20 Expected: All pass.

• Step 7: Commit

git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs
git commit -m "perf: merge graph_forward + graph_adam into single graph_step — 1 launch per step"

Task 1.3: Move spectral norm into graph

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (submit_forward_ops)
• Modify: crates/ml/src/trainers/dqn/fused_training.rs:610-616 (spectral norm call)

Spectral norm runs BEFORE the forward pass on weight matrices. It's element-wise per weight matrix with a fixed iteration count. Capturable.

• Step 1: Move spectral norm into submit_forward_ops, before cuBLAS forward

// gpu_dqn_trainer.rs, start of submit_forward_ops(), before launch_cublas_forward:
self.apply_spectral_norm_impl()?;

• Step 2: Remove from run_full_step

Delete line 615-616 from fused_training.rs.

• Step 3: Compile and test
• Step 4: Commit

git commit -m "perf: capture spectral norm in graph_step — all element-wise ops now graphed"

---

Phase 2: Dual-Stream Overlap (Week 2)

Rationale

After Phase 1, per-step is: graph_step (~3ms) + EMA + attention + IQL + IQN (~8ms) = ~11ms. The auxiliary operations (EMA, attention, IQL, IQN) are independent of graph_step's Adam update — they read save_h_s2 (output of forward pass) and have their own weight buffers. Running them on a parallel stream overlaps with graph_step.

Expected: max(graph_step ~3ms, aux_stream ~8ms) = ~8ms/step

Task 2.1: Create auxiliary CUDA stream

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs (struct + constructor)

• Modify: crates/ml/src/trainers/dqn/trainer/constructor.rs (stream fork)

• Step 1: Fork a second stream in FusedTrainingCtx

// fused_training.rs, add field:
pub(crate) aux_stream: Arc<CudaStream>,
// CUDA events for synchronization:
pub(crate) forward_done_event: cudarc::driver::CudaEvent,
pub(crate) aux_done_event: cudarc::driver::CudaEvent,

In constructor:

let aux_stream = stream.fork()
    .map_err(|e| anyhow::anyhow!("Fork aux stream: {e}"))?;
let forward_done_event = stream.create_event(false)
    .map_err(|e| anyhow::anyhow!("Forward done event: {e}"))?;
let aux_done_event = aux_stream.create_event(false)
    .map_err(|e| anyhow::anyhow!("Aux done event: {e}"))?;

• Step 2: Compile and test
• Step 3: Commit

Task 2.2: Move auxiliary ops to aux_stream

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs:640-860 (run_full_step)

• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs (stream parameter)

• Modify: crates/ml/src/cuda_pipeline/gpu_iqn_head.rs (stream parameter)

• Modify: crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs (stream parameter)

• Step 1: Record event after graph_step replay completes (main stream)

// fused_training.rs, after graph_step replay:
self.forward_done_event.record(&self.stream)?;

• Step 2: Wait for forward_done_event on aux_stream

// Before aux operations:
self.aux_stream.wait_event(&self.forward_done_event)?;

• Step 3: Move EMA, attention, IQL, IQN to aux_stream

Each auxiliary module needs to accept a stream parameter or be initialized with the aux_stream. The simplest approach: pass &self.aux_stream to each call.

For attention: self.trainer.apply_attention_forward_on_stream(attn, self.batch_size, &self.aux_stream) For IQL: iql.train_value_step_on_stream(states, rewards, &self.aux_stream) For IQN: iqn.train_iqn_step_gpu_on_stream(..., &self.aux_stream) For EMA: self.trainer.target_ema_update_on_stream(..., &self.aux_stream)

Alternative (simpler): Store stream as a field in each aux module and swap it once during init.

• Step 4: Record aux_done_event on aux_stream, wait on main before next upload

// After all aux ops:
self.aux_done_event.record(&self.aux_stream)?;

// Before next step's upload_batch_gpu (in the step loop):
self.stream.wait_event(&self.aux_done_event)?;

This ensures the next step's batch upload doesn't overwrite buffers that the aux ops are still reading.

• Step 5: Compile and test
• Step 6: Deploy and measure epoch time

Expected: 2000 steps x 8ms = 16s/epoch.

• Step 7: Commit

git commit -m "perf: dual-stream overlap — aux ops (EMA+attn+IQL+IQN) parallel with graph_step"

---

Phase 3: Kernel Fusion (Week 3)

Rationale

The aux_stream bottleneck is IQN (~5ms: trunk forward + fwd+loss + backward + grad_norm + Adam = 5 kernel launches) and attention (~3ms: fwd + bwd + Adam = 3 kernel launches). Fusing reduces launch overhead and improves data locality.

Expected: aux_stream drops from ~8ms to ~3ms → max(3ms, 3ms) = 3ms/step = 6s/epoch

Task 3.1: Fuse IQN forward+loss+backward into single warp kernel

Files:

• Create: crates/ml/src/cuda_pipeline/iqn_fused_kernel.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_iqn_head.rs:389-575 (execute_training_pipeline)

Currently IQN runs 5 sequential kernels. The forward+loss+backward kernels all use the same online_h_s2 input and operate per-sample with one warp (32 threads). Fuse into a single kernel:

extern "C" __global__ void iqn_fused_train_kernel(
    /* forward inputs */
    const __nv_bfloat16* __restrict__ online_h_s2,
    const __nv_bfloat16* __restrict__ target_h_s2,
    const float* __restrict__ online_taus,
    const float* __restrict__ target_taus,
    const int* __restrict__ branch_actions,
    const float* __restrict__ rewards,
    const float* __restrict__ dones,
    float gamma,
    /* weights */
    const float* __restrict__ online_params,
    const float* __restrict__ target_params,
    /* outputs */
    float* __restrict__ per_sample_loss,
    float* __restrict__ total_loss,
    float* __restrict__ grad_buf,
    float* __restrict__ d_h_s2_buf,
    /* dims */
    int batch_size, int shared_h1, int hidden_dim, int embed_dim
) {
    // Each block = one sample. Warp executes forward → loss → backward sequentially.
    // All intermediate activations in registers/shared memory — no global memory round-trip.
    int sample = blockIdx.x;
    if (sample >= batch_size) return;

    // Forward pass (registers)
    // ... tau embedding → combined → Q-values ...

    // Loss computation (Huber, in registers)
    // ... quantile Huber loss ...

    // Backward pass (registers → atomicAdd to grad_buf)
    // ... chain rule through all layers ...
}

• Step 1: Write the fused CUDA kernel

Combine logic from iqn_dual_head_kernel.cu forward+loss and backward sections. Keep all intermediates in registers or shared memory. Single launch: grid=(batch_size), block=(32).

• Step 2: Add to cubin build in build.rs
• Step 3: Load from cubin in gpu_iqn_head.rs
• Step 4: Replace 3-kernel pipeline with single launch

fn execute_training_pipeline(&mut self, ...) -> Result<f32, MLError> {
    // Zero buffers
    self.stream.memset_zeros(&mut self.total_loss)?;
    self.stream.memset_zeros(&mut self.grad_buf)?;
    self.stream.memset_zeros(&mut self.d_h_s2_buf)?;

    // Single fused kernel: forward + loss + backward
    unsafe {
        self.stream.launch_builder(&self.fused_train_kernel)
            .arg(online_h_s2).arg(&self.target_h_s2)
            // ... all args ...
            .launch(LaunchConfig { grid_dim: (b as u32, 1, 1), block_dim: (32, 1, 1), shared_mem_bytes: 0 })
            .map_err(|e| MLError::ModelError(format!("IQN fused train: {e}")))?;
    }

    // Grad norm + Adam (keep separate — different grid sizes)
    // ... 2 kernel launches instead of 5 total ...
    Ok(0.0)
}

• Step 5: Compile, test, benchmark
• Step 6: Commit

Task 3.2: Fuse attention forward + backward + Adam

Files:

• Create: addition to crates/ml/src/cuda_pipeline/attention_kernel.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs

Attention is 4-head with 263K params. Fuse fwd+bwd+Adam into single kernel with shared memory for Q/K/V projections.

• Step 1: Write fused attention train kernel
• Step 2: Add to cubin build
• Step 3: Replace 3 method calls with single launch
• Step 4: Compile, test, benchmark
• Step 5: Commit

Task 3.3: Fuse EMA + spectral norm into single weight pass

Files:

• Create: crates/ml/src/cuda_pipeline/ema_spectral_fused_kernel.cu
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (target_ema_update, apply_spectral_norm)

Both EMA and spectral norm iterate over all weight matrices. Fuse into single kernel that applies spectral norm and EMA in one pass per weight matrix.

• Step 1: Write fused kernel

extern "C" __global__ void ema_spectral_fused(
    float* __restrict__ target,
    float* __restrict__ online,
    float tau, float sigma_max, int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    // Spectral norm: clamp online weight
    // (simplified — full spectral norm needs power iteration, keep separate for now)
    // EMA: target = (1-tau)*target + tau*online
    target[i] = (1.0f - tau) * target[i] + tau * online[i];
}

Note: Full spectral norm with power iteration may not be fusable — it requires a reduction across the matrix. If so, keep spectral norm separate and only fuse EMA.

• Step 2-5: Build, integrate, test, commit

---

Phase 4: Upload Overlap + Zero-Copy (Week 4)

Rationale

After Phase 3, both streams are ~3ms. The upload DtoD copies (~0.1ms) are negligible but the real win is eliminating the upload entirely by using CUDA graph node update.

Task 4.1: Double-buffered batch staging

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (add buffer B)

Pre-allocate two sets of staging buffers (A and B). While step N computes on buffer A, step N+1's batch is uploaded into buffer B. Swap each step.

• Step 1: Add second set of staging buffers

// gpu_dqn_trainer.rs, new fields:
states_buf_b: CudaSlice<half::bf16>,
next_states_buf_b: CudaSlice<half::bf16>,
rewards_buf_b: CudaSlice<f32>,
// ... etc for all batch buffers ...
active_buf: bool, // false = A, true = B

• Step 2: Swap active buffer each step
• Step 3: Upload next batch while current graph executes
• Step 4: Compile, test, benchmark
• Step 5: Commit

Task 4.2: cuGraphExecUpdate for zero-copy batch swap

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

Use CUDA 12's cuGraphExecUpdate API to update the graph's memory pointers without re-instantiation. The graph reads directly from the PER buffer's ring — no staging copies needed.

• Step 1: Check cudarc bindings for cuGraphExecUpdate

Run: grep -r "cuGraphExecUpdate\|graph_exec_update\|GraphExecUpdate" ~/.cargo/registry/src/*/cudarc-*/

If not available, use cuGraphExecKernelNodeSetParams to update individual kernel node parameters.

• Step 2: Implement graph node update for batch pointers
• Step 3: Remove upload_batch_gpu entirely
• Step 4: Compile, test, benchmark
• Step 5: Commit

---

Measurement Protocol

After each phase, deploy to H100 and measure:

# Deploy
git push origin main
argo submit -n foxhunt --from workflowtemplate/compile-and-train \
  -p commit-sha=$(git rev-parse --short HEAD) \
  -p model=dqn -p gpu-pool=ci-training-h100 \
  -p hyperopt-trials=0 -p train-epochs=5

# Measure (from training logs)
kubectl logs <pod> -n foxhunt -c main | grep "Training step breakdown"
# Expected format: sample=Xms fused=Yms guard=Zms (per-step: ...)

Expected Results

Phase	Per-step	Epoch (2000 steps)	Speedup vs baseline (614s)
Baseline (pre-fix)	307ms	614s	1x
Sync removal (today)	~25ms	~50s	12x
Phase 1: mega-graph	~11ms	~22s	28x
Phase 2: dual-stream	~8ms	~16s	38x
Phase 3: kernel fusion	~3ms	~6s	102x
Phase 4: zero-copy	~2ms	~4s	154x

---

Risk Register

Risk	Impact	Mitigation
cudarc CUDA Graph API limitations	Can't merge graphs	Use cudarc's raw FFI bindings (cudarc::driver::sys::*)
CQL gradient has hidden state dependency	Graph captures wrong state	Verify CQL reads only d_value_logits_buf + d_adv_logits_buf (both in graph)
Dual-stream race condition on shared buffers	Incorrect gradients	save_h_s2 is read-only by aux stream, written by graph. CUDA event sync ensures ordering
IQN fused kernel register pressure	Low occupancy on H100	323K params × batch=1024 fits in L2 cache. Register pressure is for intermediates only — profile with ncu
cuGraphExecUpdate not in cudarc	Can't do zero-copy	Fall back to double-buffered staging (Task 4.1) — still eliminates upload latency