# Unified Single-Graph Architecture

## Goal

Eliminate ALL ungraphed kernel launches from the training step. Every GPU operation
runs inside a single set of child graphs captured once at step 0 and replayed
unconditionally every step. Zero CUfunction sharing between graphed and ungraphed
execution — eliminates the 3100ms adam_child corruption on Hopper.

## Root Cause

On Hopper (sm_90), a CUfunction captured in a CUDA graph cannot be launched
ungraphed on the same stream without corrupting the graph's kernel state.
The current architecture captures 5 child graphs but runs selectivity, denoise,
causal intervention, and gradient vaccine UNGRAPHED after each graph replay.
These ungraphed ops share CUfunction handles (adam_update_kernel, saxpy_f32_kernel,
grad_norm_finalize_kernel, scale_f32_kernel) with the captured children, causing
3100ms adam_child replays instead of ~30ms.

## Design

### Architecture: One CUmodule, All Graphed

Load all utility kernels ONCE from ONE CUmodule (current behavior). Since
every kernel launch is inside a child graph capture, there are no ungraphed
launches — no CUfunction conflict. Cross-child sharing is safe on Hopper
(confirmed by measurement: separate CUmodule per child didn't change timing).

### Child Graph Layout

```
Training step = replay all children sequentially:

  1. spectral_child      — spectral norm (12 kernels)
  2. forward_child       — trunk forward + temporal + ISV + loss + backward (60 kernels)
                           Internal: 4 branch streams (forward + backward)
  3. ddqn_child          — Double DQN pass 3 (15 kernels)
  4. aux_child           — HER + EMA + IQL×2 + IQN + Attention + ensemble + CQL (80 kernels)
                           Internal: IQN + Attention on parallel streams
  5. post_aux_child      — selectivity fwd+bwd+adam + denoise fwd+bwd+adam + Q-stats (10 kernels)
  6. adam_child           — mamba2 bwd + grad_norm + Adam + grad_snapshot + ISV (7 kernels)
  7. maintenance_child   — causal intervention + vaccine (6 kernels)
```

Total: ~190 kernel nodes. 7 children launched sequentially. 0 ungraphed launches.

### What Moves Into the Graph

| Operation | Current | New | Child |
|-----------|---------|-----|-------|
| Selectivity fwd+bwd+adam | ungraphed, conditional (non-fatal) | graphed, unconditional | post_aux_child |
| Denoise target+bwd+adam | ungraphed, conditional (non-fatal) | graphed, unconditional | post_aux_child |
| Causal intervention | ungraphed, every N steps | graphed, every step | maintenance_child |
| Gradient vaccine | ungraphed, every 20 steps, separate batch | graphed, every step, prev_grad_buf | maintenance_child |
| Q-stats reduce | ungraphed, epoch boundary | graphed, every step | post_aux_child |
| IQN readiness update | host-side pinned read | unchanged (no kernel) | N/A |

### Dependency Chain

```
spectral → forward → ddqn → aux → post_aux → adam → maintenance
```

- `post_aux_child` runs AFTER aux because selectivity/denoise read Q-values
  produced by aux. Also writes to grad_buf which adam reads. Q-stats reads
  q_out_buf here (before Adam updates params — current step's Q-values).
- `adam_child` runs AFTER post_aux because Adam reads the final grad_buf
  (C51 + IQL + IQN + attention + ensemble + selectivity + denoise gradients).
  After Adam update, snapshots grad_buf → prev_grad_buf for vaccine comparison.
- `maintenance_child` runs AFTER adam because:
  - Causal intervention reads updated params (after Adam)
  - Vaccine compares current grad_buf against prev_grad_buf (snapshot from
    previous step's adam_child). Meaningful gradient consistency signal.

### Conditional → Unconditional Migration

**Causal intervention** currently runs every `causal_intervention_interval` steps.
Now runs every step. The kernels are tiny (~2ms) — running unconditionally adds
<1% overhead vs the 190+ kernel graph.

**Gradient vaccine** currently runs every 20 steps with a separate sampled batch.
Now runs every step. The vaccine kernels (gradient dot product + project) use
the SAME batch as the main training step. The vaccine_batch sampling is eliminated.

**Selectivity** currently wrapped in `if let Err ... non-fatal`. Now unconditional.
If selectivity forward fails on step 0 capture, the graph won't capture — fatal.
Must ensure selectivity is properly initialized before capture.

**Denoise** same as selectivity — unconditional. Must be initialized before capture.

### What's NOT in the Graph

- **Batch upload** (`upload_batch_gpu`) — different data each step, must run before replay
- **PER sampling** — host-side or async GPU
- **Experience collection** — runs between epochs
- **Validation backtest** — runs at epoch boundary
- **Curiosity training** — runs between epochs on collector stream
- **Adam step counter increment** — pinned host write (`*t_pinned = step`)
- **Stochastic depth mask update** — host-side RNG → pinned write

### Vaccine: Previous-Step Gradient Comparison

Vaccine compares current gradients against PREVIOUS step's gradients to measure
gradient direction consistency. Implementation:

- `prev_grad_buf: CudaSlice<f32>` [total_params] — allocated once, ~2.6MB
- After Adam update in `adam_child`, a DtoD copy snapshots `grad_buf → prev_grad_buf`
  (captured in the graph as a memcpy node)
- `maintenance_child` runs vaccine kernels: `dot(grad_current, prev_grad)`,
  `norm(grad_current)`, `norm(prev_grad)` → cosine similarity → project correction
- On step 0 (capture), `prev_grad_buf` is zero → vaccine produces zero correction
- From step 1+, `prev_grad_buf` has meaningful gradients from the previous step

This gives a real gradient consistency signal without needing a separate batch.
Consistent gradient direction across steps → no correction needed.
Sudden direction change → vaccine applies a stabilizing projection.

### VRAM Impact

`prev_grad_buf`: 660K × 4 bytes = 2.6MB additional. Negligible on H100.

### Capture Mode

`CU_STREAM_CAPTURE_MODE_RELAXED` (already set). Required for the aux_child's
internal multi-stream (IQN + Attention on parallel streams).

### Event Timing

Add per-child events for the new children:
- `post_aux_end` — after post_aux_child
- `maintenance_end` — after maintenance_child

Log breakdown at epoch end:
```
GPU child graph breakdown: spectral=Xms forward=Yms ddqn=Zms aux=Wms post_aux=Pms adam=Ams maintenance=Mms
```

## Success Criteria

- adam_child < 50ms (was 3100ms)
- Zero ungraphed kernel launches in training step
- All existing features (selectivity, denoise, causal, vaccine) still execute
- Epoch time < 100s (from 690s)
- 899 unit tests pass
- No CUfunction conflicts on Hopper

### CUfunction Sharing Audit

Before implementation, audit ALL CUfunction fields for cross-child sharing.
With everything graphed and no ungraphed launches, cross-child sharing is
the ONLY remaining risk. Confirmed safe by measurement (separate CUmodule
per child didn't change timing), but audit as defense-in-depth.

Rule: each CUfunction should be used in AT MOST one child graph. If shared,
either move to one child or load from separate CUmodule. The unified graph
architecture makes this easier — fewer children, clearer boundaries.

### Ensemble CUfunction Verification

`run_ensemble_step` was moved into aux_child. It uses:
- `forward_value_head_for_ensemble` (cuBLAS — handle is stateless, safe)
- `ensemble_diversity_kernel` (unique to ensemble — safe)
- `ensemble_kl_grad_kernel` (unique to ensemble — safe)
- `apply_ensemble_diversity_backward` — uses `saxpy_f32_kernel` for trunk SAXPY

`saxpy_f32_kernel` is used in aux_child (ensemble SAXPY + IQN trunk) AND
forward_child (MSE/C51 blend). Since both are graphed in different children
and cross-child sharing was measured as safe, this is acceptable. But should
be noted in the audit.

## Files Modified

- `crates/ml/src/trainers/dqn/fused_training.rs` — new children, capture everything, remove outside-graph code
- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — remove conditional guards, add prev_grad_buf, expose selectivity/denoise/causal/vaccine as public submit methods
- No new .cu files (existing kernels, just captured differently)

---

## Phase 3: Performance Optimizations (after unified graph validates)

Priority order based on impact/risk ratio.

### 3.1 Fuse Small Element-wise Kernels (HIGHEST IMPACT)

Between every cuBLAS GEMM there's a separate bias_add + ReLU kernel launch.
Each adds ~5µs dispatch overhead. With ~190 kernel nodes, ~60 are bias/ReLU.
`CUBLASLT_EPILOGUE_RELU_BIAS` fuses bias+ReLU into the GEMM epilogue — zero
extra kernel launch. `sgemm_f32_fused_relu_bias` already tries this but falls
back on failure. Making the fused path reliable halves the kernel node count
(~190 → ~130) and eliminates SM idle gaps between GEMMs.

**Impact:** ~20-30% reduction in forward+backward time.
**Risk:** Low — cuBLAS epilogue is well-supported on H100.
**Files:** `batched_forward.rs`, `batched_backward.rs`

### 3.2 Reduce IQN Quantiles 64 → 32

IQN allocates 8.6GB for tiled intermediates (`B × Q × hidden` where Q=64).
Academic literature shows diminishing returns past 32 quantiles. Halving to 32:
- Saves 4.3GB VRAM (allows larger replay buffer or headroom)
- Halves IQN's ~13 GEMMs in aux_child
- IQN is a significant chunk of aux_child's 233ms

**Impact:** ~15-20% reduction in aux_child time + 4.3GB VRAM savings.
**Risk:** Low — 32 quantiles is standard. Validate with 5-epoch train.
**Files:** `config/gpu/h100.toml`, `gpu_iqn_head.rs`

### 3.3 Double-Buffered Experience Collection

Currently: collect (2s) → train (60s) → collect → train. GPU idles during
collection. With ping-pong buffers, epoch N+1 collection runs on a separate
stream WHILE epoch N trains.

**Impact:** Saves 2s/epoch (~3% of target 60s).
**Risk:** Low — existing Phase 2 spec section 2c has the full design.
**Files:** `gpu_experience_collector.rs`, `training_loop.rs`

### 3.4 cuGraphExecUpdate for Hyperparameter Changes

When LR, c51_alpha, or epsilon change, we re-capture the entire graph (3.5s
ungraphed step 0). `cuGraphExecUpdate` can patch individual kernel node
parameters in an instantiated graph without re-capture.

**Impact:** Eliminates 3.5s at fold boundaries + epoch-level LR changes.
**Risk:** Medium — cuGraphExecUpdate has strict topology matching rules.
**Files:** `fused_training.rs`

### 3.5 Tensor Core Dimension Padding

cuBLAS WGMMA tiles on H100 are 128×128. Current dimensions:
- `shared_h2=256` — good (2 tiles)
- `adv_h=32` — suboptimal (quarter tile, wasted HMMA lanes)
- `value_h=128` — good (1 tile)

Padding `adv_h` from 32 to 128 improves WGMMA utilization for branch GEMMs.
Increases branch parameter count 4× but branch params are tiny (~5K each).

**Impact:** ~10% speedup on branch forward+backward GEMMs.
**Risk:** Low — more parameters but branch capacity isn't the bottleneck.
**Files:** `config/gpu/h100.toml`, `gpu_weights.rs`, kernel dimension constants

### 3.6 Batch Size 16384 → 8192

Counter-intuitive: at 16384, each GEMM saturates 97% of 132 SMs — no room for
multi-stream parallelism. At 8192, GEMMs use ~50% SMs — branch streams and aux
parallelism can actually fill idle SMs. 2× more steps/epoch but each ~2× faster.
Net: similar epoch time, better utilization, more frequent updates (helps learning).

**Impact:** Unlocks real multi-stream parallelism benefit. Training dynamics change.
**Risk:** Medium — more frequent updates changes convergence. Needs validation.
**Files:** `config/gpu/h100.toml`

### 3.7 Persistent Adam Kernel

Adam reads/writes 7 arrays × 660K × 4 bytes = 18.5MB. At 3.35 TB/s HBM = 5.5µs.
If measured at 30ms, the overhead is kernel launch + graph node scheduling.
A persistent kernel that loops over parameter chunks in a single launch reduces
dispatch overhead.

**Impact:** Small (adam is ~30ms after CUfunction fix, target 5µs theoretical).
**Risk:** Low — single kernel, simple loop.
**Files:** `dqn_utility_kernels.cu`, `gpu_dqn_trainer.rs`

### 3.8 Memory Pool (cuMemAllocAsync)

~200 buffer allocations via `cuMemAlloc` fragment GPU memory. `cuMemAllocAsync`
with a memory pool lets the driver reuse physical memory more efficiently and
reduces allocation time in the constructor.

**Impact:** Faster initialization, less memory fragmentation.
**Risk:** Low — pool management is driver-handled.
**Files:** `gpu_dqn_trainer.rs` constructor, cudarc extensions
