SPEED.md

# Training speed / MFU — state & open levers

Honest current numbers on **8× B300 (Blackwell Ultra, sm_103, CUDA 13 / torch cu130)**, so you
don't re-walk dead ends. MFU is measured against the **achievable bf16 matmul roofline = ~1500
TFLOP/s** (a pure `n=16384` bf16 matmul on an idle B300, power-capped ~1650 MHz/1070 W). Nominal
spec is ~2500 but the cu130 Blackwell kernels only reach ~1500 in a pure matmul, so 1500 is the
honest denominator — a real workload cannot beat it. `src/mxf/mfu.py` computes MFU (numerator counts
only real non-pad tokens, so padding lowers it honestly).

## Pretrain (`scripts/pretrain.py`) — the main training loop

| config | TFLOP/s | MFU |
|---|---|---|
| bs64, no compile | 520 | 35% |
| bs64 + `torch.compile` | 817 | 54% |
| bs96 + compile + fixed-shape (round L to mult-of-64) | 860 | **57%** ← best clean single-GPU |
| sequence packing (2048-block, compile) | 708 | 47% ← **regressed, don't use** |
| **8-GPU DDP (the actual big run): no compile** | ~540 | **~36%** |

### What's been tried / known
- **`torch.compile` fuses the 152k-vocab cross-entropy** — the single biggest win (35→54%). Do keep it.
- **Length-bucketing + rounding padded L to a multiple of 64** bounds compile to ~2–3 static shapes (otherwise it recompiles every batch — a dynamic-shape trap).
- **Sequence packing REGRESSES on this box** (57→47%): no flash-attn build for sm_103, so packed attention uses a dense 4D mask and sdpa computes the full block² score matrix, wasting ~90% of it. Correct (bit-exact no-leak proof exists) but slower. It's behind `--pack-len` (default 0 = off). It'd help *only* with a block-sparse attention backend.
- **⚠️ THE BIG OPEN BUG: `torch.compile` + DDP crashes** (single-GPU+compile ✓, 8-GPU no-compile ✓, 8-GPU+compile ✗ — all ranks, likely the layer-1 injection forward-hook's graph-break fighting DDP's reducer). So the 8-GPU run currently runs **no-compile at ~36% MFU** instead of the ~57% compile path. **Fixing this is the highest-leverage speed win** (compile-after-DDP-wrap, `static_graph=True`, `torch._dynamo` DDPOptimizer settings, or moving the injection out of a Python hook into an in-graph module so there's no graph-break).
- **No gradient checkpointing** (it silently corrupts grads: the inject-hook context exits before the checkpointed backward recompute). Don't enable it without moving injection in-graph. This caps batch at ~bs96–128.

### Untried levers (candidates for hill-climbing)
1. **Fix compile+DDP** → recovers 57% at 8-GPU (biggest single win). Try in-graph injection to kill the graph-break.
2. **FP8** (torchao float8 / TE) — Blackwell's design point, ~2× the bf16 roofline. Real code + correctness surface (LoRA + injection), and changes the roofline denominator.
3. **FlexAttention** block-sparse — would make sequence packing actually help (long fixed blocks → high arithmetic intensity, zero padding); HF-Qwen3 flex integration is the fiddly part.
4. **Target-only loss / chunked CE** — lm_head is applied to all positions incl. prompt+pad; only target positions need loss. compile already fuses most of it, but computing the head on target rows only saves memory (→ bigger batch) + FLOPs.
5. Reuse-KV for the shared 98-token prompt prefix — but injection differs per row, so prefill activations differ; not shareable as-is.

## Embed / clustering (`scripts/embed_cluster_acts.py`) — forward-only, layer-27 early-exit
Streaming out-of-core; the 8-way sharded corpus download keeps GPUs fed. Forward-only, short (64-tok)
seqs → lower arithmetic intensity than training. This is the 200M-doc prefill pass. Live MFU
measured from doc throughput; expect it to be memory/IO-influenced (short seq + mfs writes).

## Injection (`src/mxf/inject.py`) — the thing that makes compile/DDP/checkpointing hard
Norm-matched additive: `h_p += coeff·‖h_p‖·v/‖v‖` at the layer-1 marker, `.detach()`ed. It's a
`register_forward_hook` → **graph-break under compile**. Making it an in-graph module op (subclass
the layer, do the injection inside `forward` with per-row marker positions/vecs passed as tensors)
would likely unlock compile+DDP *and* gradient checkpointing at once — probably the highest-value
refactor for speed.

## Roofline reproducer
```python
import torch, time
n=16384; a=torch.randn(n,n,device="cuda",dtype=torch.bfloat16); b=a.clone()
for _ in range(50): a@b
torch.cuda.synchronize(); t=time.time()
for _ in range(50): c=a@b
torch.cuda.synchronize(); print(2*n**3*50/(time.time()-t)/1e12, "TFLOP/s")  # ~1500 on B300/cu130
```