foxhunt/docs/superpowers/plans/2026-03-15-fused-cuda-training.md

Fused CUDA Training Kernel — Implementation Plan

For agentic workers: REQUIRED: Use superpowers:subagent-driven-development (if subagents available) or superpowers:executing-plans to implement this plan. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Replace 2,100+ Candle kernel dispatches per DQN training batch with 3 fused CUDA kernels captured in a CUDA Graph, achieving ~15-20x epoch speedup on H100.

Architecture: Fused forward+loss, backward, and Adam kernels bypass Candle entirely. Weight pointers extracted from existing DuelingWeightSet/BranchingWeightSet. CUDA Graph captures the fixed-shape kernel sequence for zero-overhead replay.

Tech Stack: CUDA (NVRTC), cudarc 0.17.3, Candle (weight storage only), existing common_device_functions.cuh infrastructure.

Spec: docs/superpowers/specs/2026-03-15-fused-cuda-training-design.md

---

Chunk 1: Forward + Loss Kernel

Task 1: Write the CUDA forward+loss kernel

Files:

• Create: crates/ml/src/cuda_pipeline/dqn_training_kernel.cu

Dependencies: Uses macros/functions from common_device_functions.cuh (TILE_LAYER_WARP_CLEAN, q_forward_dueling_warp_shmem pattern, cooperative_load_tile, warp_matvec_leaky_relu_shmem).

• Step 1: Write kernel entry point and data loading

Kernel loads batch data (states, next_states, actions, rewards, dones, IS weights) into registers. One block (32 threads = 1 warp) per sample. Reuses existing warp-cooperative pattern.

• Step 2: Implement 3 forward passes

Extend existing q_forward_dueling_warp_shmem pattern:

• Online forward on states (save activations for backward: h_s1, h_s2, h_v1, h_bd)
• Target forward on next_states (no saves)
• Online forward on next_states for Double DQN action selection (no saves)

For distributional mode: output is [n_d × num_atoms] per branch, apply log_softmax per action.

• Step 3: Implement C51 distributional loss

Per-branch cross-entropy:

1. Decompose factored action: exp_idx = action / 9, ord_idx = (action % 9) / 3, urg_idx = action % 3
2. Gather current log-probs for taken action: index into log_softmax output
3. Select best next action from online forward on next_states (argmax)
4. Gather target probs for best next action
5. Bellman projection: T_z = r + γ×z×(1-done), clip to [v_min, v_max], linear interpolation scatter
6. Cross-entropy: -Σ projected × current_log_probs
7. Average over 3 branches, multiply by IS weight

• Step 4: Write outputs

Write per-sample loss, td_errors to global memory. Lane 0 atomicAdd to total_loss. Save activations to pre-allocated buffers for backward kernel.

Task 2: Write the Rust host code for forward+loss kernel

Files:

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

• Modify: crates/ml/src/cuda_pipeline/mod.rs (add module)

• Step 1: Define GpuDqnTrainer struct with pre-allocated buffers

All buffers allocated once at construction (fixed shapes for CUDA Graph compatibility):

• Batch input buffers (states, next_states, actions, rewards, dones, is_weights)

• Activation save buffers (h_s1, h_s2, h_v1, h_bd, logits, target_probs)

• Output buffers (per_sample_loss, td_errors, total_loss)

• Step 2: Implement NVRTC compilation with #define injection

Follow existing pattern from compile_forward_kernel() in gpu_backtest_evaluator.rs:

• dim_overrides → common_device_functions.cuh → dqn_training_kernel.cu

• Inject: STATE_DIM, SHARED_H1/H2, VALUE_H, ADV_H, NUM_ATOMS, V_MIN, V_MAX, branch sizes, BATCH_SIZE

• Step 3: Implement kernel launch

Extract weight pointers from DuelingWeightSet + BranchingWeightSet (online + target). Launch with grid=(batch_size, 1, 1), block=(32, 1, 1), shared_mem=tile_size.

• Step 4: Test forward+loss numerical correctness

Run same batch through Candle path and fused kernel, compare:

• Q-values within 1e-5 relative error
• Per-sample loss within 1e-4
• td_errors within 1e-4

---

Chunk 2: Backward Kernel

Task 3: Write the CUDA backward kernel

Files:

• Modify: crates/ml/src/cuda_pipeline/dqn_training_kernel.cu

• Step 1: Implement gradient through C51 cross-entropy + log_softmax

For each branch d, for taken action a_d:

• ∂CE/∂log_p = -projected_target (shape: [num_atoms])

• Through log_softmax: ∂L/∂z = ∂L/∂log_p - softmax(z) × Σ(∂L/∂log_p)

• Step 2: Implement backward through linear layers

For each linear layer (reverse order):

