foxhunt/docs/superpowers/plans/2026-03-18-gpu-resident-pipeline.md

GPU-Resident Training Pipeline Refactor

For agentic workers: REQUIRED: Use superpowers:subagent-driven-development (if subagents available) or superpowers:executing-plans to implement this plan. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Eliminate ALL GPU→CPU→GPU roundtrips from the ML training pipeline. Data stays GPU-resident from collection through forward/backward/optimizer. Only scalar metrics come to CPU.

Architecture: Experience batches use CudaSlice<f32> / GpuTensor fields (not Vec<f32>). model.update() accepts GPU-resident data directly. Delete all CPU-side batch conversion functions.

Tech Stack: cudarc 0.19, CudaSlice, GpuTensor, cuBLAS

---

Current Roundtrips (what to eliminate)

#	Path	Data/epoch	Fix
1	PPO: GPU kernel → Vec<f32> (PpoExperienceBatch) → Vec<Vec<f32>> (TrajectoryBatch) → GPU re-upload	123 MB	CudaSlice fields, delete conversion
2	DQN: GPU kernel → Vec<f32> (ExperienceBatch) — parallel GpuExperienceBatch exists but may not be primary	1.95 GB	Unify on GpuExperienceBatch
3	PPO: GPU softmax → to_vec1() → CPU rand sampling per step	16K syncs	GPU sampling kernel or batch

---

Task 1: PPO Experience Batch GPU-Resident

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs
• Modify: crates/ml/src/trainers/ppo.rs
• Modify: crates/ml-ppo/src/ppo.rs

Subtask 1a: Change PpoExperienceBatch to GPU-resident

• Step 1: In gpu_ppo_collector.rs, change PpoExperienceBatch fields:

// BEFORE:
pub struct PpoExperienceBatch {
    pub states: Vec<f32>,        // downloaded from GPU
    pub actions: Vec<i32>,
    pub log_probs: Vec<f32>,
    pub advantages: Vec<f32>,
    pub returns: Vec<f32>,
    pub done_flags: Vec<i32>,
    pub n_episodes: usize,
    pub timesteps: usize,
}

// AFTER:
pub struct PpoExperienceBatch {
    pub states: CudaSlice<f32>,      // stays on GPU
    pub actions: CudaSlice<i32>,
    pub log_probs: CudaSlice<f32>,
    pub advantages: CudaSlice<f32>,
    pub returns: CudaSlice<f32>,
    pub done_flags: CudaSlice<i32>,
    pub n_episodes: usize,
    pub timesteps: usize,
    pub state_dim: usize,
    pub stream: Arc<CudaStream>,
}

• Step 2: In collect_experiences(), remove ALL memcpy_dtoh calls. Return the GPU buffers directly.

• Step 3: Add PpoExperienceBatch::download() method for debugging/checkpoint only:

impl PpoExperienceBatch {
    pub fn download_states(&self) -> Result<Vec<f32>, MLError> {
        self.stream.clone_dtoh(&self.states).map_err(...)
    }
}

Subtask 1b: Delete cpu_batch_to_trajectory_batch

• Step 4: In ppo.rs, DELETE gpu_batch_to_trajectory_batch() function entirely.

• Step 5: Change train_gpu() to pass PpoExperienceBatch directly to model.update():

// BEFORE:
let batch = collector.collect_experiences(...)?;
let training_batch = gpu_batch_to_trajectory_batch(&batch);  // GPU→CPU conversion
let (policy_loss, value_loss) = model.update(&mut training_batch)?;  // CPU→GPU re-upload

// AFTER:
let batch = collector.collect_experiences(...)?;  // stays on GPU
let (policy_loss, value_loss) = model.update_gpu(&batch)?;  // pure GPU

Subtask 1c: Add PPOAgent::update_gpu()

• Step 6: In crates/ml-ppo/src/ppo.rs, add update_gpu() method:

pub fn update_gpu(&mut self, batch: &PpoExperienceBatch) -> Result<(f32, f32), MLError> {
    // Forward pass directly on CudaSlice data — no upload needed
    // Loss computation on GPU
    // Backward + optimizer step on GPU
    // Only readback: 2 scalar losses
}

• Step 7: Compile check: SQLX_OFFLINE=true cargo check -p ml-ppo -p ml

---

Task 2: DQN Experience Batch Unification

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
• Modify: crates/ml/src/trainers/dqn/trainer/train_step.rs

Subtask 2a: Audit which batch type is primary

• Step 1: Search for all ExperienceBatch vs GpuExperienceBatch usage:

rg "ExperienceBatch\b" crates/ml/ --type rust -n | grep -v "Gpu"
rg "GpuExperienceBatch" crates/ml/ --type rust -n

• Step 2: If ExperienceBatch (CPU) is primary in training loop:

  • Change training loop to use GpuExperienceBatch instead
  • Delete ExperienceBatch struct OR keep only for checkpoint/debugging

• Step 3: If GpuExperienceBatch is already primary:

  • Delete ExperienceBatch struct
  • Remove collect_experiences() CPU download path
  • Keep only collect_experiences_gpu() path

Subtask 2b: Verify fused CUDA trainer uses GPU batch

• Step 4: Check that GpuDqnTrainer::train_step() takes CudaSlice/Tensor directly (not Vec)
• Step 5: Compile check

---

Task 3: PPO Action Sampling on GPU

Files:

• Modify: crates/ml/src/trainers/ppo.rs (CPU fallback path)

• Modify: crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs (GPU path — already correct)

• Step 1: Verify GPU collection path already samples actions on-device (should be in kernel)

• Step 2: For CPU fallback path: batch-sample N steps at once instead of per-step sync:

// BEFORE (per-step):
for step in 0..num_steps {
    let probs = actor.forward(&state)?;
    let probs_vec = probs.to_vec1::<f32>()?;  // GPU SYNC per step
    let action = sample_from_probs(&probs_vec);
}

// AFTER (batched):
let all_probs = actor.forward_batch(&all_states)?;  // single forward
let all_probs_vec = all_probs.to_vec2::<f32>()?;    // single download
for (step, probs) in all_probs_vec.iter().enumerate() {
    let action = sample_from_probs(probs);
}

• Step 3: Compile check

---

Task 4: Verify Supervised Models Are Clean

• Step 1: Confirm TFT, Mamba2, Liquid, TLOB, KAN, xLSTM, Diffusion training loops only do scalar readbacks
• Step 2: No changes expected — just audit and document

---

Execution Strategy

Task 1 (PPO) is the biggest win (123 MB/epoch eliminated). Tasks 2-3 can follow.

Recommended: 2 agents

• Agent A: Task 1 (PPO GPU-resident batch)
• Agent B: Task 2 (DQN batch unification)
• Task 3 after both complete
• Task 4 is read-only audit

Estimated: 2-4 hours with 2 agents.