perf(qwen3-next): nested scan over hybrid attention with per-layer remat by NuojCheng · Pull Request #4964 · AI-Hypercomputer/maxtext

NuojCheng · 2026-08-21T21:56:38Z

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.

Qwen3-Next's scannable block instantiated its four heterogeneous sub-layers flat and ran them in a Python loop, so the whole block was rematerialized as a unit and all four sub-layers' activations were live together. Restructure it along the same lines as Gemma4: stack the three linear-attention (GatedDeltaNet) layers and run them through nnx_scan.apply_scanned_layers, and run the single full-attention layer inside a trip-count-one jax.lax.scan that acts as an XLA scheduling barrier. Each sub-layer is now rematerialized on its own, so the outer apply skips block-level remat. Also stop apply_scanned_layers from returning parameters out of its scan body: lax.scan stacks every carry output, so returning full state made XLA materialize a second copy of the stacked layer weights on every call. Add full_attention_layer_offset to place the full-attention layer within each cycle; it defaults to -1 (last in the cycle), reproducing the previous (layer_idx + 1) % cycle == 0 schedule.

Asserts 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 sequential unroll of the same weights, and that a block whose full-attention layer is not last is rejected.

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.

gemini-code-assist · 2026-08-21T21:59:45Z

+        elif self.is_qwen3_next and kv_caches is None:
+          y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs)


When kv_caches is not None (e.g., during inference or evaluation), the else block is executed, which calls _apply_layers_sequentially with skip_block_remat=False (the default). This causes double rematerialization for Qwen3-Next because Qwen3NextScannableBlock already handles its own internal rematerialization. We should update _apply_qwen3_next_scanned_blocks to accept kv_caches and call it whenever self.is_qwen3_next is True.

Suggested change

elif self.is_qwen3_next and kv_caches is None:

y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs)

elif self.is_qwen3_next:

y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs, kv_caches=kv_caches)

Good catch on the double-remat — fixed in 0857e11.

I didn't take the suggestion verbatim, though: the scan runs over blocks, not layers, while kv_caches is a flat list with one entry per decoder layer. Passing it straight through as kv_caches_stacked would hand block i only kv_caches[i] (a single layer's cache) instead of its four-cache group, and the block would then index into that single object.

So _apply_qwen3_next_scanned_blocks now does what _apply_gemma4_scanned_blocks does:

grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, block_length, stack=False) y, self.layers, _ = self._apply_layers_sequentially( self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, skip_block_remat=True, **layer_kwargs, ) maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped_kv_caches, scan_length, block_length, stacked=False)

and the dispatch is now elif self.is_qwen3_next: with kv_caches=kv_caches, as suggested.

Two tests cover this. TestNNXDecoderQwen3Next.test_decoder_regroups_flat_kv_caches_per_block passes a distinct sentinel per layer and checks each comes back on its own layer — it fails on the pre-fix code. TestQwen3NextScannableBlock.test_external_kv_cache_matches_scanned_path checks the block's external-KV path matches the scanned path numerically and returns one cache per sub-layer in local…global order.

Training is unaffected: with kv_caches=None both helpers pass None through, so that path is byte-identical to before and the memory numbers in the description still stand.

gemini-code-assist · 2026-08-21T21:59:45Z


    return y

+  def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs):


Update _apply_qwen3_next_scanned_blocks to accept kv_caches and pass it to _apply_layers_sequentially as kv_caches_stacked to ensure that block-level rematerialization is skipped even when kv_caches is provided.

Suggested change

def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs):

def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs, kv_caches=None):

Addressed in 0857e11 — see the reply on the dispatch thread. Short version: kv_caches is flat per-layer while the scan is per-block, so it is regrouped with prepare_kv_caches_for_scan / update_kv_caches_after_scan (as Gemma4 does) rather than passed through directly, and skip_block_remat=True now holds on this path too.

gemini-code-assist · 2026-08-21T21:59:45Z

+    y, self.layers, _ = self._apply_layers_sequentially(
+        self.layers,
+        y,
+        *layer_args,
+        length=scan_length,
+        skip_block_remat=True,
+        **layer_kwargs,
+    )


Pass kv_caches as kv_caches_stacked to _apply_layers_sequentially to support inference with external KV caches while skipping block-level rematerialization.

Suggested change

y, self.layers, _ = self._apply_layers_sequentially(

self.layers,

y,

*layer_args,

length=scan_length,

skip_block_remat=True,

**layer_kwargs,

)

y, self.layers, _ = self._apply_layers_sequentially(

self.layers,

y,

*layer_args,

length=scan_length,

skip_block_remat=True,

kv_caches_stacked=kv_caches,

**layer_kwargs,

)

