# Unified Single-Graph Architecture Implementation Plan

> **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:** Move ALL ungraphed kernel launches into child graphs — eliminates CUfunction corruption on Hopper that causes 3100ms adam_child replays.

**Architecture:** 7 child graphs captured once at step 0, replayed unconditionally every step. Zero ungraphed kernel launches. Previous outside-graph ops (selectivity, denoise, causal, vaccine, risk_sgd, multi_horizon, IQL modulate, Q-stats) all become new child graphs. `prev_grad_buf` snapshot enables meaningful vaccine gradient comparison.

**Tech Stack:** CUDA 13, cudarc (vendored), cublasLtMatmul, CU_STREAM_CAPTURE_MODE_RELAXED

---

### Task 1: CUfunction Audit — Map Every Kernel to Its Child Graph

**Files:**
- Read-only audit of `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
- Read-only audit of `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: List every CUfunction field and every child graph it's captured in**

For each of the ~104 CUfunction fields in `GpuDqnTrainer`, determine which child graph(s) it's used in AND whether it's launched ungraphed. Output a table:

```
| CUfunction field | forward_child | ddqn_child | aux_child | adam_child | ungraphed |
|------------------|---------------|------------|-----------|-----------|-----------|
| saxpy_f32_kernel | YES (blend)   | no         | YES (IQN) | YES (mamba2) | YES (vaccine) |
```

Focus on the CONFLICTS — any kernel that appears in both a child graph AND ungraphed column.

- [ ] **Step 2: Document the conflict list**

Create a list of all conflicting CUfunctions that MUST be moved from ungraphed to graphed. This informs which operations go into which new child graphs.

- [ ] **Step 3: Commit the audit as a markdown file**

Save to `docs/superpowers/audit/2026-04-18-cufunction-sharing-audit.md`.

```bash
git commit -m "audit: CUfunction sharing map across child graphs and ungraphed ops"
```

---

### Task 2: Expose Submit Methods for Outside-Graph Ops

