# 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):

```cuda
// 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)