perf(qwen3.5): reuse Qwen3-Next's nested block scan

[NuojCheng]

What

Extends the Qwen3-Next scanned-block work from #4964 to Qwen3.5.

Qwen3.5 repeats the exact same hybrid attention period as Qwen3-Next — inhomogeneous_layer_cycle_interval = 4, i.e. three GatedDeltaNet layers followed by one full-attention layer — so rather than porting the logic, this makes Qwen3.5 reuse it.

Qwen3NextScannableBlock's layer construction moves into an overridable _make_decoder_layer() hook, and Qwen3_5ScannableBlock becomes a subclass that only swaps in Qwen3_5DecoderLayer. The ~50-line hand-rolled Python-loop block it replaces is deleted, so Qwen3.5 picks up the inner scan over the homogeneous linear-attention layers, the trip-count-one lax.scan scheduling barrier around the full-attention layer, and per-sub-layer remat, for free.

Changes

• models/qwen3.py — extract layer construction from Qwen3NextScannableBlock.__init__ into _make_decoder_layer(). No behavior change for Qwen3-Next.
• models/qwen3_5.py — Qwen3_5ScannableBlock subclasses Qwen3NextScannableBlock and overrides the hook. Qwen3_5DecoderLayer gains an explicit is_full_attention_layer kwarg, because inside a scan a sub-layer's position in the cycle is not recoverable from layer_idx; it still falls back to the (layer_idx + 1) % cycle == 0 derivation when unscanned. The MoE load-balance loss is now assigned rather than sow-n — sow appends to a tuple, so the layer's Intermediate structure grows on every call, which a lax.scan body cannot express.
• layers/decoders.py, layers/nnx_decoders.py — route QWEN3_5 through the shared hybrid path. _init_scanned_qwen3_next / _apply_qwen3_next_scanned_blocks are renamed to ..._qwen3_hybrid_... and pick the block class by decoder_block, including for layers_remainder.
• checkpoint_conversion/utils/param_mapping.py — remap Qwen3.5's scanned checkpoint mapping from the old per-cycle-position layers-layer_{i} prefixes to the local_layers / global_layer tree, mirroring what perf(qwen3-next): nested scan over hybrid attention with per-layer remat #4964 did for Qwen3-Next. Local layers nest [block][local], the global layer is flat [block].
• checkpoint_conversion/utils/tensor_handling.py — stacked_axes() was guarded on isinstance(mt_key, str), so Qwen3.5's composite (wi_0, wi_1) key (fed by one fused HF gate_up_proj tensor) fell through to the leading-axis MoE layout instead of the nested-scan layout and would have silently transposed those weights. Tuple keys now resolve via their first component.

Results

AOT temp_size_in_bytes per device, using the config from train_compile_test.py::test_qwen3_5 (per_device_batch_size=1, max_target_length=1024, sparse_matmul, megablox, attention=flash, use_tokamax_splash):

model	topology	before	after	delta
qwen3.5-397b-a17b	v5p-512	75.62 GB	46.01 GB	−39%
qwen3.5-397b-a17b	v5p-128	128.05 GB (OOM)	116.25 GB (OOM)	−9%
qwen3.5-35b-a3b	v5p-512	11.99 GB	6.74 GB	−44%
qwen3.5-35b-a3b	v5p-128	16.16 GB	11.03 GB	−32%
qwen3.5-35b-a3b	v5p-8	148.11 GB (OOM)	150.02 GB (OOM)	+1%

argument_size_in_bytes is unchanged (18.61 GB at 397b/v5p-512); generated code size drops 320 MB → 163 MB.

Real training on a v7-8 (qwen3.5-35b-a3b, synthetic data, ici_fsdp_parallelism=8, per_device_batch_size=1, max_target_length=1024, 6 steps) converges on both sides with matching loss curves (step 5: 7.056 before, 7.072 after — RNG-fork differences only). Memory is a wash at that scale, 40.6 → 41.1 GB estimated temp, which is the same regime as the v5p-8 row above.

Caveat: the win depends on weight_dtype

Isolating that single variable at 397b / v5p-512:

weight_dtype	before	after
float32 (MaxText default)	75.62 GB	46.01 GB
bfloat16	69.37 GB	218.19 GB (OOM)

This is not introduced here — it is a property of #4964 itself. Measured against the pre-#4964 commit, qwen3-next-80b-a3b at v5p-512 goes 21.44 → 10.68 GB with fp32 weights (the reported win, reproduced) but 20.15 → 53.46 GB with bf16 weights. The mechanism is consistent with the design: with fp32 storage the per-layer bf16 weight cast is a temporary, and block-level remat kept all four sub-layers' casts live at once, which is exactly what per-sub-layer remat fixes; with bf16 storage there is no cast to save and only the nested scan's overhead remains. Left as-is here to keep the port faithful to Qwen3-Next — worth tracking separately.

Known gap

integration/vllm/weight_converter.py:696 still assumes the old layers/layer_{slot} tree and will not read scanned Qwen3.5 checkpoints. #4964 did not touch this file either, so Qwen3-Next has the identical gap; this PR matches that rather than fixing it.

Test plan

• tests/unit/nnx_decoders_test.py -k Qwen3 — 21 passed. The three Qwen3-Next test classes are parameterized by class attributes so Qwen3.5 subclasses reuse them, giving Qwen3.5 the same coverage: nested-scan output matches a sequential unroll numerically, and the Linen and NNX parameter trees match, both for a whole number of periods and with a remainder.
• tests/unit/param_mapping_test.py — 28 passed, including two new cases covering the local_layers / global_layer mapping shape and the composite-key axis fix.

[gemini-code-assist]

Code Review

This pull request implements hierarchical nested scans for Qwen3-Next and Qwen3.5 hybrid decoders, enabling per-layer rematerialization to optimize memory usage. It introduces scannable blocks that structure inhomogeneous attention cycles and updates checkpoint conversion utilities to support multi-axis stacked tensors with arbitrary nesting depths. Comprehensive unit tests are added to verify decoder parity and state restoration. The review feedback suggests updating the type annotations for mt_key from str to str | tuple[str, ...] in stacked_axes and _build_multi_axis_stacked_tensor to accurately reflect that composite keys can be represented as tuples.

[gemini-code-assist]

The mt_key parameter is annotated as str, but the function explicitly handles tuple inputs (lines 69-70). To ensure type safety and prevent static analysis warnings, the type annotation should be updated to str | tuple[str, ...].

Suggested change

	def stacked_axes(mt_key: str, config, depth: int) -> tuple:
	def stacked_axes(mt_key: str | tuple[str, ...], config, depth: int) -> tuple:

[gemini-code-assist]

Update the type annotation of mt_key to support tuple inputs as well, aligning with the updated stacked_axes signature.

Suggested change

	mt_key: str = "",
	mt_key: str | tuple[str, ...] = "",

[gemini-code-assist]

Update the type annotation of mt_key to support tuple inputs as well, aligning with the updated stacked_axes signature.

Suggested change

	mt_key: str = "",
	mt_key: str | tuple[str, ...] = "",

[codecov]

Codecov Report

❌ Patch coverage is 87.03704% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/utils/param_mapping.py	83.33%	4 Missing ⚠️
src/maxtext/models/qwen3_5.py	57.14%	2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!