**Files:**
- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`

Currently `run_causal_intervention`, `apply_gradient_vaccine`, `step_selectivity_adam`, `step_denoise_adam`, `step_risk_sgd`, `multi_horizon_value_forward` are called from `fused_training.rs` as individual ops with conditional guards. They need to be callable unconditionally from child graph capture closures.

- [ ] **Step 1: Create `submit_post_aux_ops` method**

New public method on `GpuDqnTrainer` that unconditionally runs:
1. `launch_selectivity_forward(batch_size)`
2. `launch_selectivity_backward(batch_size, 1.0)` — mean_loss=1.0 constant
3. `step_selectivity_adam()` — remove the conditional guard
4. `compute_denoise_target_q(batch_size)`
5. `launch_q_denoise_backward(batch_size)`
6. `step_denoise_adam()` — remove the conditional guard
7. `launch_q_stats_kernel()` — the Q-stats reduce (already written)

Remove all `if let Err ... non-fatal` wrappers. If any fails, it's fatal (graph capture requires success).

- [ ] **Step 2: Create `submit_maintenance_ops` method**

New public method that unconditionally runs:
1. `run_causal_intervention_unconditional()` — remove the step % interval check
2. `apply_gradient_vaccine_from_prev()` — new method that compares grad_buf against prev_grad_buf

- [ ] **Step 3: Add `prev_grad_buf` field and grad snapshot to adam ops**

Add `prev_grad_buf: CudaSlice<f32>` to `GpuDqnTrainer`. Allocate in constructor (same size as `grad_buf` = total_params floats).

In `submit_adam_ops`, after `launch_adam_update()`, add a DtoD copy:
```rust
// Snapshot grad_buf → prev_grad_buf for vaccine comparison
unsafe {
    cudarc::driver::result::memcpy_dtod_async(
        self.prev_grad_buf.raw_ptr(),
        self.ptrs.grad_buf,
        self.total_params * std::mem::size_of::<f32>(),
        self.stream.cu_stream(),
    )?;
}
```

- [ ] **Step 4: Create `apply_gradient_vaccine_from_prev` method**

New method that runs the vaccine kernels comparing `grad_buf` against `prev_grad_buf`:
```rust
pub fn apply_gradient_vaccine_from_prev(&self) -> Result<(), MLError> {
    // vaccine_dot: dot(grad_buf, prev_grad_buf) → vaccine_dot_buf
    // vaccine_project: project grad_buf using cosine sim → correction
    // vaccine_finalize: apply correction to grad_buf
    // (same kernels as apply_gradient_vaccine but using prev_grad_buf instead of vaccine_batch)
}
```

- [ ] **Step 5: Make `run_causal_intervention` unconditional**

Remove the `if !self.graphs_captured || step % self.causal_config.interval != 0` guard. The method should always execute its kernels.

- [ ] **Step 6: Move risk_sgd and multi_horizon into submit_post_aux_ops**

Add `step_risk_sgd()` and `multi_horizon_value_forward(batch_size)` to the end of `submit_post_aux_ops`. These currently run ungraphed and use `scale_f32_kernel`.

- [ ] **Step 7: Build and test**

```bash
SQLX_OFFLINE=true cargo check -p ml
SQLX_OFFLINE=true cargo test -p ml --lib
```

- [ ] **Step 8: Commit**

```bash
git commit -m "feat: expose unconditional submit methods for post_aux + maintenance ops"
```

---

### Task 3: Add New Child Graphs to Capture and Launch

**Files:**
- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Add child graph fields**

Add to `FusedTrainingCtx`:
```rust
post_aux_child: Option<ChildGraph>,
maintenance_child: Option<ChildGraph>,
```

Add timing events to `PhaseEvents`:
```rust
post_aux_end: cuda_sys::CUevent,
maintenance_end: cuda_sys::CUevent,
```

Initialize in constructor (None), reset_for_fold (None), event allocation, event destroy.

- [ ] **Step 2: Capture post_aux_child in capture_training_graph**

After aux_child capture, before adam_child:
```rust
let post_aux = self.capture_child_graph("post_aux", |s| {
    s.trainer.submit_post_aux_ops(s.batch_size)
        .map_err(|e| anyhow::anyhow!("{e}"))
})?;
```

- [ ] **Step 3: Capture maintenance_child in capture_training_graph**

After adam_child capture:
```rust
let maintenance = self.capture_child_graph("maintenance", |s| {
    s.trainer.submit_maintenance_ops()
        .map_err(|e| anyhow::anyhow!("{e}"))
})?;
```

- [ ] **Step 4: Launch new children in run_full_step**

Update the graphed path to launch 7 children with events:
```rust
self.spectral_child...launch(cu_stream)?;
// event: spectral_end
self.forward_child...launch(cu_stream)?;
// event: forward_end
self.ddqn_child...launch(cu_stream)?;
// event: loss_end
self.aux_child...launch(cu_stream)?;
// event: backward_end
self.post_aux_child...launch(cu_stream)?;
// event: post_aux_end
self.adam_child...launch(cu_stream)?;
// event: adam_child_end
self.maintenance_child...launch(cu_stream)?;
// event: maintenance_end
```

- [ ] **Step 5: Update ungraphed step 0 path**

In the `else` branch (ungraphed step 0), add the new ops in order:
```rust
// Phase 5 (existing): aux ops
self.submit_aux_ops(agent, gpu_batch)?;
// Phase 5b (NEW): post-aux ops
self.trainer.submit_post_aux_ops(self.batch_size)?;
// Phase 6 (existing): adam ops
// ... mamba2_backward, grad_norm, adam, ISV ...
// Phase 7 (NEW): maintenance ops
self.trainer.submit_maintenance_ops()?;
```

- [ ] **Step 6: Remove ALL outside-graph code**

Delete from `run_full_step` (after the graph launch / ungraphed step):
- Causal intervention block (lines ~1106-1112)
- Vaccine block (lines ~1113-1122)
- `self.pending_vaccine_batch = None`
- Risk SGD (line ~1138)
- Multi-horizon (line ~1142)
- Selectivity+denoise disabled comment block (lines ~1156-1161)

These operations now run INSIDE the graph. The only remaining outside-graph code:
- IQN readiness update (pinned read, no kernel)
- Readback training scalars (pinned read, no kernel)
- PER priority update (uses its own unique kernel — check if shared)
- IQL modulate_td_errors (uses its own unique kernel — check if shared)

- [ ] **Step 7: Update log_phase_timing for 7-child breakdown**

Add `post_aux_ms` and `maintenance_ms` to the breakdown log:
```rust
let post_aux_ms = elapsed(pe.backward_end, pe.post_aux_end);
let adam_child_ms = elapsed(pe.post_aux_end, pe.adam_child_end);
let maintenance_ms = elapsed(pe.adam_child_end, pe.maintenance_end);

