# Multi-Stream Pipeline + Per-Push Rewrite 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:** Double throughput from 51 sps to 100+ sps at b=256 by rewriting rl_per_push as coalesced multi-block kernel and overlapping env-step with training via two CUDA streams.

**Architecture:** Part 1 splits rl_per_push into `rl_per_push_ring` (per-batch n-step ring write with 128-thread coalesced copy) + `rl_per_push_flush` (coordinated replay write with block-0 prefix-sum). Part 2 adds a second CUDA stream for PER sample + step_synthetic + priority update, running concurrently with the next step's env pipeline.

**Tech Stack:** Rust 1.85, cudarc 0.19, CUDA 12.4, pre-compiled cubins

---

## File Map

| File | Responsibility |
|------|---------------|
| `crates/ml-alpha/cuda/rl_per_push_ring.cu` | NEW: n-step ring write, Grid=(B), Block=(128) |
| `crates/ml-alpha/cuda/rl_per_push_flush.cu` | NEW: coordinated replay write, Grid=(B), Block=(128) |
| `crates/ml-alpha/cuda/rl_per_push.cu` | DELETE (replaced by ring + flush) |
| `crates/ml-alpha/build.rs` | Register new kernels, remove old |
| `crates/ml-alpha/src/rl/gpu_replay.rs` | Add flush_flags_d, write_offsets_d buffers |
| `crates/ml-alpha/src/trainer/integrated.rs` | Wire new kernels + multi-stream split |

---

## Task 1: Write rl_per_push_ring kernel

**Files:**
- Create: `crates/ml-alpha/cuda/rl_per_push_ring.cu`

- [ ] **Step 1: Write the kernel**

Grid=(B,1,1), Block=(128,1,1). One block per batch element. 128 threads cooperate on h_t copy.

```cuda
extern "C" __global__ void rl_per_push_ring(
    const float* __restrict__ h_t_current,      // [B × 128]
    const float* __restrict__ rewards_scaled,   // [B]
    const float* __restrict__ rewards_raw,      // [B]
    const float* __restrict__ dones,            // [B]
    const float* __restrict__ log_pi_old,       // [B]
    const int*   __restrict__ actions,          // [B]
    const float* __restrict__ isv,
    // N-step ring
    float* __restrict__ nstep_scalars,          // [B × N_STEP_MAX × 3]
    float* __restrict__ nstep_h_t,             // [B × N_STEP_MAX × 128]
    unsigned int* __restrict__ nstep_write_idx, // [B]
    unsigned int* __restrict__ nstep_count,     // [B]
    // Output: per-batch flush decision
    int* __restrict__ flush_flags,              // [B] — 1 if this batch needs to flush
    int b_size
)
```

Logic:
- Thread 0: reads ISV n_step + gamma, writes scalars to ring, advances idx/count, computes should_flush, writes flush_flags[b]
- All 128 threads: coalesced copy `h_t_current[b*128 + tid]` → `nstep_h_t[(b*N_STEP_MAX + ring_idx)*128 + tid]`
- `__syncthreads()` ensures copy completes before thread 0 advances

- [ ] **Step 2: Verify nvcc compilation**

```bash
nvcc -cubin -arch sm_86 -O3 crates/ml-alpha/cuda/rl_per_push_ring.cu -o /tmp/test.cubin
```

- [ ] **Step 3: Commit**

---

## Task 2: Write rl_per_push_flush kernel

**Files:**
- Create: `crates/ml-alpha/cuda/rl_per_push_flush.cu`

- [ ] **Step 1: Write the kernel**

Grid=(B,1,1), Block=(128,1,1). Block 0 does the prefix-sum coordination. All blocks do coalesced replay writes.

```cuda
extern "C" __global__ void rl_per_push_flush(
    const int* __restrict__ flush_flags,         // [B] from rl_per_push_ring
    int* __restrict__ write_offsets,             // [B] — assigned write slots
    // Source: n-step ring (to read oldest h_t)
    const float* __restrict__ nstep_scalars,     // [B × N_STEP_MAX × 3]
    const float* __restrict__ nstep_h_t,        // [B × N_STEP_MAX × 128]
    const unsigned int* __restrict__ nstep_write_idx, // [B]
    const unsigned int* __restrict__ nstep_count,// [B]
    const float* __restrict__ h_tp1_current,    // [B × 128]
    const float* __restrict__ log_pi_old,       // [B]
    const int*   __restrict__ actions,          // [B]
    const float* __restrict__ isv,
    // Replay storage
    float* __restrict__ replay_h_t,              // [cap × 128]
    float* __restrict__ replay_h_tp1,            // [cap × 128]
    float* __restrict__ replay_scalars,          // [cap × 7]
    float* __restrict__ priority_tree,           // [2 × cap]
    unsigned int* __restrict__ write_head,       // [1]
    unsigned int* __restrict__ replay_len,       // [1]
    float* __restrict__ max_priority,            // [1]
    unsigned int* __restrict__ nstep_count_out,  // [B] — reset to 0 on flush
    int b_size,
    int capacity
)
```

