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. Looking up a name that was never registered raises a ValueError listing the ones that were, rather than a bare KeyError.

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.

Unit tests pin the binding itself rather than only the training run that exercises it:

• MoEQuantizedEinsumTest in quantizations_test.py covers create_fp8_einsum for both fp8 flavors, the fnuz dtypes NANOO asks for, and both branches of apply_einsum_in_nnx — the early return for a plain callable, and the AQT bridge, which has to reuse the wrapper it built rather than rebuild one per call.
• GetEinsumTest in moe_test.py covers the fp8 and AQT paths through get_einsum, plus the unregistered-name error.
• Correctness for both fp8 schemes, also in moe_test.py: on a freshly built layer the scaling factors are still 1 and only move on the backward pass, so the bridged einsum is pinned exactly against the plain einsum with both operands cast to the scheme's e4m3 (float8_e4m3fn for fp8, float8_e4m3fnuz for nanoo_fp8). At layer level a quantized MoE is compared against the same layer with quant=None and identical weights, within the 0.22 relative-norm threshold the qwix MoE test already uses, and asserted not to match it exactly, which is what would catch a silently unquantized fallback.

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.

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[gulsumgudukbay]

The sparse_matmul gap noted at the end of the description is now covered by #4957, so the two together fix quantized MoE on both paths. They are independent apart from a small textual overlap in tests/integration/train_tests.py, which I will rebase for whichever lands second.

[Shuwen-Fang]

looks good, thanks for catching this!

[Shuwen-Fang]

can you also add a correctness test for the fp8/nanoo_fp8 path in this file?

[gulsumgudukbay]

Added in 6cf99fc, two per scheme.

test_fp8_einsum_quantizes_both_operands pins the numerics exactly rather than with a tolerance. On a freshly built layer the scaling factors are still 1 and only move on the backward pass, so the forward pass is the plain einsum with both operands cast to the scheme's e4m3, float8_e4m3fn for fp8 and float8_e4m3fnuz for nanoo_fp8, and it also asserts the result differs from the unquantized einsum.

test_quantized_dense_matmul_tracks_unquantized runs a full MoE layer against the same layer with quant=None and identical weights. Three mantissa bits across five einsums makes that agreement loose, so it uses the 0.22 relative-norm threshold the qwix MoE test in this file already uses, plus a check that the two do not agree exactly, which is what would catch a silently unquantized fallback.

GetEinsumTest now builds a small MoE instead of a full mixtral-8x7b so the forward passes stay cheap; the class runs in about 18s on CPU.

[gobbleturk]

Note we are deprecating AQT, we will rely on qwix for quantization moving forward