tracing::info!(
    "GPU child graph breakdown: spectral={:.1}ms forward={:.1}ms ddqn={:.1}ms aux={:.1}ms post_aux={:.1}ms adam={:.1}ms maintenance={:.1}ms",
    spectral_ms, forward_ms, ddqn_ms, aux_ms, post_aux_ms, adam_child_ms, maintenance_ms,
);
```

- [ ] **Step 8: Build and test**

```bash
SQLX_OFFLINE=true cargo check -p ml
SQLX_OFFLINE=true cargo test -p ml --lib
```

- [ ] **Step 9: Commit**

```bash
git commit -m "feat: unified single-graph — 7 children, zero ungraphed kernel launches"
```

---

### Task 4: Verify Remaining Outside-Graph Ops Are Kernel-Free

**Files:**
- Read: `crates/ml/src/trainers/dqn/fused_training.rs`

- [ ] **Step 1: Audit all code after graph launch for kernel launches**

After the 7-child graph launch, the only remaining code should be:
- `iqn.read_total_loss()` — pinned memory read, no kernel
- `trainer.readback_training_scalars()` — pinned memory read, no kernel
- `trainer.update_iqn_readiness()` — host-side computation, no kernel
- PER priority update — verify `update_priorities_from_td` kernel is NOT shared
- IQL `modulate_td_errors` — verify `modulate_td_kernel` is NOT shared

If PER or IQL kernels ARE shared with graphed children, move them into a child graph too.

- [ ] **Step 2: Verify increment_warmup and adam_step_async are kernel-free**

These run before the graph launch. Verify they only write to pinned memory.

- [ ] **Step 3: Document the final outside-graph code surface**

List every line of code that runs outside the graph and confirm none launch kernels.

---

### Task 5: Deploy and Validate

- [ ] **Step 1: Push and deploy**

```bash
git push origin main
./scripts/argo-train.sh dqn --baseline --epochs 5
```

- [ ] **Step 2: Check child graph breakdown**

Expected:
```
GPU child graph breakdown:
  spectral     =   ~0.2ms
  forward      =  ~81ms
  ddqn         =   ~1ms
  aux          = ~233ms
  post_aux     =  ~10ms   (selectivity + denoise + Q-stats + risk_sgd)
  adam         =  ~30ms   (was 3100ms — NO MORE CUfunction corruption)
  maintenance  =   ~5ms   (causal + vaccine)
  total        = ~360ms
```

- [ ] **Step 3: Verify epoch time**

Expected: 178 × 360ms = **64s per epoch** (was 690s).

- [ ] **Step 4: Check training metrics**

Verify loss is dropping, grad_norm is stable, Q-values are moving. All features (selectivity, denoise, causal, vaccine) should be active.

- [ ] **Step 5: Update spec with measured results**

Update `docs/superpowers/specs/2026-04-18-unified-single-graph-design.md` with actual timing.

- [ ] **Step 6: Commit**

```bash
git commit -m "docs: unified single-graph validation results"
```
