Skip to content

Round mxfp8 block scales up to avoid saturation - #4353

Open
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:mxfp8-mx-scale-no-clip
Open

Round mxfp8 block scales up to avoid saturation#4353
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:mxfp8-mx-scale-no-clip

Conversation

@dhiltgen

Copy link
Copy Markdown
Contributor

Proposed changes

MX block scales are amax / max_element, converted to E8M0. Metal and the CPU
fallback round to nearest, which lands below the request about half the time;
the block's largest elements then fall outside the element format and saturate.

CUDA already rounds up — cutlass::float_ue8m0_t uses
cvt.rp.satfinite.ue8m0x2.f32 on device and an explicit exponent bump on host,
for both 4- and 8-bit. So mx.quantize on an MX mode gives different results
per backend today. This aligns Metal and CPU with CUDA; no CUDA code changes.

For mxfp8, an in-range element loses at most half a step, 6.25%; anything
larger is saturation. On 512x32 synthetic blocks:

worst loss on a block max blocks over 6.25%
before 29.27% 243/512
after 5.84% 0/512

29.27% is 1 - 1/sqrt(2), the worst case for round-to-nearest in log2 space.
CUTLASS v4.3.5 on the same inputs: max amax/scale 447.88, nothing over 448.

Real weights behave the same. On a Gemma 4 26B expert tensor, amax/scale
spans 318-632 against a 448 ceiling, and rel RMSE goes 7.11% -> 2.68%.

mxfp4 improves too, by less: on four Gemma 4 26B tensors, rel RMSE drops
0.8-2.5% relative and block-maximum clipping goes from about half of all blocks
to none. Round-up costs some underflow there, but it is the smaller effect.

Fixed in all three implementations: the ops.cpp fallback (the path CPU takes
for mx.quantize), fp_quantize_dequantize, and both Metal quantize kernels.

The test asserts the 6.25% property rather than a golden number and pins no
device. It fails on main and on the 0.31.0 wheel, and passes here.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@zcbenz

zcbenz commented Aug 19, 2026

Copy link
Copy Markdown
Member

Can you rebase on main? Your branch seems to be ancient.

@dhiltgen
dhiltgen force-pushed the mxfp8-mx-scale-no-clip branch from 51e8524 to e73e87d Compare August 19, 2026 20:26
@dhiltgen

Copy link
Copy Markdown
Contributor Author

Oops, cherry-picked on the wrong base, fixed.

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.

2 participants