RMSNorm + residual + SwiGLU fused into one epilogue. Fuse the forward, then choose: store activations or recompute rrms in backward to save HBM.
HBM model: bf16 (2 B/elem) · H = round(2.6875·D) (SwiGLU up-proj, matches receipt H=2752 @ D=1024) · N = B·T tokens. Same analytic accountant as epi_norm/core.py:hbm_accountant.
| design | forward | stashed | backward | total | passes |
|---|
baseline fwd = N·(7D+2H) — 3 passes: add (2D→D) + norm (D→D) + 2 gate GEMMs (2×D reads, 2H out)
fused fwd = N·(2D+2H) — 1 pass: read x,res · write gate/up out (same for store & recompute)
stored: baseline N·(2D+2H) · fused-store N·(D+2H) · fused-recompute N·(D+1) — recompute keeps only fused input + rrms scalar
bwd: baseline N·(5D+4H) · fused-store N·(3D+4H) · fused-recompute N·(3D+2H+1) — recompute trades ~N·(D+D·H) FLOPs for the stash saving