Fix MX block quantization: row-local blocking, OCP shared exponent, element rounding - #321
Fix MX block quantization: row-local blocking, OCP shared exponent, element rounding#321Shreyas8612 wants to merge 5 commits into
Conversation
flatten().reshape(-1, B) let a block span two logical rows whenever the block axis was not a multiple of the block size, so the tail of one attention head shared a scale with the start of the next. Permuting the block dimension last does not prevent this; only per-row padding does. Rows are now padded independently to a multiple of the block size and the padding is stripped on the way back. Reproduction: quantise a (4, 80) tensor with block_size 32, perturb row 0 by 100x, and row 1 changes on the old path but is bit-identical on the new one. Also ties element_is_finite to the exponent width, since a 1-bit exponent has no code left over for inf or nan.
MXFP computed the shared scale as ceil(log2(block_max)), which places the block maximum above the element format's exponent range so small elements normalise into the subnormal region and flush to zero. OCP Microscaling specifies floor(log2(block_max)) - emax_element. MXINT used ceil(log2(max)), leaving the block maximum in (0.5, 1] and clipping it by one LSB against the symmetric sign-magnitude range; it now uses ceil(log2(max/qmax)). Verified by idempotency, since a KV cache entry is re-read many times after a quantisation round trip: quantising an already-quantised tensor must be a no-op. Every element format fails this before the change and passes after, with relative error roughly halving. Adds zero-block protection, the true E8M0 signed range, canonical sign-magnitude encoding for negative zero, and the 8-bit E3M4 format. Also replaces quantile at percentile 1.0 with amax, which is mathematically identical there, avoids sorting every block, and sidesteps quantile's input-size limit.
Element quantisation round-tripped every value through a uint16 encode and decode. Fake quantisation needs the rounded value, not the bit pattern, and the encoder never emits an inf or nan exponent code, so the decoder's special-value branches were unreachable. Also corrects two format boundaries that the shared-exponent fix depends on: a 1-bit exponent has no room for the standard bias formula and uses a bias of 1, and the smallest normal biased exponent is 1 - bias rather than -bias, so the subnormal boundary was off by one exponent step. Without these, MXFP re-quantisation is not idempotent at any element format.
There was a problem hiding this comment.
Pull request overview
This PR fixes correctness issues in the MX-format quantizers (MXFP/MXINT) by making block formation row-local (preventing cross-row/head scale sharing), updating shared-exponent computation to match the OCP microscaling formulas, and simplifying/improving element rounding for idempotent re-quantization.
Changes:
- Replaced global flatten/blocking with row-local tail padding via shared shape helpers (
block_rows_for_quantize/restore_quantized_rows). - Updated MXFP shared exponent selection to the OCP-conformant
floor(log2(max)) - emax_elementbehavior and expanded supported element formats (incl. E3M4). - Refactored minifloat quantization to round directly to the value grid (skipping uint16 encode/decode) and clarified
is_finitesemantics.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/chop/nn/quantizers/mxint/mxint.py | Switches MXINT sim path to row-local blocking/restoration and updates shared-exponent logic. |
| src/chop/nn/quantizers/mxint/fake.py | Updates MXINT component extraction (amax for p=1.0, new shared exponent helper) and adds sign-magnitude encoding helper. |
| src/chop/nn/quantizers/mxfp/mxfp.py | Uses row-local blocking/restoration in MXFP sim path and threads padding metadata through quantile search. |
| src/chop/nn/quantizers/mxfp/meta.py | Expands legal element formats and adds element_max_exponent to support OCP shared-exponent computation. |
| src/chop/nn/quantizers/mxfp/helpers.py | Introduces row-local block padding/restoration helpers shared across MX quantizers. |
| src/chop/nn/quantizers/mxfp/fake.py | Implements OCP-style shared exponent selection and replaces quantile(1.0) with amax. |
| src/chop/nn/quantizers/_minifloat_mx/minifloat.py | Simplifies sim path to use direct value-quantization rather than encode/decode. |
| src/chop/nn/quantizers/_minifloat_mx/meta.py | Clarifies is_finite docstring semantics. |
| src/chop/nn/quantizers/_minifloat_mx/fake.py | Refactors minifloat rounding/field extraction and adds quantize_minifloat_value. |
| src/chop/nn/quantizers/_minifloat_mx/init.py | Exposes quantize_minifloat_value from the internal minifloat module. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| @@ -71,18 +68,23 @@ def mxint_quantizer_sim( | |||
| .to(tem_dtype) | |||
There was a problem hiding this comment.
Same reasoning as the fake.py comment below. x_max is derived from x, which is already in tem_dtype, so the per-block maximum is by construction exactly representable in that dtype and the cast back is lossless. Underflow would need the max of a set of dtype-D values to be smaller than the smallest representable dtype-D value.
Tested at float16, bfloat16 and float32 with block maxima down to 2.4e-07 — below the float16 normal minimum — with no block spuriously flushed to zero. Covered by test_small_block_maxima_are_not_flushed_to_zero in test/nn/quantizers/test_mx_block_quantization.py
| x_max = ( | ||
| x.abs() | ||
| .to(torch.float32) | ||
| .quantile(percentile, dim=1, keepdim=True) | ||
| .to(ori_dtype) | ||
| ) | ||
| magnitude.amax(dim=1, keepdim=True) | ||
| if percentile == 1.0 | ||
| else magnitude.quantile(percentile, dim=1, keepdim=True) | ||
| ).to(ori_dtype) |
There was a problem hiding this comment.
I checked this empirically and it doesn't reproduce. magnitude is derived from x, which is already in ori_dtype, so the per-block max is by construction exactly representable in that dtype and the cast back is lossless. Underflow would require the max of a set of dtype-D values to be smaller than the smallest representable dtype-D value.
Tested at float16, bfloat16 and float32 with block maxima down to 2.4e-07 — below the float16 normal minimum — with no block spuriously flushed to zero. For percentile < 1.0 the interpolated quantile lies between two block values and is bounded below by the block minimum, so it has the same property.
This is now covered by test_small_block_maxima_are_not_flushed_to_zero in test/nn/quantizers/test_mx_block_quantization.py
|
I checked this empirically and it doesn't reproduce. magnitude is derived from x, which is already in ori_dtype, so the per-block max is by construction exactly representable in that dtype and the cast back is lossless. Underflow would require the max of a set of dtype-D values to be smaller than the smallest representable dtype-D value. Tested with float16, bfloat16 and float32 at block maxima down to 2.4e-07 (below the float16 normal minimum): no block is spuriously flushed to zero in any case. For percentile < 1.0 the interpolated quantile lies between two block values and so is bounded below by the block minimum, which has the same property. |
Covers the three defects this branch fixes: blocks straddling row boundaries when the block axis is not a multiple of the block size, non-idempotent re-quantization from the shared exponent, and small block maxima being flushed to zero at reduced precision. None of these paths had test coverage, which is why the defects survived.
Formats the files this branch touches with the repository's declared formatter. mxint.py also picks up a small number of pre-existing reformats in code this branch does not otherwise modify.
|
The failing python-format check is a CI environment problem, not a problem with this branch. The step aborts with: /usr/bin/python3: No module named black I've run black over the files this branch touches so the check passes once the runner is fixed. Note src/chop/nn/quantizers/mxint/mxint.py was already non-compliant on main, so that commit also reformats a few pre-existing lines there. |
|
Following up with the precise mechanism, since this affects every PR and not just this one. python-format is a step inside the software-test job, not a separate check. It resolves black from the container image deepwok/mase-docker-cpu:latest — there's no pip install step before it — and the image no longer provides it: /usr/bin/python3: No module named black Nothing in a PR can fix this — it needs black restored to the image, or a pip install black step added before line 87. Happy to open a separate PR for the latter if that's useful. |
Fixes correctness bugs in the MX-format quantizers (MXFP and MXINT) so that blocks no longer share scale factors across attention heads, the shared exponent follows the OCP Microscaling spec, and re-quantization becomes idempotent.
What was wrong
ceil(log2(block_max)), pushing small elements into the subnormal range and flushing them to zero. MXINT usedceil(log2(max)), clipping the block max by one bit. Both now follow the OCP-conformant formulas (floor(log2(block_max)) - emax_elementfor MXFP,ceil(log2(max/qmax))for MXINT).Other fixes
element_is_finitenow tied to exponent width (a 1-bit exponent has no code for inf/nan).quantileat percentile 1.0 withamax(mathematically identical, faster, no size limit).Verification
Idempotency (quantise an already-quantised tensor; must be a no-op)
MXINT idempotency: MXINT2, MXINT4 and MXINT8 all fail on main and pass on this branch
Row-straddle repro: quantise a (4, 80) tensor with block_size=32, perturb row 0 by 100x. Before: row 1 changes. After: rows 1-3 bit-identical.