Repository navigation
[MLPerf RL 397B] Final 4,611 tok/s/chip trainer kernel & config stack (supersedes #5799, includes #5571) - #5885
Merged
Merged
Conversation
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.
…wiglu) 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.
… 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.
…s 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.
…ery_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
20c877a:
- 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.
…'s config Since "Skip the GDN kernel's segment path for unpacked training rows" (in #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.
… 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/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.
…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).
…t 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.
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.
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.
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.
…n 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.
…j token-minor layout Combine Mohit Khatwani's three GDN kernel and layout optimizations (commits 906e157, 9c48e65, and 80c1495) adapted onto 1bf840c: 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.
…ken_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).
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.
…8_bwd_dispatch, moe_combine_per_chunk) Ported from the MLPerf v6.1 DeepSeek-V3 code (6639e31, "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.
…gged_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).
…sort'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.
…ack_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).
…ault 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.
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.
…_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 818caa3, before the appended argument.
…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.
…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.
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.
…e 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.
…d_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.
uwlei
requested review from
Rohan-Bierneni,
khatwanimohit and
ultrons
and removed request for
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
jshin1394,
khatwanimohit,
michelle-yooh,
richjames0,
shuningjin,
vipannalla and
zxhe-sean
October 21, 2026 21:46
Collaborator
Author
|
the check errors are not caused by this PR |
Set `last_fwd_bwd_token` (the microbatch loss array) in `fwd_bwd` and `last_update_token` (the gradient norm) in `update`. Tunix's TrainerWorker step timer blocks on these off the serving thread to stamp when each microbatch and the optimizer step finished on device (TRAINER_STEP_TIMING, google/tunix#2713). Without them the timeline carries dispatch times only. The loss rather than the running sums, which the next microbatch's accumulate donates; the gradient norm is the same signal the throttler already waits on. On a dropless rerun of an overflowed microbatch (`_resolve_overflow_check`), the rerun's loss replaces the token, since it is the last device work dispatched.
uwlei
changed the title
Rohan-Bierneni
approved these changes
Oct 23, 2026
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.
Description
Clean, lint-passing (
pyink,mdformat,codespell,yamllint,pylint 10.00/10) PR superseding #5799 (parambole/rl-4600-final-maxtext) and including #5571 (sivaibhav/engine-step-timing-tokens) ontoatwigg/mlperf(also closes already-landed #5758 and #5783, whose changes are already present in #5721 onatwigg/mlperf).Brings the full Qwen3.5-397B-A17B RL trainer kernel, sharding, FP8 MoE, Gated DeltaNet (GDN), host-offloaded FP32 master-residual storage, and
TRAINER_STEP_TIMINGreadiness token stack ontoatwigg/mlperf, reducing 128-chip TPU v7x (32x2x4,GBS=1152) trainer step time from80.835 s/step(3,649tok/s/chip) to63.963 s/step(4,611tok/s/chip,-20.9%step time /+26.4%TPS):fp8_fake_quant_hoist=gather8,wi_tile_fwd_mlp_dim=2048+wi_tile_fwd_rhs_buffer_count=2,moe_fp8_token_all_gather=true,moe_fp8_bwd_dispatch=true,ragged_unsort_spread_masked_indices=true,moe_fallback_topk_capacity=128,moe_routing_index_math=count,moe_unsplit_prefused_wi=true,moe_roe_packed_routing_gather=true,moe_wi_gmm_fused_swiglu=true,moe_wo_gmm_fused_dswiglu=true,moe_router_grad_from_wo_input=true.gdn_conv1d_natural_layout=true,gdn_fwd_kernel_head_group_size=4,gdn_in_proj_kernel_order=true,gdn_out_proj_token_minor=true,gdn_gated_norm_layout=token_minor,gdn_cp_pass1_states_and_bf16_qkv=true,gdn_cp_fast_state_passing=true,gdn_cp_m_scan=true,gdn_cp_dh_scan=true,gdn_bwd_natural_layout=states.context_parallel_attention_load_balance_order=dual_chunk_swap_local_self,dense_grad_reduce=all_reduce,fp32_master_storage=residual16,optimizer_device_resident_gib=6.c1bd747628): Exposeslast_fwd_bwd_tokenandlast_update_tokenonMaxTextTrainingEngineso TunixTrainerWorker(StepTimer/TRAINER_STEP_TIMINGandstage_fwd_bwd) can record device completion timestamps and confirm optimizer updates without an extraget_metrics()wait.498efb5b70,c7131f8be6):.github/workflows/code_quality.ymlfailures (pyink --all-files,mdformat --all-files,codespell, andpylint 10.00/10).optimizer_device_resident_gibfloat validation intypes.py,ragged_unsort_spread_masked_indicesguard inmoe.py,gdn_cp_dh_scanvalidation inqwen3.py, and_FakeConfig/_StubConfigtest attributes). All new optimization keys remain default-off unless explicitly enabled in recipe configs.Companion Tunix PRs:
atwigg/mlperf. google/tunix#2979uwlei/rl-4600-final-tunix@54f488d8); **retarget base branch toatwigg/mlperfonce #2979 merges**. google/tunix#2980Tests
32x2x4,bodaborg-tpu7x-gscnamespacer-bf6f4) Trainer Bench A/B:atwigg/mlperf(bbd0f6f10+2d3d00ab8):80.835 s/step(3,649tok/s/chip)#5885 + #2979):63.963 s/step(4,611tok/s/chip) (7.163 smedianfwd_bwd,6.759 supdate,63.16 GiBpeak HBM /25.84 GiBheadroom, 0 MoE drops atragged_buffer_factor=2.0, matching Param'snx-m6p464.271 s/step/4,589tok/s/chip within 0.5%).pb-r40c35-s3onbodaborg-tpu7x-gsc, 128-chip trainer + 128x4-chip vLLM rollout workers):1,040s(10 steps completed cleanly, 91 micro-batches with 0 MoE drops and 0 replays underragged_buffer_factor=2.0, trainer off the critical path with ~20–25s headroom before weight sync).#5571), andtraining_engineunit test suites, plusCode Quality Check / Pre-commit Linters(PASS),Build MaxText Package / Build Wheel(PASS), andDocumentation Build Check / Build Sphinx Docs(PASS).Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.