diff --git a/docs/superpowers/specs/2026-05-26-multistream-per-push-rewrite.md b/docs/superpowers/specs/2026-05-26-multistream-per-push-rewrite.md new file mode 100644 index 000000000..36a80da15 --- /dev/null +++ b/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.