| """Honest MFU meter. Numerator counts only REAL (non-pad) tokens → padding lowers MFU, as it should. | |
| MFU = achieved_model_FLOPs / (wall_time * ROOFLINE). ROOFLINE is the *measured* B300 bf16 matmul | |
| ceiling (~1500 TFLOP/s at the power-capped 1650 MHz with torch-cu130), not a marketing spec — a | |
| pure matmul cannot beat it, so this is the honest denominator for "how close to the metal are we". | |
| """ | |
| ROOFLINE_TFLOPS = 1500.0 # measured: `n=16384 bf16 a@b` on an idle B300 (override via env if kernels improve) | |
| def flops_per_token(n_params: int, fwd_bwd: bool, layer_frac: float = 1.0) -> float: | |
| """6ND (fwd+bwd) or 2ND (fwd only). layer_frac<1 for early-exit forwards (cache_resids @ L27).""" | |
| return (6.0 if fwd_bwd else 2.0) * n_params * layer_frac | |
| def mfu(n_real_tokens: int, seconds: float, n_params: int, fwd_bwd: bool = True, | |
| layer_frac: float = 1.0, roofline_tflops: float = ROOFLINE_TFLOPS) -> tuple[float, float]: | |
| """Returns (achieved_TFLOP_s, mfu_fraction) vs the measured roofline.""" | |
| tflops = flops_per_token(n_params, fwd_bwd, layer_frac) * n_real_tokens / seconds / 1e12 | |
| return tflops, tflops / roofline_tflops | |