Fix quantized MoE on the dense_matmul path

[gulsumgudukbay]

Problem

RoutedMoE.get_einsum constructs a Linen einsum and calls it inline. That worked while the MoE layer was itself a Linen module, but since the move to NNX there is no Linen scope for it to bind to, so every quantization fails on the dense_matmul path. Tiny Mixtral, one step, CPU:

quantization	sparse_matmul=True	sparse_matmul=False
none	passes	passes
int8	passes	CallCompactUnboundModuleError
fp8	AttributeError: 'Fp8Quantization' object has no attribute 'quant_dg'	CallCompactUnboundModuleError
nanoo_fp8	AttributeError: 'NANOOFp8Quantization' object has no attribute 'quant_dg'	TypeError: Quantization.einsum() got an unexpected keyword argument 'mesh_axes'

This PR fixes the sparse_matmul=False column. The quant_dg failures in the other column are a separate bug (get_quantization_dtypes reads an attribute the fp8 classes do not define) and are left alone here.

The dense path is what runs wherever the megablox and ragged kernels are unavailable, which is why this has gone unnoticed: TPU and NVIDIA runs take the sparse path.

Fix

Bridge the einsums into NNX rather than calling them unbound.

• An fp8 einsum keeps its scaling factors and amax histories in Linen variables, so it is bridged when the parent module is built rather than on the first call. Creating that state during __call__ would grow the module graph inside the scanned layer loop, which NNX rejects, and would allocate it under a trace. The state has a fixed shape (scales (1,), amax histories (1024,)), so a canonical operand pair materializes it and the bridged einsum still accepts operands of any shape.
• AQT is bridged on first use instead, since its state is shaped after the operands.
• Each call site passes a stable einsum_name, so no two share quantization state.

NANOOFp8Quantization additionally lost its einsum and its place in the isinstance check during the same migration. Both are restored, which also gives the now-unreferenced Fp8Einsum class in quantizations.py its purpose back. Both are still present on release/v26.3.

Tests

Four tests in train_tests.py, all of which fail before this change:

• test_moe_int8, test_moe_fp8, test_moe_nanoo_fp8
• test_moe_fp8_token_dropping, which adds the dispatch and combine einsums via capacity_factor > 0

They deliberately carry no hardware marker. This is a binding bug that breaks identically on every backend, both fp8 flavors are emulated in XLA rather than needing hardware support, and the whole set runs on CPU in about 30 seconds.

Verification

Tiny Mixtral and tiny Gemma 4 26B, one step on CPU, all training after the change: int8 / fp8 / nanoo_fp8, with and without token dropping, scanned and unrolled layers, and under both the pure-NNX and the Linen decoder. pylint reports no new findings and pyink is clean.

[gemini-code-assist]

Code Review

This pull request introduces support for quantized einsums in Mixture of Experts (MoE) layers, bridging Linen-based FP8 and NANOO FP8 einsums into NNX. It adds helper functions create_fp8_einsum and apply_einsum_in_nnx to manage quantization state, updates the MoE layer to utilize these quantized einsums during dense matrix multiplications, and adds integration tests. Feedback was provided to add a defensive check in get_einsum to prevent a cryptic KeyError if an unregistered einsum_name is accessed in self.quant_einsums.

[gemini-code-assist]

If get_einsum is called with an unexpected or default einsum_name (which defaults to None, resulting in op_id = "einsum"), looking up op_id in self.quant_einsums will raise a cryptic KeyError since "einsum" is not registered in quant_einsums. Adding a defensive check with a clear error message will make debugging much easier if this method is called with an unregistered name.

Suggested change

	if self.quant_einsums is not None:
	return self.quant_einsums[op_id](*args, mutable=["_overwrite_with_gradient"])
	if self.quant_einsums is not None:
	if op_id not in self.quant_einsums:
	raise ValueError(
	f"Einsum name '{op_id}' is not registered in quant_einsums. "
	f"Available names: {list(self.quant_einsums.keys())}"
	)
	return self.quant_einsums[op_id](*args, mutable=["_overwrite_with_gradient"])

References

1. Ensure appropriate checks or guards exist before accessing dictionary keys to handle invalid inputs or states safely.

[codecov]

Codecov Report

❌ Patch coverage is 92.50000% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/quantizations.py	83.33%	1 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!