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:
```rust
// 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:
```rust
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():
```rust
// 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:
```rust
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:
```bash
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<f32>)
- [ ] **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:
```rust
// 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.**