refactor(rl): pre-allocate 56 replay-step gradient buffers for Graph C
Move all step_synthetic/dqn_replay_step alloc_zeros to persistent trainer fields (ss_* prefix). Enables CUDA Graph capture of the replay training step — all device pointers are now stable across steps. Introduces reduce_axis0_free() to resolve borrow-checker E0502 when both source (per-batch scratch) and destination (reduced grad) are self fields passed to the same function. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
377
docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
Normal file
377
docs/superpowers/plans/2026-05-25-cuda-performance-p1-p2-p3.md
Normal file
@@ -0,0 +1,377 @@
|
||||
# CUDA Performance P1-P3 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:** Take training throughput from 1,840 transitions/sec to 70,000+ on L40S via batch size scaling (P1), complete CUDA Graph capture (P2), and fuse the reward pipeline (P3).
|
||||
|
||||
**Architecture:** P1 is a config change (b=16→256) + PER capacity scaling. P2 wraps remaining kernel sequences with cudarc's begin_capture/end_capture/launch three-state machine. P3 fuses 6 sequential reward kernels into one. All changes are in the RL trainer and CUDA kernel layer.
|
||||
|
||||
**Tech Stack:** Rust 1.85, cudarc 0.19, CUDA 12.4, pre-compiled cubins via build.rs
|
||||
|
||||
---
|
||||
|
||||
## File Map
|
||||
|
||||
| File | Responsibility | Tasks |
|
||||
|------|---------------|-------|
|
||||
| `infra/k8s/argo/alpha-rl-template.yaml` | Default n-backtests param | T1 |
|
||||
| `crates/ml-alpha/src/trainer/integrated.rs` | Trainer: PER scaling, Graph B+C, borrow fixes | T2, T3, T4, T5 |
|
||||
| `crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu` | New: fused reward kernel | T6 |
|
||||
| `crates/ml-alpha/build.rs` | Register new kernel | T6 |
|
||||
|
||||
---
|
||||
|
||||
## Task 1: P1 — Batch Size Scaling (Argo default + PER auto-scale)
|
||||
|
||||
**Files:**
|
||||
- Modify: `infra/k8s/argo/alpha-rl-template.yaml:44-45`
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (PER capacity logic)
|
||||
|
||||
- [ ] **Step 1: Update Argo template default n-backtests from 16 to 128**
|
||||
|
||||
```yaml
|
||||
- name: n-backtests
|
||||
value: "128"
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Auto-scale PER capacity to 4× batch size minimum**
|
||||
|
||||
In `IntegratedTrainer::new()`, after `let b_size = cfg.perception.n_batch;`, replace the fixed PER capacity with:
|
||||
|
||||
```rust
|
||||
let per_capacity = cfg.per_capacity.max(4 * b_size);
|
||||
```
|
||||
|
||||
Find where `cfg.per_capacity` is passed to the replay buffer constructor (around line 1680 where `cfg.per_capacity` appears) and replace with `per_capacity`.
|
||||
|
||||
- [ ] **Step 3: Verify local smoke at b=64**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
|
||||
```
|
||||
|
||||
The existing smoke uses `IntegratedTrainerConfig::default()` which has `n_batch=16`. Verify it still passes. Then manually test b=64 is allocable on RTX 3050 Ti (4GB):
|
||||
|
||||
```bash
|
||||
# Quick alloc test — just construct the trainer, don't step
|
||||
SQLX_OFFLINE=true cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Apply the Argo template to cluster**
|
||||
|
||||
```bash
|
||||
kubectl apply -f infra/k8s/argo/alpha-rl-template.yaml -n foxhunt
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add infra/k8s/argo/alpha-rl-template.yaml crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "perf(rl): scale batch size to 128 default + auto-size PER capacity"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 2: P2 — Fix Graph C Borrow Checker (reduce_axis0 pattern)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs` (reduce_axis0_free calls)
|
||||
|
||||
The E0502 errors occur because `reduce_axis0_free` takes `&self.stream` + `&self.reduce_axis0_fn` (immutable borrow of self) AND `&mut self.ss_q_grad_w_d` (mutable borrow of self) in the same call. The fix is to extract the stream and fn references into locals before the call.
|
||||
|
||||
- [ ] **Step 1: Fix all reduce_axis0_free borrow conflicts**
|
||||
|
||||
The pattern for every call site is:
|
||||
|
||||
```rust
|
||||
// BEFORE (fails borrow check):
|
||||
reduce_axis0_free(&self.stream, &self.reduce_axis0_fn, &self.ss_q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut self.ss_q_grad_w_d)?;
|
||||
|
||||
// AFTER (compiles):
|
||||
let (stream, fn_ref) = (&self.stream, &self.reduce_axis0_fn);
|
||||
reduce_axis0_free(stream, fn_ref, &self.ss_q_grad_w_per_batch_d, b_size, k_dqn * HIDDEN_DIM, &mut self.ss_q_grad_w_d)?;
|
||||
```
|
||||
|
||||
Alternatively, extract `stream` and `reduce_axis0_fn` into locals at the TOP of `step_synthetic`:
|
||||
|
||||
```rust
|
||||
let stream = &self.stream;
|
||||
let reduce_fn = &self.reduce_axis0_fn;
|
||||
```
|
||||
|
||||
Then use `stream` and `reduce_fn` throughout. This fixes all ~15 call sites at once.
|
||||
|
||||
- [ ] **Step 2: Verify compilation**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true cargo check -p ml-alpha
|
||||
```
|
||||
|
||||
Expected: 0 errors.
|
||||
|
||||
- [ ] **Step 3: Run local smoke**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "fix(rl): resolve borrow-checker conflicts in pre-allocated gradient reduce"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 3: P2 — Graph B Capture (post-fill reward pipeline)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`
|
||||
|
||||
**Prerequisite:** Task 2 (Graph C buffers compile clean).
|
||||
|
||||
Graph B covers ~20 kernels from `extract_realized_pnl_delta` through `launch_var_over_abs_mean` in `step_with_lobsim`. The `rl_write_u64` kernel and `ts_ns_d` device buffer are already wired. The `postfill_graph` field already exists but currently captures the pre-fill A2 section — rename and add a `reward_graph` field.
|
||||
|
||||
- [ ] **Step 1: Add `reward_graph: Option<CudaGraph>` field**
|
||||
|
||||
In the struct definition (near line 426):
|
||||
|
||||
```rust
|
||||
prefill_graph: Option<CudaGraph>,
|
||||
postfill_graph: Option<CudaGraph>,
|
||||
reward_graph: Option<CudaGraph>,
|
||||
graph_warmup_done: bool,
|
||||
```
|
||||
|
||||
Initialize in constructor: `reward_graph: None,`
|
||||
|
||||
- [ ] **Step 2: Wrap post-fill section with three-state machine**
|
||||
|
||||
After `lobsim.step_fill_from_market_targets(...)` and before the `extract_realized_pnl_delta` block, add:
|
||||
|
||||
```rust
|
||||
if self.reward_graph.is_some() {
|
||||
self.reward_graph.as_ref().unwrap().launch()
|
||||
.context("reward graph launch")?;
|
||||
} else {
|
||||
let capturing_reward = self.postfill_graph.is_some();
|
||||
if capturing_reward {
|
||||
self.stream
|
||||
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
|
||||
.map_err(|e| anyhow::anyhow!("reward begin_capture: {e}"))?;
|
||||
}
|
||||
|
||||
// ... existing post-fill kernels (extract_realized_pnl_delta through
|
||||
// launch_var_over_abs_mean) stay here unchanged ...
|
||||
|
||||
if capturing_reward {
|
||||
let graph = self.stream
|
||||
.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)
|
||||
.context("reward end_capture")?
|
||||
.ok_or_else(|| anyhow::anyhow!("reward end_capture returned None"))?;
|
||||
self.reward_graph = Some(graph);
|
||||
}
|
||||
} // end else (reward warmup / capture)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Verify no host ops inside capture region**
|
||||
|
||||
Audit: no `alloc_zeros`, no `stream.synchronize()`, no `lobsim.pos_bytes()` HOST calls inside the region. The `lobsim.pos_d()`, `lobsim.bid_px_d()` calls return stable device pointers — safe.
|
||||
|
||||
The `launch_rl_fused_controllers(b_size)` and `launch_var_over_abs_mean(...)` are method calls — verify they contain only kernel launches (already confirmed: pure `launch_builder` + `launch`).
|
||||
|
||||
- [ ] **Step 4: Run local smoke + compute-sanitizer**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "feat(rl): CUDA Graph B capture for post-fill reward/EMA/controller pipeline (~20 kernels)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 4: P2 — Graph C Capture (replay training step)
|
||||
|
||||
**Files:**
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`
|
||||
|
||||
**Prerequisite:** Tasks 2+3 (buffers pre-allocated, compiles clean).
|
||||
|
||||
Graph C captures `step_synthetic`'s kernel sequence. The three-state machine is the same pattern. Add `training_graph: Option<CudaGraph>`.
|
||||
|
||||
- [ ] **Step 1: Add `training_graph: Option<CudaGraph>` field + warmup counter**
|
||||
|
||||
```rust
|
||||
training_graph: Option<CudaGraph>,
|
||||
training_graph_warmup_done: bool,
|
||||
```
|
||||
|
||||
Initialize: `training_graph: None, training_graph_warmup_done: false,`
|
||||
|
||||
- [ ] **Step 2: Wrap step_synthetic with capture**
|
||||
|
||||
At the top of `step_synthetic`, after the `let b_size = ...` + `let k_dqn = ...` lines:
|
||||
|
||||
```rust
|
||||
if self.training_graph.is_some() {
|
||||
self.training_graph.as_ref().unwrap().launch()
|
||||
.context("training graph launch")?;
|
||||
// Skip the rest — graph replays all kernels.
|
||||
self.stream.synchronize().context("training graph sync")?;
|
||||
// Read back losses from mapped-pinned...
|
||||
return Ok(step_synthetic_stats_from_device(self)?);
|
||||
}
|
||||
|
||||
let capturing_training = self.training_graph_warmup_done;
|
||||
if !self.training_graph_warmup_done {
|
||||
self.training_graph_warmup_done = true;
|
||||
}
|
||||
if capturing_training {
|
||||
self.stream
|
||||
.begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED)
|
||||
.map_err(|e| anyhow::anyhow!("training begin_capture: {e}"))?;
|
||||
}
|
||||
|
||||
// ... existing step_synthetic body ...
|
||||
|
||||
if capturing_training {
|
||||
let graph = self.stream
|
||||
.end_capture(CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH)
|
||||
.context("training end_capture")?
|
||||
.ok_or_else(|| anyhow::anyhow!("training end_capture returned None"))?;
|
||||
self.training_graph = Some(graph);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Extract loss readback to separate method**
|
||||
|
||||
The graph replay skips the kernel dispatches but still needs to read loss scalars back. Extract the mapped-pinned loss reads at the end of step_synthetic into a helper that both the graph-replay path and the capture path call AFTER sync.
|
||||
|
||||
- [ ] **Step 4: Verify no host ops in capture region**
|
||||
|
||||
Check: no `alloc_zeros`, no `synchronize`, no host-side branching. The `self.launch_*` helper methods must be pure kernel launches.
|
||||
|
||||
- [ ] **Step 5: Run local smoke**
|
||||
|
||||
```bash
|
||||
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml-alpha --test integrated_trainer_smoke -- --ignored --nocapture
|
||||
```
|
||||
|
||||
- [ ] **Step 6: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "feat(rl): CUDA Graph C capture for replay training step (~25 kernels)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 5: P1 — Benchmark b=128 on L40S
|
||||
|
||||
**Files:**
|
||||
- None (infrastructure test)
|
||||
|
||||
- [ ] **Step 1: Push and submit b=128 run**
|
||||
|
||||
```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=5000 \
|
||||
-p n-backtests=128
|
||||
```
|
||||
|
||||
- [ ] **Step 2: Measure throughput**
|
||||
|
||||
Wait for completion, check `elapsed` in logs. Calculate steps/sec and transitions/sec.
|
||||
|
||||
Expected: ~700 steps/sec × 128 = ~89,600 transitions/sec (48× improvement over baseline).
|
||||
|
||||
- [ ] **Step 3: If passes, submit b=256**
|
||||
|
||||
```bash
|
||||
argo submit -n foxhunt --from=wftmpl/alpha-rl \
|
||||
-p git-branch=ml-alpha-phase-a \
|
||||
-p n-steps=5000 \
|
||||
-p n-backtests=256
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Task 6: P3 — Fused Reward Pipeline Kernel
|
||||
|
||||
**Files:**
|
||||
- Create: `crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu`
|
||||
- Modify: `crates/ml-alpha/build.rs`
|
||||
- Modify: `crates/ml-alpha/src/trainer/integrated.rs`
|
||||
|
||||
- [ ] **Step 1: Write the fused kernel**
|
||||
|
||||
```cuda
|
||||
// rl_fused_reward_pipeline.cu — Fuses 6 sequential post-fill kernels:
|
||||
// extract_realized_pnl_delta + rl_reward_shaping + abs_copy +
|
||||
// ema_update_on_done(duration) + ema_update_on_done(abs_pnl) +
|
||||
// ema_update_on_done(outcome) + rl_recent_outcome_update.
|
||||
//
|
||||
// Single launch replaces 7 kernels. Eliminates 6 L2 cache round-trips
|
||||
// on rewards_d and 6 kernel launch overheads.
|
||||
//
|
||||
// Grid=(1,1,1), Block=(b_size, 1, 1). Each thread handles one batch.
|
||||
// Shared memory for tree-reduce of EMAs.
|
||||
```
|
||||
|
||||
The kernel:
|
||||
1. Reads `pos_d[b]`, `prev_realized_pnl_d[b]` → computes reward, done
|
||||
2. Applies reward shaping (entry cost, hold bonus, min-hold penalty)
|
||||
3. Writes `raw_rewards_d[b]` (snapshot before scale)
|
||||
4. Computes `|reward|` → `reward_abs_d[b]`
|
||||
5. Updates `steps_since_done_d[b]`, emits `trade_duration_emit_d[b]`
|
||||
6. Tree-reduce: 3× EMA updates to ISV (mean_trade_duration, mean_abs_pnl, outcome)
|
||||
7. Writes `outcome_ema_d[b]` (per-batch)
|
||||
8. Writes final `rewards_d[b]`, `dones_d[b]`
|
||||
|
||||
- [ ] **Step 2: Register in build.rs**
|
||||
|
||||
```rust
|
||||
"rl_fused_reward_pipeline", // P3: 7→1 fused reward extraction + shaping + EMA update
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Wire into trainer — replace 7 kernel launches**
|
||||
|
||||
Replace the section from `extract_realized_pnl_delta` through `rl_recent_outcome_update` with a single launch of the fused kernel.
|
||||
|
||||
- [ ] **Step 4: Run local smoke + oracle test**
|
||||
|
||||
Verify loss trajectory matches the unfused version (within fp noise).
|
||||
|
||||
- [ ] **Step 5: Commit**
|
||||
|
||||
```bash
|
||||
git add crates/ml-alpha/cuda/rl_fused_reward_pipeline.cu crates/ml-alpha/build.rs crates/ml-alpha/src/trainer/integrated.rs
|
||||
git commit -m "perf(rl): fuse 7 post-fill reward kernels into one launch (P3)"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Kill Criteria
|
||||
|
||||
- P1: b=128 completes without OOM on L40S and throughput ≥ 5× vs b=16
|
||||
- P2 Graph B: local smoke passes after capture (step 1 warmup, step 2 capture, step 3+ replay)
|
||||
- P2 Graph C: same three-state machine, local smoke passes
|
||||
- P3: loss trajectory at step 1000 matches within 1% of unfused version
|
||||
|
||||
## Expected Final Throughput
|
||||
|
||||
| Config | Transitions/sec | vs Baseline |
|
||||
|--------|----------------|-------------|
|
||||
| Baseline (b=16, no graphs) | 1,840 | 1× |
|
||||
| P1 only (b=128) | ~89,600 | 48× |
|
||||
| P1 + P2 (b=128, all graphs) | ~92,000 | 50× |
|
||||
| P1 + P2 + P3 (b=128, fused) | ~95,000 | 52× |
|
||||
| P1 at b=256 | ~150,000+ | 80×+ |
|
||||
143
docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
Normal file
143
docs/superpowers/specs/2026-05-25-cuda-performance-roadmap.md
Normal file
@@ -0,0 +1,143 @@
|
||||
# CUDA Performance Roadmap
|
||||
|
||||
**Goal:** Maximize training throughput (steps/sec × batch_size = transitions/sec) on L40S (48GB) and H100 (80GB).
|
||||
|
||||
**Current baseline:** 115 steps/sec at b=16 on L40S = **1,840 transitions/sec**.
|
||||
|
||||
---
|
||||
|
||||
## P1: Batch Size Scaling (highest impact, minimal code)
|
||||
|
||||
**Problem:** At b=16, kernels complete in microseconds — GPU SMs are idle 90%+ of the time waiting for the next launch. Launch overhead (~10μs per kernel) dominates.
|
||||
|
||||
**Memory budget per batch:** 238 KB (dominated by `q_grad_w_per_batch` at 115 KB for the C51 backward).
|
||||
|
||||
**Available VRAM:**
|
||||
|
||||
| GPU | Total | Fixed overhead | Available | Max b (theoretical) | Recommended b |
|
||||
|-----|-------|----------------|-----------|---------------------|---------------|
|
||||
| RTX 3050 Ti | 4 GB | ~90 MB | 3.4 GB | 15,000 | 64 (smoke) |
|
||||
| L40S | 48 GB | ~90 MB | 47 GB | 208,000 | 256 |
|
||||
| H100 | 80 GB | ~90 MB | 79 GB | 350,000 | 512 |
|
||||
|
||||
**Expected throughput at higher batch sizes (L40S):**
|
||||
|
||||
| b_size | Transitions/sec | Speedup vs b=16 | Notes |
|
||||
|--------|----------------|------------------|-------|
|
||||
| 16 | 1,840 | 1.0× | Current — launch-limited |
|
||||
| 64 | ~25,600 | ~14× | GPU starts saturating |
|
||||
| 128 | ~44,800 | ~24× | Good SM utilization |
|
||||
| 256 | ~71,000 | ~39× | Near peak for these kernels |
|
||||
| 512 | ~92,000 | ~50× | Diminishing returns |
|
||||
|
||||
**Implementation:**
|
||||
- `n_backtests` CLI param already controls b_size
|
||||
- PER capacity must scale: `per_capacity = max(32768, 4 * b_size)`
|
||||
- Encoder K-loop seq_len is independent of batch size
|
||||
- Verify: `--n-backtests 256` on L40S, measure actual throughput
|
||||
|
||||
**Blockers:** None. Ship today.
|
||||
|
||||
---
|
||||
|
||||
## P2: CUDA Graph Capture (done for A+A2, remaining B+C)
|
||||
|
||||
**Status:**
|
||||
- Graph A (pre-snapshot, 20 kernels): ✅ captured
|
||||
- Graph A2 (post-snapshot/pre-fill, 7 kernels): ✅ captured
|
||||
- Graph B (post-fill reward/EMA/controllers, ~20 kernels): in progress
|
||||
- Blocker resolved: `ts_ns` moved to device-resident u64 via `rl_write_u64` kernel
|
||||
- 51 gradient buffers being pre-allocated for Graph C
|
||||
- Graph C (replay training step, ~25 kernels × K iterations): in progress
|
||||
- Blocker: 51 per-step `alloc_zeros` → persistent fields (agent working)
|
||||
|
||||
**Expected gain:** At b=16, ~5% (launch overhead is small fraction). At b=256, negligible (<1%). **Graph capture is insurance for when we increase K (replay-to-env ratio).**
|
||||
|
||||
---
|
||||
|
||||
## P3: Kernel Fusion — Reward Pipeline
|
||||
|
||||
**Problem:** The post-fill reward pipeline launches 6+ small sequential kernels on the same data: `extract_realized_pnl_delta` → `rl_reward_shaping` → `abs_copy` → `apply_reward_scale` → 3× `ema_update_on_done`. Each reads/writes rewards_d, dones_d.
|
||||
|
||||
**Fix:** Fuse into `rl_fused_reward_pipeline.cu`:
|
||||
- Single kernel, one block per batch
|
||||
- Reads pos_d, prev_realized_pnl_d once
|
||||
- Writes rewards_d, dones_d, reward_abs_d, raw_rewards_d
|
||||
- Updates 3 ISV EMA slots inline
|
||||
- Saves 5 kernel launches + 5× L2 cache round-trips on rewards_d
|
||||
|
||||
**Expected gain:** ~50μs/step at b=16, ~200μs at b=256 (memory bandwidth limited).
|
||||
|
||||
---
|
||||
|
||||
## P4: FP16 Gradient Accumulation
|
||||
|
||||
**Problem:** `reduce_axis0` and Adam steps are memory-bandwidth bound. Reading/writing f32 gradients at full precision wastes half the bandwidth.
|
||||
|
||||
**Fix:** Mixed-precision gradient pipeline:
|
||||
- Forward/backward compute stays f32 (numerical stability)
|
||||
- Per-batch gradient scratch (`*_per_batch_d`) stored as f16
|
||||
- `reduce_axis0` reads f16, accumulates f32, writes f32 reduced gradient
|
||||
- Adam reads f32 gradient, updates f32 weights
|
||||
|
||||
**Expected gain:** 2× bandwidth on gradient reduce + Adam = ~30% speedup on the training step (reduce_axis0 is ~40% of step_synthetic).
|
||||
|
||||
---
|
||||
|
||||
## P5: Multi-Stream Overlap
|
||||
|
||||
**Problem:** The step pipeline is strictly sequential: env-step → PER push/sample (host) → replay training. The host-side PER operations block the GPU.
|
||||
|
||||
**Fix:** Two CUDA streams:
|
||||
- Stream 0: env-step (Graphs A → fill → B) + PER push
|
||||
- Stream 1: replay training (Graph C × K) from PREVIOUS step's PER sample
|
||||
|
||||
Pipeline: while stream 1 trains on step N's data, stream 0 runs step N+1's env-step. PER sample for step N+1 happens during stream 1's training.
|
||||
|
||||
**Expected gain:** Hides PER latency (~200μs) + env-step overlap. ~1.3× at K=1, ~1.1× at K=4 (training dominates).
|
||||
|
||||
**Prerequisite:** Graph C captured (P2).
|
||||
|
||||
---
|
||||
|
||||
## P6: Persistent Kernel for K-loop
|
||||
|
||||
**Problem:** The replay training step runs K times per env step. Each iteration launches Graph C + syncs. At K=4, that's 4 graph launches + 4 syncs.
|
||||
|
||||
**Fix:** Single persistent kernel that:
|
||||
- Stays resident on SMs
|
||||
- Reads a "work counter" from device memory
|
||||
- For each K: reads PER sample indices from a ring buffer, runs the full training step inline
|
||||
- No host interaction until all K iterations complete
|
||||
|
||||
**Expected gain:** Eliminates K-1 sync points. At K=4: ~300μs saved/step. At K=8: ~700μs.
|
||||
|
||||
**Prerequisite:** Graph C working (P2), multi-stream (P5).
|
||||
|
||||
---
|
||||
|
||||
## P7: Warp-Cooperative Action Selection
|
||||
|
||||
**Problem:** `rl_pi_action_kernel` uses 1 thread per batch (sequential CDF walk). At b=256, that's 256 blocks × 1 thread = 256 SMs used at 0.4% occupancy each.
|
||||
|
||||
**Fix:** Warp-cooperative softmax + CDF walk:
|
||||
- 32 threads per batch (one warp)
|
||||
- Warp-shuffle for parallel softmax reduction
|
||||
- Parallel prefix-sum for CDF
|
||||
- Single-warp ballot for multinomial threshold crossing
|
||||
|
||||
**Expected gain:** ~10× faster action selection kernel. Negligible at b=16, ~50μs at b=256.
|
||||
|
||||
---
|
||||
|
||||
## Priority Order
|
||||
|
||||
1. **P1 (batch size)** — Ship immediately, 14-50× throughput increase
|
||||
2. **P2 (Graph B+C)** — In progress, enables P5/P6
|
||||
3. **P3 (fused reward)** — Medium effort, good constant-factor win
|
||||
4. **P4 (FP16 grads)** — 30% training step speedup
|
||||
5. **P5 (multi-stream)** — 1.3× overlap, needs P2
|
||||
6. **P6 (persistent K-loop)** — Eliminates sync overhead, needs P5
|
||||
7. **P7 (warp-coop action)** — Polish, only matters at large b
|
||||
|
||||
**Target:** P1 alone takes us from 1,840 to ~70,000 transitions/sec on L40S (38×). Combined with P2-P4: **~100,000 transitions/sec** — 1M steps in 10 seconds.
|
||||
Reference in New Issue
Block a user