foxhunt/docs/superpowers/specs/2026-04-18-true-single-graph-design.md

True Single-Graph Architecture — Zero CPU on Hot Path

Goal

Eliminate ALL CPU involvement from the per-step training loop. One cuGraphExecLaunch per step. Zero ungraphed kernel launches. Zero DtoD copies. Zero host-side computation. The CPU's only job is calling cuGraphExecLaunch(parent_exec, stream) once per step.

Root Cause

The current architecture has 18 ungraphed GPU operations per step on the same stream as the graphed children:

• PER sampling: 10 kernels (prefix scan, sample, gather×6, IS weights)
• Batch upload: 2 pad kernels + 4 DtoD copies
• Post-graph: IQL modulate (1 kernel) + PER priority update (1 kernel)

On Hopper (sm_90), ungraphed CUfunction launches on the same stream as captured graph children corrupt the graph's kernel state, causing 100x replay slowdown (3100ms instead of ~30ms). This explains why CUfunction isolation between children didn't fix the issue — the corruption comes from the ungraphed PER/IQL/upload operations, not from cross-child sharing.

Additionally, the CPU orchestrates 18 kernel launches per step via cudarc's safe API, putting host code in the critical path between every GPU operation.

Design

Architecture: One Parent Graph, Zero Ungraphed Launches

Per step (CPU does ONE thing):
  cuGraphExecLaunch(parent_exec, stream)

Parent graph nodes (captured once at step 0, replayed every step):

  ┌─ per_sample_node ────────────────────────────────────────────┐
  │  prefix_scan → binary_search → i64_to_u32 → IS_weights      │
  │  gather_states_padded → gather_next_states_padded            │
  │  gather_actions → gather_rewards → gather_dones              │
  │  gather_episode_ids                                          │
  └──────────────────────────────────────────────────────────────┘
  │
  ├─ stochastic_depth_node (Philox RNG, GPU-side counter)
  │
  ├─ spectral_node
  ├─ forward_node (cuBLAS trunk + ISV + temporal + loss + backward)
  ├─ ddqn_node
  ├─ aux_node (HER + EMA + IQL×2 + IQN + Attention + CQL)
  ├─ post_aux_node (selectivity + denoise + Q-stats + risk_sgd)
  ├─ adam_grad_node (mamba2_bwd + pruning + grad_norm)
  ├─ adam_update_node (Adam + grad_snapshot + unflatten + ISV)
  ├─ maintenance_node (causal + vaccine)
  │
  ├─ iql_modulate_node (advantage-weighted TD error modulation)
  └─ per_priority_node (priority writeback + max_priority update)

What Changes

Operation	Current	New
PER sampling (10 kernels)	Ungraphed, CPU-orchestrated	Child graph node in parent
Batch upload (2 pad + 4 DtoD)	Ungraphed DtoD + pad kernels	ELIMINATED — gather writes directly to padded trainer buffers
Stochastic depth mask	Ungraphed kernel + host RNG	GPU-side Philox counter (already a kernel)
Adam step counter	*t_pinned = step host write	GPU kernel increments pinned counter
IQL modulate_td_errors	Ungraphed kernel after graph	Child graph node in parent
PER priority update	Ungraphed kernel after graph	Child graph node in parent
HER donor generation	Ungraphed kernel before graph	Child graph node in parent
Sub-trainer adam_step increments	Host pinned writes	GPU kernel increments all counters
IQN tau annealing	Host cosine computation	GPU kernel reads step counter, computes tau
PhaseEvents timing	cuEventRecord per child (BROKEN)	REMOVED — use nsys for profiling
7+ cuGraphExecLaunch	Sequential launches	Single parent via cuGraphAddChildGraphNode
cudarc memset_zeros	device_ptr_mut() overhead	Raw cuMemsetD8Async

Direct-to-Trainer Gather (Eliminates upload_batch_gpu)

Current flow:

replay_buffer.states[idx] → sample_states → DtoD → pad_kernel → trainer.states_buf

New flow:

replay_buffer.states[idx] → gather_padded_kernel → trainer.states_buf

New gather_f32_rows_padded kernel:

// Gathers rows from src[indices[i]] into dst with padding: dst row width = dst_stride
// src row width = src_dim, extra columns zero-filled.
__global__ void gather_f32_rows_padded(
    float* dst, const float* src, const long long* indices,
    int src_dim, int dst_stride, int batch_size)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int row = tid / dst_stride;
    int col = tid % dst_stride;
    if (row >= batch_size) return;
    long long src_row = indices[row];
    dst[row * dst_stride + col] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f;
}

