plan: unified single-graph — 5 tasks, zero ungraphed kernel launches

Task 1: CUfunction audit (map every kernel to its child graph) Task 2: Expose submit methods (post_aux + maintenance + prev_grad_buf) Task 3: Add new child graphs (capture + launch 7 children) Task 4: Verify remaining outside-graph code is kernel-free Task 5: Deploy and validate (target: adam ~30ms, epoch ~64s) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-18 21:26:40 +02:00
parent 55f26e0f51
commit c054b4e1d8
1 changed files with 309 additions and 0 deletions
--- a/docs/superpowers/plans/2026-04-18-unified-single-graph.md
+++ b/docs/superpowers/plans/2026-04-18-unified-single-graph.md
@@ -0,0 +1,309 @@
+# 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"
+```