Skip to content

Commit ad964df

Browse files
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

0 commit comments

Comments
 (0)