plan: multi-stream + per-push rewrite — 5 tasks
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<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)
|
||||
Reference in New Issue
Block a user