spec: multi-stream pipeline + per-push rewrite — target 100 sps at b=256

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

docs/superpowers/specs/2026-05-26-multistream-per-push-rewrite.md

@@ -0,0 +1,112 @@
+# 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.