Share one RNG stream across the bridged NVFP4 quantizers

[ecnal-cienet]

Description

PR #5027 stopped the Linen to NNX bridge keeping a forked Rngs in every wrapper, which the unrolled decoder pays for once per layer. NVFP4 was left out: TransformerEngine draws sr_rng for the DGRAD quantizer, so dropping the fork there raised InvalidRngError. It still carries the whole cost, 1351 u32 entry parameters on llama3-8b te_nvfp4_no_rht unrolled against 7 for te_fp8_currentscaling.

TE folds a per-quantizer hash into whatever it draws, so the wrappers never needed streams of their own, and NNX stores a shared Rngs once however many modules reference it.

• needs_apply_rngs becomes ApplyRngs. A boolean could only say whether a backend draws. The enum separates one that never draws, NONE, from one that draws but does not need its own stream, SHARED. PRIVATE stays the default, and AQT and Qwix stay on it.
• ToNNX.share_rngs swaps the fork for the caller's Rngs. __init__ still forks first, so the caller's streams advance as before and parameter initialization is unchanged.

Tests

llama3-8b unrolled, 8xH100, entry parameters from --xla_dump_to. NVFP4 kernels need Blackwell, but compilation succeeds on H100.

	entry params	of which u32
main	2231	1351
this PR	887	7

887 / 7 is what te_fp8_currentscaling already gets, so nothing scales with the layer count. Throughput for NVFP4 itself needs GB300.

te_fp8_currentscaling forced onto the shared path, to exercise it where it can run: 21 steps, loss 12.258 to 1.977 to 0.153, median 0.6570s, its usual numbers.

scan_layers=False is expected to be the faster of the two. Both configs of a recipe, back to back on one node:

recipe	scan on	scan off
te_fp8_currentscaling	0.7360	0.6600	11.5% faster
te_fp8_delayedscaling	0.7300	0.6960	4.9% faster
unquantized	0.8750	0.8350	4.8% faster
fp8	0.6960	0.7050	1.3% slower

fp8 stayed on that side over two more repeats, so removing the per-layer state does not on its own reorder every recipe. te_mxfp8 fails on H100 in cublas_gemm, on main as well.

Nine configurations against main across llama3-8b, llama2-7b, gemma2-2b, deepseek2-16b and gpt-oss-20b, covering each enum value plus the moe.py call site and the attention-sinks path: every loss lands inside main's own run-to-run range, which is 0.002 wide when one configuration is repeated on one image.

JAX_PLATFORMS=cpu pytest tests/unit/nnx_quant_bridge_rng_test.py

SharedRngStateTest pins what this is for: a shared backend costs the same at 1, 2 or 8 layers, with a private backend alongside so it cannot pass by measuring nothing.

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 adds Gemma2-2b tokenizer configurations and refactors RNG handling in quantization backends by replacing the boolean needs_apply_rngs with an enum ApplyRngs (NONE, SHARED, PRIVATE). This change allows sharing RNG streams across unrolled layers to optimize state usage. Feedback on the changes suggests updating the type signature of the new share_rngs method in nnx_wrappers.py to accept Rngs | jax.Array | None to prevent potential static type-checking failures.

[gemini-code-assist]

The share_rngs method is annotated to accept only Rngs. However, ToNNX.__init__ accepts Rngs | jax.Array | None. To prevent static type-checking failures (e.g., with pytype or mypy) when rngs is None or a jax.Array, the type signature of share_rngs should be updated to match.

Suggested change

	def share_rngs(self, rngs: Rngs):
	def share_rngs(self, rngs: Rngs | jax.Array | None):

[codecov]

Codecov Report

❌ Patch coverage is 80.95238% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	33.33%	1 Missing and 1 partial ⚠️
src/maxtext/layers/quantizations.py	84.61%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!