foxhunt/docs/superpowers/plans/2026-04-19-true-single-graph.md

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

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

• Step 4: Commit

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

// 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)

// 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)

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

// 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:

nvcc -cubin -arch=sm_90 -o graph_utility_kernels.cubin graph_utility_kernels.cu

• Step 6: Commit

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:

/// 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:

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

SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 8: Commit

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):

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

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

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

• Step 6: Commit

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).

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):

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

// 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

SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 7: Commit

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:

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:

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:

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

SQLX_OFFLINE=true cargo check -p ml -p ml-dqn
SQLX_OFFLINE=true cargo test -p ml --lib

• Step 7: Commit

git commit -m "feat: true single-graph — ONE cuGraphExecLaunch per step, zero ungraphed launches"

---

Task 7: Deploy and Validate

• Step 1: Push and deploy

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

./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

git commit -m "docs: true single-graph validation — step time Xms, epoch Ys"