Logic:
- Block 0, thread 0: scan flush_flags[0..B], prefix-sum → write_offsets[b]. Advance write_head by total. Update replay_len. Write total to `write_offsets[B]` (sentinel for other blocks to read). `__threadfence()`.
- All blocks: if flush_flags[blockIdx.x] == 0, return early
- 128 threads coalesced: copy oldest h_t from ring → replay_h_t[slot*128 + tid]
- 128 threads coalesced: copy h_tp1_current[b*128 + tid] → replay_h_tp1[slot*128 + tid]
- Thread 0: write scalars (n-step return, action, etc.), set priority leaf, reset nstep_count

Blocks > 0 need to wait for block 0's prefix-sum. Use a global flag:
- Block 0 writes `write_offsets[B] = 1` after scan + threadfence
- Other blocks spin on `write_offsets[B]` (simple spin-wait — B blocks, only block 0 does the scan, all others spin for <1μs)

- [ ] **Step 2: Verify nvcc compilation**
- [ ] **Step 3: Commit**

---

## Task 3: Wire new kernels into trainer (replace old rl_per_push)

**Files:**
- Delete: `crates/ml-alpha/cuda/rl_per_push.cu`
- Modify: `crates/ml-alpha/build.rs` (replace `rl_per_push` with `rl_per_push_ring` + `rl_per_push_flush`)
- Modify: `crates/ml-alpha/src/rl/gpu_replay.rs` (add flush_flags_d, write_offsets_d)
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (new cubin consts + fn handles, replace launch)

- [ ] **Step 1: Add flush_flags_d and write_offsets_d to GpuReplayBuffer**

```rust
pub flush_flags_d: CudaSlice<i32>,      // [B]
pub write_offsets_d: CudaSlice<i32>,    // [B + 1] (extra slot for ready-flag)
```

- [ ] **Step 2: Update build.rs** — remove `rl_per_push`, add `rl_per_push_ring` + `rl_per_push_flush`

- [ ] **Step 3: Replace the single launch in step_with_lobsim with two launches**

Replace the rl_per_push block (~line 5652) with:
```rust
// Phase 1: ring write + flush decision
{
    let cfg = LaunchConfig {
        grid_dim: (b_size as u32, 1, 1),
        block_dim: (128, 1, 1),
        shared_mem_bytes: 0,
    };
    // launch rl_per_push_ring...
}
// Phase 2: coordinated flush (coalesced replay write)
{
    let cfg = LaunchConfig {
        grid_dim: (b_size as u32, 1, 1),
        block_dim: (128, 1, 1),
        shared_mem_bytes: 0,
    };
    // launch rl_per_push_flush...
}
```

- [ ] **Step 4: Verify compilation + smoke test**
- [ ] **Step 5: Commit**

---

## Task 4: Multi-stream — add train_stream + event sync

**Files:**
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`

- [ ] **Step 1: Add second stream field**

```rust
pub train_stream: Arc<CudaStream>,
push_done_event: CudaEvent,  // signals push completed on env stream
```

Init in constructor:
```rust
let train_stream = dev.cuda_stream().context("train stream")?.clone();
let push_done_event = ... // cudarc event creation
```

- [ ] **Step 2: Move PER sample + step_synthetic + priority update to train_stream**

In the K-loop (line ~5734), change all `self.stream.launch_builder(...)` calls for:
- rl_per_sample
- step_synthetic's internal kernels
- rl_per_update_priority
- rl_per_tree_rebuild

To use `self.train_stream.launch_builder(...)` instead.

This requires step_synthetic to accept a stream parameter (or use the train_stream internally). The simplest: change `self.stream` references inside step_synthetic to `self.train_stream`.

Actually simpler: DON'T split step_synthetic. Just move the WHOLE K-loop (sample + train + priority) to train_stream, and keep everything before it (encoder + action + fill + reward + push) on self.stream.

- [ ] **Step 3: Add event synchronization**

After PER push on env stream:
```rust
self.push_done_event.record(&self.stream)?;
```

Before PER sample on train stream:
```rust
self.train_stream.wait_event(&self.push_done_event)?;
```

This ensures the sample reads data that push has finished writing.

- [ ] **Step 4: Move train_stream sync to end of step**

The current `self.stream.synchronize()` at the end of step_with_lobsim needs to sync BOTH streams:
```rust
self.stream.synchronize()?;
self.train_stream.synchronize()?;
```

Or better: have step N's train_stream sync happen at the START of step N+1 (overlapping with step N+1's env work). This is the pipeline overlap.

- [ ] **Step 5: Smoke test + verify**
- [ ] **Step 6: Commit**

---

## Task 5: Benchmark — verify 100+ sps at b=256

- [ ] **Step 1: Push + submit 100-step benchmark**

```bash
git push origin ml-alpha-phase-a
argo submit -n foxhunt --from=wftmpl/alpha-rl \
  -p git-branch=ml-alpha-phase-a -p n-steps=100 -p n-backtests=256
```

- [ ] **Step 2: Verify sps ≥ 80**
- [ ] **Step 3: If pass, submit nsys profile run to confirm overlap**

```bash
argo submit -n foxhunt --from=wftmpl/alpha-rl \
  -p git-branch=ml-alpha-phase-a -p n-steps=100 -p n-backtests=256 -p nsys-profile=true
```

---

## Kill Criteria

- rl_per_push_ring + flush combined: < 30μs/step (down from 124μs)
- Multi-stream: encoder forward on stream 0 overlaps with step_synthetic on stream 1 (visible in nsys timeline)
- End-to-end: ≥ 80 sps at b=256 (target 100)
- Smoke test passes (losses finite, replay fills)
