# Multi-Stream Pipeline + Per-Push Rewrite

**Goal:** Double throughput at b=256 from 51 sps to 100+ sps by overlapping compute and rewriting the rl_per_push kernel.

**Current bottleneck breakdown (nsys at b=256):**
- `rl_per_push`: 46% of GPU time (124μs/call, instruction-bound)
- Encoder (CfC+Mamba2): 24% of GPU time (sequential per K=32)
- RL kernels (graphs + sample + rebuild): 20%
- Everything else: 10%

---

## Part 1: Multi-Block rl_per_push Rewrite

### Problem
Current: Grid=(1), Block=(256). One thread per batch copies 128 floats via 32 float4 loads. At 256 threads × 32 float4 ops each = 8,192 instructions issued from a single SM. The SM is instruction-pipeline bound — only 1 warp scheduler issuing at a time.

### Fix: Grid=(B), Block=(128). One block per batch, 128 threads per block.

Each thread copies ONE float. The 128-float h_t memcpy becomes a single coalesced 512-byte transaction (128 threads × 4 bytes = perfect cache-line alignment).

**Write-head coordination** changes from shared-memory prefix-sum (single block) to a **two-kernel approach**:

**Kernel A: `rl_per_push_ring`** — Grid=(B), Block=(128)
- Each block handles one batch element
- 128 threads cooperate to copy h_t into the n-step ring (1 float per thread)
- Thread 0: writes scalars, advances ring, decides flush
- Thread 0: writes `flush_flag_d[b] = should_flush` to global memory
- NO coordination across blocks needed here

**Kernel B: `rl_per_push_flush`** — Grid=(B), Block=(128)
- Preceded by a single-thread coordinator that scans flush_flag_d and writes offsets
- Actually: fuse the scan into block 0 of this kernel
- Block 0, thread 0: prefix-sums flush_flag_d[0..B], writes write_offsets_d[b], advances write_head
- `__threadfence()` to publish
- All blocks: if flush_flag_d[b] == 0, return. Otherwise read write_offsets_d[b] and 128 threads cooperatively write h_t + h_tp1 to replay (1 float per thread, fully coalesced)

Two kernel launches (A + B) replace one kernel launch. But each is 10× faster because of full coalescing.

**Expected:** 124μs → ~15μs (8× faster, removes 46% of GPU time)

---

## Part 2: Multi-Stream Overlap

### Problem
The step pipeline is strictly sequential:
```
encoder_fwd → Graph A (action) → fill → Graph B (reward) → PER push → PER sample → step_synthetic → priority update + rebuild
```

The encoder forward and step_synthetic are the two heaviest phases. They're independent: encoder processes step N+1's observations while step_synthetic trains on step N's PER sample.

### Fix: Two-stream pipeline

```
Stream 0 (env):     [encoder_fwd(N)] → [Graph A] → [fill] → [Graph B] → [PER push]
Stream 1 (train):   [PER sample(N-1)] → [step_synthetic(N-1)] → [priority update] → [tree rebuild]
```

**Pipeline execution (steady state):**
```
Time →
Stream 0: |--encoder(N)--|--GraphA--|--fill--|--GraphB--|--push--|
Stream 1:                                   |--sample(N-1)--|--train(N-1)--|--prio--|--rebuild--|
```

At step N:
1. Stream 0: runs encoder forward, action selection, fill, reward pipeline, PER push for step N
2. Stream 1 (concurrent): runs PER sample, training step, priority update for step N-1's data

Dependency: Stream 1's PER sample needs step N-1's push to have completed (it has — push was on stream 0 in the PREVIOUS iteration). An inter-stream event ensures ordering:
- After PER push on stream 0: `event_push_done.record(stream_0)`
- Before PER sample on stream 1: `stream_1.wait_event(event_push_done)`

**Overlap gain:**
- Encoder forward (stream 0): ~5ms at b=256
- Step_synthetic (stream 1): ~5ms at b=256
- Without overlap: 10ms sequential
- With overlap: ~5ms (hidden behind each other)

**Expected:** 20ms/step → ~12ms/step = **83 sps** at b=256

---

## Part 3: Combined Impact

| Optimization | Before | After | Gain |
|-------------|--------|-------|------|
| Per-push rewrite | 124μs (46%) | ~15μs (6%) | -109μs/step |
| Multi-stream overlap | sequential | pipelined | -5ms/step |
| **Combined** | 20ms/step (51 sps) | **~10ms/step (100 sps)** | **2×** |

At b=256 × 100 sps = **25,600 transitions/sec** (14× original baseline of 1,840).

---

## Implementation Order

1. **rl_per_push rewrite** (2 kernels replacing 1) — pure CUDA, independent of Rust pipeline
2. **Multi-stream wiring** — add second CudaStream, split step_with_lobsim into env-stream and train-stream, add event synchronization
3. **Benchmark** — verify 100+ sps at b=256

---

## Constraints

- Per `feedback_no_atomicadd`: flush coordination via prefix-sum, not atomic increment
- Per `feedback_cpu_is_read_only`: no host involvement in the pipeline
- Per `feedback_no_stubs`: both kernels wired in same commit
- The multi-stream approach adds one step of training latency (trains on N-1 while env-stepping N) — acceptable, standard in async RL
- Graph capture: Graphs A, A2, B are on stream 0. Graph C (training) moves to stream 1. Separate capture state machines per stream.
