Skip to content

Fp8 support#65

Open
bdepyzy wants to merge 15 commits into
Dao-AILab:mainfrom
bdepyzy:fp8-support
Open

Fp8 support#65
bdepyzy wants to merge 15 commits into
Dao-AILab:mainfrom
bdepyzy:fp8-support

Conversation

@bdepyzy

@bdepyzy bdepyzy commented Jun 25, 2026

Copy link
Copy Markdown

Add plumbing for per-tensor fp8 expert GEMM

Adds per-tensor float8_e4m3fn support for the MoE expert GEMMs, forward and backward, behind use_fp8=True.

What it does

  • Forward: per-tensor amamax quantization of activations and weights to fp8; weight casts cached across steps (keyed by data_ptr + _version, auto-invalidated after optimizer steps). Runs fp8 grouped GEMMs via gemm/gemm_gated, then descales with .mul_().
  • Backward: fp8 GEMMs for dx, dw1, dw2, dh, ds via gemm/gemm_dgated, descaled with .mul_().
  • dtype mapping: extends quack's torch2cute_dtype_map to cover fp8 (the CuTe kernels already support fp8 element types, but the Python dispatch layer doesn't map it by default).

fp8 is currently slower than bf16. The fp8 GEMMs themselves are ~2x faster, but per-tensor fp8 requires descaling each GEMM output by its activation×weight scale before the nonlinear GLU, and since those separate .mul_() passes are memory-bound and dominate at these shapes. This descale is fundamental to per-tensor fp8 (it cannot be folded through swiglu).

I have a working speedup of ~1.35x, but it requires folding the descale into the GEMM epilogue, something that quack doesn't currently expose. With that, the descale passes disappear and fp8's 2x compute advantage wins. This PR is the sonicmome-side plumbing; switching from .mul_() to alpha= is a one-line-per-call change once quack supports it. If this PR into sonicmoe is good, then I will open a PR into quack to make this faster.

Correctness: fp8 fwd+bwd matches bf16 within per-tensor fp8 tolerance ~2–3% mean rel error across output and gradients (test asserts rtol=0.2).

@bdepyzy bdepyzy mentioned this pull request Jun 25, 2026
@bdepyzy bdepyzy marked this pull request as draft June 25, 2026 16:15
@bdepyzy bdepyzy marked this pull request as ready for review June 25, 2026 17:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant