# Fused CUDA Training Kernel — Design Spec

## Problem

The DQN training loop dispatches **2,100+ CUDA kernels per batch** through Candle:
- Forward: ~50 dispatches (3 network passes × ~17 ops)
- Loss: ~60 dispatches (branching C51 cross-entropy)
- Backward: ~1,000 dispatches (autograd traversal)
- Gradient clipping: ~500 dispatches (per-parameter norm + scale)
- Adam optimizer: ~500 dispatches (per-parameter update)

At 1,531 batches/epoch = **3.2M kernel launches/epoch**. Each launch carries ~5–10μs CPU overhead. On H100, the network is tiny (210K params) — pure compute per batch is <1μs, but launch overhead adds **~20ms per batch**. Total: ~30s/epoch of pure dispatch latency vs <1.5s of actual compute.

**Root cause:** Candle's `cuMemAlloc_v2` allocations during backward invalidate CUDA Graph capture. We cannot wrap the existing Candle chain in a graph.

## Solution

Replace the entire Candle training dispatch chain with **3 fused CUDA kernels**, captured in a **single CUDA Graph**:

```
Per-batch (outside graph):  PER sample → GPU batch ready
CUDA Graph replay (1 launch): [forward_loss → backward → adam_update]
Per-batch (outside graph):  PER priority update → target network EMA (every N steps)
```

Dispatch reduction: **2,100 → 1** per batch. Epoch speedup: **~20x** on launch overhead.

## Architecture

### Network: Branching Dueling DQN with C51

Production config (from BranchingConfig):
- `state_dim`: 48 (no OFI) or 56 (with OFI), tensor-core aligned
- `shared_hidden_dims`: [256, 256]
- `value_hidden_dim`: 128, outputs [1, num_atoms] = [1, 51]
- `branch_hidden_dim`: 128
- `branch_sizes`: [5, 3, 3] (exposure, order, urgency)
- `num_atoms`: 51 (C51 distributional)
- `v_min`: -25.0, `v_max`: 25.0
- `use_distributional`: true
- `use_noisy`: true (NoisyNet in heads)
- Activation: LeakyReLU(0.01) after shared/FC layers

### Weight Layout (20 tensors, ~210K params)

From DuelingWeightSet + BranchingWeightSet in `gpu_weights.rs`:

| Layer | Weight | Shape | Count |
|-------|--------|-------|-------|
| shared_0 | w_s1, b_s1 | [256, STATE_DIM], [256] | ~14.6K |
| shared_1 | w_s2, b_s2 | [256, 256], [256] | 65.8K |
| value_fc | w_v1, b_v1 | [128, 256], [128] | 32.9K |
| value_out | w_v2, b_v2 | [51, 128], [51] | 6.6K |
| branch_0_fc | w_a1, b_a1 | [128, 256], [128] | 32.9K |
| branch_0_out | w_a2, b_a2 | [255, 128], [255] | 32.9K |
| branch_1_fc | w_bo1, b_bo1 | [128, 256], [128] | 32.9K |
| branch_1_out | w_bo2, b_bo2 | [153, 128], [153] | 19.7K |
| branch_2_fc | w_bu1, b_bu1 | [128, 256], [128] | 32.9K |
| branch_2_out | w_bu2, b_bu2 | [153, 128], [153] | 19.7K |
| **Total** | | | **~291K** |

Note: Branch output dims = n_d × num_atoms (distributional). Non-distributional: n_d only.

### Three Forward Passes Per Batch

1. **Online on states** (training=true, noise+dropout active) → current Q-distribution
2. **Target on next_states** (training=false) → target Q-distribution for Bellman
3. **Online on next_states** (training=false) → action selection for Double DQN

### C51 Loss (Per Branch)

For each branch d ∈ {0,1,2}:
1. Gather current log-probs for taken action aᵈ: [batch, num_atoms]
2. Get target probs for best next action (Double DQN selects): [batch, num_atoms]
3. Bellman projection: Tᵤ = r + γz (1-done), project onto support atoms via linear interpolation
4. Cross-entropy: CEᵈ = -Σⱼ projected[j] × current_log_probs[j]

Final: loss = (1/3) Σᵈ CEᵈ × IS_weight

## Phase 1: Forward + Loss Kernel

**File:** `crates/ml/src/cuda_pipeline/dqn_training_kernel.cu`

### Kernel Signature

