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

H100 Epoch Optimization Phase 2: Kernel Occupancy & cuBLAS

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 per-training-step time from 10.7ms to <1ms by fixing kernel occupancy (1.56% → >50%) and replacing warp-matvec with cuBLAS tensor core matmul.

Architecture: Replace the 1-warp-per-sample fused kernel design with a batched cuBLAS GEMM architecture. Each layer becomes a standard matrix multiply over the full batch, enabling H100 tensor cores (3,958 TFLOPS BF16). Fuse the 20 per-step BF16 conversions and 20 EMA updates into single flat-buffer kernels.

Tech Stack: Rust, cudarc 0.19 (cuBLAS bindings), CUDA, H100 tensor cores

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

---

Kernel Audit Findings

Current per-step breakdown (47 kernel launches):

Kernel	Launches/step	Threads/block	Occupancy (H100)	Issue
forward+loss	1	32	1.56%	46KB stack spill, zero latency hiding
backward	1	32	1.56%	atomicAdd serialization on 288K grads
grad_norm	1	256	~100%	OK
adam_update	1	256	~100%	OK
f32_to_bf16	20	256	~100%	20 launches → fuse to 1
bf16_to_f32	2	256	~100%	OK
DtoD unflatten	20	—	—	20 driver memcpy → fuse
PER sample	15	256	varies	prefix_sum single-block
EMA (per update)	20	256	varies	20 launches → fuse to 1

Root cause: 1 warp/sample architecture

The forward+loss and backward kernels process each batch sample independently in 1 warp (32 threads). With 134KB shmem/block on H100, only 1 block fits per SM → 32 threads per SM → 1.56% occupancy. The remaining 98.4% of H100 compute is idle.

---

Implementation Tasks

Task 1: Restructure forward pass — batched cuBLAS GEMM

The single biggest win. Replace the 3× branching_forward_distributional() warp-matvec loops with 6 batched cuBLAS SGEMM calls per forward pass.

Current architecture (per sample, per layer):

for each sample (warp):
    for each tile of weight matrix:
        cooperative_load_tile_bf16(shmem)
        warp_matvec(input_dist, shmem, output_dist)

New architecture (batched over all samples):

// All samples processed in one GEMM call per layer
cublasSgemmStridedBatched(
    // [batch_size, in_dim] @ [in_dim, out_dim] → [batch_size, out_dim]
)
// Then batch ReLU kernel
// Then next layer GEMM

Files:

• Create: crates/ml/src/cuda_pipeline/batched_forward.rs — cuBLAS-based forward pass
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — replace graph-captured forward with cuBLAS
• Keep: existing warp kernel as fallback for sm_<70

Layers to convert (per forward pass):

1. w_s1 [256, 48] @ states [B, 48]^T → h_s1 [B, 256] + bias + ReLU
2. w_s2 [256, 256] @ h_s1 [B, 256]^T → h_s2 [B, 256] + bias + ReLU
3. w_v1 [128, 256] @ h_s2 [B, 256]^T → h_v [B, 128] + bias + ReLU
4. w_v2 [51, 128] @ h_v [B, 128]^T → value_logits [B, 51] + bias
5. Per-branch: w_b0 [5*51, 128] @ h_s2 via value head → branch logits

Expected occupancy: cuBLAS auto-tunes thread blocks for H100 — typically >60% for these matrix sizes at B=256+.

---

Task 2: Restructure backward pass — batched GEMM replaces atomicAdd

Replace per-warp atomicAdd gradient accumulation with batched GEMM (outer product):

// dW = dY^T @ X (batched outer product, sum over batch)
// dX = dY @ W (batched matmul for upstream gradients)
cublasSgemm(CUBLAS_OP_T, CUBLAS_OP_N, ...)

This eliminates all atomicAdd contention on the 288K gradient buffer.

Files:

• Create: crates/ml/src/cuda_pipeline/batched_backward.rs
• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

---

Task 3: Fuse BF16 conversions — single flat kernel

Replace 20 separate f32_to_bf16_kernel launches with 1 launch over the flat params_buf:

// Single kernel: convert entire flat params_buf (288K elements)
f32_to_bf16_flat<<<ceil(288321/256), 256>>>(params_buf, bf16_buf, 288321);

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — sync_online_bf16()

---

Task 4: Fuse EMA updates — single flat kernel

Replace 20 separate dqn_ema_kernel launches with 1 launch:

// Single kernel: EMA over entire flat params
dqn_ema_flat<<<ceil(288321/256), 256>>>(target_flat, online_flat, tau, 288321);

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs — target_ema_update()

---

Task 5: Multi-block prefix sum (CUB DeviceScan)

Replace single-block prefix sum with multi-block parallel scan. At N=100K:

• Current: 1 block, 100K sequential adds on 1 SM
• New: 100K/256 = 391 blocks across 132 SMs

Options:

• Use CUB's DeviceScan::InclusiveSum via cudarc FFI
• Or implement 3-phase scan: local scan → scan of block totals → scatter

Files:

• Modify: crates/ml-dqn/src/prefix_sum_kernel.cu
• Modify: crates/ml-dqn/src/gpu_replay_buffer.rs — pfx_sum() launch config

---

Task 6: C51 loss kernel — reduce register spill

The forward+loss kernel spills 46KB/thread to local memory (L1/L2 cached stack). Key culprits:

• online_log_probs[561] — 561 floats per lane
• target_log_probs[561] — same
• online_next_lp_scratch[561] — same

With cuBLAS forward (Task 1), these arrays move from per-thread registers to batch buffers in global memory. The C51 loss becomes a separate fused kernel operating on the batch output tensors.

Files:

• Create: crates/ml/src/cuda_pipeline/c51_loss_kernel.cu — dedicated C51 loss over batch outputs

---

Task 7: Profile and validate on H100

Measure:

1. Per-step time breakdown: cuBLAS GEMM vs fused kernel
2. Total epoch time with 300 steps
3. GPU utilization via nvidia-smi dmon
4. Compare with Phase 1 baseline (37.56s)

Target: <5s/epoch with 300 steps, 51 atoms, batch_size=1024

---

Expected Speedup Stack

Optimization	Current	Target	Speedup
cuBLAS forward (Task 1)	1.56% occupancy	>60% occupancy	10-30×
Batched backward (Task 2)	atomicAdd 288K	GEMM outer product	5-10×
Fuse BF16 (Task 3)	20 launches	1 launch	2×
Fuse EMA (Task 4)	20 launches	1 launch	2×
Multi-block prefix sum (Task 5)	1 SM	132 SMs	10× on PER
C51 loss kernel (Task 6)	46KB spill	batched global	5×

Net expected: 10.7ms/step → 0.3-1.0ms/step 300 steps/epoch: 3.2s → 90-300ms Full epoch (including experience + validation): <2s target