Fix NNX throughput regression with scan_layers=False

[ecnal-cienet]

Description

NNX modules that hold an nnx.Rngs put it in the model state: a (key, count) pair per stream, which the host marshals into the jitted step on every call. Scanned, the decoder body is traced once and it is paid once. Unrolled, it is paid per layer. On llama3-8b with te_fp8_currentscaling that took the entry parameter count from 881 to 1277 and cost ~8% step time. All 396 extra parameters are RNG state.

Three places hold RNGs they never draw from:

• The quantization bridge. ToNNX forks the caller's Rngs and keeps it, one wrapper per quantized DenseGeneral, so llama3-8b ends up with 672 streams and 673 extra scalar counter kernels per step. Kernel time is unchanged, but GPU utilization falls from 86.0% to 76.7% as those launches serialize into the dependency chain. Adds Quantization.needs_apply_rngs, default True so a backend only opts out deliberately, and ToNNX.release_rngs() to drop the fork after init. AQT keeps its RNGs: its config sets rng_type="jax.uniform", which does draw at apply time.
• Dropout. It forked even at rate 0, where nnx.Dropout.__call__ returns its input before touching self.rngs. It now keeps the fork only when the rate can draw. It still forks either way, because fork() advances the caller's streams and skipping it would change parameter initialization.
• AttentionOp. It held rngs only for the cudnn_flash_te bridge, which draws only with attention dropout on. It now holds them only in that case.

Also drops the lazy_init in cudnn_flash_attention, which ran a full max_target_length dummy forward pass on every call, except on the attention-sinks path. That is the one case where TE declares a variable, softmax_offset, which _inject_te_softmax_offset then grafts sinks into. Sinks therefore still lazy_init, and get a call-local Rngs when AttentionOp holds none, since the drawn value is overwritten by the graft and the wrapper never reaches the model state. Everywhere else TE declares nothing and the output is bit-identical without priming.

The moe.py hunk is unreachable today, since MoE with TransformerEngine fails earlier on 'TransformerEngineQuantization' object has no attribute 'quant_dg'. It is there so GateLogit is not left behind when that is fixed.

NVFP4 needs its RNGs. Opting the whole TransformerEngine backend out was too broad, and te_nvfp4 / te_nvfp4_no_rht died before the first step with InvalidRngError: None needs PRNG for "sr_rng". NVFP4ScalingQuantizeConfig calls make_rng("sr_rng") for the DGRAD quantizer unless stochastic rounding is disabled, and it defaults to on. With the fork released the Linen apply has no RNG collection at all, so not even Flax's fallback to params is available. needs_apply_rngs is now derived from the recipe on that backend.

For review

• Reference HLO is regenerated. The Dropout and AttentionOp fixes apply to every NNX model. After normalizing instruction ids the only difference is the dropped u32 RNG arrays (30 llama3_8b, 54 deepseek3, 18 qwen3_1.7b) and the renumbering that follows. No operation, shape or layout changes.
• Existing checkpoints still resume. split_for_checkpoint routes RNG state to nnx_aux, so an older checkpoint carries entries the model no longer holds and nnx.replace_by_pure_dict raises on them. train_state_nnx.apply_checkpoint_aux now skips those, logging what it skips. Only nnx_aux is filtered, so a genuinely missing weight still raises.
• NVFP4 throughput needs GB300. On H100 the RNG failure is gone and the step compiles and launches, then dies in create_2D_tensor_map because the kernels need Blackwell. origin/main dies there identically.

FIXES: b/552606153

Testing

8xH100-80GB via NVIDIA's test-maxtext.sh, cudnn_flash_te, 21 steps, -b 1 --fsdp=8. Two images from the same tree differing only by this PR, run concurrently on separate nodes. Median step time over the last 15 steps:

model / quantization	run	before	after
llama3-8b te_fp8_currentscaling	enable_nnx=false	0.6530	0.6580
llama3-8b te_fp8_currentscaling	NNX	0.7050 (+8.0%)	0.6620 (+0.6%)
llama3-8b fp8	enable_nnx=false	0.6940	0.6950
llama3-8b fp8	NNX	0.7640 (+10.1%)	0.7050 (+1.4%)

The fp8 gap is the scan-on versus scan-off difference NVIDIA reported: unrolled NNX was paying the per-layer cost that scanned NNX pays once.

Logs from the first measurement, on an earlier base: before Linen, before NNX, after Linen, after NNX.

gpt-oss-20b with attention sinks, scanned, runs on both: 21 steps, loss 6.310 before and 6.318 after. Unrolled it OOMs on both, at 43.9 GiB. te_nvfp4_no_rht no longer fails on sr_rng.

JAX_PLATFORMS=cpu pytest tests/unit/nnx_quant_bridge_rng_test.py tests/unit/attention_test.py

Each case fails if the matching fix is reverted. The bridge tests use stub backends so they run without TE, and cover the TE recipe table when it is present. The attention cases pin both directions of the lazy_init split: skipped without sinks, and still run with them, with an Rngs to initialize from.

JAX_PLATFORMS=cpu pytest tests/unit matches main: same pre-existing environment failures, none new. Losses unchanged on llama2-7b (te_fp8_currentscaling, int8, unquantized), gemma2-2b, and mixtral-8x7b / qwen3-30b-a3b for the moe.py hunk.

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 addresses a throughput regression (b/552606153) by optimizing the NNX bridge wrapper. It removes a redundant lazy_init call on the TransformerEngine DotProductAttention bridge, which previously caused an extra forward trace per layer. Additionally, it introduces a release_rngs mechanism to drop forked Rngs after initialization for quantization backends that do not require RNGs at apply time (such as TransformerEngineQuantization), preventing unnecessary RNG state from being carried and incremented on device. Comprehensive unit tests have been added to verify these optimizations and guard against future regressions. I have no further feedback to provide as there are no review comments.

[codecov]

Codecov Report

❌ Patch coverage is 83.67347% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/quantizations.py	66.66%	4 Missing ⚠️
src/maxtext/layers/attention_op.py	85.71%	2 Missing ⚠️
src/maxtext/layers/linears.py	80.00%	0 Missing and 1 partial ⚠️
...trainers/diloco/utils/spmd_diloco_checkpointing.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!