# 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** |
