foxhunt/docs/superpowers/plans/2026-03-21-h100-epoch-optimization-phase3.md

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

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

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

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:

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

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

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

// 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

// 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).

// 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

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:

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:

// 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:

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:

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:

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

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

grep -rn "dqn_experience_kernel\|dqn_full_experience_kernel" crates/ml/src/

• Step 2: Delete the file

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

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

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

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

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.

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