perf: cuBLAS SGEMM pipeline + dead code elimination — 37s → 86ms/epoch (430x)

Phase 2: Replace 1-warp/sample fused kernels with cuBLAS SGEMM batched forward/backward. - batched_forward.rs: cuBLAS SGEMM forward (10 GEMM + bias/ReLU per pass) - batched_backward.rs: cuBLAS SGEMM backward (chain rule via GEMM, no atomicAdd) - c51_loss_kernel.cu: standalone C51 distributional loss (256 threads, 2KB shmem) - c51_grad_kernel: dL/d_logits with dueling routing for cuBLAS backward - BF16 alignment fix: pad offsets to even for short2 vectorized loads - Training step: 10.7ms → 0.7ms (15x) on RTX 3050 Phase 3: Unified cuBLAS Q-forward + dead code elimination (-4,400 lines net). - Rewrite experience collector: timestep loop + cuBLAS replaces monolithic 3,272-line kernel - Delete dqn_training_kernel.cu (1,385 lines) — replaced by dqn_utility_kernels.cu (118 lines) - Delete dqn_experience_kernel.cu (3,272 lines) — replaced by experience_kernels.cu (656 lines) - Remove BF16 warp-matvec helpers from common_device_functions.cuh (-159 lines) - Remove dead methods/fields from GpuDqnTrainer (-500 lines) - Experience collection: 348ms → 12ms (29x) on RTX 3050 - No fallback paths — cuBLAS is the only Q-forward implementation - All 1,514 tests pass, GPU smoke test verified with real data Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-21 15:33:00 +01:00
parent e1b8b46255
commit c2d116dfdf
14 changed files with 4762 additions and 6241 deletions
--- a/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase3.md
+++ b/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase3.md
@@ -0,0 +1,497 @@
+# H100 Epoch Optimization Phase 3: Unified cuBLAS + Dead Code Elimination
+
+> **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:** Reduce full-epoch wall time from 423ms to <150ms by unifying on one cuBLAS SGEMM Q-forward, restructuring the experience kernel, and eliminating ~1,900 lines of dead code.
+
+**Architecture:** Replace the embedded warp-matvec Q-forward in the experience kernel with timestep-batched cuBLAS SGEMM calls (reusing `CublasForward` from Phase 2). Split the 3,272-line monolithic experience kernel into 3 focused kernels (state_gather, action_select, env_step). Delete all dead BF16 warp-matvec code from training and common headers.
+
+**Tech Stack:** Rust, cudarc 0.19 (cuBLAS), CUDA (NVRTC), Philox RNG
+
+**Spec:** `docs/superpowers/specs/2026-03-21-h100-epoch-optimization-phase3-design.md`
+
+---
+
+## File Structure
+
+### Files to Create
+| File | Responsibility |
+|------|---------------|
+| `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu` | `grad_norm` + `adam_update` kernels (extracted from dqn_training_kernel.cu) |
+| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | 3 focused kernels: `state_gather`, `action_select`, `env_step` |
+
+### Files to Delete
+| File | Lines | Reason |
+|------|-------|--------|
+| `crates/ml/src/cuda_pipeline/dqn_training_kernel.cu` | 1,385 | All kernels replaced. grad_norm + adam extracted to utility file. |
+| `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu` | 3,272 | Monolithic kernel replaced by 3 focused kernels + cuBLAS |
+
+### Files to Modify
+| File | Changes |
+|------|---------|
+| `gpu_dqn_trainer.rs` (3,072 lines) | Remove ~600 lines: dead methods, dead kernel fields, BF16 sync from graph. Rewrite `compile_training_kernels` to load from `dqn_utility_kernels.cu`. Replace `forward_only_q` with cuBLAS. |
+| `gpu_experience_collector.rs` (1,554 lines) | Rewrite `launch_kernel` to timestep loop with cuBLAS + 3 small kernels. Share `CublasForward`. |
+| `common_device_functions.cuh` (1,718 lines) | Remove BF16 warp-matvec helpers (~300 lines): `cooperative_load_tile_bf16`, `cooperative_load_bias_bf16_to_f32`, `warp_matvec_bf16_shmem`, `BRANCHING_FORWARD_DISTRIBUTIONAL` macro. |
+| `mod.rs` (cuda_pipeline) | Update module exports |
+
+---
+
+### Task 1: Extract grad_norm + adam_update to utility kernel file
+
+**Files:**
+- Create: `crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — `compile_training_kernels()` function
+
+- [ ] **Step 1: Create `dqn_utility_kernels.cu`**
+
+Copy `dqn_grad_norm_kernel` (lines 1207-1241 of dqn_training_kernel.cu) and `dqn_adam_update_kernel` (lines 1253-1293) into the new file. No other kernels. Add the same `#ifndef` guards for compile-time constants.
+
+- [ ] **Step 2: Update `compile_training_kernels` in gpu_dqn_trainer.rs**
+
+Change `include_str!("dqn_training_kernel.cu")` to `include_str!("dqn_utility_kernels.cu")`. Remove function loads for `dqn_forward_loss_kernel`, `dqn_forward_only_kernel`, `dqn_backward_kernel`. Only load `dqn_grad_norm_kernel` and `dqn_adam_update_kernel`.
+
+Update the function return type from 7-tuple to 2-tuple `(grad_norm, adam_update)`.
+
+- [ ] **Step 3: Update the constructor**
+
+Change the constructor call site to destructure the 2-tuple. Remove `forward_loss_kernel`, `forward_only_kernel`, `backward_kernel` fields from the struct.
+
+- [ ] **Step 4: Verify compilation**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -3`
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+git commit -m "refactor: extract grad_norm + adam_update to dqn_utility_kernels.cu"
+```
+
+---
+
+### Task 2: Remove dead methods and fields from GpuDqnTrainer
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`
+
+- [ ] **Step 1: Remove dead struct fields**
+
+Delete these fields from `GpuDqnTrainer`:
+- `forward_loss_kernel: CudaFunction`
+- `forward_only_kernel: CudaFunction` (already removed in Task 1)
+- `backward_kernel: CudaFunction` (already removed in Task 1)
+- `shmem_bytes: usize` (only used by old kernel launch configs)
+
+- [ ] **Step 2: Remove dead methods**
+
+Delete these methods entirely:
+- `forward_only_q()` (~30 lines) — zero callers
+- `launch_forward_only()` (~50 lines) — only called by dead `forward_only_q`
+- `forward_loss()` (~80 lines) — zero callers outside file
+- `launch_forward_loss()` (~70 lines) — only called by dead `forward_loss`
+- `launch_backward()` (~65 lines) — replaced by `launch_cublas_backward`
+- `q_out_buf()` accessor — only used by dead `forward_only_q`
+
+- [ ] **Step 3: Remove BF16 sync from training graph**
+
+In `submit_training_ops()`, remove step 8 (`self.sync_online_bf16(...)`). The cuBLAS forward reads F32 from `params_buf`, not BF16 from `bf16_params_buf`.
+
+- [ ] **Step 4: Remove dead helper functions**
+
+Delete:
+- `compute_shmem_bytes()` — only used by old kernels
+- `compute_shmem_tile_rows()` — only used by old kernels
+- Any shared memory sizing constants only referenced by deleted code
+
+- [ ] **Step 5: Replace `forward_only_q` with cuBLAS version**
+
+Add a new `forward_only_q_cublas()` method that uses `CublasForward::forward_online()` to compute Q-values. Same signature as the old method but uses the cuBLAS path. This is for any future inference needs (e.g., GPU training guard Q-value monitoring).
+
+Read `launch_cublas_forward` for the pattern — extract Q-values from `on_v_logits_buf` and `on_b_logits_buf` via DtoD readback.
+
+- [ ] **Step 6: Verify compilation + tests**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -3`
+Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib 2>&1 | tail -3`
+Run: `SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -3`
+Expected: 359 + 855 passed, 0 failed
+
+- [ ] **Step 7: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
+git commit -m "refactor: remove ~500 lines of dead code from GpuDqnTrainer"
+```
+
+---
+
+### Task 3: Remove BF16 warp-matvec helpers from common_device_functions.cuh
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/common_device_functions.cuh`
+
+- [ ] **Step 1: Identify BF16 warp-matvec code to remove**
+
+Delete these functions/macros (approximate line ranges from current file):
+- `cooperative_load_tile_bf16()` (~line 412-428)
+- `cooperative_load_bias_bf16_to_f32()` (~line 434-441)
+- `warp_matvec_bf16_shmem()` and associated helpers (~line 446-550)
+- `BRANCHING_FORWARD_DISTRIBUTIONAL` macro (~line 530-551)
+- `SHMEM_TILE_ROWS_BF16` constant
+- Any other BF16-specific forward pass helpers
+
+Keep:
+- `f32_to_bf16_kernel` — used by BF16 seg sync (for supervised validation)
+- `bf16_to_f32_kernel` — used by BF16→F32 batch upload conversion
+- Block reduction helpers (`warp_reduce_sum_all`, etc.)
+- Philox RNG helpers
+- All `#define` constants (STATE_DIM, MARKET_DIM, etc.)
+
+- [ ] **Step 2: Verify compilation**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -3`
+This will catch any remaining references to deleted helpers.
+
+- [ ] **Step 3: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/common_device_functions.cuh
+git commit -m "refactor: remove BF16 warp-matvec helpers (~300 lines dead code)"
+```
+
+---
+
+### Task 4: Delete dqn_training_kernel.cu
+
+**Files:**
+- Delete: `crates/ml/src/cuda_pipeline/dqn_training_kernel.cu`
+- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — remove `include_str!("dqn_training_kernel.cu")`
+
+- [ ] **Step 1: Verify no remaining references**
+
+Search for `dqn_training_kernel` in the codebase:
+```bash
+grep -rn "dqn_training_kernel" crates/ml/src/
+```
+Only `include_str!("dqn_training_kernel.cu")` should remain (already replaced in Task 1).
+
+- [ ] **Step 2: Delete the file**
+
+```bash
+rm crates/ml/src/cuda_pipeline/dqn_training_kernel.cu
+```
+
+- [ ] **Step 3: Verify compilation + tests**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml-dqn --lib && SQLX_OFFLINE=true cargo test -p ml --lib`
+
+- [ ] **Step 4: Commit**
+
+```bash
+git add -A crates/ml/src/cuda_pipeline/
+git commit -m "refactor: delete dqn_training_kernel.cu (1,385 lines — all kernels replaced)"
+```
+
+---
+
+### Task 5: Create experience_kernels.cu — 3 focused kernels
+
+**Files:**
+- Create: `crates/ml/src/cuda_pipeline/experience_kernels.cu`
+
+This file contains the 3 small kernels that replace the monolithic experience kernel. The Q-forward is handled by cuBLAS between kernel calls.
+
+- [ ] **Step 1: Write `state_gather_kernel`**
+
+```cuda
+// Gather market features at each episode's current timestep into a batch.
+// Grid: ceil(N/256), Block: 256.
+extern "C" __global__ void experience_state_gather(
+    const float* __restrict__ market_features,  // [total_bars, MARKET_DIM]
+    const float* __restrict__ ofi_features,     // [total_bars, OFI_DIM] or NULL
+    const int* __restrict__ episode_starts,     // [N]
+    const int* __restrict__ current_timesteps,  // [N] current t for each episode
+    const float* __restrict__ portfolio_states, // [N, PORTFOLIO_DIM]
+    float* __restrict__ batch_states,           // [N, STATE_DIM] output
+    int N, int total_bars, int state_dim, int market_dim, int ofi_dim)
+```
+
+Each thread gathers one episode's state: market features at `episode_starts[i] + current_timesteps[i]`, concatenated with portfolio state and OFI features.
+
+- [ ] **Step 2: Write `action_select_kernel`**
+
+```cuda
+// Epsilon-greedy action selection on cuBLAS Q-value output.
+// Grid: ceil(N/256), Block: 256.
+extern "C" __global__ void experience_action_select(
+    const float* __restrict__ q_values,     // [N, TOTAL_ACTIONS(11)]
+    int* __restrict__ out_actions,          // [N]
+    unsigned int* __restrict__ rng_states,  // [N]
+    float epsilon,
+    int N, int total_actions,
+    int b0_size, int b1_size, int b2_size,
+    int use_branching)
+```
+
+Per-episode: with probability epsilon pick random factored action, otherwise argmax over Q-values per branch → combine into factored action.
+
+- [ ] **Step 3: Write `env_step_kernel`**
+
+This is the largest of the 3 — it handles portfolio simulation, barrier tracking, reward computation, curiosity bonus, DSR, and done detection. Extract the per-timestep body from the monolithic kernel (the part AFTER the Q-forward and action selection).
+
+```cuda
+// Environment step: execute action, compute reward, update portfolio/barriers.
+// Grid: ceil(N/256), Block: 256.
+extern "C" __global__ void experience_env_step(
+    // Market data
+    const float* __restrict__ targets,          // [total_bars, 4]
+    const int* __restrict__ episode_starts,     // [N]
+    const int* __restrict__ current_timesteps,  // [N]
+    const int* __restrict__ actions,            // [N]
+    // Episode state (read-write)
+    float* __restrict__ portfolio_states,       // [N, PORTFOLIO_STATE_SIZE]
+    float* __restrict__ barrier_states,         // [N, BARRIER_STATE_SIZE]
+    // Config (read-only)
+    const float* __restrict__ barrier_config,   // [3]
+    float max_position, float gamma,
+    float barrier_scale, float risk_weight,
+    float hold_reward, float tx_cost_multiplier,
+    // Fill simulation config
+    float fill_median_spread, float fill_median_vol,
+    float fill_ioc_fill_prob, float fill_limit_fill_min,
+    float fill_limit_fill_max, float fill_spread_cost_frac,
+    float fill_spread_capture_frac, int fill_simulation_enabled,
+    // DSR config
+    int use_dsr, float dsr_eta,
+    // Output for this timestep
+    float* __restrict__ out_rewards,            // [N]
+    int* __restrict__ out_dones,                // [N]
+    float* __restrict__ out_states,             // [N, STATE_DIM] (the state BEFORE action)
+    // Misc
+    int N, int total_bars, int state_dim,
+    int use_branching, int b0_size, int b1_size, int b2_size,
+    unsigned int* rng_states)
+```
+
+- [ ] **Step 4: Verify kernel syntax**
+
+Compile with NVRTC dry run (the Rust compilation will validate when integrated in Task 6).
+
+- [ ] **Step 5: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/experience_kernels.cu
+git commit -m "feat: 3 focused experience kernels (state_gather + action_select + env_step)"
+```
+
+---
+
+### Task 6: Rewrite GpuExperienceCollector to use timestep-loop + cuBLAS
+
+**Files:**
+- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`
+
+This is the core task. Replace the monolithic `launch_kernel()` with a timestep loop.
+
+- [ ] **Step 1: Add CublasForward as shared dependency**
+
+The experience collector needs access to the same `CublasForward` instance used by the trainer. Add a field:
+```rust
+cublas_forward: Arc<CublasForward>,  // shared with GpuDqnTrainer
+```
+
+Update the constructor to accept `Arc<CublasForward>` instead of creating its own cuBLAS handle.
+
+- [ ] **Step 2: Compile the 3 new experience kernels**
+
+Add `compile_experience_kernels()` function that compiles `experience_kernels.cu` via NVRTC with the same `#define` injection pattern as the training kernels. Store the 3 `CudaFunction` handles.
+
+- [ ] **Step 3: Pre-allocate timestep buffers**
+
+Add pre-allocated buffers for the timestep loop:
+```rust
+// Per-timestep batch buffers (reused each timestep)
+batch_states: CudaSlice<f32>,        // [N, STATE_DIM]
+batch_q_values: CudaSlice<f32>,      // [N, 11]
+batch_actions: CudaSlice<i32>,       // [N]
+current_timesteps: CudaSlice<i32>,   // [N]
+// cuBLAS scratch (reused from CublasForward — already allocated)
+```
+
+Also pre-allocate cuBLAS activation scratch buffers for the forward pass:
+```rust
+exp_h_s1: CudaSlice<f32>,   // [N, SH1]
+exp_h_s2: CudaSlice<f32>,   // [N, SH2]
+exp_h_v: CudaSlice<f32>,    // [N, VH]
+exp_h_b: CudaSlice<f32>,    // [N, AH]
+exp_v_logits: CudaSlice<f32>, // [N, NA]
+exp_b_logits: CudaSlice<f32>, // [N, total_branch_atoms]
+```
+
+- [ ] **Step 4: Rewrite `launch_kernel()` as timestep loop**
+
+Replace the monolithic kernel launch with:
+```rust
+fn launch_timestep_loop(&mut self, ...) -> Result<(usize, usize), MLError> {
+    let n = n_episodes;
+    let l = timesteps;
+
+    // Upload episode starts + initial state
+    // ...
+
+    // Initialize current_timesteps to all zeros
+    self.stream.memset_zeros(&mut self.current_timesteps)?;
+
+    for t in 0..l {
+        // 1. Gather states: market features + portfolio → batch_states[N, SD]
+        self.launch_state_gather(n)?;
+
+        // 2. cuBLAS forward: batch_states → Q-values
+        let param_sizes = compute_param_sizes(&self.forward_config);
+        let w_ptrs = f32_weight_ptrs(&self.online_params_buf, &param_sizes, &self.stream);
+        self.cublas_forward.forward_online(
+            &self.stream,
+            &self.batch_states,
+            &w_ptrs,
+            &self.exp_h_s1, &self.exp_h_s2,
+            &self.exp_h_v, &self.exp_h_b, &self.exp_h_b, &self.exp_h_b,
+            &self.exp_v_logits, &self.exp_b_logits,
+        )?;
+
+        // 3. Action selection: epsilon-greedy on Q-values
+        self.launch_action_select(n, epsilon)?;
+
+        // 4. Environment step: execute action, compute reward
+        self.launch_env_step(n, t)?;
+
+        // 5. Increment current_timesteps
+        // (simple kernel or managed by env_step_kernel internally)
+    }
+
+    Ok((n_episodes, timesteps))
+}
+```
+
+- [ ] **Step 5: Flatten weight sync to single DtoD**
+
+Replace the per-tensor weight sync (`sync_dueling_weights_branching` + `sync_branching_weights`) with a single DtoD copy of the flat `params_buf`:
+```rust
+pub fn sync_weights_flat(&mut self, params_buf: &CudaSlice<f32>) -> Result<(), MLError> {
+    let bytes = self.total_params * std::mem::size_of::<f32>();
+    dtod_copy(
+        raw_device_ptr(&self.online_params_buf, &self.stream),
+        raw_device_ptr(params_buf, &self.stream),
+        bytes, &self.stream, 0, "weight_sync_flat",
+    )
+}
+```
+
+- [ ] **Step 6: Verify compilation + tests**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml`
+Run: `SQLX_OFFLINE=true cargo test -p ml-dqn --lib`
+Run: `SQLX_OFFLINE=true cargo test -p ml --lib`
+
+- [ ] **Step 7: Run GPU smoke test**
+
+Run: `SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_train_step_produces_finite_metrics --ignored --test-threads=1`
+
+- [ ] **Step 8: Commit**
+
+```bash
+git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
+git commit -m "perf: rewrite experience collector — timestep loop + cuBLAS Q-forward"
+```
+
+---
+
+### Task 7: Delete dqn_experience_kernel.cu
+
+**Files:**
+- Delete: `crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu`
+
+- [ ] **Step 1: Verify no remaining references**
+
+```bash
+grep -rn "dqn_experience_kernel\|dqn_full_experience_kernel" crates/ml/src/
+```
+
+- [ ] **Step 2: Delete the file**
+
+```bash
+rm crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu
+```
+
+- [ ] **Step 3: Verify compilation + full test suite**
+
+Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml-dqn --lib && SQLX_OFFLINE=true cargo test -p ml --lib`
+
+- [ ] **Step 4: Commit**
+
+```bash
+git add -A crates/ml/src/cuda_pipeline/
+git commit -m "refactor: delete dqn_experience_kernel.cu (3,272 lines — replaced by cuBLAS + 3 focused kernels)"
+```
+
+---
+
+### Task 8: Profile and validate on local GPU
+
+**Files:**
+- Read: training_loop.rs phase breakdown output
+
+- [ ] **Step 1: Run profiling on RTX 3050**
+
+```bash
+SQLX_OFFLINE=true cargo run --release --example train_baseline_rl -p ml -- \
+  --model dqn --data-dir test_data/futures-baseline --symbol ES.FUT \
+  --epochs 3 --max-steps-per-epoch 50 --batch-size 64 \
+  --train-months 3 --val-months 1 --test-months 1 --step-months 3 \
+  2>&1 | grep -E "phase breakdown|step breakdown"
+```
+
+Expected: experience phase < 100ms (down from 348ms), total epoch < 200ms.
+
+- [ ] **Step 2: Run full test suite**
+
+```bash
+SQLX_OFFLINE=true cargo test -p ml-core --lib && \
+SQLX_OFFLINE=true cargo test -p ml-dqn --lib && \
+SQLX_OFFLINE=true cargo test -p ml --lib
+```
+
+Expected: 300 + 359 + 855 = 1,514 passed, 0 failed.
+
+- [ ] **Step 3: Run GPU smoke test with real data**
+
+```bash
+SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline \
+  cargo test -p ml --lib -- test_train_step_produces_finite_metrics --ignored --test-threads=1
+```
+
+Expected: PASSED with finite loss.
+
+- [ ] **Step 4: Commit profiling results**
+
+Update spec with measured Phase 3 speedup.
+
+```bash
+git add docs/superpowers/specs/2026-03-21-h100-epoch-optimization-phase3-design.md
+git commit -m "docs: Phase 3 profiling results"
+```
+
+---
+
+## Summary
+
+| Task | Lines removed | Lines added | Net |
+|------|:---:|:---:|:---:|
+| 1. Extract utility kernels | ~1,200 | ~90 | -1,110 |
+| 2. Dead methods/fields | ~500 | ~30 | -470 |
+| 3. BF16 warp-matvec helpers | ~300 | 0 | -300 |
+| 4. Delete dqn_training_kernel.cu | 1,385 | 0 | -1,385 |
+| 5. Experience kernels | 0 | ~400 | +400 |
+| 6. Rewrite experience collector | ~800 | ~400 | -400 |
+| 7. Delete dqn_experience_kernel.cu | 3,272 | 0 | -3,272 |
+| 8. Profile | 0 | 0 | 0 |
+| **Total** | **~7,457** | **~920** | **-6,537** |
--- a/docs/superpowers/specs/2026-03-21-h100-epoch-optimization-phase3-design.md
+++ b/docs/superpowers/specs/2026-03-21-h100-epoch-optimization-phase3-design.md
@@ -0,0 +1,134 @@
+# H100 Epoch Optimization Phase 3: Unified cuBLAS + Dead Code Elimination
+
+## Goal
+
+Reduce full-epoch wall time from 423ms to <150ms by:
+1. Replacing the embedded warp-matvec Q-forward in the experience kernel with cuBLAS SGEMM (fixing the 82% bottleneck)
+2. Eliminating all dead code from Phase 1/2 transitions (6,000+ lines of stale CUDA + Rust)
+3. Removing BF16 warp-matvec infrastructure that is no longer on any active code path
+4. Unifying on one Q-forward implementation (cuBLAS) across training, experience collection, and inference
+
+## Current State (Post-Phase 2)
+
+### RTX 3050 profiling (50 steps, batch_size=64, num_atoms=11):
+| Phase | Time (ms) | % | Notes |
+|-------|-----------|---|-------|
+| Init | 8 | 2% | GPU data upload |
+| **Experience** | **348** | **82%** | Warp-matvec Q-forward at 1.56% occupancy |
+| Training | 39 | 9% | cuBLAS SGEMM, 0.7ms/step |
+| Validation | 27 | 6% | Candle framework (not GPU trainer) |
+| **Total** | **423** | 100% | |
+
+### Dead code inventory:
+| Code | Lines | Status |
+|------|-------|--------|
+| `dqn_forward_loss_kernel` (BF16 fused) | ~200 | No callers in training path |
+| `dqn_forward_only_kernel` | ~100 | Zero callers anywhere |
+| `dqn_backward_kernel` (atomicAdd) | ~300 | Replaced by cuBLAS backward |
+| `forward_only_q()` + `launch_forward_only()` | ~60 | Zero callers |
+| `forward_loss()` (validation) | ~60 | Zero callers outside gpu_dqn_trainer.rs |
+| `launch_forward_loss()` | ~70 | Only called by dead `forward_loss()` |
+| `launch_backward()` | ~65 | Replaced by cuBLAS backward |
+| BF16 warp-matvec helpers in `.cuh` | ~300 | Only used by dead kernels |
+| BF16 sync in training graph | ~3 | cuBLAS reads F32, not BF16 |
+| Experience kernel warp-matvec Q-forward | ~700 | Replaced by cuBLAS (this phase) |
+| **Total dead code** | **~1,900** | |
+
+## Architecture
+
+### Single Q-Forward: cuBLAS SGEMM everywhere
+
+After Phase 3, there is ONE way to run the DQN Q-network forward pass: `CublasForward::forward_online()`. It is used in:
+1. **Training** — 3 passes per step (online, target, online-next for DDQN)
+2. **Experience collection** — 1 pass per timestep, batched across N episodes
+3. **Inference** (`forward_only_q`) — via cuBLAS (replaces BF16 kernel)
+
+### Experience Kernel Restructuring
+
+**Current**: Monolithic 3,272-line CUDA kernel. Each warp runs one episode for L timesteps. The Q-forward is embedded inside the kernel using warp-cooperative shared-memory matvec.
+
+**New**: Timestep-level loop in Rust, calling cuBLAS between environment-step kernels:
+
+```
+for t in 0..timesteps_per_episode:
+    1. state_gather_kernel:   read episode states → batch [N, SD]
+    2. CublasForward::forward_online(): [N, SD] → Q-values [N, 11]
+    3. action_select_kernel:  epsilon-greedy on Q-values → actions [N]
+    4. env_step_kernel:       portfolio sim, reward, done → new states [N, SD]
+```
+
+The monolithic experience kernel is split into 3 small focused kernels:
+- `state_gather_kernel` — reads market features at each episode's current timestep
+- `action_select_kernel` — epsilon-greedy with Philox RNG (already GPU-native)
+- `env_step_kernel` — portfolio simulation, barrier tracking, curiosity reward, TD error
+
+Episode state (position, cash, barriers) moves from thread registers to global memory buffers (pre-allocated, reused across timesteps).
+
+**Kernel launch count**: L timesteps × (1 gather + 10 GEMM + 5 bias + 1 action + 1 env) = ~18L launches. With L=100: 1,800 launches × ~3μs = 5.4ms overhead. Acceptable vs 348ms baseline.
+
+### Files to Delete
+
+| File | Lines | Reason |
+|------|-------|--------|
+| `dqn_training_kernel.cu` | 1,385 | All kernels replaced: forward→cuBLAS, backward→cuBLAS, loss→c51_loss_kernel.cu. Only `grad_norm` and `adam_update` survive — move them to a new small file. |
+| `dqn_experience_kernel.cu` | 3,272 | Monolithic kernel replaced by 3 focused kernels + cuBLAS |
+
+### Files to Create
+
+| File | Purpose |
+|------|---------|
+| `dqn_utility_kernels.cu` | grad_norm + adam_update (extracted from dqn_training_kernel.cu) |
+| `experience_gather_kernel.cu` | state_gather + action_select + env_step (3 small kernels) |
+
+### Files to Modify
+
+| File | Changes |
+|------|---------|
+| `gpu_dqn_trainer.rs` | Remove: `forward_only_q`, `launch_forward_only`, `forward_loss`, `launch_forward_loss`, `launch_backward`, `forward_only_kernel` field, BF16 sync from submit_training_ops. Replace `forward_only_q` with cuBLAS-based version. |
+| `gpu_experience_collector.rs` | Replace monolithic `launch_kernel` with timestep loop calling cuBLAS + 3 small kernels. Share `CublasForward` instance with trainer. |
+| `common_device_functions.cuh` | Remove: `cooperative_load_tile_bf16`, `cooperative_load_bias_bf16_to_f32`, `warp_matvec_bf16_shmem`, `BRANCHING_FORWARD_DISTRIBUTIONAL` macro, all BF16 warp-matvec helpers. Keep: f32_to_bf16_kernel, bf16_to_f32_kernel (for supervised validation), block reductions, Philox RNG. |
+| `mod.rs` (cuda_pipeline) | Update module declarations |
+| `fused_training.rs` | Remove references to old kernels |
+
+### Dead Code Removal from gpu_dqn_trainer.rs
+
+Fields to remove:
+- `forward_loss_kernel: CudaFunction`
+- `forward_only_kernel: CudaFunction`
+- `backward_kernel: CudaFunction`
+- `bf16_mirrors_initialized: bool` (move to lazy init in forward_only_q only)
+- `shmem_bytes: usize` (only used by old kernels)
+
+Methods to remove:
+- `forward_only_q()`, `launch_forward_only()`
+- `forward_loss()`, `launch_forward_loss()`
+- `launch_backward()`
+- `ensure_bf16_mirrors()` (inline into the one remaining caller if any)
+- `sync_online_bf16()` from `submit_training_ops` (keep the method, remove the call from training graph)
+
+### Weight Sync Simplification
+
+Current: `sync_gpu_weights()` → extract 20 tensors from VarStore → DtoD copy each.
+New: DtoD copy flat `params_buf` directly to experience collector. Single 1.1MB memcpy.
+
+### Pipeline Overlap (future — not in Phase 3 scope)
+
+Experience collection and training CAN overlap on separate CUDA streams since DQN is off-policy. This is noted but deferred to avoid scope creep. Phase 3 focuses on making each phase fast, not on pipelining them.
+
+## Expected Results
+
+| Phase | Before | After | Speedup |
+|-------|--------|-------|---------|
+| Experience | 348ms | ~50ms | 7x (cuBLAS occupancy) |
+| Training | 39ms | ~35ms | 1.1x (remove BF16 sync from graph) |
+| Validation | 27ms | 27ms | — (unchanged, uses Candle) |
+| Init | 8ms | 8ms | — |
+| Compile | ~3s | ~1.5s | 2x (fewer CUDA files) |
+| **Total epoch** | **423ms** | **~120ms** | **3.5x** |
+| Dead code removed | — | ~1,900 lines | cleaner codebase |
+
+## Non-Goals
+
+- Pipeline overlap (experience + training on separate streams) — separate phase
+- Supervised model cuBLAS conversion — those models have their own forward pass architecture, not DQN's branching dueling network
+- cuBLAS HGEMM (BF16 tensor cores) — F32 SGEMM is sufficient; BF16 is a follow-up