Repository navigation
Conversation
fwd_bwd differentiated every parameter, and update() then zeroed the gradients of the ones trainable_parameters_mask freezes before the norm, clipping and the optimizer saw them. Every micro-batch computed, reduce-scattered and accumulated those gradients only to throw them away. diff_wrapper now applies jax.lax.stop_gradient to the frozen leaves, inside the differentiated function: JAX emits no weight-gradient op for them and fwd_bwd returns a constant zero in their place. The gradient through a frozen weight into the activations still comes from the same ops, so the trainable gradients, the norm and the update are the same computation. Unconditional whenever a mask is set, with no config key. Nothing in the engine reads a frozen leaf's gradient: the norm, clipping, the NaN and spike skips and the optimizer all come after update() zeroes it, and no API returns gradients. update() still zeroes them, for an accumulator restored from an intra-step checkpoint written by an older engine or under another mask. The leaves keep their slots in the gradient tree, as zeros, so kernel signatures, the accumulator and checkpoints are unchanged. pre_train/train.py is untouched: it clips over the whole gradient tree, frozen leaves included, so the same stop there would change its updates (README, "Differences from pre_train/train.py"). Both kernels read the mask through one helper, _freeze_mask_for, so the leaves fwd_bwd stops are exactly the leaves update() zeroes. It reads the compile-time cache once and, when it is empty, computes the mask from the tree, as _update_kernel already did: on the eager path, and for a dropless or final kernel traced on a background thread after _disable_pure_state emptied the cache, where checking the cache alone would silently skip the stop. The MLPerf RL recipes freeze the MoE router (^(?!.*routed_experts/gate/kernel).*). Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), with the same stop behind an environment switch: - AOT with the TPU libtpu: the router's f32[4096,512] weight-gradient fusions and f32[16,512] reduce-scatters are gone from fwd_bwd (8 -> 0 each in the program, one per MoE layer when run). - Regular micro-batch 10.189 -> 10.131 s, both with moe_routing_maps=device (o275x-stk1a-a vs o275x-stk2, two jobs on one booking): -58 ms against a predicted -45 ms (-27 to -65). - Numerics A/B (nxstk2 vs nxref, nxref2): 5 of 5 gates pass; loss 4.6e-8 relative, grad_norm 4.8e-5 relative against a replicate spread of 1.9e-5. Tests (maxtext_engine_frozen_grad_test.py, new), against the same engine with the stop removed: on a tiny decoder with one MLP kernel frozen, trained by AdamW with clipping on four CPU devices, fwd_bwd lowers to one dot_general fewer, shaped like the kernel, and none added; it returns exactly zero for the kernel; and every other gradient, the running sums, the gradient norms and the trained weights are bitwise equal, under fsdp, under the deferred all-reduce and on the eager path. On the one-layer Qwen3.5 MoE rig, with the recipes' router mask and Tunix's router replay payload and loss, the same holds. A stop that also hits one trainable kernel is caught, and the failure names it. README: the config key, the test and the difference from pre_train/train.py.
…e_routing_index_math) The ring-of-experts ragged sort derives its routing index maps from the top-k expert ids with sorts: a stable argsort for the expert permutation, a gather of a repeated iota for the sorted token indices, a second argsort for the inverse permutation, a stable argsort in the SparseCore gather-reduce's validity preprocessing, and, for tokens without forced ids (every token under moe_exclude_padding_from_routing), jax.lax.top_k on the router probabilities. XLA lowers each stable argsort in the forward pass with an extra, duplicated iota operand. moe_routing_index_math="count" (default "sort", unchanged) computes the same maps with less work: - the permutation by one single-operand sort of the packed uint32 key (expert << pos_bits | slot), whose values are distinct, so the unstable sort has exactly the stable order (packed_stable_sort); - the sorted token indices as slot // topk, without a gather; - the inverse permutation by counting, on this shard's slots only, -1 elsewhere (ragged_sort._local_slot_positions): a slot's rank is the slots of lower local experts, plus the slots of its expert in earlier tokens (a blocked prefix count as a triangular bf16 matmul, exact), plus earlier slots of the same token with the same expert. Every reader of the inverse masks it to this shard's window, so nothing reads the -1s, and the unsort backward gets the forward permutation instead of argsorting the partial inverse back (it raises without it); - the router-weight gradient back to slot order by a masked gather instead of a TensorCore scatter (the table is padded to the slot count so XLA offloads the gather to the SparseCore); - the gather-reduce validity partition as one packed uint32 sort (_preprocess(packed=True)); - the top-k indices of non-forced tokens by k argmax passes (topk_indices_by_argmax), which break ties the way top_k does. That these are cheaper than top_k's sort is expected, not measured: they were timed only as part of this key. Exact: the same permutation everywhere, so the outputs and gradients are the same bits. The validator rejects "count" where no ring-of-experts ragged sort runs (it needs sparse_matmul, te_moe_block=False, use_ring_of_experts and use_ragged_sort). It starts the validator form the later keys of this series share: _reject_unless names every unmet requirement, and the fields of the other config mixins are read directly rather than through getattr with a default. With the key off the lowered TPU program of the layer is the same as before the change (checked on the forward-plus-backward lowering of a shrunk Qwen3.5 MoE layer for v7x). Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), with this change switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 10.022 -> 9.875 s (-0.147 s; arm o275x-mf-cnt-prof against o275x-mf-ctl, the flags-off control in the same pod), 99.25 -> 98.05 s per step, 0 dropped tokens in both. The same control measured in different pods spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks with this key in both measured stacks: with moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and with the fused SiLU GMM in place of the fused activation (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: that is the stack, not this key alone; and it ran on the environment-switch build, not on this branch's binary. Tests (moe_routing_index_math_test.py, new), on CPU: - _local_slot_positions equals argsort(argsort(ids)) minus the window start on every shard's slots and -1 elsewhere, across partial blocks, the padding id and repeated experts; topk_indices_by_argmax equals top_k's indices with ties; _preprocess(packed=True) equals the argsort version. The SparseCore gather-reduce is TPU-only (CPU takes its reference fallback), so its preprocessing is tested directly. - On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local routing, two chunks), "count" is bit-identical to "sort" in the output, the loss and every gradient: plain and pipelined chunk loop, padding excluded from routing, forced ids with a buffer that drops tokens, dropless, saved routing maps, bf16 router, and the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, fp8 fake-quant). - Structure, two-sided: "count" traces fewer sorts than "sort", none of them stable or with more than one operand (so a partial fallback to jnp.argsort fails), and no top_k under padding exclusion; every ragged_gather_reduce call gets packed_preprocess=True under "count" and False under "sort" (a spy, since only the TPU kernel reads it). A counted inverse shifted by one is caught, and "count" without the forward permutation raises. - The validator rejects "count" for each unmet requirement, naming it.
…ort_unique_pair_argsort, moe_prefused_swiglu_vjp) Three rewrites of MoE glue ops that leave every value unchanged, each behind its own key (default off, program unchanged): moe_unsplit_prefused_wi. With prefuse_moe_weights, the layer sliced the fused wi = [w0 | w1] into its halves at entry, so the fp8 fake-quant ran twice, the FSDP all-gather and the expert-weight prefetch barrier carried two operands, and gmm_up concatenated the gathered halves back before the up-projection. The key keeps wi whole up to the GEMM; w1_kernel=None marks the layout in sparse_matmul. The fake-quant scales are per output channel, so rounding the whole tensor gives the halves' values. moe_ragged_sort_unique_pair_argsort. With moe_routing_index_math="sort", the ring-of-experts ragged sort's two argsorts and the SparseCore gather-reduce's validity sort run as unstable sorts of the unique (key, index) pairs (argsort_unique_pairs). A stable jnp.argsort lowers in the forward pass with a third, duplicated iota operand; the pair sort has two operands and exactly one result, the stable permutation (on v7x, an s32[655360] argsort takes 0.61-0.63 ms with three operands and 0.42 ms with two). moe_prefused_swiglu_vjp. With prefuse_moe_weights, the activation on the whole [gate | up] up-projection output runs as one custom-VJP op whose backward takes jax.vjp of the same activation and returns d[gate | up] as one array. Split, the two [rows, mlp] cotangents leave their fusion separately and XLA joins them (add of two pads) in an extra pass over the expert buffer. Any gated activation works (SwiGLU, GPT-OSS's clipped form). Same values; the gradient bits can differ only in the sign of an exact zero (the add of zero pads gives -0.0 + 0.0 = +0.0, the concatenate keeps -0.0; measured on XLA:CPU). Which keys to use. With the rest of this series the recommended set is moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_roe_packed_routing_gather, moe_router_grad_from_wo_input and moe_wi_gmm_fused_swiglu. Of this commit's keys only the whole wi is in it: "count" removes the sorts the pair argsort rewrites, and the fused SiLU GMM never forms the [gate | up] output the fused activation reads, so the validators reject the pair argsort with "count" and the fused activation with the fused SiLU. The two stay as standalone alternatives for configurations the recommended keys do not cover: the pair argsort for runs that keep moe_routing_index_math="sort" (it keeps the maps as sorts, the smaller change), and the fused activation where the fused SiLU GMM does not apply (it needs tokamax gmm_v2 and plain silu + linear; the fused activation works with any GMM backend and gated activation). The validators, in the form moe_routing_index_math started, reject each key for every unmet requirement: the pair argsort off the ring-of-experts ragged sort or with "count"; the whole wi off prefused sparse-matmul training, with weight quantization, non-default expert-weight sharding, embedding chunks, chunk weight-gradient accumulation, or an expert MLP dimension that tensor, tensor_sequence or autoregressive parallelism shards (a device would then hold a slice of [w0 | w1] instead of matching slices of w0 and w1); the fused activation off prefused sparse matmul, with MLP biases or embedding chunks, or under a remat policy that saves the routed experts' gate or up half (_remat_keeps_moe_tensor_requirement, shared with the later keys of this series: the minimal* presets, or custom with mlpwi_0, moe_mlpwi_0 or moe_mlpwi_1 saved or offloaded). The fused activation never forms those halves, so such a policy would silently stop saving the MoE tensor (the shared expert's dense MLP, tagged mlpwi_0 too, would still be saved). An MLP dimension sharded under custom logical axis rules, weight quantization objects and sparsity, which only the built layer knows, raise at trace time instead of falling back. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), all three together, switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 10.022 -> 9.852 s (-0.170 s; arm o275x-mf-gb-prof against its in-pod control o275x-mf-ctl), 99.25 -> 97.70 s per step, 0 dropped tokens in both. The same control measured in different pods spreads by +/-0.02-0.03 s per micro-batch. The parts were not timed separately, and moe_routing_index_math was "sort" in that arm. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks for moe_unsplit_prefused_wi and moe_prefused_swiglu_vjp stacked with moe_routing_index_math="count", moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and for the whole wi in the recommended set above (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: stack level, not per key; the environment-switch build, not this branch's binary; and the pair argsort is in neither run (it is inert next to "count"; a permutation, exact by construction, and bit-identical on CPU). Tests (moe_glue_rewrites_test.py, new), on CPU: - argsort_unique_pairs equals the stable argsort (ties, booleans, every axis, index dtype); _preprocess gives the same partition with it (the SparseCore gather-reduce is TPU-only, so its preprocessing is tested directly); _prefused_ffn_activation gives the split value bit for bit and the split VJP value for value, with bit differences only at zeros (SwiGLU and clipped GPT-OSS, bf16 and f32, a -0.0 cotangent row). - On the shrunk Qwen3.5 routed MoE on 8 CPU devices with prefuse_moe_weights (cp-as-ep, local routing, two chunks), each key alone and all three are bit-identical to the plain program in the output, the loss and every gradient: plain and pipelined chunk loop, fp8 fake-quant with padding excluded, forced ids with drops and dropless, with "count", and with the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75). - Structure: fewer concatenates (and fake-quant barriers) with the whole wi; exactly two two-key sorts per chunk with the pair argsort, and unique_pair_argsort=True on every gather-reduce call (a spy, since only the TPU kernel reads it); one more custom_vjp per chunk with the fused activation. A fused backward that swaps d(gate) and d(up), and a pair argsort with reversed ties under a dropping buffer, are caught. - On an expert 4 x tensor 2 mesh the layer runs without the key and raises at trace time when the key reaches it past the validator. - Each validator rejects its key for each unmet requirement, naming it.
… experts (moe_roe_packed_routing_gather) With ring_of_experts_local_routing and moe_chunk_pipeline, every token chunk routes its own tokens and all-gathers its top-k weights and expert ids over the expert-parallel axis: two all-gathers, two SparseCore offloads, per chunk, and in the backward one weight-cotangent reduce-scatter per chunk. moe_roe_packed_routing_gather routes all of the layer's local tokens once and all-gathers weights and ids as ONE float32 [batch, seq, 2 * top_k] array (pack_routing_for_all_gather: ids are < num_experts or -1, exact in float32), which each chunk slices: one gather per layer and one reduce-scatter in the backward. Routing is per token, so each slice holds what the chunk would have computed: bit-identical with float32 routing weights (float32_weight_sum, the default). With bfloat16 weights the values round-trip exactly and their cotangent is reduce-scattered in float32 instead of bfloat16. The all-gather is the ring of experts' own (_ep_all_gather, now defined once for roe_ag_and_route and this gather), so a SparseCore pin applies to both alike. The per-chunk load-balance loss is computed as without the key, from this shard's own ids of the chunk, so the loss and the output are bit-identical; the router gradient then agrees to float32 reduction order (the top-k and the loss terms reach the logits in another order; 1.2e-7 relative on CPU). The gathered array is a routing map: tagged moe_routing_maps, so moe_routing_maps=device saves it and the rematerialized forward skips the gather; untagged under remat. The validator (the shared form) requires sparse_matmul, te_moe_block=False, ring_of_experts_local_routing and moe_chunk_pipeline (the packed gather replaces the pipelined loop's per-chunk gathers only), and rejects use_random_routing (random routing draws per call, so one draw for the layer differs from one per chunk). If XLA strips optimization barriers the pipeline is skipped with a warning, and the packed gather too. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, load_balance_loss_weight 0, moe_routing_maps=device), switched by an environment variable in a build that a TPU lowering check shows compiling to the same program as this commit: regular micro-batch 9.938 s (arm o275x-mf-pra-prof), 98.68 s per step, 0 dropped tokens. That pod ran no flags-off control; against the flags-off control of another pod (o275x-gdn-ctl, 10.033 s, 99.50 s per step) that is -0.095 s per micro-batch, with a spread of the same control across pods of +/-0.02-0.03 s per micro-batch. The environment-switch build dropped the load-balance loss (it returned None for it); this change computes it, which the measured run, at weight 0, does not exercise. The timing arms run with the benchmark's log-prob gates on (loss and gradient norm read 0), so they say nothing about numerics. TPU numerics: a 3-step A/B with the gates off, against a reference run and a replicate of it, passes all six checks with this key in both measured stacks: with moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp and moe_router_grad_from_wo_input (run nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6 against a replicate spread of 1.9e-5), and with the fused SiLU GMM in place of the fused activation (run nxs6silu: gradient-norm difference 2.5e-6). Caveats: stack level, not this key alone; the environment-switch build, not this branch's binary. Tests (moe_roe_packed_routing_gather_test.py, new): - pack_routing_for_all_gather / unpack_routing round-trip bit for bit with ids up to 4095 and -1 (bfloat16 holds every integer only up to 256), float32 and bfloat16 weights, and the packed array is float32. - On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local routing, pipelined loop, two chunks), bit-identical output, loss and every gradient against the per-chunk gathers: real routing, padding excluded from routing, forced ids with a dropping buffer, saved routing maps, with moe_routing_index_math=count, and with the benchmark's MoE settings (expert-weight prefetch, row-major reduce-scatter, ragged_buffer_factor 2.75, prefused weights with fp8 fake-quant). With the load-balance loss on: output and loss bit-identical, gradients within 1e-5 of max (measured 1.2e-7). bf16 routing weights: within relative L2 1e-2 (bit-identical on CPU). - Structure: the plain loop all-gathers two [..., top_k] arrays per chunk and the packed one exactly one [..., 2 * top_k] array, tagged moe_routing_maps only under moe_routing_maps=device; nothing is tagged moe_route. - A packed gather shifted along the sequence, and the packed path with its load-balance loss dropped, are caught. - The validator rejects the key for each unmet requirement, naming it, and accepts the load-balance loss. The SparseCore-pinned all-gather is a TPU-only placement and is not reached on CPU.
The up projection wrote the [buffer, 2 * mlp] GMM output and a separate
elementwise pass read it back to apply the SwiGLU, over every buffer row,
the ragged buffer's padding included.
moe_wi_gmm_fused_swiglu calls megablox ops.gmm_silu instead: the tokamax
gmm_v2 kernel with fuse_act="silu" applies the activation to the float32
accumulator of the rows it computes and writes the [buffer, mlp]
activation. When differentiated it also writes the gate and up halves
([buffer, 3 * mlp] in all), rounded to the output dtype exactly as the
unfused GMM rounds its output (the new return_pre_act option of gmm_v2:
both halves go through the store path of the 2D output, with its
zero-init DMAs and partial-row carry), and the custom VJP is the unfused
backward: the SwiGLU VJP on those bf16 halves, then gmm's own DLHS / TGMM
on [d gate | d up]. ops._gmm_silu_bwd calls gmm's backward with keyword
arguments, as gmm's custom VJP would for a plain bf16 gmm_v2 GEMM. Works
with split weights and with the whole [w0 | w1] of moe_unsplit_prefused_wi.
The layer hands the fused GMM the routing operands explicitly
(get_gmm_operands_for_local_experts, which get_gmm_for_local_experts now
wraps; the key-off program is unchanged).
Numerics change: silu reads the f32 accumulator instead of its bf16
rounding (the more precise value), so the activation, and through it
everything downstream, differs by rounding.
The validator (the shared form) requires prefused sparse-matmul training,
use_tokamax_gmm and use_gmm_v2, silu + linear activations without
clipping (and not GPT-OSS), no weight quantization, MLP bias, embedding
chunks, flat FSDP weights, 3D TC GMM, token all-gather quantization or
chunk weight-gradient accumulation, and a remat policy that saves neither
half of the routed experts' up projection (the rule
moe_prefused_swiglu_vjp uses: the fused GMM never forms the
[buffer, 2 * mlp] output, so such a policy would silently stop saving the
MoE tensor; remat_policy=full and the shared expert's own names are
fine). moe_prefused_swiglu_vjp is rejected next to it, as it would do
nothing. Weight quantization objects and an explicit weight all-gather,
which only the built layer knows, raise at trace time.
Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build whose kernel files differ from these
only in docstrings, comments and line breaks (AST-identical) and whose
program a TPU lowering check matches to this commit. Same pod, stacked on
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather and
moe_router_grad_from_wo_input (the fused activation has nothing to act on
next to this key, and the validator rejects the two together): regular
micro-batch 9.494 -> 9.418 s (-0.0755 s), 86.970 -> 86.253 s per step
(-0.716 s), and the step's second micro-batch, a proxy for the step
boundary, 9.450 -> 9.375 s, so the boundary pays nothing extra (arms
o275x-s6silu-l3-prof against o275x-s6c-l3-prof, one pod, on a build with
the staged step boundary, which is why its steps are shorter than the
other arms of this series). The same control measured in different pods
spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with the
benchmark's log-prob gates on (loss and gradient norm read 0), so they say
nothing about numerics.
TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks for this key stacked with
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run
nxs6silu: N1 1.4e-6, gradient-norm difference 2.5e-6, within the
replicate spread of 1.9e-5 measured with the same reference). Caveats:
stack level, not this key alone; the environment-switch build, not this
branch's binary. megablox_gmm_silu_test.py, the kernel's own TPU test,
has not been run yet, so that stack run is the kernel's only TPU
evidence.
Tests:
- megablox_gmm_silu_test.py (new, TPU only, skipped on CPU): gmm_v2's
gate and up equal the unfused output bit for bit, rows outside the
groups are zero, act is silu of the f32 accumulator, the act-only call
matches, with an EP shard (group_offset); gmm_silu's VJP equals the
unfused path's value for value.
- moe_wi_gmm_fused_swiglu_test.py (new, CPU):
- ops.gmm_silu with stand-ins for the Pallas kernel and gmm's TPU
backward only: the forward tile is tiling[2] halved, act alone is
written undifferentiated and act + gate + up differentiated, the
backward hands gmm's backward [d gate | d up] of the SwiGLU on the
bf16 halves bit for bit, with this op's operands as residuals and its
DLHS options on the right keywords (each option set alone, both
group offsets); an odd forward tile is rejected.
- The shrunk Qwen3.5 MoE on 8 CPU devices, with jnp stand-ins for gmm
and gmm_silu (gmm_silu as the unfused composition of the same
stand-in): the layer with the key is bit-identical to the layer
without it (output, loss, every gradient): plain and pipelined loop,
dropless (non-zero group offsets), padding excluded, the whole
prefused wi, with the other MoE keys of this series, and with those
and the benchmark's MoE settings. One gmm_silu call per chunk, with
the plain path's wi tiling and the config's DLHS options
(moe_gmm_v2_dlhs_transpose_rhs True in the benchmark case). Stand-ins
that swap the gate and up halves, or ignore the group offset, are
caught.
- Key off: the dropless ring-of-experts ragged sort matches the ring of
experts without it (use_ragged_sort=False), and fails to once its GMM
ignores the group offset, so the routing operands that both paths now
take from get_gmm_operands_for_local_experts are checked against a
reference this change does not touch.
- The validator rejects the key for each unmet requirement, naming it.
…ward that reads the GEMM input (moe_router_grad_from_wo_input)
On the ring-of-experts ragged path the down projection y = h @ wo (per
expert, over the sorted rows) feeds ring_ragged_unsort, whose backward
needs y for the routing-weight gradient <g, y>. Under remat the layer's
backward therefore recomputes the down-projection GEMM only to form y,
and the combine backward makes a pass over two [rows, hidden] tensors that
also writes dy = w * g.
moe_router_grad_from_wo_input runs the activation, the down projection and
the unsort as one custom-VJP op (ragged_sort.ring_ragged_unsort_down_projection).
Its forward is the GEMM followed by ring_ragged_unsort's forward, so the
output is bit-identical. Its backward never reads y: row j of the GEMM is
linear in h_j and the weight is a scalar per row, so
<g, h wo> = <g wo^T, h>, (w g) wo^T = w (g wo^T) and h^T (w g) = (w h)^T g.
It takes the GEMM's own VJP at the input w * h with the unweighted
cotangent g (u = g wo^T, dwo = (w h)^T g), then dh = w u and the routing
gradient <u, h>. The GEMMs keep their shapes, the recomputed forward GEMM
has no reader and XLA drops it, and the weighting and the dot run over
[rows, mlp] tensors. In the pipelined chunk loop, tie 2 moves into the op
as a forward-only barrier on the GEMM output. The cotangent rows outside
this shard's groups are left unmasked: the GEMM backward never mixes them
into group rows and every output is masked.
The op shares its code with ring_ragged_unsort: the unsort's forward and
the fan-out half of its backward are now helpers (_ring_unsort_window,
_ring_unsort_combine, _ring_unsort_fan_out, with _unsort_kwargs for their
kernel options) that both call, so the fused forward is the unsort's by
construction. The helpers keep the unsort's ops in their order: lowered
for TPU v7x (forward plus backward of a shrunk Qwen3.5 MoE layer, the
benchmark's MoE settings, also the dropless retry program), the key-off
program is the same as the parent commit's and every key-on program the
same as before this refactor.
Numerics change: each gradient keeps one rounding to the activation dtype,
at a different point (w * h instead of w * g; dh after the product with w
instead of before the GEMM; the routing gradient from g wo^T instead of
y), so the gradients agree with the two ops apart to bf16 rounding, not bit
for bit. On the CPU layer below they differ by 1e-3 to 4.2e-3 relative L2.
On the op alone, against the same function in float32, the fused
gradients are as far as the two ops' are: the ratio of the two errors is
0.99-1.01 for dh and dwo and 0.94-1.07 for the routing gradient, over 6
seeds with a full and a dropping buffer.
The validator (the shared form) requires the ring-of-experts ragged sort
(sparse_matmul, te_moe_block=False, use_ring_of_experts and
use_ragged_sort) and a down projection that is the bare GEMM of the
sorted rows: no tensor parallelism (also checked on the mesh at trace
time), MLP bias, weight quantization, TC ragged sort, chunk weight-gradient
accumulation, pre-quantized unsort backward, token all-gather quantization
or embedding chunks; and a remat policy that saves no name of the routed
experts' down-projection output (mlpwo, moe_mlpwo: the minimal* and
save_dot_except_mlpwi presets, or custom with either saved or offloaded),
the rule of the series: the op never forms that output, so such a policy
would silently stop saving the MoE tensor. It composes with
moe_routing_index_math, the prefused-weight keys and
moe_wi_gmm_fused_swiglu. The layer raises if it is asked to return the
weights for chunk weight-gradient accumulation; _moe_body's three exits
now share one construction of its outputs.
Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build that a TPU lowering check shows
compiling to the same program as this commit: regular micro-batch 10.022
-> 9.765 s (-0.257 s; arm o275x-mf-a14-prof against its in-pod control
o275x-mf-ctl), 99.25 -> 97.05 s per step, 0 dropped tokens in both. All
five MoE keys of this series together (moe_routing_index_math=count,
moe_unsplit_prefused_wi, moe_roe_packed_routing_gather, this key and
moe_wi_gmm_fused_swiglu), in one pod on a build with the staged step
boundary: 10.0445 -> 9.510 s per micro-batch (-0.535 s), 91.87 -> 87.04 s
per step (o275x-s5moe5-l3-prof against o275x-s5c-l3-prof). The same
control measured in different pods spreads by +/-0.02-0.03 s per
micro-batch. The timing arms run with the benchmark's log-prob gates on
(loss and gradient norm read 0), so they say nothing about numerics.
TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks with this key in both measured
stacks: with moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp and moe_roe_packed_routing_gather (run nxs6c: loss
relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6
against a replicate spread of 1.9e-5), and with the fused SiLU GMM in
place of the fused activation (run nxs6silu: gradient-norm difference
2.5e-6). Caveats: stack level, not this key alone; the environment-switch
build, not this branch's binary.
Tests (moe_router_grad_from_wo_input_test.py, new), on 8 CPU devices:
- The op on a 4-device EP mesh with a jnp grouped matmul as the down
projection, the full buffer and a packed buffer that drops tokens,
"sort" and "count": output bit-identical to the two ops apart, every
gradient within relative L2 1e-2 of theirs; a forward tie comes back
unchanged with its cotangent passed through and is tied to the GEMM
output by a barrier; "count" without the forward permutation raises.
- Precision of the op alone (its input h given, so only the backward
rounds), against the same function written densely in float32 (the
ragged gather-reduce returns bf16, so the two ops cannot be the
reference): each fused gradient at most 1.10x as far as the two ops',
on 6 seeds x {full buffer, dropping buffer}, with routing weights where
bf16 rounds them worst. A backward that rounds the routing weight to
bf16 measures 1.31-1.34 and one that shrinks the routing gradient by
0.5% 3.0-3.3; both are caught.
- The shrunk Qwen3.5 routed MoE: output and loss bit-identical to the key
off, gradients within relative L2 1e-2: plain and pipelined loop,
dropless, "count" with padding excluded, saved routing maps, prefused
weights with the whole wi and the fused SwiGLU VJP, and the benchmark's
MoE settings with every compatible key of the series on against all
off. The compiled rematerialized backward has exactly
num_moe_token_chunks GEMM loops fewer (19 -> 17 for 2 chunks), the
pipelined forward keeps its four barriers, and the next chunk's routed
tokens are read from the op's tie output.
- On an expert 4 x tensor 2 mesh the layer runs without the key and
raises at trace time when the key reaches it past the validator.
- A fused backward that drops the routing weight is caught at the op and
in the layer; the validator rejects the key for each unmet requirement,
naming it.
…kv_dtype) The GDN kernel path's Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) builds its halo-extended input [conv_state | qkv] in float32 and slices the gdn_conv_kernel_dim taps out of it. On the Qwen3.5-397B RL trainer that is an f32[1, 16387, 12288] = 805.5 MB copy per GDN layer in the backward pass, written once and read once per tap. With gdn_conv1d_input_in_qkv_dtype=True the input stays in qkv's dtype and each tap is cast to float32 where it is read, so XLA fuses the cast into the consumer and the float32 copy is never written. No optimization barrier. The option reaches the Conv1D through gdn_decoupled_conv1d's new static argument conv_input_in_qkv_dtype, which scopes it around the op's primal, forward and backward rules (the backward rule is traced after the model's forward returns, so only the custom_vjp's static arguments reach it). Numerics: the same float32 values in the same order, so op by op every output is bit-identical. Compiled, the forward outputs and the qkv, conv-state and bias gradients stay bit-identical; the conv weight gradient (a float32 reduction over batch x sequence) can be summed in a different order once XLA fuses the cast into it (CPU: 11 of 64 elements, <= 5.1e-6 relative, in 2 of 16 configurations). Default off. Requires use_gdn_kernel=True (validated).
![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 introduces several advanced performance optimizations and configuration keys for Mixture of Experts (MoE) and GatedDeltaNet (GDN) layers in MaxText. Key MoE enhancements include keeping prefused weights unsplit during quantization and FSDP all-gather (moe_unsplit_prefused_wi), fusing gated activations (moe_prefused_swiglu_vjp and moe_wi_gmm_fused_swiglu), packing routing gathers (moe_roe_packed_routing_gather), and optimizing backward passes to avoid down-projection GEMM recomputation (moe_router_grad_from_wo_input). For GDN, it adds support for a 128 chunk size, a new triangular-inverse algorithm (xor_dnc), and an option to keep halo-extended inputs in their native activation dtype (gdn_conv1d_input_in_qkv_dtype) to save HBM. Additionally, it integrates support for freezing parameters via trainable_parameters_mask in the training engine. Extensive unit tests are added to validate these exact mathematical rewrites and optimizations. I have no feedback to provide as there are no review comments.
Follow-up to the six commits of this PR (moe_routing_index_math, moe_unsplit_prefused_wi / moe_ragged_sort_unique_pair_argsort / moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather, moe_wi_gmm_fused_swiglu, moe_router_grad_from_wo_input, gdn_conv1d_input_in_qkv_dtype). No new keys and no behaviour change. Refactors: - gmm_v2 kernel: inner_kernel's 2D store is now the helper _store_group_rows, which the activation output and the return_pre_act gate and up halves all call, instead of the halves using a copy of the store; the 3D stores share its row mask and partial-row carry (_mask_group_rows, _carry_partial_rows). - GDN api.py: gdn_decoupled_conv1d and its forward and backward rules keep their bodies at the original indentation. A small decorator enters conv_input_in_qkv_dtype_scope from the static argument for the primal and forward rule; the backward rule runs its argument shim, then its body (_gdn_decoupled_conv1d_bwd_impl) inside the scope. The api.py diff against the parent shrinks from 848 lines to the real changes. - moe.py: _use_prefused_swiglu_vjp inlined (it only returned the config flag); two pointer comments dropped. - types.py: the MoE validators' requirement helpers come before the validators. Validators and guards: - Both chunk loops now refuse moe_router_grad_from_wo_input with chunk weight-gradient accumulation at trace time, before any work: the pipelined loop had no such check, and the plain loop's came after the up projection, which can raise a different error first (the config validator rejects the pair in both cases). - gdn_conv1d_input_in_qkv_dtype is read directly instead of through getattr with a default, in its validator and in model_runner (gdn_cp_head_sharded_test.py's SimpleNamespace config gains the field). Tests: - moe_router_grad_from_wo_input_test.py: chunk_wgrad_guard, both chunk loops raise with moe_accumulate_chunk_wgrad past the validator. The op precision check also bounds both paths' output and gradients against the float32 reference (relative L2 1e-2, measured at most 2.4e-3): the two paths share the unsort helpers, so a defect there leaves the fused/two-ops ratio at 1, and the forward on the dropping buffer was unchecked. - moe_wi_gmm_fused_swiglu_test.py: GmmSiluVjpTest also with heuristic tiling (use_gmm_v2_heuristic_tiling must reach gmm's backward). - moe_glue_rewrites_test.py and moe_routing_index_math_test.py: validator cases for tensor_sequence and autoregressive parallelism, for every minimal* remat preset, and an error that names two unmet requirements. - Test docstrings no longer cite history or TPU runs that did not happen (megablox_gmm_silu_test.py is TPU-only and not yet run on TPU); each _BENCH_FLAGS is commented as the subset of the benchmark's MoE settings that file needs; one-line docstring summaries. - gdn_conv1d_input_in_qkv_dtype_test.py: dict literals, a docstring for _conv_inputs, the repo's license header. Docs and comments: - base.yml, types.py and the compute_conv1d_bwd.py comment describe gdn_conv1d_input_in_qkv_dtype without model-specific numbers, and say the compiled bit identity was measured on CPU. - nnx_decoders.py and types.py point at each other for the MLP checkpoint names of the remat presets (_REMAT_PRESET_MLP_NAMES), so the two stay in sync. Lint (pylint 3.3.8, pyink 24.10.1, codespell 2.4.1, yamllint 1.35.0, repo config): the changed lines add no messages against the base; the 3 pyink hunks and 5 pylint messages of the GDN test are gone. Program identity against the previous head: unchanged, by normalized TPU v7x lowering, except where intentionally changed (no change was meant to alter a program).
…ews (gdn_conv1d_dense_token_stack) The Pallas GDN forward kernel (kernels/gdn/wrapper.py fused_conv1d_gdn) convolves a chunk one token at a time and stacks the per-token [1, 1, dim] float32 rows with jnp.stack(axis=1) (kernels/gdn/compute_conv1d.py causal_conv1d). The stacked [1, chunk, 1, dim] value has a second-minor dimension of 1, and the TPU compiler of the October 2026 Pathways build (and of the libtpu 0.0.51 nightlies) lays it out in (8, 128) tiles padded from 1 to 8 sublanes. The SiLU applied next inherits that layout, so its exp and reciprocal run on 8x as many vector registers as the data fills. Older compilers keep the value dense. With gdn_conv1d_dense_token_stack=True each row is viewed as [dim / 1024, 8, 128], which is free because the row already fills dim / 1024 float32 registers; the views are concatenated along the leading axis and the result is reshaped to [1, chunk, 1, dim] (compute_conv1d.stack_conv1d_tokens). The same values in the same order; the compiler keeps the value dense on old and new builds alike. The option reaches the kernel through gdn_decoupled_conv1d's new static argument conv_dense_token_stack. Its primal and forward rule pass it to every forward-kernel call: the local call, both sequence-sharded context-parallel passes, and the forward rule's replay under remat (the backward rule never runs the forward kernel). fused_conv1d_gdn hands it to the per-sequence kernel only: its batched (decode) kernel, which convolves one token per sequence, keeps jnp.stack for those one-token rows, and the tokamax Op wrapper (kernels/gdn/pallas_mosaic_tpu.py) does not take the option. MaxText's GDN layer runs neither, since gdn_decoupled_conv1d calls fused_conv1d_gdn prefill-only. The option is left out of the kernel name and metadata, so outside the kernel body the lowered program is unchanged and profiles keep the kernel name. Numerics: every element lands where jnp.stack puts it, so the kernel's outputs do not change. On CPU, causal_conv1d and the SiLU after it give the same bits with and without the option (the kernel itself does not run on CPU). The TPU numerics A/B of the measured build below passed all six gates (nxs6gr). The validator requires a decoder with Gated DeltaNet layers (qwen3_next or qwen3_5), use_gdn_kernel=True, and a width that each forward-kernel call convolves (2 * key heads * key head dim + value heads * value head dim; one rank's heads under head-sharded GDN context parallelism) that is a multiple of 1024. Like the existing sequence-length check, it counts a context parallelism left at -1 (filled in when the mesh is built) as 1; fused_conv1d_gdn checks the per-call width again when the layer is traced, which catches such a config. Where the width is not a multiple of 1024 the measured build below silently kept jnp.stack; this change rejects the config instead. Default off. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, ragged_buffer_factor 2.75, staged step boundary), with the same change behind an environment switch (MAXTEXT_GDN_CONV_DENSE_STACK=1) on the stack-6 base: 85.731 s per step, regular micro-batch 9.3552 s (o275x-s6gr-l3-prof), against 86.926 s / 9.4910 s for the same-pod arm o275x-s6n2-l3-prof: -1.195 s per step. That arm's only change, a 2D GDN out_proj, is a no-op (XLA already canonicalises the 4D out_proj into the same 2D dot; TensorCore busy time -2 ms per micro-batch). The stack-6 control in another pod measured 86.970 s / 9.494 s (o275x-s6c-l3). Equivalence to the measured build, on CPU: lowered for a TPU7x target (no TPU needed), the GDN layer's forward and backward program with the key on is the measured build's program with its switch on, and with the key off is the base program (StableHLO and every Mosaic kernel body, compared without source locations). 18 configurations: no context parallelism, head-sharded context parallelism 2, sequence-sharded context parallelism 4 (also at the 397B's 16 key and 64 value heads), with and without packing and remat, and gdn_conv1d_input_in_qkv_dtype off and on. In each, the key changes only the forward-kernel bodies. Outside that layer the two differ: the measured build's switch also regrouped the batched kernel's one-token rows (at a decode tile of 1) and reached the tokamax Op wrapper. Tests (gdn_conv1d_dense_token_stack_test.py, new), on CPU: - stack_conv1d_tokens(dense=True) equals jnp.stack(axis=1) bit for bit, and dense=False is jnp.stack itself (the same jaxpr); causal_conv1d and the SiLU after it give the same bits with the option at widths 1024 and 12288, packed or not, with and without conv bias, for full and partial chunks; a token-minor regrouping is caught. - Traced and lowered for a TPU7x target on the CPU host, every forward-kernel call assembles its conv output densely exactly when the option is on: fused_conv1d_gdn; the custom VJP's primal, forward rule and remat replay; both sequence-sharded CP passes, under the forward rule and, forward only, under the primal; head-sharded CP; and model_runner reading the config key. The batched (decode) kernel is the same program with and without the option. The Mosaic lowering keeps the kernel name and changes only the kernel body. The option composes with gdn_conv1d_input_in_qkv_dtype, and the backward rule accepts every positional form. An option the conv ignores, one the forward rule drops and one the primal drops each fail every check they break. - The width check and the validator reject the option where it would do nothing, including one rank's width under head-sharded CP. The validator's copy of model_runner's head-vs-sequence CP rule agrees with model_runner at every GDN CP mode, degrees 1-8 on each context axis, nine head-count pairs and head dims 128 and 256. The kernel name and metadata do not change.
parambole
added a commit
to google/tunix
that referenced
this pull request
Oct 15, 2026
… join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now selects the trainer cell that measured fastest on the RL trainer bench with every numerics gate passing. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains sixteen keys and loses one: - wi_tile_dlhs_mlp_dim=2048, decoder_layer_input=device, moe_routing_maps=device and moe_gmm_v2_dlhs_transpose_rhs=true: the K tile of wi's input-gradient GMM matched to the fused gate and up weights, each decoder layer's input kept in HBM, the backward reusing the forward's MoE routing maps, and the input-gradient GMMs reading the expert weights through the kernel's transpose_rhs. - moe_routing_index_math=count, moe_unsplit_prefused_wi=true, moe_roe_packed_routing_gather=true, moe_router_grad_from_wo_input=true and moe_wi_gmm_fused_swiglu=true: the MoE keys of AI-Hypercomputer/maxtext#5758 in the set its commits recommend (the routing index maps by counting instead of argsorts, the fused wi kept whole up to the GEMM, one packed routing all-gather per layer, the routing-weight gradient from the down projection's input, the SwiGLU in the wi GEMM's epilogue). - gdn_conv1d_input_in_qkv_dtype=true and gdn_conv1d_dense_token_stack=true: the GDN Conv1D backward without its f32 copy of the input, and the forward kernel's Conv1D output built from dense per-token views. - wi_tile_fwd_batch_seq, wi_tile_drhs_batch_seq, wo_tile_fwd_batch_seq, wo_tile_dlhs_batch_seq and wo_tile_drhs_batch_seq=512 (from 256): row tiles of 512 for the five MoE GMMs with 1024 tiles. wi_tile_dlhs_batch_seq stays 256 next to its 2048 K tile. - sa_use_base2_exp=false leaves the block, so the trainer keeps the recipe's sa_use_base2_exp=true: base-2 exponentials measured 0.025 s faster per micro-batch (o275x-gx-b2e), and their numerics run passes (nxgx-b2e). In mlperf_gbs1024_submission.sh, which turns TRAINER_TUNED_FLAGS on, the switch also sets three defaults, each still overridable: - COMPUTE_LOGPS_CHUNK_SIZE=4096 (was 1024), and 1024 when TUNIX_LOGPS_SHARDED_HEAD turns the sharded head off (0, false, off or no, in any case): chunks of 4096 were measured with the head only, and without it the log-prob program failed to compile on TPU at chunks of 2048 and 4096. - TUNIX_LOGPS_SHARDED_HEAD=fused (was local). TRAINER_TUNED_FLAGS=false keeps local. - PACK_PAD_POSITION=start: each packed row puts its padding first. Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step, mean of steps 1 to 4, no token dropped in any arm): - before these keys, the fused head and pad first: 106.30 s per step (o275x-hd), 2,774 tokens/s/chip; - this cell: 82.671 and 82.670 s per step in two pods (o275x-c9pfxhv-l3), 3,567 tokens/s/chip. Part of the difference is the staged step boundary (AI-Hypercomputer/maxtext#5718), which this recipe already turns on and o275x-hd did not run; per regular micro-batch (the median of micro-batches 3 to 9), which the boundary does not change, about 10.77 -> 9.01 s; - within it, each against a control in the same pod: padding first 83.285 -> 82.043 s per step (-1.242 s) and the fused head 83.253 -> 82.708 s per step (-0.545 s). The recipe lists the measurement of each key. The measured build switched the MaxText changes with environment variables, and it also had one attention change on that neither MaxText PR carries (attention's load-balancing permute hoisted; alone 99.504 -> 99.267 s per step). The keys themselves, on the MaxText change's own code (2d82c5199) without gdn_conv1d_dense_token_stack, the fused head and pad first, measured 85.295 and 85.251 s per step (o275x-prc-l3-prof, o275x-prc-l3-r), within 0.23 s of what the environment-switch arms predicted. Both arms set moe_route=device, which nothing in #5698 or #5758 tags, so moe_route stays out of the block as before; and the stop of the frozen router's gradient, an environment switch in the measured build, is unconditional in #5758, with no key. Numerics (loss gates off, random advantages, three steps, against two reference runs whose gradient norms differ by 1.9e-5 relative): all six gates pass for this cell (nxc9pfxhv) and for the keys on the MaxText change's own code (nxprc). moe_wo_gmm_fused_dswiglu stays out: next to the fused SiLU GMM at these tiles it computed a wrong SwiGLU gradient, and its fix is still being verified on TPU. moe_routing_maps and moe_gmm_v2_dlhs_transpose_rhs exist in AI-Hypercomputer/maxtext#5698, and the seven other new MaxText keys in the MaxText change of AI-Hypercomputer/maxtext#5758 (on top of #5698); none of the nine is defined at this tree's pin (e526ea02), and MaxText refuses a key it does not define when the trainer starts. TRAINER_TUNED_FLAGS=true, the default of mlperf_gbs1024_submission.sh, therefore needs requirements/maxtext_requirements.txt moved to a MaxText with #5758 before an image built from this tree runs it. TRAINER_TUNED_FLAGS=false and the other 397B v7x recipes pass none of the nine. Tests (mlperf_submission_launcher_test.py): the exact-set test expects 23 keys, _MAXTEXT_DEFAULTS has the nine new keys' defaults, and the switch-off and other-recipe tests check that none of the nine is passed; the default launch has base-2 exponentials on, chunks of 4096, the fused head and --pack_pad_position=start; TRAINER_TUNED_FLAGS=false keeps the local head and the padding at the end; an explicit pad position, head or chunk size wins; TUNIX_LOGPS_SHARDED_HEAD=0, false, OFF or No gives chunks of 1024; the other 397B v7x recipes pass no pad position. Fifteen deliberately wrong copies of the recipes each fail at least one of these tests: a key dropped from the block or given the wrong value, one tile left out, wi_tile_dlhs_batch_seq raised to 512, sa_use_base2_exp=false kept, moe_wo_gmm_fused_dswiglu or moe_route=device added, a new key passed outside the block at its default, the local head under the switch, the chunk fallback removed or made case-sensitive, chunks of 1024 by default, pad first outside the switch, pad_align passed, and both recipes as they were before this change.
parambole
marked this pull request as ready for review
October 15, 2026 17:42
parambole
requested review from
NuojCheng,
RissyRan,
aireenmei,
bvandermoon,
darisoy,
gagika,
gobbleturk,
huytransformer,
igorts-git,
jiangjy1982,
michelle-yooh,
richjames0,
shralex,
shuningjin,
xibinliu and
zxhe-sean
as code owners
October 15, 2026 17:42
parambole
requested review from
A9isha,
SurbhiJainUSC,
abhinavclemson,
dipannita08,
hengtaoguo,
khatwanimohit and
vipannalla
as code owners
October 15, 2026 17:42
parambole
added a commit
that referenced
this pull request
Oct 16, 2026
At d66a820, the MoE, GDN and model files of this branch come from #5758, which is built on all of #5698. But this branch ports only part of #5698 onto atwigg/mlperf. Several files it needs are still the atwigg/mlperf versions, so the trainer does not start: - utils/sharding.py lacks CompoundLogicalAxis, so `import maxtext.layers.moe` raises ImportError; - utils/max_utils.py lacks is_eval (moe.py); - models/weaver.py is missing, and common_types.py lacks DecoderBlockType.WEAVER, which nnx_decoders.py maps for every decoder; - kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py builds GmmConfigs without transpose_rhs, which the gmm_v2 kernel requires: TypeError when the MoE backward compiles. This brings src/maxtext to the tree it was developed and measured on: #5758 (20c877a, on #5698 + #5718) plus this branch's own commits (the GDN segment-path skip, the GDN forward zero-fill skip, gather_scanned_kernels_once and the cp-as-ep-dense rules), which apply to it cleanly.
4 tasks done
uwlei
added a commit
that referenced
this pull request
Oct 23, 2026
… (supersedes #5799, includes #5571) (#5885) * GDN: rename a dense-stack test local that codespell flags codespell 2.4.1 with the repo's pre-commit arguments reads the local alog in gdn_conv1d_dense_token_stack_test.py as a misspelling of "along" (six times), and the hook's -w would rewrite it. It is the A_log parameter, so it is now a_log, as everywhere else in the GDN code. Test only. * MoE: SwiGLU backward in the wo DLHS GMM epilogue (moe_wo_gmm_fused_dswiglu) The DLHS GMM of the MoE down projection (wo) wrote u = g @ wo^T, the activation's cotangent, to HBM in bfloat16, and separate elementwise passes read it back: the SwiGLU VJP over [buffer, mlp] and, with moe_router_grad_from_wo_input, the routing-weight product w * u and the dot <u, h> that is the routing weight's gradient. Only then did d[gate | up] reach the up projection's backward. moe_wo_gmm_fused_dswiglu runs the SwiGLU backward in that kernel's epilogue. The tokamax gmm_v2 kernel's new fuse_dact="swiglu" option turns the float32 accumulator u of each row into d[gate | up] (apply_dact_fn: d gate = s u up silu'(gate), d up = s u silu(gate)) and writes that [buffer, 2 * mlp] array instead of u. It reads [gate | up] as an input, an optional per-row scale s (the routing weight), and with return_dact_dot also writes <u, silu(gate) * up> per row. megablox ops.gmm_dswiglu is the down projection with this backward: gmm's forward, and a backward that returns d[gate | up] and a zero cotangent for the activation, whose gradient the layer stops (gate_up carries it). With moe_router_grad_from_wo_input, ring_ragged_unsort_down_projection takes [gate | up] too and calls gmm_dswiglu_bwd with the routing weight as the row scale and the dot as the routing-weight gradient. The up projection hands its whole [gate | up] output on; with moe_wi_gmm_fused_swiglu the fused wi GMM writes [gate | up] as one array next to the activation (gmm_silu(fuse_bwd_dswiglu=True), gmm_v2 return_pre_act="concat") and its backward takes the epilogue's d[gate | up] straight to gmm's DLHS / TGMM. The forward is unchanged. Numerics change: the SwiGLU derivative reads the float32 accumulator u where the unfused path reads its bfloat16 rounding (u is a kernel output there), and runs in float32 with one rounding of d[gate | up] (unfused: the bf16 VJP). With moe_router_grad_from_wo_input the routing weight multiplies the float32 u (unfused: bf16(w * bf16(u))), and the routing weight's gradient is <u, silu(gate) * up> in float32 from the bf16 gate and up (unfused: <bf16(u), h> with h the bf16 activation the forward fed to the down projection; with moe_wi_gmm_fused_swiglu, silu of the wi accumulator). The output, the loss and the wo gradient are bit-identical; the other gradients differ by bf16 rounding: 1.6e-3 to 5.5e-3 relative L2 on the CPU layer below, with the kernel's own epilogue. Fused SiLU with a forward tile of 1024, and stack 9. The "concat" form of return_pre_act stores each n tile's accumulator [gate_j | up_j] as one block, which is [gate | up] only with one tile per half, and the measured build accepted more. Stack 9 (moe_wi_gmm_fused_swiglu with this key) ran wi_tile_fwd_mlp_dim=1024 against Qwen3.5-397B's MoE MLP dim of 1024: two tiles per half, so the fused wi GMM wrote [gate_0 | up_0 | gate_1 | up_1] and the epilogue read [gate_0 | up_0] as the gate. On the CPU layer, with the kernel's block layout modelled, the gradients came out 0.89 to 1.07 relative L2 from the unfused ones, against 2.3e-3 to 5.4e-3 with one tile. So the stack-9 numerics results are not valid: nxs9p and nxs9pf ran with that tile and passed the 6 numerics gates anyway, because those scalar gates do not see gradient direction. Stack 9's speed numbers stand. Here make_gmm_configs raises for "concat" unless one n tile spans each half, as the fused DLHS raises unless one n tile spans its own [d gate | d up], and the config validator rejects both tile settings before the model is built (below). Only the float32 epilogue that was measured is ported. The validator requires prefuse_moe_weights, use_tokamax_gmm and use_gmm_v2, sparse matmul, silu + linear activations without clipping, no MLP bias, no quantization, and none of: embedding chunks, flat FSDP weights, the TC ragged sort, the 3D TC GMM, token all-gather quantization, chunk weight-gradient accumulation. The remat policy must not save the routed experts' gate or up half (mlpwi_0, moe_mlpwi_0, moe_mlpwi_1; the minimal* policies save mlpwi_0): the whole [gate | up] goes on, so those names are never formed. moe_prefused_swiglu_vjp next to it would do nothing and is rejected. With explicit tiles and a MoE MLP dim no mesh axis shards (no tensor, tensor_sequence or autoregressive parallelism), it also checks the kernels' one-tile rules: wo_tile_dlhs_mlp_dim must be the MoE MLP dim, and with moe_wi_gmm_fused_swiglu wi_tile_fwd_mlp_dim at least twice it. The stack-9 argv therefore fails at config time, not at the first train-step trace; with heuristic tiling or a sharded MLP dim the kernels still check at trace time. Weight quantization objects and an explicit weight all-gather, which only the built layer knows, raise at trace time. The program is the measured build's: on 8 CPU devices, for seven layer configurations (with and without moe_router_grad_from_wo_input and moe_wi_gmm_fused_swiglu, plain and pipelined loop, dropless, dropping, padding excluded), the key gives the environment switch's outputs and gradients bit for bit (jnp models of the two Pallas kernels), the same lowered StableHLO, and the same TPU v7x lowering with the Mosaic kernel bodies once source locations are stripped; with the key off, the parent commit's. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, staged step boundary), with the same change behind an environment switch (PARAMBOLE_GMM_FUSE_DSWIGLU=1, the float32 epilogue) on stack 6 (moe_router_grad_from_wo_input, "count" index math, the packed routing gather and the glue rewrites, no fused wi SwiGLU): 87.061 -> 86.528 s per step, the mean of steps 1-4 (o275x-s6n3c-l3 vs o275x-s6n3-l3, same pod, booking r-7dbc0), 0 dropped tokens in both. Both arms logged loss 0 and grad_norm 0 at every step (sample_mask/kept_frac 0), so the kernels were timed on zero cotangents. Numerics, TPU A/B nxs6n3 against the reference pair nxref / nxref2 (parambole-bn6c-7dbc0, kept_frac 1.0, real gradients): PASS on all 6 gates. V1 kept 1.0; V2 loss max_rel 4.60e-8; N1 3.39e-6 and N1-last 6.70e-6 (tolerance 5e-5); N2 1.26e-8 (1e-4); N3, the relative grad_norm difference, 1.03e-4 against the replicate's 1.92e-5 (2e-3). Tests: - megablox_gmm_dswiglu_test.py (new). CPU: apply_dact_fn equals jax.vjp of silu(gate) * up in float32 (row scale, dot) and the check rejects a derivative that drops silu's second term and swapped halves; make_gmm_configs, for a TPU v7x target, accepts fuse_dact and "concat" with one n tile and rejects them with two, with a tile wider than the dim, wrong shapes or another activation; the fused DLHS (row scale, dot, transposed rhs) and gmm_silu's [gate | up] forward lower to one Mosaic kernel each for a v7x on the CPU host; with pallas_call replaced by a recorder, the wrapper hands the kernel [gate | up] as given and the row scale in float32 for each accepted scale shape (a bf16 rounding of it is caught). TPU only, not run: the fused DLHS against the epilogue of the unfused float32 accumulator with an EP shard; gmm_silu's [gate | up] equal to the unfused output bit for bit; gmm_dswiglu's VJP against gmm and the SwiGLU VJP, with the DLHS reading wo^T or wo as stored (transpose_rhs). Run on CPU against the jnp stand-ins of the layer test below, those 10 cases pass. - moe_wo_gmm_fused_dswiglu_test.py (new, CPU): on the shrunk Qwen3.5 MoE on 8 CPU devices, with jnp stand-ins for gmm_v2 and tgmm_v2 and everything above them real. With a stand-in epilogue that rounds where the unfused backward rounds, the layer with the key is bit-identical to the layer without it (output, loss, every gradient): with and without moe_router_grad_from_wo_input and moe_wi_gmm_fused_swiglu, plain and pipelined loop, dropless, padding excluded, the measured stack's keys with the bench's transpose_rhs DLHS (with both of those keys the router and input gradients within relative L2 1e-2, the activation rounding above). With the kernel's own epilogue, output and loss bit-identical and gradients within 1e-2. The fused DLHS is traced once per chunk, and the activation's VJP is never built with the key (the unfused reference builds it). Epilogues that swap the halves or drop the routing weight are caught. The validator rejects the key off its supported path, including quantization, clipping, a minimal* remat policy and the stack-9 tiles, and accepts the measured stack, any wi forward tile without moe_wi_gmm_fused_swiglu, heuristic tiling and a sharded MLP dim. * MoE: read the fused wi GMM's per-tile [gate | up] in the fused SwiGLU backward (moe_wo_gmm_fused_dswiglu) return_pre_act="concat" stores each n tile's accumulator [gate_j | up_j] as one block, so with moe_wi_gmm_fused_swiglu and more than one forward n tile per half (wi_tile_fwd_mlp_dim below twice the per-shard MLP dim: 1024 against Qwen3.5-397B's 1024) the [gate | up] that moe_wo_gmm_fused_dswiglu's epilogue reads is laid out [gate_0 | up_0 | gate_1 | up_1 | ...]. The parent commit refuses that combination (make_gmm_configs, and the config validator for explicit tiles); the measured build accepted it and took [gate_0 | up_0] as the gate, i.e. wrong gradients for the wi weights and everything upstream of them, with the output, the loss and the wo gradient unchanged. The epilogue now reads the forward's layout. gmm_v2's fuse_dact takes dact_pre_act_tile_n, the forward's tile_n per half, and split_dact_gate_up de-interleaves the [tile_m, 2 * tile_n] gate | up block the kernel already loads (the fused DLHS keeps its single n tile, so the block is the whole row): the gate is the even pieces, up the odd ones, with static lane-aligned slices in VMEM. There is no extra HBM pass, and [d_gate | d_up] stays the contiguous layout the wi GMM's DLHS / TGMM read. megablox ops.gmm_silu_pre_act_tile_n is the forward's tile (gmm_silu tiles with it), gmm_dswiglu and gmm_dswiglu_bwd take dact_pre_act_tile_n, and the MoE layer passes it with moe_wi_gmm_fused_swiglu and None otherwise (the GMM output is contiguous). make_gmm_configs accepts "concat" with any tile_n that divides each half, requires dact_pre_act_tile_n to be a multiple of 128 dividing size_n (the single n tile spans size_n, so a wider forward tile cannot exist), and normalises a value equal to size_n to None: the contiguous reader keeps its configs, metadata, kernel name and Mosaic body. With use_gmm_v2_heuristic_tiling the reader cannot know the forward's tile, so the two keys together raise there. The config validator follows the kernels: with moe_wi_gmm_fused_swiglu and explicit tiles it now accepts any wi_tile_fwd_mlp_dim whose half is the MoE MLP dim or a multiple of 128 dividing it (stack 9's 1024 against Qwen3.5-397B's 1024 included), instead of requiring one forward n tile per half, and still requires the fused DLHS's single n tile. Not taken: tiling the fused DLHS like the forward (wo_tile_dlhs_mlp_dim = the forward tile) fixes the input but writes [d_gate_j | d_up_j] per tile, which the wi backward would misread; relayouting [gate | up] in XLA before the DLHS costs a [buffer, 2 * mlp] HBM read and write per layer and chunk. The forward is unchanged. Without moe_wi_gmm_fused_swiglu, or with one forward n tile per half, the program is the parent's. Tests: - megablox_gmm_dswiglu_test.py. CPU: the pre-activation layout, with the forward's store modelled from the kernel source (FusedWeightsRef concatenates the gate and up weight tiles of n tile j; "concat" stores that accumulator as block j): split_dact_gate_up recovers gate and up bit for bit at nine (size_n, forward tile_n) pairs with 1 to 8 pieces per half, the epilogue on them equals the epilogue on the true halves bit for bit, and the contiguous reader, swapped pieces and reversed tiles are each caught at every interleaved layout; None and one piece per half trace to the parent's jaxpr. End to end through ops (gmm_silu(fuse_bwd_dswiglu) -> gmm_dswiglu, jnp stand-ins for gmm_v2 / tgmm_v2 that store and read that layout): dx, d wi and d wo match the float32 autodiff reference at 1, 2 and 4 forward tiles per half, and the contiguous reader, or a caller that drops dact_pre_act_tile_n, gets dx and d wi wrong (d wo right). make_gmm_configs for a TPU v7x target: "concat" with two tiles accepted, and a non-dividing tile and one wider than each half rejected; dact_pre_act_tile_n legal values, the size_n value normalised away (no metadata key, same kernel name), wider / non-dividing / off-lane / zero values and use without fuse_dact rejected. The fused DLHS lowers to one Mosaic kernel with and without the layout, and gmm_silu's [gate | up] forward with one and two tiles per half. TPU only, not run: the real gmm_silu gate_up at 2 and 4 forward tiles per half read by the fused DLHS equals the epilogue of the unfused [gate | up], and the contiguous reader does not. - moe_wo_gmm_fused_dswiglu_test.py (CPU, 8 devices): the gmm_v2 stand-in now writes the fused wi GMM's [gate | up] per tile and its fused DLHS reads it with split_dact_gate_up at the dact_pre_act_tile_n the layer passes. With two forward tiles per half (wi_tile_fwd_mlp_dim = the MLP dim, the 397B cell's shape) the layer with the key is bit-identical to the unfused layer on the measured stack and the pipelined loop, and within bf16 rounding with the kernel's epilogue; the layer passes the forward's tile (None without moe_wi_gmm_fused_swiglu); the measured build's reader is caught on both. The validator accepts stack 9's tiles and rejects a wi forward tile that is odd or whose half is off the 128-lane grid or does not divide the MLP dim. Run in the bench image (jax 0.11.0) on a CPU host: the CPU tests above, 103 of 103 in megablox_gmm_dswiglu_test.py and the validator class, and the layer test with all 13 cases passing and all 5 deliberately wrong epilogues and readers caught (the measured build's reader at two forward tiles per half: wi gradient 0.89 / 0.90 relative L2, router 1.04, input 0.50 / 0.91). For a TPU v7x target, lowered on that host with the Mosaic bodies decoded and printed without locations: the fused DLHS with dact_pre_act_tile_n None, or equal to size_n, has the parent's Mosaic body; with 512 of 1024 (the 397B cell) it differs only in its name and in the epilogue, where each of the four splits of the loaded gate | up block into two halves becomes four lane slices and two concatenations (sixteen slices with 128 of 1024). One routed-MoE layer's value_and_grad lowers to the parent's program with moe_wo_gmm_fused_dswiglu alone and with moe_wi_gmm_fused_swiglu at one forward tile per half; at two tiles per half the parent refuses the config and this commit lowers, to the same program as the environment-switch build below. Measured with the same change behind an environment switch on the RL trainer benchmark's overlay (comb67y, PARAMBOLE_N3_PREACT_LAYOUT=fixed; "legacy" keeps the measured build's program, checked equal on that lowering): an AOT compile of the 128-chip TPU v7x cell (fsdp 32 x context 4 x expert 2, fused wi SwiGLU and fused SwiGLU backward on) builds the fixed epilogue (-dact_swiglu-preact_tn_512) with 78.00 GiB in the default memory space, byte-identical to the measured build. On TPU v7x with that build: a one-chip check of the real kernels at the 397B cell's shapes (the logic of the TPU-only kernel test above, two forward n tiles per half) against a float32 reference of the SwiGLU backward, relative L2 1.658e-3 for this reader and 1.068 for the measured build's (parambole-bck-2d4c2); the numerics A/B nxc9ypfckhv passed all 6 gates (the relative grad_norm difference, gate N3: 0.0557873 vs 0.0557790, 1.58e-4 against 2e-3). Step time, 128 TPU v7x, staged arms on that build: with moe_wi_gmm_fused_swiglu and this key, 82.148 s per step (o275x-c9ypfctl-l3) against 83.316 s with moe_wo_gmm_fused_dswiglu off (o275x-c9pfxctl-l3); with further levers on top, 81.606 s (o275x-c9ypfall-l3) against 81.64 s for the measured build's reader in the same configuration (c9pfall). The fixed reader costs nothing. * Attention: reuse the forward's load-balanced positions and segment ids in the backward (context_parallel_attention_load_balance_hoist_ids) Under context_parallel_attention_load_balance every attention layer reorders its input, positions and segment ids into DUAL_CHUNK_SWAP order (tokamax_ring_attention.reorder_for_load_balance: two collective-permutes of half a shard, then a per-rank select between the received halves and their concatenation). The positions and segment ids are the same in every layer, since they depend only on the batch, and jax.lax.scan already computes the forward's copies once, outside the layer loop. But under remat the backward replay of each attention layer reorders them again, inside the loop. With context_parallel_attention_load_balance_hoist_ids=True, Attention._to_load_balanced_order names the reordered positions and segment ids maxtext_utils.ATTENTION_LB_IDS ("attention_lb_ids") and maxtext_utils.get_save_and_offload_names adds that name to the remat_policy=custom save list, on device. The saved value depends only on loop-invariant inputs, so the scan's partial evaluation hoists it out of the loop: one copy, which the backward reads. With the key off neither the name nor the save list changes. What changes in the program: in the lowered program the backward layer scan no longer holds the int32 collective-permutes of the ids (4 with packed segment ids, 2 without). Compiled for TPU v7x, XLA already moves those permutes out of the loop without the key; what the key removes from the compiled backward loop is the rest of the reorder (the select and concatenation of the received halves), the float conversion of the positions for RoPE and, with Splash attention, the segment-id min/max reductions. AOT-compiled for the benchmark's 128-chip mesh (the 60-layer trainer program of the measured build at the benchmark's stack-1A settings, its switch off and on, the benchmark's LIBTPU_INIT_ARGS), the while loops hold 47 fewer instructions (33,135 -> 33,088; among them 4 select, 4 pad and 2 add on the [1, 16384] int32 id rows), while the collective-permutes and the memory (temp 89.61 GiB, resident 98.31 GiB) are unchanged. With the switch off that program is identical to the base compile. Numerics: exact by construction, the backward reads the integers it would otherwise recompute. On CPU the logits and every parameter gradient are bit-identical with the key on and off. The TPU numerics A/B of the measured build passed all six gates on stack 4 (nxgx-x17h) and on the recipe of the best run below (nxc9pfxhv). The validator requires context_parallel_attention_load_balance=True with an even context parallelism of at least 2 (the degree AttentionOp.attention_load_balance_context_size balances over; attention keeps natural order otherwise), remat_policy='custom' (the only policy whose save list takes the name: 'full' recomputes everything, 'none' keeps everything, the presets save fixed names) and scan_layers=True (the hoist out of the layer loop is the scan's). It is a model validator of the HardwareAndMesh section, so it runs on MaxTextConfig and RLConfig. Default off. Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step, ragged_buffer_factor 2.75), with the same change behind an environment switch (PARAMBOLE_ATTN_HOIST_PERMUTE=1): on the stack-3 base, 99.267 s per step (o275x-s3-x17h, steps 1-4) against 99.504 s for the stack-3 control in another pod of the same booking (o275x-gdn-ctl): -0.237 s per step. There is no same-pod control. The switch is on in every later stack, including the best verified run: 82.671 s per step (o275x-c9pfxhv-l3, steps 1-4; its log shows the switch active), i.e. 3,567 tokens/s/chip for 9 x 64 x 65,536 tokens per step on 128 chips. Its share of that run was not measured separately. Equivalence to the measured build, on CPU: lowered for TPU v7x (compile-only devices, no TPU), the training VJP of eight tiny models with the key on is the measured build's program with its switch on, applied to this tree (the attentions.py and maxtext_utils.py that the best run imported), and with the key off is the program of this commit's parent (StableHLO and every Mosaic kernel body, compared without source locations). Configurations: hybrid Qwen3.5 (GatedDeltaNet, gated attention, MoE) on a context 4 x expert 2 mesh with dot-product attention, packed and unpacked; the same with Splash attention; Splash with the Pallas GDN kernel, sequence-sharded GDN context parallelism and context and gdn_cp_state saved; the hybrid on context 2 x expert 4; a dense decoder on fsdp 2 x context 4 with dot-product and with Splash attention; the dense decoder on context 8. Negative controls differ in all eight: the key on against the parent, the key off against the measured build, and the measured build with its second switch (SAVE_LB_INPUT, set in no measured run and not ported) also on. Compiled for TPU v7x (three of the configurations), the optimized HLO with the key off is the parent's and with the key on the measured build's. Differences from the measured build: its environment switch becomes a config key with a validator (the switch silently did nothing off the custom-remat path), and its log line is dropped (the key is in the config dump). Tests (context_parallel_attention_load_balance_hoist_ids_test.py, new), on CPU: - On 8 devices, for the hybrid Qwen3.5 with packed segments and without segment ids and for the dense decoder, all with remat_policy=custom and scan_layers: the logits and every parameter gradient are bit-identical with the key on and off; in the jaxpr the backward layer scan loses its int32 ppermutes (4 packed, 2 unpacked) while the int32 ppermutes outside every scan and all float ppermutes stay; the save list gains the name, and the program names it only with the key. Saving the natural-order positions fails the values check (every array differs) and a save list without the name fails the program check. - get_save_and_offload_names adds the name after the GDN names, only with the key and only under remat_policy=custom, without mutating the config's lists. - The validator accepts context parallelism 2 and 4 (ICI, DCN or both, and on the expert axis) and rejects, one case per requirement, attention load balancing off, context parallelism 1, 3 and unset, remat_policy full, minimal, none and minimal_offloaded, and unscanned layers; an error names every unmet requirement; it is registered on MaxTextConfig and RLConfig; pyconfig defaults the key to off and rejects it through the validator. - 13 deliberately broken copies each fail at least one of these tests: the reorder applied twice, the natural-order positions saved, the option not reaching the naming site, the save list without the name, the name offloaded instead of saved, either reader taking the load balancing switch for the key, the key ignored entirely, the segment ids not named, and the validator accepting any remat policy, unscanned layers, an odd context parallelism or no attention load balancing. The parent with the new test file fails 21 of its 23 tests. * Attention: gather the query kernel once per layer in the backward (query_proj_kernel_gather_once) The attention query kernel (with Qwen3-Next / Qwen3.5's output gate it is [emb, heads, 2 * head_dim]) is sharded on its embed dimension ('embed_attn'; fsdp x context x expert under cp-as-ep), so the partitioner all-gathers it for every dot that reads it. Under remat the backward of a scanned layer reads it twice: in the replay of q = x @ wq and in the input gradient dx = dq @ wq^T. On the Qwen3.5-397B-A17B RL trainer cell the TPU compiler gathers the 128 MiB kernel twice per attention layer in the backward, 256-way, from the same source slice. The second gather, for the input gradient, is queued behind splash's dK/dV reduce-scatters on the same SparseCore core and issued 18 compute instructions ahead of its use, so the input-gradient matmul can stall on it. Whether that stall is on the step's critical path depends on the chip; the step-time effect of this change has not been measured on TPU. With query_proj_kernel_gather_once=True the query projection is a custom VJP (linears._dot_general_reusing_gathered_kernel). Its forward gathers the kernel with one explicit sharding constraint and keeps the gathered kernel as the residual. Its backward computes the input gradient from that residual and the kernel gradient with the stock dot_general transposes (jax.linear_transpose). Under remat the residual is recomputed with the rest of the replayed forward, so the replay's gathered kernel serves the input gradient: one gather per layer in the backward instead of two (on this cell 15 x 128 MiB = 1.875 GiB less SparseCore gather traffic per micro-batch), and the input gradient no longer depends on a gather issued late in the iteration. The kernel gradient carries no sharding constraint, so it is reduce-scattered to the kernel's sharding as before. A plain with_sharding_constraint would not do: its transpose constrains the kernel gradient to the gathered sharding (an all-reduce), which also changes the gradient bits (a test variant below shows both). DenseGeneral takes the gathered layout as a new keyword-only argument, gathered_kernel_axes (the kernel's logical axes with the gathered ones set to None); the attention layer passes the query kernel's axes with 'embed_attn' set to None, for training-mode layers only. Default off; off, the code path and the compiled program are unchanged. Exactness: the forward and the replay contract the same gathered values as before, the input gradient contracts over heads and head_dim, which are not sharded, so its sums are formed the same way, and the kernel gradient path is unchanged. The loss, the input gradient and every parameter gradient are the same bits with the key on as off, on CPU and on TPU (v5p, executing the test's two meshes below). The validator (shared by MaxTextConfig and RLConfig) rejects the key where it would not remove a gather or where the result has not been verified: attention_type mla or compressed (their layers build their own query projections), fused_qkv, quantization, explicit shard mode, ici_context_autoregressive_parallelism > 1 (replicated kernel), an 'embed_attn' that maps to no mesh axis of size > 1, q_heads or kv sharded (tensor parallelism on the query kernel), and a remat policy that saves the query projection or does not rematerialize it (it is replayed under full, save_out_proj, or custom with query_proj 'remat'). Evidence: AOT compiles of the trainer's fwd_bwd at the exact cell of PR #5758 (pdev arm o275x-prf-l3-prof: fsdp 32 x context 4 x expert 2, 256 devices of tpu7x-256, cp-as-ep, all_gather context parallelism, logps chunk 4096; libtpu 0.0.51 n1010; no TPU used), against the baseline compile of 20c877aa1: - query-kernel gathers per backward iteration 2 -> 1 (forward 1 -> 1), 3 -> 2 in the module. The remaining backward gather is the replay's, issued at the top of the iteration; the input-gradient dot reads it through a bitcast, 566 compute instructions after it lands, and no layout copy is added. The input-gradient gather that sat behind the dK/dV reduce-scatters is gone, and no other collective changes. - memory: device peak 104.46 GiB both ways. fwd_bwd temporaries -130,944 B; the token-dropless retry kernel's -133,824,384 B. - key off: the after-optimizations fwd_bwd HLO equals the baseline's line for line once metadata is stripped. Tests (query_proj_kernel_gather_once_test.py, new; 33 cases, all on CPU, one of them compiles for a TPU target and needs libtpu): the custom VJP equals lax.dot_general (value and both gradients, bit for bit); the attention layer passes the gathered axes to the query projection only, for training only; DenseGeneral rejects quantized, sliced, explicit-mode and mesh-less use; the validator accepts and rejects as above, in pyconfig too, and RLConfig runs it. In a child process with 8 CPU devices, a scanned, rematerialized stack of Qwen3NextFullAttention layers (each followed by an RMS norm, as in the decoder layer), on cp-as-ep rules and on the default rules: with the key every dot that reads the query kernel (forward, replay, input gradient) reads the gathered kernel and none does without it; values are bit-identical; the backward has one gather of the kernel; the output kernel's gathers and every other collective are unchanged. When libtpu is installed the same program checks also run compiled for tpu7x-8, where the backward gathers the query kernel twice without the key and once with it. Ten deliberately wrong variants each fail at least one case: residual not used, gather keeping the kernel's sharding, wrong dimension gathered, gate half not reusing the gather, key name misread, validator accepting a policy that saves the projection, plain sharding constraint without the VJP, key applied outside training, RLConfig not calling the validator, and a perturbed input gradient. * Tests: set packing in the gdn_conv1d_dense_token_stack validator test's config Since "Skip the GDN kernel's segment path for unpacked training rows" (in AI-Hypercomputer/maxtext#5721), validate_gdn_config reads cfg.packing, and gdn_conv1d_dense_token_stack_test.py's _validator_cfg builds the config as a SimpleNamespace without that field. So 9 of the file's 41 cases (the ValidatorTest cases that reach the GDN validator) fail with AttributeError: 'types.SimpleNamespace' object has no attribute 'packing' (types.py:5548). _validator_cfg now sets packing=True: the validator then skips its new packing=False checks, so every case checks what it checked before. Test only. This is the gdn_conv1d_dense_token_stack_test.py hunk of 6314a6b05 ("Tests: bring the tests in line with the completed #5698 port"), which was written for #5721 and not pushed before #5721 merged. * MoE: hoist the fp8 fake-quant absmax out of the layer scan and gather fp8 codes (fp8_fake_quant_hoist) With fp8_moe_fake_quant=true every routed-expert weight is rounded, in every layer and in both the forward and the rematerialized forward, to the per-channel float8_e4m3fn grid of quantize_weight_per_channel. The absmax that fixes each channel's scale reduces over the weight's input dimension, which is FSDP-sharded, so each layer pays a max all-reduce in each pass (on the Qwen3.5-397B-A17B RL trainer cell, 30 per micro-batch), and the FSDP all-gather then moves the bf16 fake-quantized weight, 2 bytes per element. fp8_fake_quant_hoist (default off) changes that inside NNXDecoder._apply_layers_sequentially: - absmax: moe.stacked_fp8_fake_quant_absmax computes the per-channel absmax of every eligible RoutedMoE weight for the whole scanned stack before the scan (one max all-reduce per weight instead of one per layer and pass). The stack rides in the scan's xs, and moe.fp8_hoisted_absmax_attached hands each layer its slice for the duration of the layer call (the attribute is removed before the scan splits the layer's state, so the graphdef and the scanned state are as without it). The layer fake-quantizes with it (fake_quantize_weight_per_channel(absmax=...)). - gather8: absmax, plus sparse_matmul quantizes each weight on its FSDP shard with the hoisted absmax (quantize_weight_codes_per_channel), all-gathers the float8_e4m3fn codes and dequantizes after the gather and the expert-weight prefetch barrier (dequantize_weight_codes_ste, a custom VJP whose backward hands the cotangent to the shard through the gather's sharding constraint, as the straight-through estimator did). The gather moves 1 byte per element instead of 2. The codes stay float8_e4m3fn because on TPU v7x the SparseCore all-gather offload takes 8-bit floats but not uint8, which falls back to the TensorCore. The gathered scale's layout is pinned at its first use in the layer, so the wo scale is gathered once for the stack instead of in every layer. A whole prefused wi (moe_unsplit_prefused_wi, w1_kernel=None) is quantized, gathered and dequantized as one tensor. The mode comes from the new config key fp8_fake_quant_hoist (off | absmax | gather8) or, while the key is off, from the env var PARAMBOLE_FQ_HOIST (0, absmax, 1 = gather8), so a run can A/B it without changing its argv. Both set to different modes raises, as does any other env value. A mode that the config cannot honour raises at trace time with the unmet requirement (moe.fp8_fake_quant_hoist_enabled) instead of quietly running the original code. A bare `off` on the command line or in YAML parses as False; a before-validator maps it to 'off'. Every enabled trace logs one [FQHOIST] WARNING line, the liveness check for an arm. Exactness: a max is exact in any order and over any sharding, so every layer gets the absmax, scale and codes it computed before; the dequantization is the same elementwise expression with a per-channel scale, so it commutes with the gather; -0.0 codes are stored as +0.0, because the original w_q + (w - stop_gradient(w)) turns every -0.0 of w_q into +0.0. On CPU, on a shrunk Qwen3.5 through the production model path (scan, remat, ring of experts, chunk pipeline, prefetch barrier, forced routing), the loss, the logits and every parameter gradient are bit-identical to the unpatched tree in both modes, set by key or by env, with moe_unsplit_prefused_wi on and off; deliberately wrong variants (scale without the bf16 rounding, the neighbouring layer's absmax, a scaled straight-through gradient, the neighbouring channel's scale after the gather, a scaled custom-VJP cotangent, codes without the -0.0 canonicalization) each change the bits. The step-time effect has not been measured on TPU. * Tests: fp8_fake_quant_hoist off vs absmax vs gather8 tests/unit/fp8_fake_quant_hoist_test.py, built from the CPU proofs of the fp8 fake-quant hoist: - In a child process with 8 CPU devices, a shrunk Qwen3.5 through the production model path (decoder layer scan under remat, sparse_matmul with prefused unsplit wi, ring of experts, chunk pipeline, expert-weight prefetch barrier, forced routing) on the cp-as-ep rules: with fp8_fake_quant_hoist=absmax and =gather8 the loss, the logits and every parameter gradient are bit-identical to =off. The compiled programs show that the hoist ran: no fake-quant max all-reduce is left inside the layer loop, and with gather8 the loop all-gathers 8-bit-float expert weights. - In process: the hoisted and stacked absmax, the fp8 codes with -0.0 stored as +0.0 and their straight-through dequantization (with a control that keeps -0.0 and differs), the mode resolution from the key and PARAMBOLE_FQ_HOIST, and the raises for a bad env value, a key/env conflict, a bad key value and a mode the config cannot honour. * MoE: wi_tile_fwd_rhs_buffer_count, the fused wi forward GMM's weight pipeline depth gmm_v2 takes rhs_buffer_count (default None: today's 3, or 2 for a chunked 3D lhs); the metadata key and the scope suffix -rhsbuf_N appear only when it is set, so every default call compiles to the same kernel. ops.gmm_silu passes it to the forward kernel only (fwd_rhs_buffer_count). The layer passes wi_tile_fwd_rhs_buffer_count (default 0) only when that GMM runs one K tile; the validator rejects it without moe_wi_gmm_fused_swiglu or with heuristic tiling. With 2, wi_tile_fwd_mlp_dim=2048 at wi_tile_fwd_batch_seq=512 fits VMEM: the fused SwiGLU wi forward kernel (tm 512, tk 4096, tn 1024 per half) needs 50.95 MiB of the 57.60 MiB kernel limit instead of 66.95 MiB (kernel-only AOT for TPU v7x, two-sided). * Tests: wi_tile_fwd_rhs_buffer_count, the fused wi forward GMM's weight pipeline depth tests/unit/wi_tile_fwd_rhs_buffer_count_test.py, built from the CPU checks of the key: - The config key: 0 by default; accepted at 2 and 3 next to moe_wi_gmm_fused_swiglu with explicit tiles; rejected, naming the key, at 1 and below, without moe_wi_gmm_fused_swiglu, and with heuristic tiling. - make_gmm_configs under a TPU v7x target: without the option the config, the metadata, the kernel name and the weight BlockSpec's three buffers are unchanged; a depth of N sets the field, the metadata, the -rhsbuf_N name suffix and Buffered(buffer_count=N); a depth below 2 raises. - ops.gmm_silu lowered for a TPU v7x on a CPU host, with and without fuse_bwd_dswiglu: only the forward kernel takes the depth. - In a child process with 8 CPU devices, the MoE layer with jnp stand-ins for the two Pallas kernels: the key reaches every fused wi forward call, and no other gmm_v2 call, only when that GMM runs one K tile; the output, loss and gradients are unchanged. - On TPU only: gmm_v2's fused forward writes the same bits at depths 2 and 4 as at the default. * GDN CP forward: fold the rank's chunk transitions in a Pallas scan With gdn_cp_matmul_precision=high on TPU, the sequence-sharded CP forward computes m_local (the product of the rank's 256 chunk transitions M_c = g_last I - k_bwd^T (t_inv k_beta_g)) with gdn_cp_m_scan, a Pallas emit_pipeline over (batch, head group, chunk) that keeps the running product in VMEM, instead of compose_local_from_t_inv's balanced tree of HBM-bound f32 [chunks, heads, 128, 128] XLA fusions. s_ext_local is pass 1's end state, as the tree returns it. CPU and 'highest' keep the XLA tree. Same terms in sequential instead of tree order, every matmul at 3-pass bf16 like the tree's Precision.HIGH, the segment-masked gating cumsum at HIGHEST. The scan reads k / t_inv / b / a with pallas_gdn_bwd_kernel's block layout and runs all 64 heads of the 397B layer in one group (head_tile 64). Numerics (v6e, 397B per-device geometry: 16 kq / 64 v heads, dims 128, 16384 tokens, chunk 64, real pass-1 t_inv): against the XLA tree at HIGHEST, the scan's m_local is within 1.7e-6 (relative to max) and the old HIGH tree within 9.9e-4 (long-memory gates; with model-like gates m_local is exactly 0 in all three). The whole CP layer on 4 devices (S = 65536, packed and single long documents) gives a bit-identical output. v6e per call: 3.3 ms against 8.3 ms for the XLA tree. On the v7x heavy device the tree costs 3.2 ms per GDN layer (144 ms per micro-batch). 397B v7x AOT: fwd_bwd temp 93.31 GiB (93.37 before); the forward f32 [256,1,64,128,128] tree fusions are gone. Test: gdn_cp_dh_scan_test compares the interpreted scan with the tree (no segments, one long document, packed documents; one and two head groups), with chunk-order and segment-start negative controls. * GDN CP backward: dS_ext_local as a reverse Pallas scan With gdn_cp_matmul_precision=high on TPU, the sequence-sharded CP backward computes ds_ext_local with gdn_cp_dh_scan, a reverse Pallas emit_pipeline over (batch, 2 groups of 32 heads, chunk) in the backward kernel's form of the state recurrence (acc <- g_last acc + q_g^T do - k_beta_g^T t_inv^T (attn^T do + k_bwd acc)), instead of compose_bwd_local_from_t_inv's balanced tree of HBM-bound f32 [chunks, heads, 128, 128] XLA fusions. dm_local is the transposed forward m_local, as the tree returns it with m_local_cached. CPU and 'highest' keep the XLA tree. Numerics (v6e, 397B per-device geometry, real pass-1 t_inv): against the XLA tree at HIGHEST, the scan's ds_ext_local is within 3.1e-6 (relative to max), the old HIGH tree within 1.7e-5 to 4.8e-4. Whole CP layer on 4 devices (S = 65536): output bit-identical; bf16 input grads within 1 bf16 ulp (as HIGH vs HIGHEST); f32 param grads differ by 1e-4 to 2e-3 relative, the size of the old path's own HIGH-vs-HIGHEST difference. v6e per call: 4.8 ms against 24 ms for the XLA tree (v7x heavy device: tree 8.95 ms per GDN layer, 403 ms per micro-batch). Memory: 397B v7x AOT fwd_bwd temp +1.3 GiB. In the backward loop body XLA now materializes the f32 conv recompute (about 2.25 GiB) before the second GDN layer's gdn_bwd_kernel instead of after it, so it overlaps that layer's replayed chunk_states. The bench config still fits and runs on 128 v7x chips. * GDN CP scans: opt-in keys gdn_cp_m_scan and gdn_cp_dh_scan (default off) The two Pallas VMEM scans of the previous two commits (the forward m_local transition product and the backward ds_ext_local reverse scan) ran whenever gdn_cp_matmul_precision was 'high' on TPU. Gate each on its own key instead, both default False, so the default program is unchanged: - configs: gdn_cp_m_scan / gdn_cp_dh_scan in base.yml and types.py; the GDN validator rejects them without Gated DeltaNet layers, without use_gdn_kernel=True, or with gdn_chunk_size != 64. - kernels/gdn/model_runner.py passes them to gdn_decoupled_conv1d. - kernels/gdn/gdn_bwd/api.py: cp_m_scan / cp_dh_scan are trailing nondiff arguments of the custom VJP (now 14), forwarded to the CP forward and backward; the positional-call shim accepts 10..15 arguments. - The gdn_cp_matmul_precision description is restored: precision alone no longer selects the scans. The key tests are in the next commit. * Tests: gdn_cp_m_scan and gdn_cp_dh_scan, the GDN context-parallel scan keys tests/unit/gdn_cp_dh_scan_test.py, on top of the scan-vs-tree tests of the two scan commits: - The config validation: both keys default to False and are refused without Gated DeltaNet layers, without use_gdn_kernel=True, and with gdn_chunk_size != 64. - The gating, traced through model_runner.run_gdn_kernel_layer on an abstract TPU7x context mesh: no scan pallas_call with both keys off, only the forward m scan with gdn_cp_m_scan alone, both scans with both keys. - The CPU numerics through gdn_decoupled_conv1d, in a child process with 2 CPU devices so that it runs whatever the parent's device count: keys off bit-identical to the call without them, keys on close to keys off, forward and backward. * GDN seq-CP pass 1 states-only kernel, bf16 QKV chunk DMA, and out_proj token-minor layout Combine Mohit Khatwani's three GDN kernel and layout optimizations (commits 906e15714, 9c48e653a, and 80c149572) adapted onto 1bf840cff: 1. GDN seq-CP pass 1: states-only forward kernel (GDNMode.PER_SEQ_STATES), no output or chunk states allocated or written to HBM, called via api._run_local_gdn_states_fwd and wrapper.fused_conv1d_gdn_states. 2. GDN forward kernel: DMA bf16 qkv chunks via ChunkInBufferedRef and upcast to fp32 in VMEM via vmem_ldst.upcast_large_to_compact for prefill-only calls. 3. GDN forward out_proj: pin the gated-norm output token-minor (0, 2, 1) in training via moe._forward_layout_constraint. * GDN: opt-in keys gdn_cp_pass1_states_and_bf16_qkv and gdn_out_proj_token_minor (default off) Gate the states-only CP pass 1 kernel, VMEM bf16 QKV chunk DMA, and out_proj token-minor layout constraint behind two opt-in config keys (both default False) so default programs are unchanged: - configs: add gdn_cp_pass1_states_and_bf16_qkv and gdn_out_proj_token_minor in base.yml and types.py; validate_gdn_config rejects both without Gated DeltaNet layers and rejects gdn_cp_pass1_states_and_bf16_qkv=True without use_gdn_kernel=True, with gdn_chunk_size != 64, or with dtype != 'bfloat16'. - kernels/gdn/wrapper.py: gate use_bf16_qkv_dma on bool(bf16_qkv_dma and is_prefill_only and act_in_dtype == jnp.bfloat16), preserve public wrapper defaults (compute_precision=jnp.float32.dtype, use_qk_norm_in_gdn=True), and fix inner_kernel_without_t_inv slot_args for GDNMode.PER_SEQ_STATES. - kernels/gdn/gdn_bwd/api.py and model_runner.py: thread pass1_states_and_bf16_qkv through gdn_decoupled_conv1d and _run_cp_gdn_decoupled_fwd_impl. - models/qwen3.py: gate _forward_layout_constraint on getattr(cfg, 'gdn_out_proj_token_minor', False). * Tests: gdn_cp_pass1_states_and_bf16_qkv and gdn_out_proj_token_minor Add tests/unit/gdn_pass1_states_and_out_proj_test.py and update tests/unit/gdn_bwd_pallas_test.py: - Config validation for gdn_cp_pass1_states_and_bf16_qkv and gdn_out_proj_token_minor (decoder_block, use_gdn_kernel, gdn_chunk_size, and dtype checks). - Pallas TPU interpret-mode execution on CPU for _run_local_gdn_states_fwd and bf16 vs fp32 QKV DMA paths. - Chunked GDN PER_SEQ_STATES vs PER_SEQ bit-identity with non-underflowing gates and continuation segment IDs at kernel_size - 1. - Real Qwen3NextGatedDeltaNet out_proj token-minor lowering checks in train vs prefill mode. - 2-rank shard_map CP forward and backward bit-identity and single-rank agreement. * MoE ring of experts: opt-in fp8 combine keys (moe_fp8_combine, moe_fp8_bwd_dispatch, moe_combine_per_chunk) Ported from the MLPerf v6.1 DeepSeek-V3 code (6639e3166, "code from image integrate_v2:mlperf61v4-stackall6"), adapted to the Qwen3.5 ring-of-experts path (pipelined chunk loop, router grad from the wo input, fused SwiGLU). All keys default false; with them off the program is unchanged. - moe_fp8_combine: the forward combine reduce-scatter (_ring_reduce_scatter) becomes an fp8 e4m3 all-to-all with one f32 absmax scale per row, dequantized and summed in f32 on the receiver. The backward stays the all-gather of the cotangent. - moe_fp8_bwd_dispatch (+ moe_fp8_bwd_dispatch_qtype, default float8_e4m3fn): the backward all-gather of the combine cotangent carries an fp8 payload with a per-row f32 scale. - moe_fp8_combine_pin_sparse_core: pins that all-to-all to the EP SparseCore. - moe_combine_per_chunk: in the plain chunk loop, chunk c's unpermute waits on chunk c-1's combine reduce-scatter (forward-only data dependency on the routing weights, remat name moe_combine_order). The pipelined chunk loop already orders them so (its tie 3); there the key has no effect. Helpers quantize_rows_fp8, dequantize_rows_fp8, ep_combine_reduce_scatter and _order_after are the DeepSeek-V3 versions. The validator rejects the fp8 keys with moe_quantize_combine_bwd_method or moe_combine_bwd_direct_qarray, which they would bypass. * MoE ring unsort: spread the masked lanes of the slot-order gather (ragged_unsort_spread_masked_indices) With moe_routing_index_math=count and a packed buffer, the backward of the ring-of-experts unsort sends the routing-weight gradient back to slot order with a gather through the partial inverse permutation. Slots of other shards and dropped slots are masked: their index is 0, so at expert parallelism 8 about 7 of every 8 lookups read table entry 0. The masked lanes are discarded by the outer where. The new key (default false) gives each masked slot its own slot index, an in-bounds entry of the n-entry table, instead. The values are bit-identical; only the addresses the gather touches change. The same offloaded SparseCore gather program runs 40x faster with a permutation index than with the entry-0 fill, so a hot address is the suspected cause of the slow one. The key is read at both ring unsort call sites (ring_ragged_unsort and ring_ragged_unsort_down_projection, including the fused dswiglu backward). * Tests: ragged_unsort_spread_masked_indices, the spread of the ring unsort's masked gather lanes tests/unit/ragged_unsort_spread_masked_indices_test.py: - The key: a bool, False by default in types.py and base.yml, resolved by pyconfig, and True when passed. - ragged_sort._ring_unsort_fan_out directly, every shard of a 4-shard ring, 3 seeds x {packed buffer, packed buffer that drops, full buffer} x {"count", "sort"}: the gathered rows, weights, mask and slot-order value bit-identical with the key on; again for "count" with NaN, +Inf and 3e38 in the per-row values outside the mask, which the spread lanes now read. - Structure: with the key on the traced fan-out has exactly one iota more for "count" on a packed buffer (the spread index) and is the same program for "sort" and the full buffer, so a key that leaves the index at 0 fails. - Deliberately wrong spreads (no outer where, valid lanes one row off, one valid slot dropped) must fail the fan-out comparison. - In a child process with 8 CPU devices, the op: ring_ragged_sort and then ring_ragged_unsort_down_projection or ring_ragged_unsort, forward and VJP, on 4- and 8-device expert-parallel meshes: output, dx, dw, dwi and dwo bit-identical with the key on; a slot-order gradient one ulp off and a spread without the outer where must be detected. - In a child process with 8 CPU devices, the shrunk Qwen3.5 routed MoE with moe_routing_index_math="count" (plain unsort, fused down projection, and the benchmark's MoE settings): output, loss and gradients bit-identical with the key on, and the key's value reaches both ring unsort call sites and every fan-out call. * MoE router replay: opt-in compaction of the fallback top-k (moe_fallback_topk_capacity) Under router replay get_topk computes the router's own top-k ids for every token (10 argmax passes over the [tokens, 512] router probabilities on the 397B) and then keeps them only for tokens without valid forced ids; padding is routed nowhere under moe_exclude_padding_from_routing. With moe_fallback_topk_capacity > 0 the ids are computed only for the tokens that use them: those rows, at most that many per call, are gathered in token order, top-k'd, and written back to their tokens with a compare-select-sum (no gather or scatter over the tokens). A lax.cond computes every row instead when more tokens use them, so the result is the same for any input. top-k is row-wise, the gathered rows are exact copies, and a non-using token's ids are never read, so get_topk's weights, ids and gradients are bit-identical to the default (0 = every token, the program unchanged). * Config: fp32_master_storage, read by Tunix's Fp32MasterOptimizer (default f32) "residual16" asks Tunix's Fp32MasterOptimizer (weight_dtype=bfloat16) to store a scanned layer's float32 master as the 16 low bits its bfloat16 weight does not hold plus one packed bit per element, rebuilt bit for bit at each update. MaxText itself does not read the key; the default "f32" changes nothing. * GDN CP: opt-in 2-round cross-rank state scan (gdn_cp_fast_state_passing) * Tests: moe_fallback_topk_capacity and fp32_master_storage tests/unit/moe_fallback_topk_capacity_test.py: - The key: an int, 0 by default in types.py and base.yml, resolved by pyconfig; a negative value is refused. - fallback_topk_indices_compacted directly, for both top-k functions get_topk uses (argmax passes for moe_routing_index_math="count", lax.top_k for "sort"), random, exactly tied and duplicated score rows, and using-token counts on both sides of the capacity (0, 1, C - 1, C, C + 1, all): every using token gets the full top-k's ids; on the compacted branch every other token gets 0; above the capacity every token gets the full top-k's ids. - RoutedMoE.get_topk on the Qwen3.5 forced-routing path, forward and VJP: weights, ids and d(gate_logits) with capacities 1, 7, 64, using - 1 and using bit-identical to capacity 0, for "count" and "sort" with moe_exclude_padding_from_routing on and off, on forced ids mixing valid tokens, unset (-1) slots, duplicate ids and padding rows. - Structure: get_topk's program has a lax.cond with a capacity and none without, so a key that never reaches get_topk fails. - Deliberately wrong helpers (a write-back one row off, no lax.cond, the cond one row too permissive) are caught by the same get_topk comparison; a correct copy is not. tests/unit/fp32_master_storage_test.py: the key takes exactly the two storages Tunix's Fp32MasterOptimizer accepts ("f32", "residual16"), defaults to "f32" in types.py and base.yml, is resolved by pyconfig, and any other value is refused when the config is built. * Tests: the pass-1 test's negative controls find pass1_states_and_bf16_qkv by name gdn_pass1_states_and_out_proj_test.py drops pass1_states_and_bf16_qkv in two negative controls by overwriting gdn_decoupled_conv1d's last argument. gdn_cp_fast_state_passing appends an argument after it, so the controls overwrote that one instead, still saw the states-only kernel, and test_tpu_programs_in_a_child_process failed (control_dropped_by_forward_rule and control_dropped_by_primal). The controls now look the argument up in the primal's signature. 19 passed here, and 19 with this file at 818caa328, before the appended argument. * Tests: the 2-round test's op-level case gets weak-decay inputs and a matmul count test_gdn_decoupled_conv1d_fwd_bwd_bit_identical drew a_log, a and b from N(0, 1). Over a rank's 128 tokens that makes each rank's state transition 1e-31 to 5e-29 at most, so the states passed on two or three ranks underflow to 0 in float32, and the case stayed bit-identical with the 2-round scan broken (a wrong permutation, or the path disabled). It now uses a_log = -6 and a = b = -4 (transition entries up to about 0.96 per rank), and it also checks that the fast program has 6 fewer dot_generals than the slow one: here the 2-round scans save 3 matmuls each and the generic fast loop 2, and a disabled 2-round path gives bit-identical values. On CPU, the correct code passes both tests; the forward or the backward 2-round scan disabled, or given a wrong permutation, fails both. With the test file before this commit only the direct cp_gdn test failed under those four breaks. * MoE ring of experts: fp8 forward token all-gather (moe_fp8_token_all_gather) The dispatch's token all-gather over the expert axis carries an fp8 e4m3 payload with one float32 absmax scale per token row (quantize_rows_fp8), and the gathered rows are dequantized to the activation dtype right after the gather. Everything downstream (sort, expert GMMs, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu, moe_router_grad_from_wo_input) reads the same bfloat16 token buffer as before, so the key works with those keys, unlike moe_quantize_token_all_gather (fp8 tokens into the GMM as a QArray), which their validators reject. The backward is straight-through: the plain bf16 reduce-scatter of the cotangent, the transpose of the unquantized gather. Default false; with it off the program is unchanged. The validator requires sparse_matmul and use_ring_of_experts, and rejects moe_quantize_token_all_gather, moe_tc_ragged_3d_dispatch, num_moe_emb_chunks and the SparseCore-pinned EP all-gathers. * Tests: moe_fp8_token_all_gather on a CPU mesh The helper against the plain all-gather (forward per-row error in the e4m3 band, zero rows exact, backward bit-identical) and the routed-MoE layer on the bench stack (pipelined chunk loop, moe_router_grad_from_wo_input, moe_wi_gmm_fused_swiglu, moe_wo_gmm_fused_dswiglu) with the gmm_v2/tgmm_v2 stand-ins of moe_wo_gmm_fused_dswiglu_test.py. Broken gathers (scale dropped, local scale, backward without the reduce) must fail. Config: default off, accepted on the bench stack, rejected where it cannot run. * Tests: moe_fp8_token_all_gather, judge the broken layer gathers by the key's own gate The broken layer variants must fail the output band the real key passes. Call the validator directly for use_ring_of_experts=False (num_moe_token_chunks' validator fires first in a full config), and place the single-shard input on its mesh. * GDN backward kernel on the natural operand layouts: opt-in key gdn_bwd_natural_layout (default off) Port of the backward part of #5628 onto this tree's backward kernel (recompute_t_inv, kq-head matmuls) and the sequence-CP dh scan. With gdn_bwd_natural_layout=true the Pallas GDN backward kernel reads qkv_conv [B, S, dim] as three q / k / v channel windows, do [B, S, Hv, dv], chunk_states [B, C, Hv, dk, dv] and t_inv [B, C, Hv, cs, cs] in place and writes dy_conv [B, S, dim] directly (make_bwd_block_specs_natural), instead of the XLA slice / transpose / concatenate relayouts into and out of the [B, G, C, ...] head-group-major layout around the pallas_call. Every VMEM block holds the same values as before; only the HBM index maps change, so the gradients are the same. Differences from #5628: do is read in its 4-D natural layout (the layout the CP dh scan already reads in place) instead of a 2-D [B, S, Hv * dv] view, which is a different tiled layout and would need its own relayout; dht / dh0 keep the [1, Hv_tile, dk, dv] VMEM block through a [B, 1, Hv, dk, dv] view, so the kernel body changes only for q / k / v and dq / dk / dv; shapes whose k / v channel offsets are not whole blocks keep the relayout path (no head-tile halving); the bf16 forward input part is not included. With gdn_cp_dh_scan off the key also enables #5628's head-major glue in cp_gdn.compose_bwd_local_from_t_inv. The key is a new trailing nondiff argument of the gdn_decoupled_conv1d custom VJP (the primal and forward rule ignore it; the backward rule's positional-call shim gains one form) and is read in model_runner as a config attribute. The kernel is named gdn_bwd_kernel_nat (gdn_bwd_kernel_recompute_<algo>_nat) when the natural layout is used. Cherry-picked from 77d006a8f (d4ffdd1ad picked onto parambole/w25-mega and the tail commit) onto this branch, where gdn_decoupled_conv1d has cp_fast_state_passing (25) as its last nondiff argument and no cp_scan_cost_estimate. bwd_natural_layout is position 26: nondiff_argnums gains 26, the backward rule's comment counts 17 nondiff arguments (19 with residuals and cotangents), and its positional shim handles cp_fast_state_passing as a two-omitted case and bwd_natural_layout as the one-omitted case. tests/unit/gdn_pass1_states_and_out_proj_test.py keeps finding pass1_states_and_bf16_qkv by name in its two negative controls, and its positional-form check gains forms 18 and 19, with the 19-argument full form as the reference. * Tests: gdn_bwd_natural_layout bit identity, wrong index maps detected, config wiring CPU (Pallas interpret mode): the natural-layout backward kernel against the relayout path at three head layouts, with and without packed segment ids, dh0 and in-kernel t_inv recompute (bitwise); the unaligned shape keeps the relayout path; three deliberately wrong natural index maps are detected; the config key reaches the kernel through model_runner (gdn_bwd_kernel vs gdn_bwd_kernel_nat) and a config without the key raises; the key defaults to off. * gdn_bwd_natural_layout: per-operand-group modes; "states" avoids the XLA relayout ripple The key becomes a string: "off" (default) | "states" | "all" | "no_do" | "inputs" | "states_dy". Each mode names the operand groups the Pallas GDN backward kernel reads / writes in place: qkv (q / k / v windows of qkv_conv), states (chunk_states, t_inv, dht, dh0), do and dy (dq / dk / dv windows of dy_conv); "all" is the previous behaviour. The kernel body reads q / k / v and writes dq / dk / dv as three refs only for the qkv / dy groups (split_qkv_in / split_qkv_out); every other block keeps the VMEM shape of the relayout path, so all modes are bit-identical to it. Why: exact-cell AOTs of the Qwen3.5-397B bench config show that wiring the kernel directly to the bf16 qkv_conv, do or dy_conv tensors makes the TPU compiler re-lay out the in_proj_qkvz backward (and, for do, the gated RMSNorm), which gives back most of what the removed relayouts cost. Reading only the f32 chunk_states / t_inv in place ("states") removes the chunk_states transpose and the three t_inv copies per layer and changes nothing else in the program. The head-major composition glue (gdn_cp_dh_scan off) now follows the qkv / do groups. * GDN: opt-in natural layout for the gated-norm output (gdn_gated_norm_layout) Under gdn_out_proj_token_minor the reshaped gated-norm output feeds out_proj token-minor, and XLA propagates that layout back into the float32 RMSNorm math: every call re-tiles the bf16 GDN kernel output and transposes a float32 [tokens, value_dim] copy of it, in the forward and again in the remat replay. gdn_gated_norm_layout (default "auto" = unchanged) pins the norm output row-major: "natural" in the forward only (moe.forward_row_major), so the norm math stays in the kernel layout and one bf16 transpose to token-minor follows it; "natural_fwd_bwd" also pins its cotangent, so the norm backward stays natural and produces d_o in the layout the GDN backward kernels read. Layout only: values are unchanged. Training path only; requires gdn_out_proj_token_minor=True and a Gated DeltaNet decoder. The key is read as an attribute, so the four tests whose SimpleNamespace configs feed validate_gdn_config gain gdn_gated_norm_layout="auto". Cherry-picked from 35b6f0e2a, written on parambole/w25-mega, whose gdn_cp_pass1_states_and_bf16_qkv comment in base.yml is two lines long; on this branch it is the four-line version, so base.yml conflicted at that comment. Resolved as the union: gdn_gated_norm_layout's lines follow gdn_out_proj_token_minor as in the original, then this branch's comment unchanged; nothing else differs. * Tests: gdn_gated_norm_layout is bit-identical and pins the GDN gated-norm output row-major Two-sided CPU checks through the real Qwen3NextGatedDeltaNet, alone and as a remat'd two-layer scan: loss and every gradient bit-identical for auto, natural and natural_fwd_bwd; the lowered programs carry exactly the expected row-major LayoutConstraints on the norm output (one per forward evaluation, plus one per backward for natural_fwd_bwd) and none in prefill; config defaults and refusals. Planted variants (pin dropped, permutation (0, 1, 3, 2), heads reversed, no cotangent pin) must fail the same checks. * GDN: gdn_gated_norm_layout="token_minor" pins the bf16 GDN output into the norm On the Qwen3.5-397B training cell XLA runs the gated RMSNorm token-minor: the gate z leaves in_proj_qkvz token-minor and out_proj reads its input token-minor (gdn_out_proj_token_minor). It relayouts the GDN kernel output only after the float32 convert, so every norm call (forward and remat replay) re-tiles the bf16 output, converts it and transposes the float32 [tokens, value_dim] copy. "token_minor" pins the bf16 GDN output token-minor (0, 2, 3, 1) in the forward, before the norm: the relayout becomes one bf16 transpose ahead of the convert and the norm reads bf16. Pinning the norm output row-major instead ("natural", "natural_fwd_bwd") keeps the norm token-minor in parts and makes XLA transpose z and the norm output as well, so those modes add copies on that cell; they stay available for comparison. Layout only: values are unchanged. The test covers the new mode, its prefill exclusion and two planted variants. * Tests: token_minor reorders the gated-norm scale-gradient sum on CPU On XLA:CPU the token_minor pin changes the order of the float32 sum over the tokens that forms the gated-norm scale gradient (max |d| 9.5e-7, rel 1.3e-7); the loss, the input gradient and every other parameter gradient stay bit-identical. The test now requires exactly that: bit-identity everywhere except that one leaf, which must match to a relative 1e-6, and the other modes stay exact. Planted controls: token_minor with reversed heads, the token_minor arrays under the exact rule, and the scale gradient moved 1e-5 relative must all fail. * Attention load balance: opt-in reorder without the self-copy context_parallel_attention_load_balance moves each attention layer's tokens into DUAL_CHUNK_SWAP order and back with two collective-permutes of half a shard. In each permute some rank sends its half to itself (N=4: rank 0 for the first halves, rank 2 for the second), which runs as a local copy inside the collective-permute. New key context_parallel_attention_load_balance_order: - dual_chunk_swap (default): today's code, unchanged. - dual_chunk_swap_local_self: the (r, r) pairs are dropped from both permutes and the self rank takes its own half with a select. Same token order and values; forward and gradients are bit-identical. A perm made only of self pairs (N=2) is no permute at all. Requires context_parallel_attention_load_balance=True with an even context parallelism of at least 2 (so the Ulysses path is refused). The reorder logs one [attn-lb-order] line naming the order. The existing test's reorder patch forwards the new keyword. * Tests: context_parallel_attention_load_balance_order is bit-identical and drops the self pairs tests/unit/attention_load_balance_order_test.py, for the opt-in reorder without the self-copy. In a subprocess with 8 CPU devices: - the reorder under dual_chunk_swap_local_self (both directions and their VJPs) equals dual_chunk_swap bit for bit at context 2, 4 and 8, for float32, bfloat16 and int32 inputs; - its compiled collective-permutes carry dual_chunk_swap's source-target pairs minus every self pair (at context 2 the all-self first-half permute is gone), and dual_chunk_swap's program has exactly the self pairs load_balance_self_ranks names; - a tiny hybrid Qwen3.5 (GatedDeltaNet and gated attention, MoE) on a context=4 x expert=2 mesh with the attention lo…
ultrons
pushed a commit
to ultrons/tunix
that referenced
this pull request
Oct 25, 2026
…her_once join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now also turns on the two remaining keys of the MaxText change on top of AI-Hypercomputer/maxtext#5721 (branch parambole/rl-4600-final-maxtext, pinned two commits ago): the SwiGLU backward fused into the down projection's input-gradient GMM, and the query kernel gathered once per attention layer in the backward. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains: - moe_wo_gmm_fused_dswiglu=true, on every trainer mesh: the input-gradient GMM of the MoE down projection turns its float32 accumulator into d[gate | up] in its epilogue (with moe_router_grad_from_wo_input, also the routing weight's gradient) instead of writing the activation's gradient in bfloat16 for separate passes to read back. The gradients of wi and of everything before it change by bfloat16 rounding. MaxText refuses the key unless that GMM has one n tile (wo_tile_dlhs_mlp_dim = the MoE MLP dim, 1024 in this recipe) and, with moe_wi_gmm_fused_swiglu, wi_tile_fwd_mlp_dim halves to that dim or to a multiple of 128 dividing it (this recipe's 1024 gives two n tiles of 512 per half); it also needs explicit tiles, prefuse_moe_weights, use_tokamax_gmm, use_gmm_v2 and a remat policy that does not save the routed experts' gate or up half. - query_proj_kernel_gather_once=true, at every trainer mesh except the 256-chip 32x2x8 and 64x1x8: the attention query projection keeps the kernel it gathered in the remat replay for its input gradient, so the backward gathers the query kernel once per attention layer instead of twice. The values are the same. MaxText refuses the key unless the backward replays the query projection (remat_policy=custom with query_proj=remat). google#2884's block sets query_proj=device at 32x2x8 and 64x1x8, where the backward does not replay the projection, so the key goes in a new default branch of that mesh case instead of on the tuned line: passed there, MaxText would refuse the 256-chip trainer when it starts. Measured on TPU together (F2): the bench cell of this block at 32x2x4 with these two keys, context_parallel_attention_load_balance_hoist_ids and cp-as-ep-dense-no-cp, on the MaxText of AI-Hypercomputer/maxtext#5758 with the change on top of it (4ed82d793, whose src/maxtext is the pinned MaxText's without #5721's three commits of Rohan Bierneni and two model YAMLs). Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, padding first, 9 micro-batches per step, mean of steps 1 to 4, no token dropped), booking r-7a4cc: 80.398 and 80.342 s per step in two pods (o275x-f2-l3-prof in parambole-bf2-book, o275x-f2-l3-r in parambole-bf2r-book), 3,668 and 3,671 tokens/s/chip, regular micro-batch 8.748 and 8.749 s, against 82.373 s per step, 3,580 tokens/s/chip and 8.966 s for the code of #5758 with the keys before these four (o275x-prfc-f1 in parambole-bf1-book). The F2 pods' own control arms did not finish, so the comparison is across pods: -1.974 and -2.030 s per step. Numerics in both pods (nxf2: loss gates off, random advantages, three steps): all six gates pass against that code under the same protocol (nxprf-l3) and against the reference pair (nxref, nxref2); the loss is 0.0304453575 in all three steps, nxprf-l3's to the ten digits logged, and the gradient norm (gate N3) is within 1.97e-4 and 1.23e-4 relative of nxprf-l3's (tolerance 2e-3). Each key's own measurement: - moe_wo_gmm_fused_dswiglu, behind an environment switch on an overlay build with the keys of google#2913's recipe, the hoist and padding first, without the fused log-prob head, against a control in another pod: 83.316 -> 82.148 s per step (o275x-c9pfxctl-l3 -> o275x-c9ypfctl-l3), 9.0845 -> 8.9545 s per regular micro-batch. On one TPU v7x chip its kernels at the cell's shapes match a float32 SwiGLU backward to 1.658e-3 relative L2 (parambole-bck-2d4c2), and its numerics run passes all six gates (nxc9ypfckhv). - query_proj_kernel_gather_once: not timed on TPU alone. An AOT compile of the cell with cp-as-ep gathers the query kernel once per attention layer in the backward instead of twice at the same device peak (104.46 GiB), and the MaxText change's test, executed on TPU v5p, gives bit-identical values and gradients with and without it. Not run: with #5721's three commits (#5721's own content with the rule ran at 81.230 and 81.168 s per step in booking r-67724, o275x-rbd-v16 and o275x-rbd-r-v16), at the 256-chip meshes, or end to end. This is b7f7a6570 (branch parambole/rl-trainer-4600-next of google/tunix) without its third value, custom_mesh_and_rule=cp-as-ep-dense-no-cp, which google#2884 already passes, ported onto google#2884 and replayed onto its head b16ee52. The port keeps google#2884's structure: the header (google#2884 rewrote it for the 256-chip meshes; it gains one line naming the rule under TRAINER_TUNED_FLAGS), the end of the block's comment, the tuned line (google#2884's, with moe_wo_gmm_fused_dswiglu after the hoist key), the mesh case (the new default branch) and, in the test, the defaults (the two keys join the MaxText-change defaults ahead of google#2884's per-tensor remat keys). On replay the same two comments as in the previous commit conflicted (the recipe's sentence on where the keys are defined and the test's comment on _BENCH_MAXTEXT_KEYS), resolved the same way, now naming all three keys. The block's comment on the whole cell replaces "has not run on TPU" with the F2 measurement above. Tests, in the bench image (dbdcf5f8dcba; CPU, no network, the tree read-only, stubs for kubectl, gcloud and gsutil first on PATH): - mlperf_submission_launcher_test.py: 48 of 48 pass at this commit, at each of the two commits before it, and at google#2884's b16ee52 and 8647d69. - Five deliberately wrong copies of the recipe each fail exactly the tests predicted before the run: moe_wo_gmm_fused_dswiglu or context_parallel_attention_load_balance_hoist_ids dropped from the block (the 3 exact-set tests); the default branch not passing query_proj_kernel_gather_once (those 3, the default-cell test and the two 256-chip mesh tests); query_proj_kernel_gather_once also on the tuned line (the 4 tests of the 256-chip meshes); moe_wo_gmm_fused_dswiglu passed only with TRAINER_TUNED_FLAGS=false (the 3 and the two switch-off tests). This test fails 6 tests on b16ee52's recipe, and b16ee52's test fails 5 on this recipe, both as predicted. - Rendered, the submission launch differs from the parent's only in the trainer's MaxText flags: moe_wo_gmm_fused_dswiglu=true and query_proj_kernel_gather_once=true at 32x2x4, moe_wo_gmm_fused_dswiglu=true at 32x2x8 and 64x1x8. With TRAINER_TUNED_FLAGS=false, and the other 397B v7x recipes at their defaults, the renders are byte-identical. - MaxText config check (the rendered trainer JobSet through run_trainer_node's argument parsing and build_maxtext_config to pyconfig.initialize, which runs the validators, with mocked devices): the pinned MaxText (94f11c09d) accepts the trainer at 32x2x4 (256 devices; all three keys on, query_proj remat), at 32x2x8 and 64x1x8 (512 devices; query_proj_kernel_gather_once off, query_proj device) and with TRAINER_TUNED_FLAGS=false. It refuses the 32x2x4 trainer on the previous pin bbd0f6f10, which does not define context_parallel_attention_load_balance_hoist_ids, and refuses query_proj_kernel_gather_once at 32x2x8 next to the block's query_proj=device ("query_proj_kernel_gather_once=True requires a remat policy that replays the query projection").
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on
parambole/rl-trainer-4600-base(#5698's head3792fdef8, the MaxText the RL trainer bench image runs, merged withatwigg/mlperfe526ea0), so the diff shows only this work. Retarget toatwigg/mlperfonce #5698 lands. Companion tunix PR: google/tunix#2913.Nine opt-in config keys, all off by default (seven recommended, plus two alternatives in
048d61701), that speed up the Qwen3.5-397B-A17B RL trainer on 128 TPU v7x. On the RL trainer bench (#5701), this PR's head with the companion tunix recipe, run with no environment switches, takes 82.18 s per step against 106.30 s before this work: 2,774 -> 3,589 tok/s/chip (two runs, 82.200 and 82.158 s; row 4 of Results). The six-gate numerics A/B passes for the same changes on this PR's pre-review code (85.30 s, row 2) and on an overlay build of the full recipe (82.67 s = 3,567 tok/s/chip, row 3); it has not yet run on the exact head build. The 106.30 s cell also lacked the already-merged staged step boundary (#5718); per regular micro-batch, which the boundary does not change, the time goes from 10.77 to 8.95 s (row 4). The 82.67 s cell ran on an overlay build of the bench image's MaxText that also had one attention change, which follows in a separate PR, and lacked the merged #5568 (see Results).Where each change acts
One decoder layer: full attention in 15 of the 60 layers, GDN in the other 45, and the MoE block in every layer. Blue: changed by this PR, labelled with the key that turns it on. Gray: unchanged.
flowchart TB classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 subgraph MIX["token mixer: attention or GDN"] ATT["full attention"]:::same CONV["GDN Conv1D + SiLU<br/>gdn_conv1d_dense_token_stack<br/>gdn_conv1d_input_in_qkv_dtype"]:::pr DR["GDN gated delta rule<br/>xor_dnc, always on"]:::pr CONV --> DR end subgraph MOE["MoE block"] RT["routing<br/>moe_routing_index_math=count<br/>moe_roe_packed_routing_gather"]:::pr WI["wi GMM + SwiGLU<br/>moe_unsplit_prefused_wi<br/>moe_wi_gmm_fused_swiglu"]:::pr WO["wo GMM + unsort<br/>moe_router_grad_from_wo_input"]:::pr RT --> WI --> WO end ATT ---> RT DR --> RT style MIX fill:none,stroke:#9ca3af style MOE fill:none,stroke:#9ca3afChanges
moe_routing_index_math=countmoe_roe_packed_routing_gathero275x-mf-ctl; -0.096 againsto275x-gdn-ctl)moe_unsplit_prefused_wimoe_wi_gmm_fused_swiglufuse_act="silu")moe_prefused_swiglu_vjpin its place)moe_router_grad_from_wo_inputgdn_conv1d_input_in_qkv_dtypegdn_conv1d_dense_token_stackjnp.stackrows, which the October 2026 TPU compiler pads from 1 to 8 sublanes (older compilers keep them dense)jnp.stackGains are against a control in the same pod unless marked (for
gdn_conv1d_dense_token_stack, a same-pod arm whose only change, a 2D GDN out_proj, is a no-op); mb = micro-batch, 9 per step. Also in the branch, with no key:dede90a9f: the GDN triangular inverse with xor_dnc, Rohan's Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 cherry-picked as a dependency (drop it when Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 merges). -0.191 s/mb (cross-pod), TPU A/B PASS 6/6.7f5dcc7ce:stop_gradienton frozen parameters insidefwd_bwd; the RL recipe freezes the router. -0.058 s/mb, both arms withmoe_routing_maps=device(cross-pod), TPU A/B PASS 6/6.8cc954eb2: review feedback. It lowers the MoE layer, the gmm_v2 kernel and the GDN op to the same TPU programs as the commit before it (see Testing).Results
RL trainer bench (#5701), 128 TPU v7x. Seconds per optimizer step, mean of steps 1-4.
o275x-hd)o275x-prc-l3-profand its repeato275x-prc-l3-r, one pod)nxprc)o275x-c9pfxhv-l3, two pods)nxc9pfxhv)o275x-prf-l3-profand its repeato275x-prf-l3-r, one pod)2d82c5199(same MoE-layer program as the published commits, GDN and model sources unchanged; see Testing), with the keys set as config: the companion recipe's key set withoutgdn_conv1d_dense_token_stack, pad first and the fused log-prob head. It includes the six recommended earlier keys of this PR (048d61701's two alternative keys were off), the GMM tiles at 512 and log-prob chunk 4096.gdn_conv1d_input_in_qkv_dtypewas process-wide in that build and is scoped here, but the GDN layer program is the same (see Testing). The stop_gradient on frozen parameters has no such check. That build also hoisted attention's load-balancing permute, which is in neither PR (alone, against another pod's control: 99.504 -> 99.267 s per step, regular micro-batch 10.033 -> 10.020 s).o275x-c9ypfxhv-l3), took 82.945 s per step (3,556 tok/s/chip; regular micro-batch 9.04 s). That pod's overlay build also contained the follow-up dSwiGLU fix, with its key off.moe_route=device, which the recipe leaves at its default; no code on this branch tags a tensor with that name.20c877aa1, Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe google/tunix#2913bee63532f) installed in the bench image and run with the recipe's settings as config, with no environment switches: the recipe's 16 MaxText keys, the fused log-prob head at chunks of 4096, padding first (the bench driver's equivalent ofpack_pad_position=start),ragged_buffer_factor2.75 with 0 dropped tokens, and the same Pathways server image as row 3. It does not have the attention hoist that row 3 had, and it does have the merged perf(training_engine): buffer microbatch metrics in Python lists and gate checkpoint metrics on intra-step state #5568, which row 3 lacked. The six-gate numerics A/B has not been run on this build; rows 2 and 3 pass it.How to turn it on
Each MoE key's validator refuses a config in which the key would do nothing and names every unmet requirement.
gdn_conv1d_dense_token_stackis refused without Gated DeltaNet layers, withoutuse_gdn_kernel, or when the Conv1D width each kernel call convolves is not a multiple of 1024, one requirement at a time;gdn_conv1d_input_in_qkv_dtypeis refused only foruse_gdn_kernel=Falseon a GDN decoder. The companion recipe,TRAINER_TUNED_FLAGS=truein google/tunix#2913, sets these seven together with its other MaxText keys; the 82.67 s cell measured those settings on the overlay build (see Results).Testing
8cc954eb2; the dense-stack test, and the Conv1D dtype test again, on the tip20c877aa1. The gmm SiLU kernel test (megablox_gmm_silu_test.py) needs TPU and has not run there yet.7f5dcc7ce).cb709c93clowers to the same program as the pre-review tree that ran row 2 andnxprc, in 45 of 45 cases (9 key sets in 5 variants), with the GDN and model sources unchanged between the two;8cc954eb2lowers it identically (16 of 16) and the tip changes no MoE source. Withgdn_conv1d_dense_token_stackon, the GDN layer equals the overlay's switch applied to8cc954eb2in 18 of 18 configurations. Against the measured overlay tree itself, with its switches, the tip lowers the MoE layer (the recipe's five MoE keys) and the GDN layer (both GDN keys, andgdn_conv1d_input_in_qkv_dtypealone, at the 397B head counts with sequence-sharded context parallelism 4; both keys also without context parallelism) to the same programs. Negative controls differ, as they must.8cc954eb2;gdn_cpalso on the tip): GDN backward tests 14 passed, 1 skipped;gdn_cp7 passed; in a 31-file regression subset, run oncb709c93c(before the review-feedback and dense-stack commits), every failure also fails at7f5dcc7ce(the tree before the keyed commits).nxprc) and for the full recipe (nxc9pfxhv). The gates compare scalars of the first of the run's 3 steps (the later two are logged, not gated), so they catch gross errors, not a wrong gradient direction; exactness rests on the CPU bit-identity tests, and the two rounding-level keys have finer checks (details below).alogsix times ingdn_conv1d_dense_token_stack_test.py(read as "along"); the rename follows in a separate PR. pyink reformats 4 hunks, all in lines of the cherry-picked Optimize GDN triangular inverse with xor_dnc and support gdn_chunk_size=128 #5667 (dede90a9f); the other pyink hunks in the changed files are already at base.Not in this PR (follow-up PRs)
moe_wo_gmm_fused_dswiglu(SwiGLU backward in the wo DLHS GMM epilogue) and its layout fix, in a follow-up PR.wi_tile_fwd_mlp_dim1024 the forward kernel halves its N tile to 512, so the array was laid out per tile as [g0|u0|g1|u1], and the DLHS epilogue split it in half: a wrong gradient. The key's own commit refuses that tile setting at config time; the fix commit reads the per-tile layout. This PR's fused SiLU writes gate and up as separate arrays and is not affected.nxc9ypfckhv, PASS 6/6 (these gates also passed the unfixed version, see Numerics, so the kernel check is the evidence for the fix).o275x-c9pfxctl-l3-> 82.148 ino275x-c9ypfctl-l3, two pods).How it was measured: bench, timing, logs, profiles
ragged_buffer_factor2.75, with 0 dropped tokens in every arm.r-4e5a9/parambole-bxa-4e5a9; rows 2-3:r-5f2bf/parambole-bpr-5f2bf,parambole-bvx2-5f2bfandparambole-bvx4-5f2bf; the intermediate arm:parambole-bvx-5f2bf). The same control measured in different pods spreads by at most +/-0.01 s per micro-batch:o275x-hdin four pods took 10.7735-10.7845 s per regular micro-batch and 106.04-106.30 s per step.o275x-c9pfxctl-l3, same overlay build): 83.32 s per step. Per regular micro-batch the chain is 10.77 -> 9.30 -> 9.08 -> 9.01 s. Between row 2 and this arm,gdn_conv1d_dense_token_stackand tunix pad first are added and the build changes too (overlay build: the hoist on, perf(training_engine): buffer microbatch metrics in Python lists and gate checkpoint metrics on intra-step state #5568 absent).gs://parambole-397b-replay-nap/win-tpu/sched/<booking>/<job>/<arm>.log, for example:r-4e5a9/parambole-bxa-4e5a9/o275x-hd.logr-5f2bf/parambole-bpr-5f2bf/o275x-prc-l3-prof.logando275x-prc-l3-r.log; numericsr-5f2bf/parambole-bnpr-5f2bf/nxprc.logr-5f2bf/parambole-bvx2-5f2bf/o275x-c9pfxhv-l3.log, replicater-5f2bf/parambole-bvx4-5f2bf/o275x-c9pfxhv-l3-r.log, third podr-2d4c2/parambole-bn3s-2d4c2/o275x-c9ypfxhv-l3.log; numericsr-5f2bf/parambole-bvx3-5f2bf/nxc9pfxhv.logo275x-prc-l3-prof, under.../r-5f2bf/parambole-bpr-5f2bf/o275x-prc-l3-prof/.../tensorboard/plugins/profile/.o275x-c9pfxall-l3-prof, 82.787 s per step; the recipe arm itself ran without a profiler):.../r-5f2bf/parambole-bvx-5f2bf/o275x-c9pfxall-l3-prof/parambole-bvx-5f2bf-o275x-c9pfxall-l3-prof/tensorboard/plugins/profile/2026_10_15_14_43_36/server-gke-tpu-174635e9-sj65.0.0.xplane.pb.o275x-hd-prof.o275x-*/nx*arm ID maps to a log at the path above. Arms ending in-profalso have an xplane profile under.../<job>/<arm>/<job>-<arm>/tensorboard/plugins/profile/.Numerics: the six checks, per-run results, and what they cannot see
The bench's gates-off A/B, run for every kept change: log-prob gates off, random advantages, 3 steps, against a reference run and a replicate of it. Six checks: every arm keeps its tokens; loss, relative difference; mean log-prob difference; last-micro-batch log-prob difference; entropy, relative difference; gradient norm, relative difference.
nxprc2d82c5199) with the six MoE/GDN keys of stack 8T as confignxc9pfxhvnxs6grgdn_conv1d_dense_token_stackon the environment-switch build, stacked as in its timing arm (withmoe_prefused_swiglu_vjp, not the fused SiLU); steps 0-1 of 3, cut in step 2 by a dropped proxy connectionnxs6silumoe_wi_gmm_fused_swigluon the environment-switch build, stacked as in its timing armnxgx2-c1b1gdn_conv1d_input_in_qkv_dtypealonenxt67,nxstk2stop_gradientWhat these checks cannot see. The six gates compare scalars of the first optimizer step (the run has three steps; the later two are logged, not gated), so they rule out gross errors only, not a wrong gradient direction.
moe_wo_gmm_fused_dswiglu) next to the fused SiLU at a 512-wide forward tile. That key is not in this PR.moe_router_grad_from_wo_input: an op-level precision test against a float32 reference, both paths within relative L2 1e-2, measured at most 2.4e-3.moe_wi_gmm_fused_swiglu: kernel testmegablox_gmm_silu_test.py, which needs TPU and has not yet run on TPU. Runsnxprcandnxs6siluare that kernel's TPU evidence.Tests: per-file counts, program identity, mutation tests, lint
All tests run on CPU in the bench image. Unit tests on the tree of
8cc954eb2; the dense-stack test, and the Conv1D dtype test again, on the tip20c877aa1.add4a1e74tests/unit/moe_routing_index_math_test.py048d61701tests/unit/moe_glue_rewrites_test.pyc88674bbctests/unit/moe_roe_packed_routing_gather_test.py4ab9f123dtests/unit/moe_wi_gmm_fused_swiglu_test.py;tests/unit/megablox_gmm_silu_test.py(TPU only)63cba2632tests/unit/moe_router_grad_from_wo_input_test.py, including the float32-reference precision testcb709c93ctests/unit/gdn_conv1d_input_in_qkv_dtype_test.py20c877aa1tests/unit/gdn_conv1d_dense_token_stack_test.py, including a program check in a child processRelated suites (tree of
8cc954eb2; the threegdn_cpsuites also on the tip, same results):gdn_cp: 7 passed.gdn_cp_head_shardedandgdn_cp_precision: one failing test each, the same at base and at tip. It is a pre-existingSimpleNamespaceerror atqwen3.py:897.cb709c93c(before the review-feedback and dense-stack commits): every failure there also fails at7f5dcc7ce(the tree before the keyed commits).Program identity (normalised TPU v7x StableHLO, Mosaic kernel bodies decoded, source locations stripped):
7f5dcc7ce), for the MoE layer and the GDN op on the local and context-parallel paths.8cc954eb2vs the publishedcb709c93c: MoE layer identical for 8 key sets in 2 variants (bench, retry), 16 of 16; gmm_v2 kernel identical for 2D, group offset, transpose_rhs, 3D, windowed 3D and SiLU; GDN op identical. The tip20c877aa1changes no MoE source.cb709c93cvs the pre-review tree that ran theo275x-prc-l3arms andnxprc: the MoE layer is identical for all 9 key sets in 5 variants (45 of 45), and the GDN and model sources are unchanged, so those measurements apply to this branch.gdn_conv1d_dense_token_stack: key off equals8cc954eb2, and key on equals the overlay's switch applied to8cc954eb2, in 18 of 18 configurations (162 probes). Negative controls differ, as they must.gdn_conv1d_input_in_qkv_dtypealone matches at the first. Negative controls differ, as they must.Mutation tests (tests shown to fail on deliberately broken copies):
8cc954eb2: R1, R4, R6, R8, R8p, R9, R10 and R11.gdn_conv1d_dense_token_stack: 11 of 11 mutants from the test reviewer are killed; each fails at least one test that passes on the clean tip (41 passed). They cover each context-parallel kernel call site dropping the option, the backward shim using the wrong slot, the batched kernel getting the option, a concat off the token path, the runner reading the wrong key, a bf16 round trip and a dtype guard that accepts bf16, a token shift by one, and a validator threshold.Lint (the CI pins), tip vs base: pylint 3.3.8, 0 new messages; yamllint clean; codespell 2.4.1 flags
alogsix times intests/unit/gdn_conv1d_dense_token_stack_test.py(read as "along"; the rename toa_logfollows in a separate PR); pyink 24.10.1 reformats 4 hunks, all in lines the cherry-picked #5667 adds (dede90a9f: 3 ingdn_bwd/api.py, 1 ingdn_bwd_pallas_test.py), and the other pyink hunks in the changed files are already at base. Each tool fires on a deliberately broken file.Tried and not kept (same bench)
ragged_buffer_factor2.25: not allowed; 2.75 stays, or 2.0 only with zero drops.moe_pin_sparse_core_all_gathers: +0.363.gdn_cp_mode=head: +0.294.o275x-gdn-ctl).Corrections to the published commit messages
The six key commits (
add4a1e74tocb709c93c) were published before the last review pass, so their messages cannot be changed. The follow-up commit "Address review feedback on the MoE and GDN keys" carries the code changes; the corrected wording for the messages is here.TPU numerics of this PR's own code (replaces "it ran on the environment-switch build, not on this branch's binary"). Applies to
moe_routing_index_math="count",moe_unsplit_prefused_wi,moe_roe_packed_routing_gather,moe_wi_gmm_fused_swiglu,moe_router_grad_from_wo_inputandgdn_conv1d_input_in_qkv_dtype.nxprcran this PR's code on TPU with those six keys set as config, not as environment switches. It ran the pre-review revision of these commits; a TPU v7x lowering check shows that revision lowers the MoE layer to the same program as the published commits for all 9 key sets checked, in 5 variants (45 of 45), with the GDN and model sources unchanged between the two, and the follow-up commit keeps that identity.o275x-prc-l3-prof, and 85.251 s in its repeato275x-prc-l3-r.nxprc:moe_prefused_swiglu_vjp(TPU numerics only from environment-switch-build runs:nxs6c,nxs6lc,nxs6t512andnxs6gr, each PASS 6/6) andmoe_ragged_sort_unique_pair_argsort(in no TPU run; a permutation, exact by construction, bit-identical on CPU).What the A/B can catch (add wherever the A/B is cited). These gates compare scalars of the first of 3 steps, and they have passed a run with known-wrong MoE gradients, so they rule out gross errors only. Exactness of the exact rewrites rests on the CPU bit-identity tests. The two numerics changes need more care:
moe_wi_gmm_fused_swiglu, whose kernel test (megablox_gmm_silu_test.py) is not yet run on TPU, andmoe_router_grad_from_wo_input, a rounding-level gradient change whose finer check is the op-level precision test against a float32 reference.moe_ragged_sort_unique_pair_argsortandmoe_prefused_swiglu_vjp(048d61701). Neither key is in the recommended set, and neither was timed alone. Both were timed only together with the whole wi (moe_unsplit_prefused_wi): all three keys in one arm, withmoe_routing_index_math="sort", -0.170 s per micro-batch (o275x-mf-gb-profagainst its in-pod controlo275x-mf-ctl). They are kept as alternatives for configurations the recommended set does not cover. The pair argsort is for runs that keepmoe_routing_index_math="sort". The fused activation is for runs withoutmoe_wi_gmm_fused_swiglu, which needs tokamax gmm_v2 and plain silu + linear, while the fused activation works with any GMM backend and gated activation.moe_wi_gmm_fused_swiglu(4ab9f123d).megablox_gmm_silu_test.py, the kernel's own TPU test, is not yet run on TPU. Runsnxprcandnxs6siluare the kernel's only TPU evidence._store_group_rows), not a copy of it. Withreturn_pre_actoff, the kernel lowers to the same Mosaic body as before.moe_router_grad_from_wo_input(63cba2632).Test additions from the mutant re-review.
moe_wi_gmm_fused_swiglu_test.py: the gmm_silu VJP test also runs with heuristic tiling, so the option must reach gmm's backward.moe_glue_rewrites_test.py: the whole wi is rejected under tensor_sequence and autoregressive parallelism, and the fused activation under every minimal* remat preset.moe_routing_index_math_test.py: the validator error names every unmet requirement when several are unmet.8cc954eb2: these additions kill the re-review's surviving mutants R1 (forward and backward error bound, dropping buffer), R4 (forward error 0.585 against the 1e-2 bound), R6, R9 (both parallelism cases), R10 and R11, plus the trace-time guard removed from either chunk loop. Not addressed: R5 (its killing test,moe_megatron_aux_loss_test, still needs 8 host devices to run) and R14 (needs a down-projection stand-in that writes NaN outside its groups).gdn_conv1d_input_in_qkv_dtype(cb709c93c): measurement and numerics.o275x-gx2-c1b1-prof). That pod ran no flags-off control. Against the flags-off control of another pod (o275x-gdn-ctl, 10.033 s, 99.504 s per step), that is -0.050 s per micro-batch and -0.45 s per step. The same control measured in different pods spreads by at most +/-0.01 s per micro-batch (o275x-hdin four pods: 10.7735-10.7845 s).conv1d_siluforward and backward in the process through an environment variable. This commit scopes the option togdn_decoupled_conv1d. A lowering check of one GDN layer (16 key / 64 value heads, sequence-sharded context parallelism 4, segment ids, remat) gives the same program for the build's switch and this commit's key, alone and together with the dense stack; with both on, also without context parallelism.nxgx2-c1b1, environment-switch build): all six checks pass. Loss relative difference 4.6e-8, mean log-prob difference 2.6e-6, gradient-norm relative difference 9.3e-5. That is 4.9x the replicate spread of 1.9e-5, under the 2e-3 tolerance. Runnxprcran this commit's scoped option together with the five recommended MoE keys and passes all six checks. The option is also in runnxs6c's stack.dede90a9f). Its api.py import hunk does not apply without it.