# 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"
```