```c
extern "C" __global__ void dqn_forward_loss_kernel(
    // Batch data (from GPU PER)
    const float* states,         // [B, STATE_DIM]
    const float* next_states,    // [B, STATE_DIM]
    const int*   actions,        // [B] factored action indices (0-44)
    const float* rewards,        // [B]
    const float* dones,          // [B]
    const float* is_weights,     // [B] PER importance-sampling weights
    // Online network weights (20 tensors via #define pointers)
    // Target network weights (20 tensors)
    // Saved activations output (for backward kernel)
    float* saved_h_s1,           // [B, SHARED_H1]
    float* saved_h_s2,           // [B, SHARED_H2]
    float* saved_h_v1,           // [B, VALUE_H]
    float* saved_h_bd,           // [B, 3, ADV_H] (3 branches)
    float* saved_logits,         // [B, total_branch_output_atoms]
    float* saved_target_probs,   // [B, 3, NUM_ATOMS] projected targets (detached)
    // Outputs
    float* out_loss,             // [B] per-sample weighted loss
    float* out_td_errors,        // [B] for PER priority update
    float* out_total_loss,       // [1] mean loss (for monitoring)
    // Config
    int batch_size,
    float gamma,
    float v_min, float v_max, int num_atoms
);
```

### Thread Layout

- grid = (batch_size, 1, 1) — one block per sample
- block = (32, 1, 1) — one warp per block (matches existing pattern)
- Shared memory: SHMEM_TILE_ROWS × SHMEM_MAX_IN_DIM for weight tiling

### Algorithm (per block/sample)

1. Load state[i] into distributed registers (stride-32)
2. **Online forward (training):**
   - shared_0 → LeakyReLU → shared_1 → LeakyReLU → save h_s2
   - value_fc → LeakyReLU → save h_v1 → value_out → value_logits [51]
   - For each branch d: branch_fc → LeakyReLU → save h_bd → branch_out → branch_logits [n_d×51]
   - Compute dueling: Q_d = V + A_d - mean(A_d) per atom
   - log_softmax per action within branch → save logits for backward
3. **Target forward (inference):** Same architecture with target weights, no saves
4. **Online forward on next_states (inference):** For Double DQN action selection
5. **Action decomposition:** factored_action → (exposure_idx, order_idx, urgency_idx)
6. **C51 loss per branch:**
   - Gather current log-probs for taken action
   - Select best next action (from online forward on next_states)
   - Gather target probs for best next action
   - Bellman projection: T_z = r + γ×z×(1-done), clip to [v_min, v_max]
   - Linear interpolation onto support atoms (scatter-add)
   - Cross-entropy: -Σ projected × current_log_probs
7. **Average over branches, multiply by IS weight**
8. **Write:** per-sample loss, td_error, atomicAdd to total_loss

### NoisyNet Handling

NoisyNet layers use: y = (μ_w + σ_w ⊙ ε_w) x + (μ_b + σ_b ⊙ ε_b)

For the fused kernel, noise tensors ε are pre-generated on GPU before the graph replay. The kernel reads them as additional input buffers. This allows different noise per batch while keeping the graph structure fixed.

## Phase 2: Backward Kernel

```c
extern "C" __global__ void dqn_backward_kernel(
    // Saved activations (from forward kernel)
    const float* saved_h_s1, saved_h_s2, saved_h_v1, saved_h_bd, saved_logits,
    const float* saved_target_probs,
    const float* states,         // [B, STATE_DIM] for input gradients
    const float* is_weights,     // [B]
    // Online network weights (for W^T in backward)
    // ... (same weight pointers as forward)
    // Gradient output buffers (zero-initialized before launch)
    float* grad_w_s1, float* grad_b_s1,  // [256, STATE_DIM], [256]
    float* grad_w_s2, float* grad_b_s2,  // ...
    // ... all 20 gradient tensors
    int batch_size
);
```

### Backward Math

For each sample i (one block per sample):

1. **∂L/∂logits** from C51 cross-entropy:
   - For each branch d, for the taken action aᵈ:
   - ∂CE/∂log_p = -projected_target  (shape: [num_atoms])
   - Through log_softmax: ∂L/∂z = ∂L/∂log_p - softmax(z) × Σ(∂L/∂log_p)

2. **Backward through branch output layers** (linear, no activation):
   - ∂L/∂W_out = ∂L/∂z × h_bd^T → atomicAdd to grad buffer
   - ∂L/∂b_out = ∂L/∂z → atomicAdd
   - ∂L/∂h_bd = W_out^T × ∂L/∂z

