plan: true single-graph — 7 tasks, zero CPU on hot path

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-19 00:06:05 +02:00
parent e8d382fc2f
commit 4f7b1f4804
1 changed files with 535 additions and 0 deletions
--- a/docs/superpowers/plans/2026-04-19-true-single-graph.md
+++ b/docs/superpowers/plans/2026-04-19-true-single-graph.md
@@ -0,0 +1,535 @@
+# True Single-Graph 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:** Single `cuGraphExecLaunch` per step — zero ungraphed kernel launches, zero DtoD copies, zero host computation on GPU stream.
+
+**Architecture:** Compose all child graphs (PER sampling, training pipeline, priority update) into one parent graph via `cuGraphAddChildGraphNode`. GPU-side counters replace host writes. Direct-to-trainer gather eliminates batch upload DtoD copies.
+
+**Tech Stack:** CUDA 13, cudarc (vendored), cublasLtMatmul, `cuGraphAddChildGraphNode`, `CU_STREAM_CAPTURE_MODE_RELAXED`
+
+---
+
+### Task 1: Isolate ALL Ungraphed CUfunctions into Separate CUmodule
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
+
+Every CUfunction launched ungraphed on the training stream MUST come from a different CUmodule than any CUfunction captured in a child graph. Currently `pad_states_kernel` and `stochastic_depth_rng_kernel` are from the main utility CUmodule — same as graphed `adam_update_kernel`, `grad_norm_kernel`, etc.
+
+- [ ] **Step 1: Audit every ungraphed kernel launch**
+
+Search for all `launch_builder` calls in functions that run OUTSIDE graph capture: `upload_batch_gpu`, `update_stochastic_depth_mask`, `sample_proportional` (in gpu_replay_buffer.rs), `modulate_td_errors` (in gpu_iql_trainer.rs), `update_priorities_gpu`. List each CUfunction and which CUmodule it comes from.
+
+- [ ] **Step 2: Load ungraphed utility cubin from separate CUmodule**
+
+Already partially done — `pad_states_kernel` and `stochastic_depth_rng_kernel` moved to `ungraphed_module` in previous commit. Verify no other utility kernel from the main CUmodule is launched ungraphed. Check `relu_mask_kernel` (used in backward — graphed, OK), `clip_grad_kernel` (used where?), `spectral_norm_batched_kernel` (graphed in spectral_child, OK).
+
+- [ ] **Step 3: Build and test**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+- [ ] **Step 4: Commit**
+
+```bash
+git commit -m "fix: complete CUmodule isolation — zero ungraphed launches from graphed CUmodules"
+```
+
+---
+
+### Task 2: Write `gather_padded` and `increment_step_counters` CUDA Kernels
+
+**Files:**
+- Create: `crates/ml/src/cuda_pipeline/graph_utility_kernels.cu`
+
+- [ ] **Step 1: Write `gather_f32_rows_padded` kernel**
+
+```c
+// gather_f32_rows_padded: gathers rows from src[indices[i]] into dst with zero-padding.
+// dst row width = dst_stride (padded), src row width = src_dim (unpadded).
+// Eliminates separate gather + DtoD + pad_states pipeline.
+extern "C" __global__ void gather_f32_rows_padded(
+    float* __restrict__ dst,
+    const float* __restrict__ src,
+    const long long* __restrict__ indices,
+    int src_dim, int dst_stride, int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    int total = batch_size * dst_stride;
+    if (tid >= total) return;
+    int row = tid / dst_stride;
+    int col = tid % dst_stride;
+    long long src_row = indices[row];
+    dst[tid] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f;
+}
+```
+
+- [ ] **Step 2: Write `gather_f32_scalar` kernel (rewards, dones)**
+
+```c
+// gather_f32_scalar: gathers scalar values from src[indices[i]] into dst.
+extern "C" __global__ void gather_f32_scalar(
+    float* __restrict__ dst,
+    const float* __restrict__ src,
+    const long long* __restrict__ indices,
+    int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= batch_size) return;
+    dst[tid] = src[indices[tid]];
+}
+```
+
+- [ ] **Step 3: Write `gather_i32_scalar` kernel (actions)**
+
+```c
+extern "C" __global__ void gather_i32_scalar(
+    int* __restrict__ dst,
+    const int* __restrict__ src,
+    const long long* __restrict__ indices,
+    int batch_size)
+{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= batch_size) return;
+    dst[tid] = src[indices[tid]];
+}
+```
+
+- [ ] **Step 4: Write `increment_step_counters` kernel**
+
+```c
+// Single-thread kernel: increments all per-step counters and computes cosine-annealed tau.
+// All pointers are pinned device-mapped — GPU writes, host reads.
+extern "C" __global__ void increment_step_counters(
+    int* adam_t, int* sel_t, int* denoise_t, int* iql_t,
+    int* iql_low_t, int* iqn_t, int* attn_t, int* rng_step,
+    float* tau_out, float tau_init, float tau_final, int anneal_steps)
+{
+    int step = atomicAdd(adam_t, 1) + 1;
+    atomicAdd(sel_t, 1);
+    atomicAdd(denoise_t, 1);
+    atomicAdd(iql_t, 1);
+    atomicAdd(iql_low_t, 1);
+    atomicAdd(iqn_t, 1);
+    atomicAdd(attn_t, 1);
+    atomicAdd(rng_step, 1);
+    float progress = fminf((float)step / (float)anneal_steps, 1.0f);
+    float cosine = 0.5f * (1.0f + cosf(progress * 3.14159265f));
+    *tau_out = tau_final + (tau_init - tau_final) * cosine;
+}
+```
+
+- [ ] **Step 5: Compile to cubin**
+
+Add to the cubin compilation in `gpu_dqn_trainer.rs` constructor (same pattern as other `.cu` files — `include_bytes!` of precompiled cubin, or NVRTC compile at init). Precompile:
+
+```bash
+nvcc -cubin -arch=sm_90 -o graph_utility_kernels.cubin graph_utility_kernels.cu
+```
+
+- [ ] **Step 6: Commit**
+
+```bash
+git commit -m "feat: gather_padded + increment_step_counters CUDA kernels for true single-graph"
+```
+
+---
+
+### Task 3: Direct-to-Trainer Gather — Eliminate upload_batch_gpu
+
+**Files:**
+- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
+- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
+
+- [ ] **Step 1: Add trainer buffer pointers to GpuReplayBuffer**
+
+Add fields to `GpuReplayBuffer` for the trainer's destination buffers:
+
+```rust
+/// Trainer destination buffer pointers (set once at init, stable for graph capture).
+/// PER gather writes directly to these — eliminates DtoD copy + pad_states.
+trainer_states_ptr: u64,        // [B, state_dim_padded] f32
+trainer_next_states_ptr: u64,   // [B, state_dim_padded] f32
+trainer_actions_ptr: u64,       // [B] i32
+trainer_rewards_ptr: u64,       // [B] f32
+trainer_dones_ptr: u64,         // [B] f32
+trainer_is_weights_ptr: u64,    // [B] f32
+trainer_state_dim_padded: usize,
+```
+
+Add `pub fn set_trainer_buffers(&mut self, ...)` method to wire these from fused_training.rs init.
+
+- [ ] **Step 2: Replace gather kernels with gather_padded in sample_proportional**
+
+Replace the 6 gather calls + IS weights with direct-to-trainer gather calls using the new kernels from Task 2. States and next_states use `gather_f32_rows_padded` (handles padding). Scalars use `gather_f32_scalar` / `gather_i32_scalar`.
+
+Remove the intermediate `sample_states`, `sample_next_states`, `sample_actions`, `sample_rewards`, `sample_dones`, `sample_weights` buffers — they're no longer needed.
+
+Update the return type: `GpuBatchPtrs` now points to the trainer's buffers instead of the replay buffer's sample buffers.
+
+- [ ] **Step 3: Remove `upload_batch_gpu` from gpu_dqn_trainer.rs**
+
+Delete the `upload_batch_gpu` method entirely. The trainer's buffers are now populated directly by the PER gather.
+
+- [ ] **Step 4: Remove `upload_batch_gpu` call from fused_training.rs**
+
+In `run_full_step`, remove the `self.trainer.upload_batch_gpu(gpu_batch)` call. The PER sampling child graph now writes directly to trainer buffers.
+
+- [ ] **Step 5: Wire trainer buffer pointers at init**
+
+In `FusedTrainingCtx::new`, after creating the trainer and replay buffer:
+
+```rust
+agent.primary_dqn_mut().memory.set_trainer_buffers(
+    trainer.states_buf_ptr(),
+    trainer.next_states_buf_ptr(),
+    trainer.actions_buf_ptr(),
+    trainer.rewards_buf_ptr(),
+    trainer.dones_buf_ptr(),
+    trainer.is_weights_buf_ptr(),
+    trainer.state_dim_padded(),
+);
+```
+
+- [ ] **Step 6: Convert ALL memset_zeros in gpu_replay_buffer.rs to raw cuMemsetD8Async**
+
+Replace `self.stream.memset_zeros(&mut self.scan_tile_state)` etc. with raw `cudarc::driver::sys::cuMemsetD8Async(ptr, 0, bytes, cu_stream)`.
+
+- [ ] **Step 7: Build and test**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+- [ ] **Step 8: Commit**
+
+```bash
+git commit -m "feat: direct-to-trainer gather — eliminates upload_batch_gpu + 6 DtoD copies"
+```
+
+---
+
+### Task 4: GPU-Side Counters — Eliminate Host Writes
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
+- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_attention.rs`
+
+- [ ] **Step 1: Load `increment_step_counters` kernel from cubin**
+
+In `gpu_dqn_trainer.rs`, load from a NEW CUmodule (graph_utility cubin — NOT the main utility cubin):
+
+```rust
+let graph_util_module = context.load_cubin(GRAPH_UTILITY_CUBIN.to_vec())?;
+let increment_counters_kernel = graph_util_module.load_function("increment_step_counters")?;
+```
+
+Add field `increment_counters_kernel: CudaFunction` to `GpuDqnTrainer`.
+
+- [ ] **Step 2: Create `submit_counter_increments` method**
+
+```rust
+pub(crate) fn submit_counter_increments(&self) -> Result<(), MLError> {
+    unsafe {
+        self.stream.launch_builder(&self.increment_counters_kernel)
+            .arg(&self.t_dev_ptr)           // adam_t
+            .arg(&self.sel_t_dev_ptr)       // sel_t
+            .arg(&self.denoise_t_dev_ptr)   // denoise_t
+            .arg(&self.iql_t_dev_ptr)       // iql_t
+            .arg(&self.iql_low_t_dev_ptr)   // iql_low_t
+            .arg(&self.iqn_t_dev_ptr)       // iqn_t
+            .arg(&self.attn_t_dev_ptr)      // attn_t
+            .arg(&self.rng_step_dev_ptr)    // rng_step
+            .arg(&self.tau_dev_ptr)         // tau_out
+            .arg(&self.tau_init)
+            .arg(&self.tau_final)
+            .arg(&self.tau_anneal_steps)
+            .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 })
+            .map_err(|e| MLError::ModelError(format!("increment_counters: {e}")))?;
+    }
+    Ok(())
+}
+```
+
+- [ ] **Step 3: Remove host-side counter increments from per-step path**
+
+In `fused_training.rs`, remove:
+- `self.trainer.adam_step_async()` — now GPU-side
+- `self.pre_replay_state_update(agent)` — counter increments now GPU-side, tau now GPU-side
+- `self.increment_warmup()` — move to GPU or keep as host (non-critical)
+
+In sub-trainers, remove `increment_adam_step()` calls from per-step path — counters incremented by the GPU kernel.
+
+- [ ] **Step 4: Expose sub-trainer t_dev_ptr fields**
+
+Add public accessors for the pinned device-mapped step counter pointers in IQL, IQN, Attention trainers so `submit_counter_increments` can pass them as kernel args.
+
+- [ ] **Step 5: Build and test**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+- [ ] **Step 6: Commit**
+
+```bash
+git commit -m "feat: GPU-side counter increments + tau annealing — zero host writes per step"
+```
+
+---
+
+### Task 5: Capture PER Sampling + IQL Modulate + PER Priority as Child Graphs
+
+**Files:**
+- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
+- Modify: `crates/ml-dqn/src/gpu_replay_buffer.rs`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs`
+
+- [ ] **Step 1: Expose `submit_per_sample` on GpuReplayBuffer**
+
+New method that runs ALL PER sampling ops (prefix scan, sample, gather_padded, IS weights) as stream operations suitable for graph capture. No host-side `rng_step` increment (handled by GPU counter kernel).
+
+```rust
+pub fn submit_per_sample(&mut self, batch_size: usize) -> Result<(), MLError> {
+    // All the kernel launches from sample_proportional, but using raw cuMemsetD8Async
+    // and the gather_padded kernels that write directly to trainer buffers.
+}
+```
+
+- [ ] **Step 2: Expose `submit_modulate_td_errors` on GpuIqlTrainer**
+
+New method that runs `modulate_td_errors` using raw u64 pointers (not `&mut CudaSlice`):
+
+```rust
+pub fn submit_modulate_td_errors(&self, td_errors_ptr: u64, indices_ptr: u64, write_pos: i32, capacity: i32) -> Result<(), MLError>
+```
+
+- [ ] **Step 3: Expose `submit_per_priority_update` on GpuReplayBuffer**
+
+New method that runs `update_priorities_gpu` using raw pointers for graph capture.
+
+- [ ] **Step 4: Capture new child graphs in `capture_training_graph`**
+
+```rust
+// Before training children:
+let per_sample = self.capture_child_graph("per_sample", |s| {
+    let agent = ...;  // need mutable access to replay buffer
+    agent.primary_dqn_mut().memory.submit_per_sample(s.batch_size)
+        .map_err(|e| anyhow::anyhow!("{e}"))
+})?;
+
+let counters = self.capture_child_graph("counters", |s| {
+    s.trainer.submit_counter_increments()
+        .map_err(|e| anyhow::anyhow!("{e}"))?;
+    s.trainer.update_stochastic_depth_mask()
+        .map_err(|e| anyhow::anyhow!("{e}"))
+})?;
+
+// After training children:
+let iql_mod = self.capture_child_graph("iql_modulate", |s| {
+    s.gpu_iql.submit_modulate_td_errors(...)
+        .map_err(|e| anyhow::anyhow!("{e}"))
+})?;
+
+let per_priority = self.capture_child_graph("per_priority", |s| {
+    agent.primary_dqn_mut().memory.submit_per_priority_update(...)
+        .map_err(|e| anyhow::anyhow!("{e}"))
+})?;
+```
+
+- [ ] **Step 5: Remove ungraphed calls from run_full_step**
+
+Remove from the pre-graph section:
+- `upload_batch_gpu` (already removed in Task 3)
+- `update_stochastic_depth_mask` (now in counters child)
+- `adam_step_async` (now GPU-side in counters child)
+- `pre_replay_state_update` (now GPU-side)
+
+Remove from the post-graph section:
+- `gpu_iql.modulate_td_errors` (now in iql_mod child)
+- `agent.update_priorities_from_td` (now in per_priority child)
+
+- [ ] **Step 6: Build and test**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+- [ ] **Step 7: Commit**
+
+```bash
+git commit -m "feat: capture PER sampling + IQL modulate + priority update as child graphs"
+```
+
+---
+
+### Task 6: Compose Parent Graph — Single cuGraphExecLaunch
+
+**Files:**
+- Modify: `crates/ml/src/trainers/dqn/fused_training.rs`
+
+- [ ] **Step 1: Implement `compose_parent_graph`**
+
+After all children are captured, compose into a single parent graph:
+
+```rust
+fn compose_parent_graph(&mut self) -> Result<()> {
+    let mut parent: cuda_sys::CUgraph = std::ptr::null_mut();
+    unsafe { cuda_sys::cuGraphCreate(&mut parent, 0); }
+
+    // Sequential dependency chain:
+    // per_sample → counters → spectral → forward → ddqn → aux →
+    // post_aux → adam_grad → adam → maintenance → iql_mod → per_priority
+
+    let children = [
+        &self.per_sample_child,
+        &self.counters_child,
+        &self.spectral_child,
+        &self.forward_child,
+        &self.ddqn_child,
+        &self.aux_child,
+        &self.post_aux_child,
+        &self.adam_grad_child,
+        &self.adam_child,
+        &self.maintenance_child,
+        &self.iql_modulate_child,
+        &self.per_priority_child,
+    ];
+
+    let mut prev_node: cuda_sys::CUgraphNode = std::ptr::null_mut();
+    for child_opt in children {
+        if let Some(child) = child_opt {
+            let mut node: cuda_sys::CUgraphNode = std::ptr::null_mut();
+            let deps = if prev_node.is_null() { &[] } else { std::slice::from_ref(&prev_node) };
+            unsafe {
+                cuda_sys::cuGraphAddChildGraphNode(
+                    &mut node, parent,
+                    deps.as_ptr(), deps.len(),
+                    child.graph,
+                );
+            }
+            prev_node = node;
+        }
+    }
+
+    let mut exec: cuda_sys::CUgraphExec = std::ptr::null_mut();
+    unsafe {
+        cuda_sys::cuGraphInstantiate(
+            &mut exec, parent, 0,
+        );
+    }
+    self.parent_exec = Some(exec);
+    self.parent_graph = Some(parent);
+    Ok(())
+}
+```
+
+- [ ] **Step 2: Replace sequential child launches with single parent launch**
+
+In `run_full_step`, replace the entire block of 12 `child.launch(cu_stream)` calls with:
+
+```rust
+if let Some(exec) = self.parent_exec {
+    unsafe { cuda_sys::cuGraphLaunch(exec, cu_stream); }
+}
+```
+
+- [ ] **Step 3: Remove PhaseEvents**
+
+Delete the `PhaseEvents` struct, all `cuEventRecord` calls in `run_full_step`, the `log_phase_timing` method, and all event creation/destruction code. Replace with a single wall-clock timer per epoch. Use nsys for detailed profiling when needed.
+
+- [ ] **Step 4: Remove HER donor generation from pre-graph path**
+
+Move HER `generate_random_donors_gpu` into the aux_child capture (it's already called from `submit_aux_ops`). If not, capture it as part of the per_sample child or a separate her_child.
+
+- [ ] **Step 5: Verify zero ungraphed launches**
+
+After the refactor, the per-step path should be:
+
+```rust
+pub fn run_full_step(...) -> Result<FusedStepResult> {
+    // ONE graph launch — everything inside
+    if let Some(exec) = self.parent_exec {
+        unsafe { cuda_sys::cuGraphLaunch(exec, self.stream.cu_stream()); }
+    } else {
+        // Ungraphed step 0: run all ops, then capture + compose parent
+        self.run_ungraphed_step_0(agent, gpu_batch)?;
+        self.capture_training_graph(agent, gpu_batch)?;
+        self.compose_parent_graph()?;
+    }
+
+    // Host-only: pinned scalar reads (nanoseconds, no GPU ops)
+    let fused_result = self.trainer.readback_training_scalars()?;
+    if let Some(ref iqn) = self.gpu_iqn {
+        if let Ok(loss) = iqn.read_total_loss() {
+            self.trainer.update_iqn_readiness(loss);
+        }
+    }
+    self.steps_since_varmap_sync += 1;
+
+    Ok(FusedStepResult { ... })
+}
+```
+
+- [ ] **Step 6: Build and test**
+
+```bash
+SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+- [ ] **Step 7: Commit**
+
+```bash
+git commit -m "feat: true single-graph — ONE cuGraphExecLaunch per step, zero ungraphed launches"
+```
+
+---
+
+### Task 7: Deploy and Validate
+
+- [ ] **Step 1: Push and deploy**
+
+```bash
+git push origin main
+./scripts/argo-train.sh dqn --baseline --epochs 5
+```
+
+- [ ] **Step 2: Check step timing**
+
+Expected:
+```
+STEP_TIMING step=50 step_ms=250-400
+```
+(was 3500ms)
+
+- [ ] **Step 3: Check epoch time**
+
+Expected: < 80s per epoch (was 670s).
+
+- [ ] **Step 4: Run nsys profiled epoch**
+
+```bash
+./scripts/argo-train.sh dqn --baseline --epochs 1 --sanitizer nsys
+```
+
+Download nsys_profile.nsys-rep and verify:
+- ONE cuGraphExecLaunch per step
+- Zero ungraphed kernel launches
+- All ~210 kernel nodes execute within the parent graph
+
+- [ ] **Step 5: Commit validation results**
+
+```bash
+git commit -m "docs: true single-graph validation — step time Xms, epoch Ys"
+```