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

55c368d already decided what should happen for schemes that define no gmm quantization rule: they "execute
unquantized GMM". That change handled the qwix rule and left this read untouched, so fp8 crashes instead of taking the
fallback it was given. Reading quant_dg defensively is enough to get it there.

Expert matmuls silently running in the compute dtype while the dense layers run fp8 is easy to miss from the config
alone, so the layer logs that once when it is built, pointing at sparse_matmul=False for anyone who wants the experts
quantized too.

Not addressed here: actually quantizing the gmm for fp8/nanoo_fp8. That needs kernel support, and fp8_full
remains the scheme that quantizes the expert matmuls.

Tests

Two tests on a tiny Mixtral, sparse_matmul=True with fp8 and nanoo_fp8. No hardware marker, since the fp8
schemes are emulated in XLA and the failure is backend independent; they run on CPU in about 14s together, and both
fail before this change.

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; 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

❌ Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	75.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!