d4c9c9c34e
Based on kernel deep dive: 1.56% occupancy on H100 from 1-warp/sample architecture, 46KB stack spill, 47 kernel launches per step. 7 tasks: cuBLAS forward, batched backward, fuse BF16/EMA, multi-block prefix sum, C51 loss kernel, H100 profiling. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
179 lines
6.6 KiB
Markdown
179 lines
6.6 KiB
Markdown
# 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**
|