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

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

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