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

H100 Epoch Optimization: 37s → <5s per epoch

Date: 2026-03-21 Target: <5s/epoch on H100 (80GB), <10s acceptable, <3s stretch goal Current: 37.56s/epoch (H100, 51 atoms, 300 training steps, CUDA Graph)

Profiling Results

RTX 3050 phase breakdown (proportions apply to H100):

Phase	Time	%	Per-step
Init (data upload)	10ms	0.1%	once/epoch
Experience collection	348ms	2.3%	1 kernel launch
Training steps	14,858ms	97.2%	300 steps
Validation	27ms	0.2%	1 backtest pass

H100 epoch time: 37.56s × 97.2% training = 36.5s in training steps. At 300 steps/epoch: 122ms per training step.

Root Cause Analysis

Primary bottleneck: PER sampling CPU roundtrips (300×/epoch)

Each sample_proportional() call does:

1. DtoH: cs_total() reads prefix sum total (1 scalar) — forces stream sync
2. CPU RNG: rand::thread_rng().gen::<f32>() × batch_size — serial CPU work
3. HtoD: uploads random thresholds (batch_size × 4 bytes)
4. HtoD: zeros batch_max (4 bytes)

300 calls × ~100μs per sync barrier = 30ms just for sync, but the CPU RNG + alloc + memcpy overhead adds 100-400ms per call.

Secondary: kernel underutilization

Training kernel: block_dim=(32,1,1), grid_dim=(batch_size,1,1) = 256 blocks × 32 threads = 8,192 threads on 132-SM H100 (3.0% occupancy). Each SM can run 2048 threads — we're using 1 warp per SM on a fraction of 132 SMs.

Tertiary: sequential phase execution

Experience collection → Training → Validation runs serially. No overlap.

Design: Three-Phase Optimization

Phase 1: Eliminate CPU Roundtrips (target: 10-15s/epoch)

1a. GPU-native PER threshold generation

Replace CPU rand::thread_rng() with Philox CUDA RNG kernel (already have this pattern in IQN tau sampler):

// New kernel: per_generate_thresholds_kernel
__global__ void per_generate_thresholds(
    float* __restrict__ thresholds,   // [batch_size] output
    float total_sum,                   // from prefix sum (stays on GPU)
    unsigned long long seed,           // counter-based, no state
    int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    // Philox-based PRNG (same as iqn_sample_taus_kernel)
    unsigned int r = philox_uint(seed, i);
    float u = __uint_as_float(0x3f800000u | (r >> 9)) - 1.0f; // uniform [0,1)
    thresholds[i] = u * total_sum;
}

Key change: total_sum stays on GPU — read it via a single DtoD copy from cs_buf[n-1] to a scalar buffer, then pass as kernel arg. Zero DtoH.

1b. Fuse PER kernels into single launch

Current: 10+ kernel launches per sample (pow_alpha, prefix_sum, searchsorted, 5 gathers, IS weights, reduce).

Fuse into 2 kernels:

1. per_sample_and_gather_kernel: prefix sum total (already computed) → thresholds → searchsorted → gather all fields
2. per_is_weights_kernel: compute IS weights + normalize

Or capture PER sampling into a second CUDA Graph.

1c. Remove adam_step HtoD

Store adam_step as a GPU scalar, increment via 1-element kernel or atomicAdd.

1d. Audit and remove all hot-path to_host() calls

Remove GpuTensor::to_host() from the public API. Replace with to_host_checkpoint() (clearly marked for cold-path use only). Compile-time guard via #[deprecated] or feature gate.

Phase 2: Increase Kernel Occupancy (target: 5-8s/epoch)

2a. Batched training: multiple samples per warp

Current: 1 warp (32 threads) per sample, 64 samples per graph replay. Target: 4-8 samples per block (128-256 threads), or increase batch_size to 512-1024.

With batch_size=512: grid=(512,1,1), block=(32,1,1) = 16K threads = 7.5% occupancy. With batch_size=1024 + 4 samples/block: 8K threads at 128 threads/block = better SM utilization.

2b. cuBLAS for network layers

Replace warp-cooperative matvec with cuBLAS cublasSgemmStridedBatched for the shared trunk layers (256×256 matmul). For batch_size=1024, this is a [1024, 256] × [256, 256] GEMM — cuBLAS tensor cores will dominate.

Trade-off: cuBLAS adds launch overhead but uses tensor cores (3.9 PFLOPS BF16 on H100 vs ~100 TFLOPS custom kernel).

Profile both and pick winner. Custom kernel may win for batch<256.

2c. BF16 throughout (not just storage)

Current: BF16 weight storage → F32 accumulation in kernel. Target: BF16 weights + BF16 activations + F32 accumulation (tensor core matmul).

H100 tensor cores: BF16 GEMM at 989 TFLOPS vs F32 at 67 TFLOPS = 14.8× potential.

Phase 3: Pipeline Overlap (target: <5s/epoch)

3a. Dual-stream: overlap experience + training

