perf(qwen3-next): nested scan over hybrid attention with per-layer remat

[NuojCheng]

Description

Qwen3-Next has a hybrid decoder stack: each period of inhomogeneous_layer_cycle_interval (4) layers is three linear-attention (GatedDeltaNet) layers followed by one full-attention layer. Qwen3NextScannableBlock instantiated those four heterogeneous sub-layers flat (layer_0 … layer_3) and ran them in a Python loop, so the whole block was rematerialized as a single unit and all four sub-layers' activations were live at once.

This PR restructures the block along the same lines as Gemma4ScannableBlock:

• the three linear-attention layers are homogeneous, so they are stacked and run through nnx_scan.apply_scanned_layers;
• the single full-attention layer runs inside a trip-count-one jax.lax.scan, which acts as an XLA scheduling barrier (skip-simplify-while-loops_trip-count-one);
• each sub-layer is rematerialized on its own (apply_internal_remat=True), so the outer apply passes skip_block_remat=True and nothing is rematerialized twice.

A second, independent fix in nnx_scan.apply_scanned_layers: the scan body returned nnx.state(current_layer) — the full state, parameters included. jax.lax.scan stacks every per-iteration output, so returning parameters made XLA materialize a second copy of the stacked layer weights on every call. The body now returns only the state that was not fed in as a scan input.

Dropping every nnx.Param would be too blunt. Qwix materializes LoRA adapters while tracing the scan body, and nnx.LoRAParam is an nnx.Param subclass, so they would go out with the carried weights and LoRA setup would fail with LoRA module path matched target modules, but nnx.LoRAParam is still missing. Diffing against the parameter paths that went in keeps the memory saving and still lets anything the body creates out.

External (vLLM) KV caches arrive as a flat list with one entry per decoder layer, but the NNX scan runs over blocks, so _apply_qwen3_next_scanned_blocks regroups them into per-block tuples and writes them back afterwards — the same prepare_kv_caches_for_scan / update_kv_caches_after_scan pair Gemma4 uses. That keeps skip_block_remat=True on the KV-cache path instead of falling back to the generic, block-rematerialized branch. Qwen3NextScannableBlock now also threads one cache per sub-layer and returns the updated ones; previously the block handed the same cache object to all four sub-layers and dropped the updates. The Linen Decoder already did the equivalent regrouping and is unchanged in this respect.

Layer counts that are not a multiple of the cycle are handled on both decoders. The pure-NNX decoder built num_decoder_layers // cycle blocks and stopped there, silently dropping the trailing layers; the Linen Decoder did keep them, but spelled them out as individual Qwen3NextDecoderLayerToLinen layers, which named their parameters differently from anything the NNX side produces. Both now put the leftovers into one short Qwen3NextScannableBlock named layers_remainder, so no layer is dropped and a single checkpoint mapping still covers both trees. Gemma4 and Qwen3-Next share the code that applies it (_apply_remainder_block) rather than carrying two copies of the same split / jax.checkpoint / merge dance.

Checkpoint conversion

The new layout needs matching conversion support, so that HF Qwen3-Next checkpoints still load into MaxText and still save back out. This follows #4530, which did the same for Gemma4.

The scanned branch of QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING, and its hook map, now describe the nested layout instead of the old layer_{0..3} keys. Writing psa for param_scan_axis, a parameter lands in one of four shapes:

MaxText key	HF keys nest as	stacked axes
global_layer-*	[block]	psa
local_layers-*	[block][local]	psa, psa + 1
global_layer-…-routed_experts-w*	[expert][block]	0, 1
local_layers-…-routed_experts-w*	[expert][block][local]	0, psa, psa + 1

Only the last row needed new machinery. Routed experts are expert-stacked inside the nested local scan, so they take a third stacked axis, and the conversion helpers only handled two. _build_multi_axis_stacked_tensor (to_maxtext) and its inverse in process_maxtext_param (to_huggingface) are generalized from two axes to N, and the axis placement moves into a shared stacked_axes helper in tensor_handling.py, next to nesting_depth and slice_shape. A third copy of the same helper lived in utils/utils.py, unreachable — nothing calls that module's _get_hf_loading_function — so it is deleted here rather than updated. Nested-scan detection widens from scanned_blocks-local_layers, which is Gemma4's module name, to -local_layers, so Qwen3-Next's layers-local_layers matches as well; mt_key is threaded through the load_dynamic lazy path so it makes the same choice.

The Linen Decoder needed a fix first. _apply_qwen3_next_scanned_blocks still had a broadcast-arg spec copied from Gemma4, which no longer matched Qwen3NextScannableBlock.__call__ and passed 12 positional args to a 10-argument signature, and it named the scanned module scanned_blocks where the pure-NNX decoder uses layers. With both corrected, and with the remainder unified as described above, the two decoders emit byte-identical parameter names and shapes, so a single mapping serves both.

Results

Activation memory drops by roughly 45–60%. Every row below is the same command on both sides, changing only the checkout.

AOT compile, real qwen3-next-80b-a3b, compile_topology=v5p-256, per_device_batch_size=1, max_target_length=2048, remat_policy=full, attention=flash, dtype=bfloat16:

CompiledMemoryStats	main	this PR	change
temp_size_in_bytes	30.25 GB	16.10 GB	−46.8%
argument_size_in_bytes	7.49 GB	7.49 GB	—
generated_code_size_in_bytes	312 MB	153 MB	−50.9%
total (args + temp)	37.74 GB	23.59 GB	−37.5%

Real training runs on a v5p-8 (4 chips, ici_fsdp_parallelism=4), scaled-down qwen3-next, max_target_length=4096:

	main total / temp	this PR total / temp	temp change
per_device_batch_size=2	26.7 / 22.1 GB	17.0 / 12.3 GB	−44.3%
per_device_batch_size=12	116.9 / 112.2 GB	48.2 / 43.6 GB	−61.1%

argument_size is byte-identical on both sides in all three experiments, confirming the two trees are the same architecture with the same parameter count — the saving is purely activations, which is why the gap widens with batch size.

Step time is unchanged: 1.544 s (main) vs 1.538 s (this PR) at batch 2, and 8.181 s vs 8.130 s at batch 12.

Shortcomings and follow-ups

• Parameter layout changes. A block's weights move from layer_0…layer_3 to a stacked local_layers plus a global_layer. HF checkpoints round-trip through the updated conversion, but a MaxText checkpoint saved from the old tree has to be re-converted.
• The full-attention layer must be last in the cycle, because the local scan runs before the global layer. That is where stock Qwen3-Next puts it, and the block's constructor raises rather than silently reordering the model if a block straddles a cycle boundary or covers more than one full-attention layer. Supporting a hybrid stack that interleaves differently means splitting the stacked local scan in two; no shipped config needs it, so the position stays derived from inhomogeneous_layer_cycle_interval rather than exposed as a config flag.
• No dense-prefix support. first_num_dense_layers is 0 for stock Qwen3-Next, so a scanned dense prefix is left as a follow-up.
• Gemma4 and Qwen3-Next still have near-identical block implementations. The remainder path is now shared (_apply_remainder_block), but factoring out a shared hybrid-block base, and unifying nnx_scan.apply_scanned_layers with NNXDecoder._apply_layers_sequentially (two appliers with two different state write-back conventions), is worth doing separately.

Tests

Block and decoder structure

Ten new tests in tests/unit/nnx_decoders_test.py, next to the analogous TestGemma4ScannableBlock.

TestQwen3NextScannableBlock asserts that the block splits a cycle into a stacked local scan plus one global layer, that the local params really are stacked along param_scan_axis, that the nested scans reproduce a plain sequential unroll of the same weights to rtol=atol=1e-5, that the external-KV path matches the scanned path and returns one cache per sub-layer in local…global order, and that a block starting off a cycle boundary -- which would straddle two periods and put the full-attention layer third of four -- is rejected rather than silently reordered.