Addressed in 0857e11 — see the reply on the dispatch thread. Short version: kv_caches is flat per-layer while the scan is per-block, so it is regrouped with prepare_kv_caches_for_scan / update_kv_caches_after_scan (as Gemma4 does) rather than passed through directly, and skip_block_remat=True now holds on this path too.

codecov · 2026-08-21T22:02:28Z

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!

The scanned-block apply was guarded on `kv_caches is None`, so an inference call fell through to the generic branch and rematerialized the whole block on top of the block's own per-layer remat. Route Qwen3-Next through `_apply_qwen3_next_scanned_blocks` unconditionally and regroup the flat per-layer cache list into per-block tuples before the scan, writing it back afterwards -- the `prepare_kv_caches_for_scan` / `update_kv_caches_after_scan` pair Gemma4 already uses. The scan runs over blocks, so the flat list would otherwise hand block i only `kv_caches[i]`. Training (kv_caches=None) is unaffected: the grouping helpers pass None through, so that path is byte-identical to before. Move the scannable-block tests out of their own file into tests/unit/nnx_decoders_test.py alongside TestGemma4ScannableBlock, and add coverage for both halves of the KV path.

NuojCheng · 2026-08-21T22:22:49Z

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.

gemini-code-assist · 2026-08-21T22:25:13Z

+from maxtext.layers import nnx_scan
+from jax.experimental import xla_metadata


According to PEP 8, imports should be grouped in the following order:

Standard library imports.

Related third party imports.

Local application/library specific imports.

The third-party import from jax.experimental import xla_metadata is currently placed within the local maxtext imports group. It should be moved to the top of the file and grouped with the other jax imports (around lines 24-29).

Suggested change

from maxtext.layers import nnx_scan

from jax.experimental import xla_metadata

from maxtext.layers import nnx_scan

References

PEP 8: Imports should be grouped: 1. Standard library imports. 2. Related third party imports. 3. Local application/library specific imports. You should put a blank line between each group of imports. ^(link)

Done in 99cbb62 — moved up next to from jax.ad_checkpoint import checkpoint_name so it sits in the jax group rather than at the top of the file, keeping that group alphabetical.

gemini-code-assist · 2026-08-21T22:25:13Z

+        per_layer_states.append(nnx.state(layer))
+
+      stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states)
+      if scan_axis != 0:
+        stacked_params, stacked_other = stacked_state.split(nnx.Param, ...)
+        stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params)
+        stacked_state = nnx.State.merge(stacked_params, stacked_other)
+      nnx.update(self.local_layers, stacked_state)


Similar to the optimization introduced in nnx_scan.py, we can avoid returning, stacking, and moving the axes of the read-only Param variables in the external KV cache path as well. Since parameters are never modified during the forward pass, we only need to collect and update the non-Param state (rest) back into self.local_layers.

_, _, current_rest = nnx.split(layer, nnx.Param, ...) per_layer_states.append(current_rest) stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) nnx.update(self.local_layers, stacked_state)

Agreed, and applied in 99cbb62.

Worth spelling out why the param_scan_axis round-trip disappears along with the stacking, since it looks like a separate change: nnx_scan.create_scanned_layers only puts Params on param_scan_axis (add_scan_metadata(stacked_params, param_scan_axis)); non-Param state gets axis 0 (add_scan_metadata(stacked_rest, 0)). The state read at the top of this method is therefore already axis-0 and sliced as such, so once Params are out of per_layer_states the stack goes straight back on axis 0 — matching what apply_scanned_layers does with scanned_rest. Net effect here:

_, _, updated_state = nnx.split(layer, nnx.Param, ...) per_layer_states.append(updated_state) nnx.update(self.local_layers, jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states))

Qwen3NextScannableBlock tests still pass (6/6), and the four NNX suites are at 71 passed / 2 skipped.

One note for anyone following up: Gemma4ScannableBlock._forward_with_external_kv_cache (gemma4.py:633) still has the original nnx.state(layer) + moveaxis version, so it keeps the extra param copy. I left it alone to keep this PR scoped to Qwen3-Next, but the same fix applies verbatim.

…path Parameters are read-only in the forward pass, so stacking them back into self.local_layers allocated a second copy of every layer weight. Collect only the non-Param state, which also removes the param_scan_axis round-trip: unlike Params, non-Param state is stacked on axis 0 by nnx_scan.create_scanned_layers. Same reasoning as the scan-body change in nnx_scan.apply_scanned_layers, applied to the static-unroll path. Also move the xla_metadata import into the jax group (PEP 8).

`apply_scanned_layers` dropped every `nnx.Param` from the scan output to avoid materializing a second copy of the stacked layer weights. That also discarded parameters created *while tracing the body*, most notably the `nnx.LoRAParam` adapters Qwix materializes, which are an `nnx.Param` subclass -- LoRA setup then failed with "LoRA module path matched target modules, but nnx.LoRAParam is still missing". Only drop the params that were fed in as scan inputs; anything else the body produces still leaves the scan.

