Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
NuojCheng wants to merge 8 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e463229
perf(qwen3-next): nested scan over hybrid attention with per-layer remat
NuojCheng Aug 21, 2026
8a4d252
test(qwen3-next): cover the scannable block's nested scans
NuojCheng Aug 21, 2026
0857e11
fix(qwen3-next): keep per-layer remat on the external KV cache path
NuojCheng Aug 21, 2026
99cbb62
refactor(qwen3-next): drop param stacking from the external KV cache …
NuojCheng Aug 21, 2026
449a81e
fix(nnx): keep params created inside the layer scan
NuojCheng Aug 22, 2026
ef4f304
feat(qwen3-next): support checkpoint conversion for the scanned block…
NuojCheng Aug 22, 2026
2968036
Keep trailing Qwen3-Next layers, unify remainder handling, add decode…
NuojCheng Aug 24, 2026
ceedbea
Drop full_attention_layer_offset
NuojCheng Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 33 additions & 40 deletions src/maxtext/checkpoint_conversion/to_maxtext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@
from maxtext.common.common_types import MODEL_MODE_TRAIN
from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS
from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns
from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns, nesting_depth, slice_shape, stacked_axes
from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys
from maxtext.inference.inference_utils import str2bool
from maxtext.layers import quantizations
Expand Down Expand Up @@ -344,66 +344,59 @@ def get_maxtext_model_info(config):


def _build_multi_axis_stacked_tensor(
hf_source_keys: List[List[str]],
hf_source_keys: List[Any],
tensor_getter_fn: Callable[[str], np.ndarray],
hook_fns: Any,
target_shape: tuple,
config,
mt_key: str = "",
) -> np.ndarray:
"""Builds a MaxText tensor by stacking HF weights along two axes.
"""Builds a MaxText tensor by stacking HF weights along several axes.

Two cases share this helper:
Three layouts share this helper, distinguished by how deeply ``hf_source_keys``
is nested and by whether ``mt_key`` names a nested block scan:
* MoE expert stacking: outer=experts, inner=layers, placed at the LEADING two
axes -> shape (num_experts, num_layers, ...).
* Gemma4 nested block scan (``scanned_blocks-local_layers``): the block's local
layers are an inner scan nested inside the outer block scan, so the two
stacked axes live at ``(param_scan_axis, param_scan_axis + 1)`` rather than
the leading axes, e.g. (emb, blocks, local, ...).
* Nested block scan (``...-local_layers-...``, gemma4 and qwen3-next): the
block's local layers are an inner scan nested inside the outer block scan,
so the two stacked axes live at ``(param_scan_axis, param_scan_axis + 1)``
rather than the leading axes, e.g. (emb, blocks, local, ...).
* Both at once (qwen3-next's routed experts, which are expert-stacked inside
a nested block scan): three axes at ``(0, param_scan_axis,
param_scan_axis + 1)``, e.g. (experts, blocks, local, ...).

Args:
hf_source_keys: A nested (2D) list of Hugging Face parameter names.
The outer list is the outer stack axis, the inner list the
inner stack axis.
hf_source_keys: A nested list of Hugging Face parameter names, one level of
nesting per stacked axis, outermost axis first. A ``tuple`` of names is a
composite HF source that the hook fuses into a single leaf, not an axis.
tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array).
hook_fns: The hook function(s) to apply to each individual weight.
target_shape: The final shape of the target MaxText tensor.
config: The MaxText pyconfig object.
mt_key: The MaxText parameter key, used to detect the gemma4 nested-scan case.
mt_key: The MaxText parameter key, used to detect the nested-scan case.

