From 4f7b1f48044b444a546804e751408d43f01099da Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 00:06:05 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20true=20single-graph=20=E2=80=94=207=20t?= =?UTF-8?q?asks,=20zero=20CPU=20on=20hot=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.6 (1M context) --- .../plans/2026-04-19-true-single-graph.md | 535 ++++++++++++++++++ 1 file changed, 535 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-19-true-single-graph.md diff --git a/docs/superpowers/plans/2026-04-19-true-single-graph.md b/docs/superpowers/plans/2026-04-19-true-single-graph.md new file mode 100644 index 000000000..49681c767 --- /dev/null +++ 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 { + // 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" +```