Repository navigation
Qwen3.5-397B RL trainer on 128 v7x: fp8 fake-quant hoist, fused SwiGLU backward, attention backward savings, the wi forward GMM's weight pipeline depth and the GDN context-parallel m scan as config keys - #5799
Draft
parambole wants to merge 14 commits into
Draft
parambole wants to merge 14 commits into
parambole wants to merge 14 commits into
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.
![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 optimization features, including fusing the MoE SwiGLU backward into the down projection's DLHS GMM (moe_wo_gmm_fused_dswiglu), gathering the attention query projection kernel once instead of twice in the backward pass (query_proj_kernel_gather_once), and hoisting the load-balanced positions and segment IDs out of the layer loop (context_parallel_attention_load_balance_hoist_ids). The review identified two critical NameError issues: a missing import of checkpoint_name in attentions.py and an undefined tpu_info variable in pallas_mosaic_tpu_v2_gmm_kernel.py.
Comment on lines
+1339
to
+1341
| reordered_positions = checkpoint_name(reordered_positions, maxtext_utils.ATTENTION_LB_IDS) | ||
| if reordered_segment_ids is not None: | ||
| reordered_segment_ids = checkpoint_name(reordered_segment_ids, maxtext_utils.ATTENTION_LB_IDS) |
There was a problem hiding this comment.
The function checkpoint_name is not imported in src/maxtext/layers/attentions.py, which will cause a NameError at runtime when context_parallel_attention_load_balance_hoist_ids is enabled. Use jax.ad_checkpoint.checkpoint_name instead to avoid the missing import issue.
Suggested change
| reordered_positions = checkpoint_name(reordered_positions, maxtext_utils.ATTENTION_LB_IDS) | |
| if reordered_segment_ids is not None: | |
| reordered_segment_ids = checkpoint_name(reordered_segment_ids, maxtext_utils.ATTENTION_LB_IDS) | |
| reordered_positions = jax.ad_checkpoint.checkpoint_name(reordered_positions, maxtext_utils.ATTENTION_LB_IDS) | |
| if reordered_segment_ids is not None: | |
| reordered_segment_ids = jax.ad_checkpoint.checkpoint_name(reordered_segment_ids, maxtext_utils.ATTENTION_LB_IDS) |
| if cfgs.return_dact_dot: | ||
| # The dot is one float32 per row: broadcast to a lane row and stored through the extra output. | ||
| tiled_dot_ref, partial_dot_ref = pre_act_refs[0] | ||
| dot_tile_padded = jnp.broadcast_to(dot_tile, (cfgs.tiles.tile_m, tpu_info.num_lanes)) |
There was a problem hiding this comment.
The variable tpu_info is not defined in the scope of inner_kernel or kernel_main, which will cause a NameError at runtime when cfgs.return_dact_dot is enabled. Use pltpu.get_tpu_info().num_lanes instead, consistent with how it is accessed elsewhere in this file.
Suggested change
| dot_tile_padded = jnp.broadcast_to(dot_tile, (cfgs.tiles.tile_m, tpu_info.num_lanes)) | |
| dot_tile_padded = jnp.broadcast_to(dot_tile, (cfgs.tiles.tile_m, pltpu.get_tpu_info().num_lanes)) |
… 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.
parambole
changed the title
parambole
changed the title
…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.
parambole
changed the title
parambole
changed the title
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
2 participants
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
atwigg/mlperfatbbd0f6f10, the squash merge of #5721 (same tree as #5721's head22fa0c559). Companion tunix PR: google/tunix#2948, which turns these keys on in the tuned recipe and pins this branch.Seven opt-in config keys, all off by default, for the Qwen3.5-397B-A17B RL trainer on 128 TPU v7x. On the RL trainer bench (#5701), the companion recipe's keys before the last two take 79.34-79.54 s per step. The last two levers were each measured on that code without the other: the 2048 wi tile at 78.99-79.14 s per step = 3,727-3,734 tok/s/chip (-0.32 and -0.36 s/step against a control in the same pod), and the GDN context-parallel m scan at 78.80-78.82 s per step = 3,742-3,743 tok/s/chip (-0.66 s/step in the same pod). Both together, this PR's head, measured 78.35 s per step = 3,764 tok/s/chip (-1.01 s/step against a control in the same pod; under Results). All six numerics gates pass for each lever; a numerics run of the combination was cancelled before it ran. For comparison:
Where each change acts
One decoder layer, scanned 60 times: full attention in 15 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 PRE["before the layer scan: fp8 absmax of every<br/>routed-expert weight, once for all layers<br/>fp8_fake_quant_hoist"]:::pr subgraph MIX["token mixer: attention or GDN"] ATT["full attention, backward:<br/>load-balanced positions and segment ids saved once<br/>context_parallel_attention_load_balance_hoist_ids<br/>query kernel gathered once instead of twice<br/>query_proj_kernel_gather_once"]:::pr GDN["GDN, context-parallel forward:<br/>chunk transitions folded by one Pallas scan<br/>gdn_cp_m_scan"]:::pr end subgraph MOE["MoE block"] RT["routing"]:::same W["expert-weight FSDP all-gather<br/>fp8 codes, 1 byte per element<br/>fp8_fake_quant_hoist=gather8"]:::pr WI["wi GMM + SwiGLU, forward: one N tile of 2048<br/>per half, two weight buffers instead of three<br/>wi_tile_fwd_rhs_buffer_count=2"]:::pr WO["wo GMM, backward: SwiGLU backward<br/>in the GMM epilogue (float32 accumulator)<br/>moe_wo_gmm_fused_dswiglu"]:::pr RT --> WI W --> WI --> WO end PRE --> W ATT --> RT GDN --> RT style MIX fill:none,stroke:#9ca3af style MOE fill:none,stroke:#9ca3afThe fp8 fake-quant hoist, before and after
With
fp8_moe_fake_quant=truethe trainer rounds every routed-expert weight to its per-channel float8_e4m3fn grid, in every layer's forward and again in the remat replay.flowchart TB classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 %% AFTER is declared first on purpose: mermaid draws the last-declared subgraph first. subgraph AFTER["after, fp8_fake_quant_hoist=gather8"] direction TB a1["once, before the scan:<br/>absmax for all layers"]:::pr --> a2["per layer: quantize<br/>the shard to fp8 codes"]:::pr --> a3["FSDP all-gather<br/>fp8 codes, 1 B/element"]:::pr --> a4["dequantize<br/>(same expression)"]:::pr --> a5["GMM"]:::same end subgraph BEFORE["before: every layer, forward and again in the remat replay"] direction TB b1["expert weight<br/>FSDP shard, bf16"]:::same --> b2["wi absmax:<br/>max all-reduce<br/>(32-way)"]:::same --> b3["round to the<br/>fp8 grid"]:::same --> b4["FSDP all-gather<br/>bf16, 2 B/element"]:::same --> b5["GMM"]:::same end style BEFORE fill:none,stroke:#9ca3af style AFTER fill:none,stroke:#1d4ed8absmax: computes every layer's absmax once, before the layer scan (for wi, one all-reduce for the whole stack), and each layer takes its slice.gather8: also quantizes on the shard, all-gathers the float8_e4m3fn codes at half the bytes, and dequantizes after the gather.w_q + (w - stop_gradient(w))did.Changes
gdn_cp_m_scangdn_cp_mode=seq, the GDN context-parallel forward folds each rank's chunk transitions into its local transition product with one Pallas scan in VMEM, instead of an XLA reduction tree of matmulsgdn_cp_dh_scan(not set by the companion recipe)gdn_cp_m_scan(same pod); a profile shows the SparseCore reduce-scatters running in inverted orderwi_tile_fwd_rhs_buffer_count(0= the kernel default, or at least2)2,wi_tile_fwd_mlp_dim=2048(one N tile per half of [gate | up]) fits the kernel's VMEM: 50.95 MiB instead of 66.95, of 57.60wi_tile_fwd_mlp_dim=2048(same pod); a replicate in that pod -0.357fp8_fake_quant_hoist(off/absmax/gather8)gather8also all-gathers fp8 codes (diagram above)gather8(same pod); a replicate in a new pod landed within 0.042 smoe_wo_gmm_fused_dswiglucontext_parallel_attention_load_balance_hoist_idsquery_proj_kernel_gather_onceResults
RL trainer bench (#5701), 128 TPU v7x. Seconds per optimizer step, mean of steps 1-4; tok/s/chip = 294,912 / s per step.
20c877aa1)parambole-bf2-v18/o275x-prfc-f2parambole-bf2-v18/o275x-f2-l3-prof,o275x-f2-l3-rnxf2)parambole-rb4-v16/o275x-rbd-v16,o275x-rbd-r-v16(booking r-67724)nxrbd-v16)94f11c09d)parambole-bpr-v18/o275x-pr-v18,o275x-pr-r-v18;parambole-v19a/o275x-v19a-ctlnxpr-v18)e1cd8bf1d), fp8 hoist onparambole-v19a/o275x-v19a-fq8(same pod aso275x-v19a-ctl);parambole-v19b/o275x-v19b-fq8r(new pod); booking r-b6424:parambole-w21ds/o275x-d1-ctl,parambole-w21cm/o275x-c1-ctl(two more pods)nx-v19b-fq8)parambole-w21ec/o275x-ec-ctl15285fd74):wi_tile_fwd_mlp_dim=2048+wi_tile_fwd_rhs_buffer_count=2parambole-w21ec/o275x-ec-e6k,o275x-ec-e6kr(same pod aso275x-ec-ctl);parambole-w21e6k/o275x-e6k1,parambole-w21p3/o275x-p3-e6k(two more pods)nx-e6k1)parambole-w21g2a/o275x-g2-ctlgdn_cp_m_scan=trueparambole-w21g2a/o275x-g2-m(same pod aso275x-g2-ctl);parambole-w21gr/o275x-gr-m(new pod)nx-gm-m)gdn_cp_dh_scan=truealone; both scansparambole-w21g3/o275x-g3-dh(same pod aso275x-g3-ctl, 79.403);parambole-w21g2a/o275x-g2-both(same pod aso275x-g2-ctl)This PR's head, the tile and the m scan together, ran in one pod with its own control (job
parambole-w22k1a), from an overlay whose ten files are this branch's at1bf840cff(tgz4b980f9c, md5-checked). A replicate and a numerics arm of the combination were planned inparambole-w22k1bbut cancelled before they ran (chip budget); each lever passed the numerics gates on its own:TPU combined:
parambole-w22k1a/o275x-k1-ctl(this overlay, all keys off) 79.361 s/step (3,716 tok/s/chip),o275x-k1-em(the tile and the m scan, same pod) 78.349 s/step (3,764 tok/s/chip): -1.012 s/step. The prediction written before the run was -1.00 (band -1.30 to -0.70); the two levers' separate same-pod deltas sum to -0.996. Logs: md5 a17b17b3 (control), 8d7b0411 (combined).Numerics
The bench's six-gate A/B: loss gates off, random advantages, 3 steps, compared against the reference pair
nxref/nxref2. The gates check the first optimizer step (step 0); steps 1-2 are logged, not gated. Below, step 0 before the fp8 hoist, with it, and with the tile on top (this PR's head). The last column is how far two runs of identical code land apart on TPU.nxpr-v18)nx-v19b-fq8)nx-e6k1)nxrefvsnxref2)nx-gm-mvsnx-v19b-fq8): the loss is the same (0.0304453575); log-prob gap 12.2219148625, 7.3e-7 absolute away (last micro-batch 12.2205133, 2.9e-6); entropy 11.9241102, 3e-8 relative; gradient norm 0.0557771586, 2.5e-6 relative. Against the reference pair: log-prob gap 6.4e-6 absolute (last micro-batch 2.9e-6), entropy 2.7e-7 and gradient norm 2.4e-5 relative. Over steps 1-2 the gradient norms differ from the hoist run's by up to 1.5e-4 relative, inside 2e-3.How to turn it on
TRAINER_TUNED_FLAGS=truein Qwen3.5-397B RL trainer on 128 v7x: the tuned recipe turns on the fp8 fake-quant hoist, fused SwiGLU backward, attention savings, the 2048 wi forward tile and the GDN context-parallel m scan, and pins the MaxText that defines them google/tunix#2948.query_proj_kernel_gather_onceis refused next toquery_proj=device.fp8_fake_quant_hoistis checked when the trainer traces the layer scan. It raises, for example, withoutfp8_moe_fake_quant=true, or forgather8withoutsparse_matmul.scan_layers=Falseit has no effect and is not refused.wi_tile_fwd_rhs_buffer_countis refused at 1 or below (0 is the kernel default), withoutmoe_wi_gmm_fused_swiglu=true, and withuse_gmm_v2_heuristic_tiling=true. The layer applies it only when the fused wi forward GMM runs one K tile (wi_tile_fwd_embed_dimat least the embed dim); otherwise, for example with the eval tiles, the kernel keeps its default.gdn_cp_m_scanandgdn_cp_dh_scanare refused without Gated DeltaNet layers, withoutuse_gdn_kernel=true, and with agdn_chunk_sizeother than 64.PARAMBOLE_FQ_HOIST(0,absmax,1= gather8) selects a mode while the key isoff. A key and an env var that disagree raise.Testing
dbdcf5f8dcbaat the head1bf840cff(the earlier files also at15285fd74).1bf840cffas at15285fd74.wi_tile_fwd_rhs_buffer_count_test.pyhas 22 tests: the config key; the kernel config (the weight BlockSpec's depth, the metadata and the kernel name); the fused forward lowered for TPU v7x on CPU, where only the forward kernel takes the depth; the MoE layer on 8 CPU devices with stand-ins for the two Pallas kernels, where the depth reaches only the fused wi forward calls and only with one K tile, with unchanged values; and 2 TPU-only cases (the same bits at depths 2 and 4), skipped on CPU. Those two pass in Pallas TPU interpret mode on CPU, with an interpreter patch that is not part of this PR, and fail on a deliberately wrong kernel output.gdn_cp_dh_scan_test.pyhas 16 tests: Mohit Khatwani's 13, which check the two scans against the XLA trees they replace (Pallas interpret mode), and 3 for the keys: the config validation; the gating, traced on an abstract TPU7x mesh (no scan with the keys off, only the m scan withgdn_cp_m_scanalone, both with both); and the CPU numerics in a child process with 2 CPU devices, so it never skips (keys off bit-identical to the call without them, keys on close). On Mohit's two commits without the key gating, exactly the 3 key tests fail.fp8_fake_quant_hoist_test.pyhas 39 tests. One builds a shrunk Qwen3.5 with 8 layers on 8 CPU devices. Its loss (6.76953268), logits and 66 gradients are bit-identical betweenabsmax,gather8andoff, and its compiled programs show the hoist ran: in-loop absmax all-reduces go from 2 to 0, and the in-loop expert gathers are 8-bit floats.gdn_cp_m_scan, the forward scan running whatever its key, and the validator without its chunk-size requirement.94f11c09d, and at15285fd74the same as one1cd8bf1d, together with the PR's 8 earlier test files andnative_fp8_dot_general_test, which imports the gmm kernel. The kernel's other tests do not run on CPU:megablox_gmm_silu_testistpu_only, andpallas_mosaic_tpu_v2_kernel_testneedschex, which the bench image lacks. Their 5 failures fail there too. At1bf840cffthe GDN tests that import the changed kernel files, the config tests and the PR's other test files give the same result per test as at15285fd74; three GDN tests fail at both (gdn_cp_head_sharded_test'stest_e2e_layer_zero_cross_document_bleed,gdn_cp_precision_test'stest_cp_precision_reaches_only_the_cp_composition_on_cpu_meshandremat_cp_state_and_route_test'stest_remat_cp_state_and_route_on_cpu_mesh).gather8, of 83.0. With the tile and two weight buffers on top, the memory report and the scheduler memory are those ofgather8alone. With the tile and both scan keys on, it fits the scheduler's memory limit with 52.5 MB to spare; with every key off, this branch's code gives the fp8 head's memory report and scheduler lines byte for byte.Not in this PR / known gaps
absmaxalone was not run on TPU.parambole-w22k1a). No replicate and no numerics run of the combination: both were cancelled before they ran. Each lever passed the numerics gates on its own.gdn_cp_dh_scan: in the code and off by default. On the bench it made the step slower, so the companion recipe does not set it.gdn_cp_m_scanand precision: with the key on, the forward state composition runs at 3-pass bf16 (high) even whengdn_cp_matmul_precision=highest; no validator ties the two. The bench and the recipe usehigh.o275x-p3-e6k) has a control profile only in another pod (o275x-c1-prof); both are listed under "How it was measured".6314a6b05+1c31cc9efcover 16 more test files and a test helper, none of which this PR touches. They are still unpushed and not in this PR.cp_as_ep_dense_no_cp_testaborts its child process on hosts with more than 64 cores unlessOPENBLAS_NUM_THREADSis 64 or lower; this also happens without this PR.How it was measured: bench, logs, profiles, and the check that this is the measured code
Bench
custom_mesh_and_rule=cp-as-ep-dense-no-cp). 9 micro-batches per optimizer step, 294,912 tokens per chip per step,ragged_buffer_factor2.75.fp8_fake_quant_hoist=gather8on the fp8 arms. The tile arms passwi_tile_fwd_mlp_dim=2048 wi_tile_fwd_rhs_buffer_count=2after them; MaxText keeps the last value, and their config logs read 2048 and 2. The companion recipe sets its own tile to 2048 instead, which resolves to the same config. The bench adds its ownenable_checkpointing=false.Logs:
gs://parambole-397b-replay-nap/win-tpu/sched/final/<job>/<arm>.log, md5 prefixes below.parambole-w21g2a/o275x-g2-m, controlo275x-g2-ctl, both scanso275x-g2-bothparambole-w21gr/o275x-gr-mparambole-w21g3/o275x-g3-dh, controlo275x-g3-ctlparambole-w21gm/nx-gm-mparambole-w21ec/o275x-ec-e6k, replicateo275x-ec-e6kr, controlo275x-ec-ctl(booking r-b6424)parambole-w21e6k/o275x-e6k1,parambole-w21p3/o275x-p3-e6kparambole-w21ds/o275x-d1-ctl,parambole-w21cm/o275x-c1-ctl,o275x-c1-profparambole-w21e6k/nx-e6k1parambole-v19a/o275x-v19a-fq8parambole-v19a/o275x-v19a-ctlparambole-v19b/o275x-v19b-fq8rparambole-v19b/nx-v19b-fq8parambole-bpr-v18/o275x-pr-v18,o275x-pr-r-v18,nxpr-v18parambole-bf2-v18/o275x-prfc-f2,o275x-f2-l3-prof,o275x-f2-l3-r,nxf2gs://parambole-397b-replay-nap/win-tpu/sched/v16/parambole-rb4-v16/o275x-rbd-v16,o275x-rbd-r-v16,nxrbd-v16Profiles (xplane): the head (tile on), one chip, under
gs://parambole-397b-replay-nap/win-tpu/sched/final/parambole-w21p3/o275x-p3-e6k/parambole-w21p3-o275x-p3-e6k/tensorboard/plugins/profile/2026_10_16_23_13_51/, and its control in another pod undergs://parambole-397b-replay-nap/win-tpu/sched/final/parambole-w21cm/o275x-c1-prof/parambole-w21cm-o275x-c1-prof/tensorboard/plugins/profile/2026_10_16_23_21_25/. Read per micro-batch over 8 TensorCores, control to head: busy 8,641.8 -> 8,585.2 ms, the wi GMM forward 401.8 -> 390.6 ms and its remat replay 247.8 -> 241.9 ms, exposed waits -37.0 ms. F2, one chip, undergs://parambole-397b-replay-nap/win-tpu/sched/final/parambole-bf2-v18/o275x-f2-l3-prof/parambole-bf2-v18-o275x-f2-l3-prof/tensorboard/plugins/profile/2026_10_16_07_24_19/.Is this the code that was measured?
The scan arms installed eight files over
94f11c09d(overlay tgz25804529, manifest09d13a74, read back from GCS). Its three scan files are byte-identical to this branch's at1bf840cff.base.ymlandtypes.pyhere are the clean 3-way merge (no conflict) of the overlay's and the tile's versions overe1cd8bf1d;moe.pyis the tile's; the other two are the fp8 hoist's files, unchanged sincee1cd8bf1d.src/maxtext/)1bf840cff15285fd74kernels/gdn/gdn_bwd/api.pykernels/gdn/gdn_bwd/dh_scan.pykernels/gdn/model_runner.pyconfigs/base.ymlconfigs/types.pylayers/moe.pyThe controls in the same pods (
o275x-g2-ctl,o275x-g3-ctl) installed the same overlay with both scan keys off (their config logs carrygdn_cp_m_scan: False).The tile arms installed seven files over
94f11c09d(overlay tgz3c5bc6e0, manifest042de2eb, read back from GCS). All seven are byte-identical to this branch's files at15285fd74: the fp8 hoist's five, three of which the tile commit also changes, and the two kernel files.e1cd8bf1d..15285fd74changes exactly five of them plus the new test file.src/maxtext/)15285fd74e1cd8bf1dconfigs/base.ymlconfigs/types.pykernels/megablox/ops.pykernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.pylayers/moe.pylayers/quantizations.pylayers/nnx_decoders.pyThe control in the same pod (
o275x-ec-ctl) installed the same overlay with the two keys off. The controls in other pods ran the previous head's code (the new key is absent from their config logs). With the keys off, the overlay's code compiles at this cell to the program of71194c53c(thesrc/maxtextofe1cd8bf1d), identical after source locations are removed, in an AOT compile for v7x.The fp8 arm installed five files over
94f11c09d. They are byte-identical to this branch's files ate1cd8bf1d, and94f11c09d..e1cd8bf1dchanges exactly those five files plus the new test file.src/maxtext/)e1cd8bf1d94f11c09dlayers/moe.pylayers/quantizations.pylayers/nnx_decoders.pyconfigs/types.pyconfigs/base.yml94f11c09dwas installed as one delta overlay over20c877aa1(23 files, every base md5 checked). On a CPU check that overlay reproduces this branch'ssrc/maxtextat94f11c09dexactly: 719 vs 719 files, 0 md5 differences.F2 ran
4ed82d793, a local commit that was not pushed. It is94f11c09dwithout three #5721 commits and without two cosmos3 model YAMLs that #5721 also added (other models). The three commits:3b126c073gather_scanned_kernels_once1f639cc3eGDN segment path skipped for unpacked rowspacking=False; this trainer resolvespacking=True6b00e919cGDN forward skips its output zero fillCommits
4dde1118ae08f7aae3moe_wo_gmm_fused_dswiglu)b57e4f04d83f60748econtext_parallel_attention_load_balance_hoist_ids)1aff05148query_proj_kernel_gather_once)94f11c09dpacking=Truein thegdn_conv1d_dense_token_stackvalidator test's config (fixes 9 of 41 cases that fail atbbd0f6f10)71194c53cfp8_fake_quant_hoist). The TPU arm's overlay was built from this commit. Its message predates the TPU measurement.e1cd8bf1dfp8_fake_quant_hoistoff vs absmax vs gather81c1162e96wi_tile_fwd_rhs_buffer_count, the fused wi forward GMM's weight pipeline depth. A cherry-pick of1768510ca, the commit the TPU arms' overlay was built from: the same patch and the samesrc/maxtext. Its message predates the TPU measurement.15285fd74wi_tile_fwd_rhs_buffer_countdd533934d4683514d5, cherry-picked.5fd3c95824026f5079, cherry-picked. Commits 11 and 12 select the scans withgdn_cp_matmul_precision=high, which the bench sets, so on their own they would turn both scans on; commit 13 moves them to their own keys, default off. The branch is pushed as one stack, so its tip is never at commit 11 or 12.bdc0681e8gdn_cp_m_scanandgdn_cp_dh_scan(default off). Its scan files are the TPU arms' overlay's.1bf840cffgdn_cp_m_scanandgdn_cp_dh_scanUnit test results at the head
1bf840cfftests/unit/)megablox_gmm_dswiglu_test.pytpu_only)moe_wo_gmm_fused_dswiglu_test.pycontext_parallel_attention_load_balance_hoist_ids_test.pyquery_proj_kernel_gather_once_test.pycp_as_ep_dense_no_cp_test.pygdn_conv1d_dense_token_stack_test.pybbd0f6f10(commit 6 fixes it)gdn_conv1d_input_in_qkv_dtype_test.pyfp8_fake_quant_hoist_test.py94f11c09dwi_tile_fwd_rhs_buffer_count_test.pytpu_only)e1cd8bf1dgdn_cp_dh_scan_test.py15285fd74ore1cd8bf1d(no scan module); the 3 key tests fail on commit 12The ten existing test files that cover the fp8 hoist's five changed files:
quantizations_test,moe_test,nnx_decoders_test,nnx_decoder_test,pyconfig_test,configs_test,moe_expert_weight_prefetch_test,moe_wi_gmm_fused_swiglu_test,qwen3_5_layers_test,moe_chunk_pipeline_test.94f11c09d.b241b5ae3, which has the samesrc/maxtextase1cd8bf1d.15285fd74they and the 8 earlier test files above give the same result per test as one1cd8bf1d, in one session; so doesnative_fp8_dot_general_test(24 passed, 1 skipped), which imports the gmm kernel.94f11c09d:test_create_fp8_einsum,moe_test'stest_init, twotest_sparse_matmul_repairs_batch_specs_only_without_expert_parallelismcases, andtest_explicit_sharding_gated_delta_net_unsupported_combinations.