TestNNXDecoderQwen3Next feeds the decoder a flat per-layer cache list with a distinct sentinel per layer and checks each one comes back on its own layer — this fails on the pre-fix code, where block i received kv_caches[i] instead of its four-cache group. It also covers the remainder: with 6 layers and a cycle of 4, layers_remainder has to hold the last two, and perturbing only that block's weights has to move the output, which it cannot do if the block is never applied. On the pre-fix code those two layers do not exist at all. With 8 layers there is no remainder block.

TestQwen3NextDecoderParity builds the parameter tree from get_maxtext_model_info on each decoder and asserts the two are equal, with and without a remainder. One mapping serves both decoders only for as long as that holds, and nothing was checking it.

Checkpoint conversion

tests/unit/param_mapping_test.py gains three round-trip tests, modelled on #4530's test_gemma4_local_layers_stack_unstack_roundtrip: stacking HF weights into the MaxText layout and un-stacking them back has to be the identity. They cover the two-axis [block][local] case, the three-axis [expert][block][local] case, and the global layer's plain leading-axis MoE case. test_qwen3_next_mapping_scanned is rewritten for the new layout.

The mapping is also checked against the model itself. For a scanned qwen3-next-80b-a3b, its key set matches get_maxtext_model_info's parameter tree exactly, with nothing missing and nothing extra, and every key's nesting depth matches the real tensor shape at the axes stacked_axes picks — on both decoders, which now produce identical trees.

Decode

Golden-logit tests only cover the forward pass, so nothing here caught a broken prefill → autoregressive handoff, which is exactly the kind of bug a restructured block can introduce: the forward pass can be bit-correct while the block mis-threads its sub-layers' carried state between decode steps.

tests/unit/deepseek_decode_consistency_test.py is the existing harness for that, added by the DeepSeek "decode ran the prefill attention path under NNX" fix; its module docstring notes the property holds for any model, so new families go in the parameter list. It decodes a tiny model two ways and requires the token ids to be equal:

• through MaxEngine — prefill on the padded prompt, init_decode_state, insert at slot 0, then generate per step, greedy sampling;
• teacher-forced — a plain MODEL_MODE_TRAIN forward pass recomputed over the whole growing prefix each step, taking argmax at the last real position.

The second is the definition of what the first is supposed to be replaying, so any divergence is a decode bug rather than a tolerance question. Both sides load the same freshly initialized nnx.Param state, with matmul_precision=highest, dtype=float32 and attention=dot_product so the comparison is exact rather than approximate. Prompt is 5 tokens, 5 tokens are generated.

Three Qwen3-Next cases are added, all on a scaled-down qwen3-next-80b-a3b (base_emb_dim=64, head_dim=32, 4 experts, gdn_chunk_size=4, vocab 64):

case	layers	what it covers
qwen3_next_unscanned	8	the per-layer path, as a control — isolates whether a failure is the block or the layer
qwen3_next_scanned	8	two whole periods through the nested scan, no remainder
qwen3_next_scanned_remainder	6	one period plus a two-layer layers_remainder block

Both attention kinds in the hybrid stack carry state across a decode step — the GatedDeltaNet recurrent and conv states, and the full-attention KV cache — and the scanned path threads them through two nested lax.scans plus, in the remainder case, a separate short block outside the scan. A block that regroups caches wrongly, hands one sub-layer another's state, or drops the write-back shows up here as a diverging token. The remainder case additionally fails outright on the pre-fix code, since layers 4 and 5 were not built.

Commands

pytest tests/unit/nnx_decoders_test.py tests/unit/nnx_scan_test.py \
       tests/unit/nnx_decoder_test.py tests/unit/maxtext_utils_nnx_test.py \
       tests/unit/param_mapping_test.py tests/unit/checkpointing_test.py
# 113 passed, 2 skipped

pytest tests/unit/deepseek_decode_consistency_test.py
# 7 passed