3. **Backward through branch FC** (linear + LeakyReLU):
   - ∂L/∂pre_bd = ∂L/∂h_bd ⊙ leaky_relu'(pre_bd)
   - ∂L/∂W_fc = ∂L/∂pre_bd × h_s2^T → atomicAdd
   - ∂L/∂b_fc = ∂L/∂pre_bd → atomicAdd
   - Accumulate ∂L/∂h_s2 from all 3 branches + value head

4. **Backward through value head** (similar to branches)

5. **Backward through shared layers** (accumulated gradients from value + all branches):
   - shared_1 backward
   - shared_0 backward

6. **Gradient accumulation:** Each block atomicAdds its per-sample gradient contribution to global gradient buffers.

### atomicAdd Strategy

- 291K gradient elements
- batch_size = 256 concurrent blocks
- Average ~1 atomicAdd per element (256 writes to 291K elements)
- H100 L2 cache (50MB) holds all gradient buffers — atomics are fast

## Phase 3: Adam Optimizer Kernel

```c
extern "C" __global__ void dqn_adam_kernel(
    // Parameters (read + write, flattened)
    float* params,        // [TOTAL_PARAMS] all weights flattened
    // Gradient buffer (read only)
    const float* grads,   // [TOTAL_PARAMS]
    // Adam state (read + write)
    float* m,             // [TOTAL_PARAMS] first moment
    float* v,             // [TOTAL_PARAMS] second moment
    // Config
    float lr, float beta1, float beta2, float eps,
    float weight_decay, float max_grad_norm,
    int step,             // for bias correction
    int total_params,
    // Output
    float* out_grad_norm  // [1] pre-clip gradient L2 norm
);
```

### Algorithm

1. **Gradient norm** (parallel reduction over all params):
   - Each thread accumulates partial norm_sq for its elements
   - Block-level reduction → global atomicAdd → shared grad_norm

2. **Gradient clipping:**
   - scale = min(max_grad_norm / (grad_norm + eps), 1.0)
   - Each thread: grad[i] *= scale

3. **Adam update** (per element, trivially parallel):
   - m[i] = β1 × m[i] + (1-β1) × grad[i]
   - v[i] = β2 × v[i] + (1-β2) × grad[i]²
   - m_hat = m[i] / (1 - β1^step)
   - v_hat = v[i] / (1 - β2^step)
   - params[i] -= lr × m_hat / (√v_hat + eps) + weight_decay × params[i]

Grid: (ceil(TOTAL_PARAMS / 256), 1, 1), block: (256, 1, 1)

## Phase 4: CUDA Graph Integration

### Graph Structure

```
PER sample (outside graph)
  ↓
begin_capture(THREAD_LOCAL)
  ↓
[noise_generate_kernel]  ← NoisyNet noise (if enabled)
  ↓
[dqn_forward_loss_kernel]
  ↓
[zero_grad_kernel]  ← memset gradient buffers to 0
  ↓
[dqn_backward_kernel]
  ↓
[dqn_adam_kernel]
  ↓
end_capture(AUTO_FREE_ON_LAUNCH)
  ↓
graph.launch() per batch (replay)
  ↓
PER priority update (outside graph)
Target EMA update (outside graph, every N steps)
```

### Input Tensor Updates Between Replays

CUDA Graph replays with fixed buffer pointers. Between replays:
- Write new batch data (states, next_states, actions, rewards, dones, is_weights) into **pre-allocated fixed buffers**
- The PER sampling writes directly into these buffers
- No new allocations between replays

### Graph Invalidation

Rebuild graph when:
- batch_size changes (should not happen — fixed in config)
- Network architecture changes (never during training)
- First batch of training (initial capture)

### Existing Pattern

Follow `gpu_backtest_evaluator.rs`:
- `SendSyncGraph` wrapper for `CudaGraph`
- `CU_STREAM_CAPTURE_MODE_THREAD_LOCAL`
- `CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH`

## Phase 5: Rust Host Code

**File:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

### GpuDqnTrainer Struct