The trainer passes its buffer pointers to the PER sampling child graph via pinned device-mapped pointers (stable addresses, graph-safe).

GPU-Side Counters (Eliminates Host Writes)

All per-step counter increments move to a single GPU kernel:

// Increments all step counters atomically. (1,1,1) grid, 1 thread.
__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, int total_steps, float tau_init, float tau_final, int anneal_steps)
{
    // Increment all counters
    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);
    // Cosine-annealed tau
    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;
}

All counter pointers are pinned device-mapped — GPU writes, host reads for monitoring.

PER Random Seed

PER sampling uses rng_step as Philox seed. Currently incremented on host. New: increment_step_counters kernel increments rng_step on GPU. The PER sampling kernel reads *rng_step from device-mapped memory. Same Philox sequence, zero host involvement.

Parent Graph Composition

fn compose_parent_graph(&mut self) -> Result<()> {
    let parent = unsafe { cuGraphCreate()? };

    // PER sampling + gather (writes directly to trainer buffers)
    let per_node = cuGraphAddChildGraphNode(parent, &[], per_sample.graph)?;

    // Counter increments + stochastic depth
    let counters_node = cuGraphAddChildGraphNode(parent, &[per_node], counters.graph)?;

    // Training pipeline (existing children)
    let spectral = cuGraphAddChildGraphNode(parent, &[counters_node], spectral.graph)?;
    let forward = cuGraphAddChildGraphNode(parent, &[spectral], forward.graph)?;
    let ddqn = cuGraphAddChildGraphNode(parent, &[forward], ddqn.graph)?;
    let aux = cuGraphAddChildGraphNode(parent, &[ddqn], aux.graph)?;
    let post_aux = cuGraphAddChildGraphNode(parent, &[aux], post_aux.graph)?;
    let adam_grad = cuGraphAddChildGraphNode(parent, &[post_aux], adam_grad.graph)?;
    let adam = cuGraphAddChildGraphNode(parent, &[adam_grad], adam_update.graph)?;
    let maint = cuGraphAddChildGraphNode(parent, &[adam], maintenance.graph)?;

    // Post-training bookkeeping
    let iql_mod = cuGraphAddChildGraphNode(parent, &[maint], iql_modulate.graph)?;
    let per_prio = cuGraphAddChildGraphNode(parent, &[iql_mod], per_priority.graph)?;

    let exec = cuGraphInstantiate(parent, AUTO_FREE_ON_LAUNCH)?;
    self.parent_exec = Some(exec);
    Ok(())
}

CUfunction Isolation

One CUmodule per child graph that uses utility kernels. The new children (per_sample, counters, iql_modulate, per_priority) use their OWN cubins from separate CUmodule loads. Zero cross-child or graphed-vs-ungraphed sharing.

What's NOT in the Graph

• Experience collection (runs between epochs on collector stream)
• Validation backtest (epoch boundary)
• Model checkpointing (epoch boundary)
• Enrichment cycle (epoch boundary)

Host-Side Per Step (After cuGraphExecLaunch)

ONLY pinned memory reads for monitoring — zero GPU operations:

• readback_training_scalars() — reads *total_loss_pinned, *grad_norm_pinned
• read_total_loss() — reads IQN loss from pinned
• update_iqn_readiness() — reads pinned loss, computes readiness (host math OK — not on GPU stream)
• Bookkeeping counters — steps_since_varmap_sync += 1 (pure host, no GPU)

These are nanosecond host memory reads. No GPU sync, no kernel launches, no stream ops.

Success Criteria

• ONE cuGraphExecLaunch per step
• Zero ungraphed kernel launches on training stream
• Zero DtoD copies in per-step path (gather writes directly to trainer buffers)
• Zero host-side computation on GPU stream critical path
• Step time < 400ms (from 3750ms)
• Epoch time < 80s (from 670s)
• 899 unit tests pass

Files Modified

• crates/ml/src/trainers/dqn/fused_training.rs — parent graph composition, remove PhaseEvents, remove upload_batch_gpu call
• crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — remove upload_batch_gpu, add increment_step_counters kernel
• crates/ml-dqn/src/gpu_replay_buffer.rs — expose sample ops as capturable methods, add gather_padded kernel
• crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs — expose modulate_td_errors as graph-capturable
• New .cu file: gather_padded_kernel.cu — gather_f32_rows_padded, increment_step_counters
• infra/docker/Dockerfile.ci-builder — nsight-systems-cli (DONE)
• infra/k8s/argo/train-template.yaml — nsys support (DONE)