perf(qwen3-next): nested scan over hybrid attention with per-layer remat - #4964
perf(qwen3-next): nested scan over hybrid attention with per-layer remat#4964NuojCheng wants to merge 6 commits into
Conversation
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[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| elif self.is_qwen3_next and kv_caches is None: | ||
| y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs) |
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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.
|
|
||
| return y | ||
|
|
||
| def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs): |
There was a problem hiding this comment.
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.
| 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): |
There was a problem hiding this comment.
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.
| y, self.layers, _ = self._apply_layers_sequentially( | ||
| self.layers, | ||
| y, | ||
| *layer_args, | ||
| length=scan_length, | ||
| skip_block_remat=True, | ||
| **layer_kwargs, | ||
| ) |
There was a problem hiding this comment.
Pass kv_caches as kv_caches_stacked to _apply_layers_sequentially to support inference with external KV caches while skipping block-level rematerialization.
| 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, | |
| ) |
There was a problem hiding this comment.
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 Report❌ Patch coverage is 📢 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.
|
Pushed 0857e11 addressing the review — the KV-cache path now keeps /gemini review |
There was a problem hiding this comment.
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.
| from maxtext.layers import nnx_scan | ||
| from jax.experimental import xla_metadata |
There was a problem hiding this comment.
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).
| 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)
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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)There was a problem hiding this comment.
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.
1 participant
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.
Qwen3NextScannableBlockinstantiated those four heterogeneous sub-layers flat (
layer_0…layer_3) and ranthem 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:through
nnx_scan.apply_scanned_layers;jax.lax.scan, whichacts as an XLA scheduling barrier
(
skip-simplify-while-loops_trip-count-one);apply_internal_remat=True), so theouter apply passes
skip_block_remat=Trueand nothing is rematerialized twice.A second, independent fix in
nnx_scan.apply_scanned_layers: the scan body returnednnx.state(current_layer)— the full state, parameters included.jax.lax.scanstacks 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.Paramwould be too blunt. Qwix materializes LoRA adapters whiletracing the scan body, and
nnx.LoRAParamis annnx.Paramsubclass, so they would goout 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 parameterpaths 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_blocksregroups theminto per-block tuples and writes them back afterwards — the same
prepare_kv_caches_for_scan/update_kv_caches_after_scanpair Gemma4 uses. Thatkeeps
skip_block_remat=Trueon the KV-cache path instead of falling back to thegeneric, block-rematerialized branch.
Qwen3NextScannableBlocknow also threads onecache 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
Decoderalready did the equivalent regrouping and is unchanged in this respect.
Finally,
full_attention_layer_offsetselects where the full-attention layer sitswithin each cycle. It defaults to
-1(last in the cycle), which reproduces theexisting
(layer_idx + 1) % cycle == 0schedule, so no current config changesbehavior.
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, nowdescribe the nested layout instead of the old
layer_{0..3}keys. Writingpsaforparam_scan_axis, a parameter lands in one of four shapes:global_layer-*[block]psalocal_layers-*[block][local]psa, psa + 1global_layer-…-routed_experts-w*[expert][block]0, 1local_layers-…-routed_experts-w*[expert][block][local]0, psa, psa + 1Only 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 inprocess_maxtext_param(to_huggingface) are generalized from two axes to N, and theaxis placement moves into a shared
stacked_axeshelper intensor_handling.py,next to
nesting_depthandslice_shape. Nested-scan detection widens fromscanned_blocks-local_layers, which is Gemma4's module name, to-local_layers, soQwen3-Next's
layers-local_layersmatches as well;mt_keyis threaded through theload_dynamiclazy path so it makes the same choice.The Linen
Decoderneeded a fix first._apply_qwen3_next_scanned_blocksstill had abroadcast-arg spec copied from Gemma4, which no longer matched
Qwen3NextScannableBlock.__call__and passed 12 positional args to a 10-argumentsignature, and it named the scanned module
scanned_blockswhere the pure-NNX decoderuses
layers. With both corrected the two decoders emit byte-identical parameter namesand 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:CompiledMemoryStatstemp_size_in_bytesargument_size_in_bytesgenerated_code_size_in_bytesReal training runs on a v5p-8 (4 chips,
ici_fsdp_parallelism=4), scaled-downqwen3-next,
max_target_length=4096:per_device_batch_size=2per_device_batch_size=12argument_sizeis byte-identical on both sides in all three experiments, confirmingthe 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
layer_0…layer_3to astacked
local_layersplus aglobal_layer. HF checkpoints round-trip through theupdated conversion, but a MaxText checkpoint saved from the old tree has to be
re-converted.
_build_multi_axis_stacked_tensorexists in three copies, into_maxtext.py,utils/tensor_handling.pyandutils/utils.py. The two live ones are updated here.The third is unreachable — nothing calls its
_get_hf_loading_function— so it isleft alone; deleting it belongs in its own PR.
before the global layer. The constructor now raises rather than silently reordering
the model; a block covering more than one full-attention layer is also rejected.
first_num_dense_layersis 0 for stock Qwen3-Next, soa scanned dense prefix is left as a follow-up.
num_decoder_layers % inhomogeneous_layer_cycle_intervaltrailing layers (48 % 4 == 0 for Qwen3-Next, so nothing is dropped today). That is
unchanged from the generic scanned path this replaces; the Linen
Decoderdoeshandle the remainder. Worth unifying.
shared hybrid-block base, and unifying
nnx_scan.apply_scanned_layerswithNNXDecoder._apply_layers_sequentially(two appliers with two different statewrite-back conventions), is worth doing separately.
Tests
Six new tests in
tests/unit/nnx_decoders_test.py, next to the analogousTestGemma4ScannableBlock:TestQwen3NextScannableBlockasserts that the block splits a cycle into a stackedlocal 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 thesame weights to
rtol=atol=1e-5, that the external-KV path matches the scanned pathand returns one cache per sub-layer in
local…globalorder, and that a block whosefull-attention layer is not last is rejected.
TestNNXDecoderQwen3Nextfeeds the decoder a flat per-layer cache list with a distinctsentinel 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.tests/unit/param_mapping_test.pygains three round-trip tests, modelled on #4530'stest_gemma4_local_layers_stack_unstack_roundtrip: stacking HF weights into theMaxText 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 theglobal layer's plain leading-axis MoE case.
test_qwen3_next_mapping_scannedisrewritten for the new layout.
The mapping is also checked against the model itself. For a scanned
qwen3-next-80b-a3b, its key set matchesget_maxtext_model_info's parameter treeexactly, with nothing missing and nothing extra, and every key's nesting depth matches
the real tensor shape at the axes
stacked_axespicks — on both decoders, which nowproduce identical trees.
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:
Reproduce (drop
override_model_configand the size overrides to run the real 80B):AOT numbers come from
Memory analysis: {compiled.memory_analysis()}intrain_compile.py:Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.