Returns:
The final, assembled NumPy array for the MaxText parameter.
"""
# Gemma4's local layers are an inner scan nested inside the block scan, so the
# two stacked axes go at (param_scan_axis, param_scan_axis + 1); everything else
# (MoE experts/layers) stacks at the leading axes.
if isinstance(mt_key, str) and "scanned_blocks-local_layers" in mt_key:
outer_axis, inner_axis = config.param_scan_axis, config.param_scan_axis + 1
else:
outer_axis, inner_axis = 0, 1

depth = nesting_depth(hf_source_keys)
axes = stacked_axes(mt_key, config, depth)
# The hook function needs the shape of an individual slice: target_shape with
# the two stacked axes removed.
stacked_axes = {outer_axis, inner_axis}
mt_slice_shape = tuple(d for i, d in enumerate(target_shape) if i not in stacked_axes)

all_outer_tensors = []
# Outer loop iterates the outer stack axis (experts, or blocks for gemma4 local)
for inner_keys in hf_source_keys:
inner_tensors = []
# Inner loop iterates the inner stack axis (layers, or local layers for gemma4)
for hf_key_single in inner_keys:
if isinstance(hf_key_single, (list, tuple)):
hf_tensor_numpy = tuple(tensor_getter_fn(k) for k in hf_key_single)
# the stacked axes removed.
mt_slice_shape = slice_shape(target_shape, axes)

def gather(keys, level):
if level == depth:
if isinstance(keys, (list, tuple)):
raw = tuple(tensor_getter_fn(k) for k in keys)
else:
hf_tensor_numpy = tensor_getter_fn(hf_key_single)
inner_tensors.append(apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns))
all_outer_tensors.append(np.stack(inner_tensors, axis=0))
stacked = np.stack(all_outer_tensors, axis=0) # (outer, inner, *slice)

# Move the (outer, inner) axes from leading positions to their targets.
if (outer_axis, inner_axis) != (0, 1):
stacked = np.moveaxis(stacked, (0, 1), (outer_axis, inner_axis))
raw = tensor_getter_fn(keys)
return apply_hook_fns(raw, mt_slice_shape, hook_fns)
# Stack with the axes leading, outermost first; they are moved into place below.
return np.stack([gather(sub, level + 1) for sub in keys], axis=0)

stacked = gather(hf_source_keys, 0)
if axes != tuple(range(depth)):
stacked = np.moveaxis(stacked, tuple(range(depth)), axes)
return stacked


Expand Down
1 change: 1 addition & 0 deletions src/maxtext/checkpoint_conversion/utils/load_dynamic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,7 @@ def tensor_getter(key):
hook_fn,
target_leaf,
maxtext_config,
mt_key,
)

# Execute transformation and assign to flat_restored
Expand Down
187 changes: 86 additions & 101 deletions src/maxtext/checkpoint_conversion/utils/param_mapping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1402,100 +1402,81 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F
}

if scan_layers:
# 2. Scan over block cycles
for block_idx in range(layer_cycle_interval):
hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval))
prefix = f"params-decoder-layers-layer_{block_idx}"
# 2. Scanned blocks. One block covers a single period of the hybrid attention
# pattern: `layer_cycle_interval - 1` linear-attention (GatedDeltaNet) layers
# run as an inner scan, then one full-attention layer. The resulting params are:
# layers-local_layers-* -> nested [block][local] (doubly scanned)
# layers-global_layer-* -> flat [block] (block scan only; the
# length-1 _scan_global_layer scan is a runtime memory boundary, not a
# param stack)
# Routed-expert weights carry an additional leading expert axis, so they are
# nested one level deeper: [expert][block][local] and [expert][block].
# Qwen3NextScannableBlock requires the full-attention layer to be last in the
# period, so the local positions are 0..cycle-2 and the global position is
# cycle-1.
num_blocks = num_main_layers // layer_cycle_interval
local_positions = list(range(layer_cycle_interval - 1))
global_position = layer_cycle_interval - 1

def hf_layer(block_idx, position, suffix):
return f"model.layers.{block_idx * layer_cycle_interval + position}.{suffix}"

# (maxtext subkey, hf suffix) pairs shared by both the local and global layers.
shared_specs = [
("input_layernorm-scale", "input_layernorm.weight"),
("post_attention_layernorm-scale", "post_attention_layernorm.weight"),
("mlp-routed_experts-gate-kernel", "mlp.gate.weight"),
("mlp-shared_expert-wi_0-kernel", "mlp.shared_expert.gate_proj.weight"),
("mlp-shared_expert-wi_1-kernel", "mlp.shared_expert.up_proj.weight"),
("mlp-shared_expert-wo-kernel", "mlp.shared_expert.down_proj.weight"),
("mlp-shared_expert_gate-kernel", "mlp.shared_expert_gate.weight"),
]
# Linear (GatedDeltaNet) attention: only ever on the local layers.
local_specs = shared_specs + [
("attention-in_proj_qkvz-kernel", "linear_attn.in_proj_qkvz.weight"),
("attention-in_proj_ba-kernel", "linear_attn.in_proj_ba.weight"),
("attention-conv1d-kernel", "linear_attn.conv1d.weight"),
("attention-A_log", "linear_attn.A_log"),
("attention-dt_bias", "linear_attn.dt_bias"),
("attention-norm-rms_norm-scale", "linear_attn.norm.weight"),
("attention-out_proj-kernel", "linear_attn.out_proj.weight"),
]
# Full attention: only ever on the global layer.
global_specs = shared_specs + [
("attention-attention-query-kernel", "self_attn.q_proj.weight"),
("attention-attention-key-kernel", "self_attn.k_proj.weight"),
("attention-attention-value-kernel", "self_attn.v_proj.weight"),
("attention-attention-out-kernel", "self_attn.o_proj.weight"),
("attention-attention-query_norm-scale", "self_attn.q_norm.weight"),
("attention-attention-key_norm-scale", "self_attn.k_norm.weight"),
]
expert_specs = [
("mlp-routed_experts-wi_0", "gate_proj.weight"),
("mlp-routed_experts-wi_1", "up_proj.weight"),
("mlp-routed_experts-wo", "down_proj.weight"),
]

# Layer norms
mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment]
f"model.layers.{i}.input_layernorm.weight" for i in hf_indices
] # pyrefly: ignore[bad-assignment]
mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment]
f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
local_prefix = "params-decoder-layers-local_layers"
for subkey, suffix in local_specs:
mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload]
[hf_layer(b, p, suffix) for p in local_positions] for b in range(num_blocks)
]
for subkey, suffix in expert_specs:
mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload]
[[hf_layer(b, p, f"mlp.experts.{e}.{suffix}") for p in local_positions] for b in range(num_blocks)]
for e in range(num_experts)
]

# Handle Interleaved Attention (Linear vs Full)
is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0

if is_full_attention_layer:
mapping.update( # pyrefly: ignore[no-matching-overload]
{
f"{prefix}-attention-attention-query-kernel": [
f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices
],
f"{prefix}-attention-attention-key-kernel": [
f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices
],
f"{prefix}-attention-attention-value-kernel": [
f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices
],
f"{prefix}-attention-attention-out-kernel": [
f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices
],
f"{prefix}-attention-attention-query_norm-scale": [
f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices
],
f"{prefix}-attention-attention-key_norm-scale": [
f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices
],
}
)
else:
# Linear/Hybrid Attention Block
mapping.update( # pyrefly: ignore[no-matching-overload]
{
f"{prefix}-attention-in_proj_qkvz-kernel": [
f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices
],
f"{prefix}-attention-in_proj_ba-kernel": [
f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices
],
f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices],
f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices],
f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices],
f"{prefix}-attention-norm-rms_norm-scale": [
f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices
],
f"{prefix}-attention-out_proj-kernel": [
f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices
],
}
)

# 3. Handle MLP: Gates and Shared Experts
mapping.update( # pyrefly: ignore[no-matching-overload]
{
f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices],
f"{prefix}-mlp-shared_expert-wi_0-kernel": [
f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices
],
f"{prefix}-mlp-shared_expert-wi_1-kernel": [
f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices
],
f"{prefix}-mlp-shared_expert-wo-kernel": [
f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices
],
f"{prefix}-mlp-shared_expert_gate-kernel": [
f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices
],
}
)

# 4. Handle MoE Routed Experts
mapping.update( # pyrefly: ignore[no-matching-overload]
{
f"{prefix}-mlp-routed_experts-wi_0": [
[f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts)
],
f"{prefix}-mlp-routed_experts-wi_1": [
[f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts)
],
f"{prefix}-mlp-routed_experts-wo": [
[f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts)
],
}
)
global_prefix = "params-decoder-layers-global_layer"
for subkey, suffix in global_specs:
mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload]
hf_layer(b, global_position, suffix) for b in range(num_blocks)
]
for subkey, suffix in expert_specs:
mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload]
[hf_layer(b, global_position, f"mlp.experts.{e}.{suffix}") for b in range(num_blocks)]
for e in range(num_experts)
]
else:
# Unscanned layer mapping
for i in range(num_main_layers):
Expand Down Expand Up @@ -1589,17 +1570,21 @@ def permute_conv(input_tensor, target_shape=None):

layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
num_main_layers = config["num_hidden_layers"]
loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers)

for i in loop_indices:
if scan_layers:
prefix = f"params-decoder-layers-layer_{i}"
block_idx = i
else:
prefix = f"params-decoder-layers_{i}"
block_idx = i % layer_cycle_interval
is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0
# Scanned blocks expose two prefixes -- the stacked local (linear-attention) layers
# and the single global (full-attention) layer -- rather than one prefix per position
# in the cycle. Unscanned models keep one prefix per decoder layer.
if scan_layers:
layer_prefixes = [
("params-decoder-layers-local_layers", False),
("params-decoder-layers-global_layer", True),
]
else:
layer_prefixes = [
(f"params-decoder-layers_{i}", (i % layer_cycle_interval + 1) % layer_cycle_interval == 0)
for i in range(num_main_layers)
]

for prefix, is_full_attention_layer in layer_prefixes:
if is_full_attention_layer:
for key in ["query", "key", "value", "out"]:
hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment]
Expand Down
Loading
Loading