Stream A: [experience_epoch_N+1] → [experience_epoch_N+2] → ...
Stream B: [train_epoch_N]        → [train_epoch_N+1]      → ...
           ↑ uses experiences from epoch N (already in buffer)

After epoch 1, experience collection for epoch N+1 overlaps with training for epoch N. ~2× throughput.

3b. Fuse PER sampling + training into mega-graph

Capture the entire sequence (PER sample → batch upload → forward+loss → backward → adam → EMA) into a single CUDA Graph. One cuGraphLaunch per step instead of 15+ kernel launches.

Requires: GPU-native PER (Phase 1a) and graph-compatible memory pool.

3c. Async validation

Move validation backtest to a separate stream, overlapping with the start of the next epoch's experience collection.

Implementation Order

1. Phase 1a: GPU Philox RNG for PER thresholds (eliminate sync barriers)
2. Phase 2a: Increase batch_size to 512+ on H100 (config change, immediate win, amplifies all later optimizations)
3. Phase 1b: Fuse PER kernels + pre-allocate all PER buffers (zero cuMemAlloc after warmup)
4. Phase 1c: GPU adam_step increment
5. Phase 1d: Remove to_host() from hot paths, deprecate the function
6. Phase 2b: Profile cuBLAS vs custom for training matmul
7. Phase 2c: BF16 activations (with C51 log_softmax numerical guard)
8. Phase 3a: Dual-stream overlap (requires resolving stream unification tension)
9. Phase 3b: Mega-graph (if 3a insufficient)

Expected Speedup Stack

Optimization	Expected speedup	Cumulative
GPU PER RNG (eliminate 300 sync barriers)	2-3×	12-18s
Batch_size 512+ (occupancy)	1.3-1.5×	9-12s
Fuse PER kernels + pre-alloc (reduce 14→2 launches)	1.2-1.5×	6-10s
cuBLAS/tensor core matmul	1.5-3×	3-6s
Dual-stream overlap	1.5-2×	2-4s

Conservative estimate: 5-8s/epoch. Aggressive: 3-4s/epoch.

Known Risks

CUDA Graph capture requires fixed buffer addresses

PER's sample_proportional currently allocates new CudaSlice buffers per call (lines 434, 436, 443, 450-451, 478, 490 — ~10 allocs × 300 calls = 3,000 GPU malloc/free per epoch). For mega-graph capture (Phase 3b), ALL buffers must be pre-allocated at fixed addresses. Phase 1b must pre-allocate these.

Dual-stream conflicts with stream unification

The stream_unification_plan.md explicitly advocates collapsing to a single forked stream to avoid cudarc event tracking conflicts. Phase 3a proposes re-introducing dual streams. Resolution: use separate CudaContext per stream (each context has independent event tracking), or complete cudarc event tracking fix first.

C51 log_softmax underflow in BF16

Phase 2c (BF16 activations) risks numerical issues in C51 distributional loss. log_softmax on BF16 can underflow for low-probability atoms. Mitigation: keep loss computation in F32 accumulation, only use BF16 for matmul operands. Profile numerically before deploying.

Prefix sum single-block scalability

The prefix sum kernel runs with grid_dim=(1,1,1) — single block processing up to ~1M elements. At large replay buffer capacities (1M+ on H100), this serializes. May need multi-block parallel scan (CUB DeviceScan) for Phase 2a batch_size increases.

Success Criteria

• Phase breakdown log shows zero DtoH in training steps (except 8-byte loss+grad_norm per step)
• Zero cuMemAlloc calls in sample_proportional after first warmup call (pre-allocated buffers only)
• scripts/gpu-hotpath-guard.sh catches any new CPU roundtrip in hot path
• H100 epoch time < 5s with 300 training steps, 51 atoms, batch_size ≥ 512
• RTX 3050 epoch time < 10s (proportional improvement from current ~15s)
• All 1861+ tests pass
• Training metrics regression < 5%: eval Sharpe within 5% of baseline, loss convergence within 10% of baseline epoch count
• to_host() deprecated with #[deprecated] attribute — only to_host_checkpoint() available for cold paths

Risk: Experience Collector Becomes Bottleneck

At 2.3% of current epoch time, the experience collector seems negligible. But on H100 with 256 episodes × 500 timesteps = 128K forward passes, it likely takes 5-10s. Once training drops to <5s, experience collection becomes the dominant phase.

Mitigations (included in Phase 3):

• Phase 3a (dual-stream): overlap experience with training → hides latency entirely
• Occupancy: experience kernel uses episodes_per_block=4-8 with warp kernel on H100 — already reasonably optimized
• Reduce timesteps: 500 timesteps may be excessive; 200 could suffice with higher episode count for same total experiences
• Profile independently: add per-kernel CUDA event timing inside experience collector to identify sub-kernel bottlenecks

Non-Goals

• Multi-GPU training (separate effort, needs NCCL)
• Model architecture changes (network dims stay 256×256)
• Hyperopt speed (separate binary, different bottleneck)