diff --git a/docs/superpowers/plans/2026-04-18-unified-single-graph.md b/docs/superpowers/plans/2026-04-18-unified-single-graph.md new file mode 100644 index 000000000..88fc1d89a --- /dev/null +++ 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` 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::(), + 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, +maintenance_child: Option, +``` + +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" +```