```rust
pub struct GpuDqnTrainer {
    stream: Arc<CudaStream>,
    // Compiled kernels
    forward_loss_fn: CudaFunction,
    backward_fn: CudaFunction,
    adam_fn: CudaFunction,
    // Pre-allocated buffers (fixed size for graph capture)
    batch_states: CudaSlice<f32>,      // [B, STATE_DIM]
    batch_next_states: CudaSlice<f32>, // [B, STATE_DIM]
    batch_actions: CudaSlice<i32>,     // [B]
    batch_rewards: CudaSlice<f32>,     // [B]
    batch_dones: CudaSlice<f32>,       // [B]
    batch_is_weights: CudaSlice<f32>,  // [B]
    // Activation save buffers
    saved_activations: SavedActivations,
    // Gradient buffers (zero before each backward)
    gradient_buf: CudaSlice<f32>,      // [TOTAL_PARAMS]
    // Adam state
    adam_m: CudaSlice<f32>,            // [TOTAL_PARAMS]
    adam_v: CudaSlice<f32>,            // [TOTAL_PARAMS]
    // Outputs
    per_sample_loss: CudaSlice<f32>,   // [B]
    td_errors: CudaSlice<f32>,         // [B]
    total_loss: CudaSlice<f32>,        // [1]
    grad_norm: CudaSlice<f32>,         // [1]
    // CUDA Graph
    training_graph: Option<SendSyncGraph>,
    // NoisyNet noise buffers
    noise_buffers: Option<NoiseBuffers>,
}
```

### Integration with DQN train_step()

```rust
// In dqn.rs train_step():
#[cfg(feature = "cuda")]
if let Some(ref mut gpu_trainer) = self.gpu_trainer {
    // 1. PER samples into gpu_trainer's fixed input buffers
    self.memory.sample_into(&mut gpu_trainer.batch_buffers)?;

    // 2. Generate NoisyNet noise (if enabled)
    gpu_trainer.generate_noise()?;

    // 3. Replay CUDA Graph (or capture on first batch)
    gpu_trainer.train_step()?;

    // 4. Update PER priorities from td_errors output
    self.memory.update_priorities_gpu(
        &gpu_trainer.batch_indices, &gpu_trainer.td_errors
    )?;

    // 5. Target network EMA (every N steps, outside graph)
    self.update_target_networks()?;

    return Ok(GpuTrainResult {
        loss_gpu: gpu_trainer.total_loss_tensor()?,
        grad_norm_gpu: gpu_trainer.grad_norm_tensor()?,
    });
}
```

## Dimension Injection

Follow existing pattern from `gpu_experience_collector.rs` — all dimensions injected via `#define` before NVRTC compilation:

```rust
let dim_overrides = format!(
    "#define STATE_DIM {state_dim}\n\
     #define MARKET_DIM {market_dim}\n\
     #define PORTFOLIO_DIM 3\n\
     #define SHARED_H1 {shared_h1}\n\
     #define SHARED_H2 {shared_h2}\n\
     #define VALUE_H {value_h}\n\
     #define ADV_H {adv_h}\n\
     #define NUM_ATOMS {num_atoms}\n\
     #define V_MIN {v_min:.6}\n\
     #define V_MAX {v_max:.6}\n\
     #define BRANCH_0_SIZE {branch_0}\n\
     #define BRANCH_1_SIZE {branch_1}\n\
     #define BRANCH_2_SIZE {branch_2}\n\
     #define BATCH_SIZE {batch_size}\n\
     #define TOTAL_PARAMS {total_params}\n\
     #define USE_NOISY_NETS {use_noisy}\n\
     #define USE_DOUBLE_DQN 1\n"
);
```

## Test Strategy

1. **Numerical correctness test:** Run fused kernel and Candle path on same input, compare:
   - Forward outputs (Q-values) within 1e-5
   - Loss values within 1e-4
   - Gradients within 1e-3 (accumulated atomicAdd tolerance)
   - Adam-updated params within 1e-4

2. **CUDA Graph smoke test:** Capture + 10 replays, verify loss decreases monotonically

3. **Integration test:** Full training epoch with fused path, compare learning curve against Candle baseline

## Risk Assessment

| Risk | Mitigation |
|------|-----------|
| Gradient correctness | Numerical diff test against Candle autograd |
| atomicAdd precision | Use float atomicAdd (sufficient for 256-sample batch) |
| NoisyNet noise patterns | Pre-generated noise buffers, updated between graph replays |
| Graph invalidation | Only rebuild on first batch (all dims fixed) |
| Occupancy concerns | Network is tiny — compute <1μs, bottleneck is purely launch overhead |
| Dropout in training | Compile-time disabled for fused path (NoisyNet replaces dropout for exploration) |

## Performance Targets

| Metric | Current | Target |
|--------|---------|--------|
| Kernel launches/batch | 2,100 | 1 (graph replay) |
| Launch overhead/batch | ~15ms | ~1μs |
| Epoch time (launch only) | ~23s | ~1.5ms |
| Expected epoch speedup | - | ~15-20x |
