# 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:
```rust
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**

```rust
#[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**

```bash
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"
```
