Commit ad964df
committed
feat(fp8): lever r1-e2e — wire TE DelayedScaling(E4M3) fp8 GEMM into transformer-block Linears (fwd+dgrad+wgrad)
New cppmega_mlx/nn/_tilelang/fp8_te_linear.py: an @mx.custom_function fp8_linear +
per-site .vjp wrapping a persistent te.Linear(K,N) under te.fp8_autocast(
DelayedScaling(E4M3)) — the MEASURED 1.57-1.83x R1 tensorwise route (§19/§20) — crossing
MLX<->torch zero-copy via _cuda_zerocopy DLPack. Forward = y=x@W.T fp8; .vjp re-runs the
TE fp8 fwd under torch autograd to recover dgrad (x.grad) + wgrad (weight.grad), both TE
fp8 GEMMs; wgrad lands on the bf16 MASTER weight cotangent (optimizer keeps updating bf16,
fp8 cast at the GEMM only). Persistent per-site te.Linear so DelayedScaling amax history
converges. ensure_nvrtc_builtins_path() called at import (pre-torch/TE).
maybe_fp8_linear_call(nn.Linear, x) call-site gate: routes the GEMM through fp8 iff
CPPMEGA_FP8_LINEAR is on AND (M,K,N) clears the MIN_M/K/N floor (default 512) AND the
Linear is bias-free; else the BYTE-IDENTICAL bf16 nn.Linear runs (verified: gate-OFF and
below-floor are byte-identical, never a fallback). Wired at the dense transformer-block
GEMMs the research named: attention q/k/v (_project_qkv) + dense out_proj (__call__), and
MoE FeedForwardExpert gate/up/down.
RULE #1: a SELECTED fp8 GEMM that cannot run RAISES with where+what (fp8_te_linear + RULE
#1 context) — NEVER a silent fp8->bf16 degrade (verified on this Mac: selected-fp8 raises).
New scripts/bench_fp8_e2e_step_20260604.py: e2e A/B step harness (bf16 arm gate-OFF vs fp8
arm gate-ON) over the real local_gb10_quarter fwd+bwd(+AdamW), reporting median tok/s, MLX
+ torch-CUDA peak GB, final loss (loss-parity), and the fp8/bf16 speedup ratio. The GB10
phase runs it to MEASURE the realized e2e effect (not extrapolated). Disjoint from lever 1's
fp8_matmul_path_c.py and the microbench.1 parent 464cbd5 commit ad964df
4 files changed
Lines changed: 703 additions & 7 deletions
File tree
- cppmega_mlx/nn
- _tilelang
- scripts
0 commit comments