From be3d5ceeae431496d227c1d35e01125f5d5c72ef Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 01:32:56 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20multi-stream=20+=20per-push=20rewrite?= =?UTF-8?q?=20=E2=80=94=205=20tasks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Opus 4.7 --- ...2026-05-26-multistream-per-push-rewrite.md | 254 ++++++++++++++++++ 1 file changed, 254 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-26-multistream-per-push-rewrite.md diff --git a/docs/superpowers/plans/2026-05-26-multistream-per-push-rewrite.md b/docs/superpowers/plans/2026-05-26-multistream-per-push-rewrite.md new file mode 100644 index 000000000..2f74b2f5d --- /dev/null +++ b/docs/superpowers/plans/2026-05-26-multistream-per-push-rewrite.md @@ -0,0 +1,254 @@ +# 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, // [B] +pub write_offsets_d: CudaSlice, // [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, +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)