… layout The scannable-block rewrite replaced the per-cycle-position `layer_{i}` params with a nested layout -- `layers-local_layers-*` (an inner scan over the linear-attention layers, nested in the block scan) and `layers-global_layer-*` (the block scan only). The HF param mapping still described the old layout, so conversion produced keys the model does not have. Rewrite the scanned branch of the qwen3-next mapping (and its hooks) for the new layout. Routed-expert weights are expert-stacked *inside* the nested scan, giving `[expert][block][local]` -- a third stacked axis, which the conversion helpers did not support. Generalize `_build_multi_axis_stacked_tensor` and the inverse in `process_maxtext_param` from two axes to N, with the axis placement factored into a shared `stacked_axes` helper, and broaden the nested-scan detection from `scanned_blocks-local_layers` (gemma4's module name) to `-local_layers` so qwen3-next's `layers-local_layers` matches too. Also fix the Linen `_apply_qwen3_next_scanned_blocks`: its broadcast-arg spec had been copied from gemma4 and no longer matched `Qwen3NextScannableBlock.__call__`, and it named the scanned module `scanned_blocks` where the pure-NNX decoder uses `layers`. Both decoder paths now emit byte-identical parameter names and shapes, so one mapping serves both.

… test Address the follow-ups from the scanned Qwen3-Next block PR that fit in it: - The pure-NNX decoder built `num_decoder_layers // cycle` blocks and stopped there, so a layer count that is not a multiple of the cycle silently dropped its trailing layers. They now go into a `layers_remainder` block that is built and applied like Gemma4's. - The Linen decoder spelled its remainder out as individual `Qwen3NextDecoderLayerToLinen` layers, which named parameters differently from the NNX side. It now uses a `Qwen3NextScannableBlockToLinen` named `layers_remainder` too, so one checkpoint mapping still serves both decoders. `TestQwen3NextDecoderParity` checks the two parameter trees match, with and without a remainder. - Gemma4 and Qwen3-Next now share `_apply_remainder_block` instead of keeping two copies of the same split/checkpoint/merge dance. - Deleted the unreachable third copy of `_build_multi_axis_stacked_tensor` and friends in checkpoint_conversion/utils/utils.py; nothing calls that module's `_get_hf_loading_function`. Decode coverage: `deepseek_decode_consistency_test.py` gains three Qwen3-Next cases (unscanned, scanned, scanned with a remainder). Both attention kinds in the hybrid stack carry state across a decode step, so a block that mis-threads its sub-layers' state diverges from the teacher-forced rollout.

The flag was added alongside the scanned block, but the block rejects every value it can take other than the default: the local layers are scanned before the global one, so the full-attention layer has to be last in the cycle. Any other offset raised, or -- on the unscanned path, which does not go through the block -- silently built a different architecture from the scanned path. Stock Qwen3-Next puts full attention last, no config set the flag, and nothing tested a non-default value working. Go back to deriving the position from the cycle, and keep the block's validation, which still earns its place: a block starting off a cycle boundary straddles two periods and would reorder the model.

NuojCheng added 2 commits August 21, 2026 20:43

gemini-code-assist Bot reviewed Aug 21, 2026

View reviewed changes

NuojCheng added 3 commits August 21, 2026 22:38

NuojCheng marked this pull request as ready for review August 23, 2026 02:58

NuojCheng requested review from A9isha, RissyRan, SurbhiJainUSC, abhinavclemson, aireenmei, bvandermoon, gagika, gobbleturk, hengtaoguo, huytransformer, igorts-git, jiangjy1982, khatwanimohit, parambole, richjames0, shralex, shuningjin, vipannalla and xibinliu as code owners August 23, 2026 02:58

NuojCheng requested review from darisoy and dipannita08 as code owners August 23, 2026 02:58

suhithr self-requested a review August 24, 2026 18:20

NuojCheng force-pushed the chengnuojin-bharatgen-scan branch from 9c3bae5 to 2968036 Compare August 24, 2026 20:49

		elif self.is_qwen3_next and kv_caches is None:
		y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs)


		return y

		def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs):

		from maxtext.layers import nnx_scan
		from jax.experimental import xla_metadata

Uh oh!

Conversation

NuojCheng commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checkpoint conversion

Results

Shortcomings and follow-ups

Tests

Block and decoder structure

Checkpoint conversion

Decode

Commands

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

NuojCheng commented Aug 21, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

NuojCheng commented Aug 21, 2026 •

edited

Loading

codecov Bot commented Aug 21, 2026 •

edited

Loading