• ∂L/∂W = outer_product(∂L/∂y, x) → atomicAdd to gradient buffer
• ∂L/∂b = ∂L/∂y → atomicAdd
• ∂L/∂x = matmul(W^T, ∂L/∂y)
• Through LeakyReLU: mask by sign of pre-activation

Accumulate ∂L/∂h_s2 from value head + all 3 branch heads.

• Step 3: Wire zero-init of gradient buffers

Before backward kernel: memset gradient buffers to 0. Use stream.memset_zeros() (captured in CUDA Graph).

• Step 4: Test backward numerical correctness

Compare gradients against Candle backward() within 1e-3 tolerance (atomicAdd accumulation noise).

Task 4: Extend Rust host for backward kernel

Files:

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

• Step 1: Add gradient buffer allocation

Flattened gradient buffer: [TOTAL_PARAMS] floats. Layout matches parameter flattening order. Zero-init between batches.

• Step 2: Implement backward kernel launch

Pass saved activation buffers, weight pointers (for W^T), gradient output buffers. Grid/block same as forward kernel.

---

Chunk 3: Adam Optimizer Kernel

Task 5: Write the Adam optimizer kernel

Files:

• Modify: crates/ml/src/cuda_pipeline/dqn_training_kernel.cu

• Step 1: Implement gradient norm + clipping

Two-pass approach:

• Pass 1: Each thread accumulates grad² for its elements, block reduction, atomicAdd to global norm

• __threadfence() + last-block detection

• Pass 2: scale = min(max_norm / (norm + eps), 1.0), each thread scales its grads

• Step 2: Implement Adam update

Per-element (trivially parallel, 256 threads per block, ceil(TOTAL_PARAMS/256) blocks):

• m[i] = β1×m[i] + (1-β1)×g[i]

• v[i] = β2×v[i] + (1-β2)×g[i]²

• m_hat = m[i] / (1-β1^t), v_hat = v[i] / (1-β2^t)

• param[i] -= lr × m_hat / (√v_hat + ε) + wd × param[i]

• Step 3: Output grad_norm

Write pre-clip gradient L2 norm to output buffer (for monitoring).

Task 6: Extend Rust host for Adam kernel

Files:

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

• Step 1: Add Adam state buffers (m, v)

Allocated once, persisted across batches. Same size as gradient buffer. Initialize to zeros.

• Step 2: Implement parameter flattening/unflattening

Map Candle VarMap tensors ↔ flat CUDA buffer for optimizer. After Adam update, sync back to VarMap tensors for target network EMA.

• Step 3: Test Adam correctness

Compare parameter updates against Candle Adam within 1e-4.

---

Chunk 4: CUDA Graph Capture

Task 7: Implement CUDA Graph capture and replay

Files:

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

• Step 1: Implement graph capture

On first batch:

stream.begin_capture(CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
// Launch: noise_gen (if noisy), zero_grad, forward_loss, backward, adam
stream.end_capture(CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)?;

• Step 2: Implement graph replay loop

Subsequent batches:

1. Write new batch data into pre-allocated input buffers (outside graph)
2. Generate new NoisyNet noise (outside graph)
3. graph.launch() — replays entire training step

• Step 3: Implement graph invalidation

invalidate_training_graph() — called when weights need manual update (target network sync). Rebuild graph on next batch.

• Step 4: Smoke test: capture + 100 replays

Verify loss decreases, no crashes, no memory leaks.

---

Chunk 5: Integration

Task 8: Wire into DQN train_step

Files:

• Modify: crates/ml-dqn/src/dqn.rs (train_step method)

• Modify: crates/ml/src/cuda_pipeline/mod.rs (register module)

• Step 1: Add gpu_trainer field to DQNAgent

Lazily initialized on first CUDA train_step. Requires: network dims, batch_size, config.

• Step 2: Modify train_step to use fused path

#[cfg(feature = "cuda")]
if self.gpu_trainer.is_some() {
    // Fused CUDA path — bypasses Candle entirely
    return self.train_step_fused(batch);
}
// Fallback: existing Candle path (tests, CPU builds)

• Step 3: Implement weight sync

After target network EMA update (Candle VarMap), sync changed weights to CUDA buffers. Invalidate CUDA Graph (target weights changed).

• Step 4: End-to-end integration test

Full training run (1 epoch, 100 batches) with fused path. Compare:

• Final loss within 10% of Candle path (different noise patterns expected)

• Learning curve shape similar

• Step 5: Commit

git add crates/ml/src/cuda_pipeline/dqn_training_kernel.cu \
        crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \
        crates/ml/src/cuda_pipeline/mod.rs \
        crates/ml-dqn/src/dqn.rs
git commit -m "feat(cuda): fused forward+loss+backward+adam training kernel with CUDA Graph"