Fix fp8 MoE on the sparse_matmul path

[gulsumgudukbay]

Problem

Any MoE model with sparse_matmul=True and an fp8 quantization crashes while the layer is being traced:

AttributeError: 'Fp8Quantization' object has no attribute 'quant_dg'
AttributeError: 'NANOOFp8Quantization' object has no attribute 'quant_dg'

get_quantization_dtypes reads self.quant.quant_dg whenever a quantization is configured, but quant_dg is an AQT
notion — the fp8 classes have never had one. Since sparse_matmul=True is the default, this is what a user hits first
when combining fp8 with a MoE model.

Fix

quant_dg is how AQT spells its numerics, not what the gmm needs. All the hand-crafted rule wants is a forward dtype,
which the fp8 schemes can simply name, so they do: float8_e4m3fn for fp8 and float8_e4m3fnuz for nanoo_fp8. The
read becomes "AQT describes itself through a quant_dg, everyone else names a dtype", and the gmm ends up quantized for
every scheme rather than only for AQT.

This keeps the three intended behaviours intact, per @shuningjin in the thread below:

1. use_qwix_quantization=false: every scheme quantizes the gmm through the hand-crafted rule. That now includes the
   fp8 schemes, which previously crashed here.
2. use_qwix_quantization=true with fp8_full: the gmm takes the qwix rule, untouched by this change.
3. use_qwix_quantization=true with any other scheme: configure_quantization returns None, so there is no dtype to
   hand over and the gmm is left unquantized, which is what a scheme declaring no gmm rule should get.

An earlier revision of this PR read quant_dg defensively, which stopped the crash but left fp8 with an unquantized
gmm and a warning saying so. That was the wrong side of case 1, so both are gone.

Tests

SparseMatmulQuantizationTest in moe_test.py asserts the lhs_quantize_dtype and rhs_quantize_dtype the grouped
matmul is actually handed, across all six cases above: fp8, nanoo_fp8, int8, unquantized, and fp8/int8 under
qwix, the last two of which must stay None. It intercepts mblx.gmm, so it costs no kernel time. The two fp8 cases
fail without this change.

Three integration tests on a tiny Mixtral with sparse_matmul=True cover fp8, nanoo_fp8 and int8 end to end, and
they run megablox rather than ragged_dot so the rule is built and used for real — which is also what confirms qwix
accepts the fnuz dtype AMD needs. A downscaled model needs tiles small enough to divide its operands, hence the tile
overrides. The two fp8 ones fail before this change, and the three take about 33s on CPU. No hardware marker: the fp8
schemes are emulated in XLA and the failure is backend independent.

Related

• Fix quantized MoE on the dense_matmul path #4955 fixes the same combination on the dense_matmul path, where every quantization fails for an unrelated reason
  (Linen einsums with no scope to bind to after the NNX migration). The two are independent, but together they make
  quantized MoE work on both paths. There is a small textual overlap in tests/integration/train_tests.py and
  tests/unit/moe_test.py; happy to rebase whichever lands second.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces safety checks and integration tests for Mixture of Experts (MoE) models using the sparse_matmul path with FP8 quantization. It adds a warning log when FP8 quantization is combined with sparse_matmul because the expert matmuls remain unquantized on this path. It also updates get_quantization_dtypes to safely access quant_dg using getattr to prevent errors when using FP8 schemes. Finally, it adds integration tests to ensure that the MoE layer builds and trains successfully under these configurations. There are no review comments, so I have no feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[gulsumgudukbay]

Codecov flagged one partial on the patch. The two fp8 tests only ever take the branch where there is no quant_dg, so the read that does find one, and the isinstance check that skips the warning, were never exercised from this PR's own tests. Added test_moe_int8_sparse_matmul in a46a1ad: int8 is the scheme whose dtypes actually reach the gmm, so it covers both arcs and guards the behaviour this change had to preserve. Measured locally with coverage --branch, the changed lines now have no missing lines or branches; the three tests run in about 28s on CPU.

[shuningjin]

Hi @gulsumgudukbay, Here is the intended behavior:

1. when use_qwix_quantization=false, all quantization scheme use quantized gmm with hand-crafted rule
2. when use_qwix_quantization=true, for fp8_full, it will use the quantized gmm with qwix rule
3. when use_qwix_quantization=true, for non-fp8_full (e.g., fp8, int8), the gmm path should be an no-op as it does not specify gmm qwix rule here.

I accidentally changed Behavior 3: #4451 (assert against empty gmm qwix rule) and #4842 (allow non-fp8-full with qwix to use quantized gmm with hand-crafted rule). We should restore to previous state (non-fp8-full with qwix does NOT use quantized gmm). Sorry for the confusion!

[gulsumgudukbay]

Hi @gulsumgudukbay, Here is the intended behavior:

1. when use_qwix_quantization=false, all quantization scheme use quantized gmm with hand-crafted rule
2. when use_qwix_quantization=true, for fp8_full, it will use the quantized gmm with qwix rule
3. when use_qwix_quantization=true, for non-fp8_full (e.g., fp8, int8), the gmm path should be an no-op as it does not specify gmm qwix rule here.

I accidentally changed Behavior 3: #4451 (assert against empty gmm qwix rule) and #4842 (allow non-fp8-full with qwix to use quantized gmm with hand-crafted rule). We should restore to previous state (non-fp8-full with qwix does NOT use quantized gmm). Sorry for the confusion!

Thanks @shuningjin, that's very helpful. I had this on the wrong side of behavior 1. just pushed 5bcaf7f, which changes the fix.

quant_dg is how AQT spells its numerics, not what the gmm needs. The hand-crafted rule only wants a forward dtype, which the fp8 schemes can name directly, so they now do: float8_e4m3fn for fp8, float8_e4m3fnuz for nanoo_fp8. now it becomes becomes "AQT describes itself through a quant_dg, everyone else names a dtype", so with use_qwix_quantization=false the gmm is quantized for every scheme, fp8 included, where it used to crash.

fp8_full under qwix is untouched, and any other scheme under qwix has no quantization object at all, so there is no dtype to pass and the gmm stays unquantized. My previous revision stopped the crash but left fp8 with an unquantized gmm and a warning saying so; both of these are gone now.

On restoring behavior 3, it seems to me like the gmm dtypes already satisfy it, since configure_quantization returns None whenever use_qwix_quantization=true. If what you have in mind for #4451 and #4842 goes further, I'll leave it to you

SparseMatmulQuantizationTest in moe_test.py intercepts mblx.gmm and asserts the dtypes it is handed for all six cases, including the two qwix ones that must stay None; the fp8 cases fail without the change. The integration tests now run megablox instead of ragged_dot so the rule is really built, which is also what confirms qwix accepts the fnuz dtype AMD needs.

[gobbleturk]

I'm surprised this is true - isn't train_main being invoked which will run on the backend?

[gulsumgudukbay]

sorry you are right, that comment was stale it predates 5bcaf7f, where fp8 stopped falling past the gmm quantization and started feeding it a dtype like every other scheme. i fixed this in 664b181, along with the int8 one next to it.

regarding the marker: conftest auto-marks unmarked tests cpu_only and skips those on accelerator testbeds, so these only run in the CPU job, which is fine/enough because the failure is an attribute read during tracing, and fp8 is emulated in XLA, so the quantized gmm still runs under megablox's interpret mode.

[gobbleturk]

typo conftest?

[gulsumgudukbay]

no, it is not a typo, conftest.py is the pytest file that auto-marks unmarked tests cpu_only. i just reworded now for it to say so properly, since "conftest runs these" was misleading anyway. it marks them, and the CPU job is what runs them sorry for the confusion.