pytest tests/post_training/integration/lora_e2e_nnx_test.py
# 3 passed

End-to-end on a v5p-8, 8 steps of synthetic data with fixed seeds. Losses track main to within 0.002 at every step despite the two trees drawing init RNG differently:

step	0	1	2	3	4	5	6	7
main	9.505	9.016	8.585	8.228	7.956	7.771	7.660	7.600
this PR	9.499	9.013	8.583	8.228	7.957	7.771	7.660	7.601

Reproduce (drop override_model_config and the size overrides to run the real 80B):

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  model_name=qwen3-next-80b-a3b override_model_config=True \
  base_output_directory=/tmp/run_out run_name=scan_check dataset_type=synthetic \
  enable_checkpointing=False steps=8 data_shuffle_seed=1234 init_weights_seed=1234 \
  base_emb_dim=2048 base_num_decoder_layers=12 base_num_query_heads=16 \
  base_num_kv_heads=2 head_dim=256 base_mlp_dim=4096 base_moe_mlp_dim=512 \
  num_experts=32 num_experts_per_tok=8 vocab_size=8192 \
  per_device_batch_size=2 max_target_length=4096 dtype=bfloat16 attention=flash \
  remat_policy=full sparse_matmul=True megablox=False ici_fsdp_parallelism=4

AOT numbers come from Memory analysis: {compiled.memory_analysis()} in train_compile.py:

python3 -m maxtext.trainers.pre_train.train_compile src/maxtext/configs/base.yml \
  model_name=qwen3-next-80b-a3b compile_topology=v5p-256 compile_topology_num_slices=1 \
  base_output_directory=/tmp/aot_out run_name=aot dataset_type=synthetic \
  enable_checkpointing=False skip_jax_distributed_system=True steps=3 \
  per_device_batch_size=1 max_target_length=2048 dtype=bfloat16 attention=flash \
  remat_policy=full sparse_matmul=True megablox=False

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 support for Qwen3-Next scanned decoder blocks using hierarchical nested scans, separating linear-attention layers from the global full-attention layer to optimize rematerialization. It also refactors NNX scanning to avoid stacking read-only parameters. The review feedback correctly identifies that when kv_caches is provided, the implementation falls back to the default path and causes double rematerialization. To resolve this, the reviewer suggests updating _apply_qwen3_next_scanned_blocks to accept and pass kv_caches directly to _apply_layers_sequentially while keeping block-level rematerialization disabled.

[codecov]

Codecov Report

❌ Patch coverage is 84.80565% with 43 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	77.58%	6 Missing and 7 partials ⚠️
src/maxtext/models/qwen3.py	88.54%	4 Missing and 7 partials ⚠️
src/maxtext/layers/decoders.py	76.66%	4 Missing and 3 partials ⚠️
...xtext/checkpoint_conversion/utils/param_mapping.py	82.60%	4 Missing ⚠️
...ext/checkpoint_conversion/utils/tensor_handling.py	88.88%	2 Missing and 2 partials ⚠️
src/maxtext/checkpoint_conversion/to_maxtext.py	85.71%	1 Missing and 1 partial ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	86.66%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[NuojCheng]

Pushed 0857e11 addressing the review — the KV-cache path now keeps skip_block_remat=True and regroups the flat per-layer cache list into per-block tuples (details in the thread replies above). Tests moved into tests/unit/nnx_decoders_test.py next to TestGemma4ScannableBlock, with new coverage for both halves of the KV path.

/gemini review

[gemini-code-assist]

Code Review

This pull request adds support for Qwen3-Next scanned decoder blocks in MaxText, implementing hierarchical nested scans (local linear-attention layers and a global full-attention layer) and adding corresponding unit tests. The feedback suggests grouping the jax.experimental.xla_metadata import with other third-party imports in qwen3.py to comply with PEP 8, and optimizing the external KV cache path in qwen3.py to avoid processing read-only parameters, mirroring the optimization made in nnx_scan.py.