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

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