Repository navigation
397B v7x RL trainer: all GDN, MoE and engine speedups (measured + pending), ordered by impact - #5762
Open
khatwanimohit wants to merge 19 commits into
Open
khatwanimohit wants to merge 19 commits into
khatwanimohit wants to merge 19 commits into
Conversation
…_mask freezes The 397B RL trainer freezes the MoE router gate with trainable_parameters_mask, but fwd_bwd still differentiated every parameter and _update_kernel zeroed the frozen gradients afterwards. So every micro-batch computed each layer's gate weight gradient, an f32 [4096, 512] HIGHEST-precision matmul, and reduce-scattered it 256-way on the SparseCore, only to discard it. Wrap the frozen leaves in stop_gradient inside diff_wrapper. Autodiff then emits neither the weight-gradient matmuls nor their reduce-scatters. Their gradients come out zero, as _update_kernel made them before (ahead of the grad norm and clipping), so the update and grad_norm are unchanged.
_run_local_gdn_decoupled_fwd called fused_conv1d_gdn with the default zero_initialize_out=True, which aliases a zero buffer into the kernel output. XLA hoists the zeros out of the layer scan into a loop-carried buffer and copies it before every kernel call, and the GDN recompute in the backward broadcasts it again. The prefill kernel never reads that buffer and writes every row (seq_len is a multiple of chunk_size, which both functions already check), so the fill is dead work. On the 397B v7x bench (16,384 tokens and 64 value heads per device) this removes six 256 MiB copies per forward layer-block iteration (about 1 ms per iteration) and three zero broadcasts per backward iteration, and frees the 256 MiB loop-carried buffer. Outputs are bit-identical.
… adding its halves With prefuse_moe_weights the gate and up projections are one GEMM, whose [tokens, 2 * mlp] output gmm_up split into two halves. Autodiff transposes each half's slice into a zero-padded [tokens, 2 * mlp] gradient and adds the two, so XLA writes both half gradients and reads them back in a pad-add fusion over the whole ragged buffer (bf16[225280,2048], 8 per micro-batch on the 397B bench). The chunk pipeline's tie 4 held the first half, so XLA also copied that half out of the GEMM output (bf16[225280,1024] slice) to feed the barrier. gmm_up now returns the fused output whole when there are no MoE biases, and _down_projection applies fused_glu, a custom_vjp around the split plus the activation. Its backward computes the two half gradients with jax.vjp of the same apply_ffn_activation and concatenates them; jnp.where(d == 0, 0, d) keeps the -0 -> +0 of the zero-padded add, so the gradient is bit-identical. The halves keep their remat names (mlpwi_0, moe_mlpwi_0, moe_mlpwi_1) inside the custom_vjp, so remat policies save the same residuals as before. Tie 4 now holds the GEMM output itself. Checks: the tiny Qwen3.5 A/B on v6e (prefused fake-quantized wi, ring of experts, chunk pipeline, CP/EP/FSDP=2) is bit-identical to the base in loss, logits and all 66 gradient leaves. FusedGluTest compares the forward and the gradient bit for bit against autodiff of the split, with zero cotangent rows, plain and under jax.checkpoint.
Pass 1 of sequence-sharded context parallelism (gdn_cp_mode=seq) only needs the local final states and t_inv, but it ran the full fused Conv1D-GDN prefill kernel and dropped the output and the chunk states. It now runs the kernel in a new GDNMode.PER_SEQ_STATES: - chunked_gdn skips the output path (q.state, q.k^T and (q.k^T * gating).(u - ws)) and computes ws = w.state_prev on its own instead of as the first half of the merged [w; q].state_prev matmul. Each row's dot product is unchanged, so the states and t_inv are bit-identical. - No output or chunk-states buffers are allocated, so the kernel writes neither to HBM. Per 397B call that drops the bf16[16384,64,128] output, its zero init, and the f32[256,64,128,128] chunk states (about 1.34 GB). Fixed behaviour, no new flag: fused_conv1d_gdn and the PER_SEQ kernel are unchanged. Pass 1 calls the new wrapper.fused_conv1d_gdn_states through api._run_local_gdn_states_fwd, which reuses the pure-JAX forward on CPU. The kernel input layout that both callers share moves to api._fused_kernel_inputs. Verified locally: - v6e model A/B (tiny Qwen3.5 with gdn_cp_mode=seq, gdn_cp_state=device, context=2) against the previous code: loss, logits and all 66 gradient leaves bit-identical. - The kernel on the local v6e, base tree against this tree, at two geometries including the 397B per-device heads (16 k / 64 v heads, dim 12288), with bf16 and f32 activations, with and without packed segments and initial states: the PER_SEQ outputs (80 arrays) and the states-only conv state, recurrent state and t_inv (48 arrays) are bit-identical. - The same comparison in Pallas TPU interpret mode on CPU with the v7x chip description (64 MiB VMEM): 40 + 24 arrays bit-identical, no unwritten (NaN) rows. - v7x AOT of the 397B bench program (tpu7x-256): compiles. fwd_bwd temp 89.26 GiB and dropless temp 69.08 GiB, both unchanged. In fwd_bwd, 3 of the 9 PER_SEQ calls become fused_conv1d_gdn_per_seq_states_b1_c64 (outputs: conv state, recurrent state, t_inv); in the dropless program 6 of 12. Zero-init output copies go from 6 to 3 in both. - The new test test_run_local_gdn_states_fwd_matches_full_forward passes on CPU (pure-JAX path).
…routing get_topk sorted every token's router probabilities for auto_indices, but under router replay only tokens that are neither forced nor padding read them (the rest are replaced by the forced ids or masked). When at most 256 tokens need them, gather those rows (jnp.nonzero(size=256)), run top_k on the [256, num_experts] block and scatter the ids back; a lax.cond falls back to the full top_k when more need it. Blocks of 256 tokens or fewer keep the plain top_k. top_k is per row, so every id that is read is unchanged: outputs and gradients are bit-identical. Checked against the previous get_topk (weights, ids, jax.grad) on CPU and on v6e at 8192x512 top-10, both branches and the 256/257 boundary, and with the tiny Qwen3.5 A/B run.
…jor copies pallas_gdn_bwd_kernel copied qkv_conv, do, chunk_states and t_inv into group-major (B, G, C, ...) layouts before the kernel and transposed dy_conv back after it. The kernel now reads head group g from the token-major arrays: q, k and v as column blocks of qkv_conv (v is sliced out only when q + k is not a whole number of v blocks), do / chunk_states / t_inv through group-indexed reshape views. It writes dq, dk and dv as column blocks of three token-major outputs. The kernel math is unchanged; head_tile stays 32 (64 would need ~86 MiB VMEM on v7x). 397B bench program (tpu7x-256 AOT; per device 64 v heads in 2 groups), per GDN layer backward, 3 -> 0 each: copies to f32[1,256,2,32,128,128] (chunk_states), f32[1,256,2,32,64,64] (t_inv), bf16[1,256,64,2,32,128] (do), the qkv concat pad_maximum_fusion.*.refs_0_, and the dy_conv transposes back (bf16[1,256,64,2,4096] / [1,256,64,2,1024]). jit_fwd_bwd copies 718 -> 694. Verified locally: - Kernel outputs bit-identical to the previous kernel: local v6e, 34/34 outputs over 6 geometries (397B layer with 2 groups, packed segments and dht/dh0; no segments; head_tile=64; f32 with 4 groups; head_tile=4; 32 heads); Pallas interpret on CPU with the v7x chip description (64 MiB VMEM tiling), 23/23 outputs over 4 geometries. - New tpu_only test_gdn_kernel_bwd_group_blocks_match_per_group_calls (2-group call vs per-group calls, bitwise) passes on the v6e; the CPU interpret tests around the kernel pass. - 397B AOT for tpu7x-256 in the bench image: compiles (Mosaic for v7x); same 316 collectives (types and shapes) as the base; fwd_bwd temp 89.26 -> 90.07 GiB (HBM 77.73 -> 78.13 GiB), dropless temp 69.08 -> 68.24 GiB (HBM 65.48 -> 64.56 GiB). - Tiny Qwen3.5 A/B on the v6e (gdn_cp_mode=seq): loss, logits and the embedding, norm, A_log, dt_bias, conv1d, router and routed-expert gradients are bit-identical, but 13 dense weight-gradient leaves differ by bf16-ulp amounts (max abs 3.1e-5; both programs deterministic). One of them is the last layer's attention out kernel, which no GDN backward feeds. Cause (HLO dumps of both): XLA's all-reduce combiner groups the tiny model's bf16 gradient all-reduces (cast_grads_after_all_reduce) differently, moving one bf16[2,256,8] in_proj_ba gradient into the other combined all-reduce, which changes the cross-device summation order of the leaves after it. The 397B program's collectives are unchanged (above).
Prefill-only calls (the training path through gdn_bwd/api.py) no longer materialise qkv.astype(f32).reshape(N, 1, D) in HBM before the fused Conv1D-GDN kernel. The wrapper reshapes the act_in-dtype qkv to [num_chunks, chunk, D] (a bitcast), ChunkInBufferedRef DMAs one whole chunk per tile, and the kernel upcasts it into a compact fp32 VMEM scratch with one strided store per 128-lane column (the inverse of load_as_qkv_large). The scratch then serves as the qkv slot, so Conv1D and the q/k/v loads see the same fp32 values in the same layout as before. bf16 -> fp32 is exact, so the outputs are bit-identical. Both prefill kernels share this input path through _fused_conv1d_gdn: the full PER_SEQ kernel and the states-only PER_SEQ_STATES kernel of CP pass 1. The general (decode/mixed) path keeps its HBM upcast and row DMAs. VMEM per tile for bf16 qkv: a 2 x 1.5 MiB bf16 window plus a 3 MiB fp32 scratch, the same 6 MiB as the old 2 x 3 MiB fp32 window. No new flag: the existing is_prefill_only argument selects the path. Verified locally against the previous kernel: - Tiny Qwen3.5 A/B on v6e (GDN kernel, gdn_cp_mode=seq, gdn_cp_state=device, context=2, so CP pass 1 runs the states-only kernel): loss, logits and all 66 gradient leaves bit-identical. - Old vs new kernels on v6e: packed segment_ids, initial conv/recurrent states and fp32 qkv at 2 kq / 8 v heads (2 x 256 tokens); segment_ids and initial states at the 397B per-device heads (16 kq / 64 v heads, dim 12288, 256 tokens); each through fused_conv1d_gdn and fused_conv1d_gdn_states; plus the general path (3 decode + 2 prefill sequences). All 53 arrays bit-identical, and the new states-only states and t_inv equal the new full kernel's. - The same cases in Pallas TPU interpret mode on CPU with the v7x chip description (64 MiB VMEM): all bit-identical. The interpreter needed two patches in the check script: a patched store discharge rule that can write through a ref reshape (jax 0.11 raises NotImplementedError for the strided stores into the reshaped fp32 scratch), and static-size row and state DMAs (it cannot slice with traced sizes). - New test test_prefill_bf16_qkv_matches_fp32_qkv_path (tpu_only) passes on v6e: the full and the states-only prefill kernels match the HBM-upcast general path bitwise. - gdn_bwd_pallas_test.py on CPU: 38 passed, 3 skipped (two need 4 devices, one is the new tpu_only test).
…gsorting The unsort backward argsorted the revert indices to rebuild their inverse, which is the forward sort's permutation. ring_ragged_sort now also returns that permutation and RouteOutput carries it to ring_ragged_unsort, so the backward reads it instead of sorting again. The permutation is named for the `moe_route` remat key in place of the sorted token indices (same size), which are now the permutation // topk (flat slot j belongs to token j // topk) instead of an s32 gather of the repeated token ids. Exact by construction: the argsort of a permutation is its inverse (distinct keys, so stability does not matter) and arange(n).repeat(topk)[p] == p // topk. Per MoE layer and token chunk this drops one argsort from the backward and the s32 token-index gather from the forward and its replay. The buffered branch's routing-weight gradient stays a scatter into zeros. A gather over the revert indices would be equivalent, but on v7x XLA compiles it to a TensorCore gather fusion over the f32[655360] slots, about 30 ms per call against about 2 ms for the scatter (measured: unpermute went from 589 ms to 4,085 ms per micro-batch on the 128-chip bench with the gather). Test: RingRaggedUnsortGradTest compares the custom backward with autodiff of a plain-gather forward, bitwise, unbuffered and with a buffer that drops slots, and checks that the gradient contains no sort.
…ar inverse
compute_gdn.invert_triangular_matrix solved each 16-row diagonal block
left-looking: every row re-stacked the rows solved so far and moved its row of
t from lanes to sublanes to form sum_j t[i, j] * x_j. The right-looking form
keeps the block in place as two sublane-aligned 8-row halves and, once row j is
final, subtracts t[i, j] * x_j from every later row. Each row subtracts the
same products, one at a time in column order, so only the summation order
changes. The inter-block matmuls are unchanged. Both prefill kernels (the full
PER_SEQ kernel and the CP pass-1 states-only kernel) use this function, as does
the pure-JAX path.
tpu7x LLO at the 397B per-device shape (16,384 tokens, 16 kq / 64 v heads,
chunk 64), loop bundles per chunk:
- full kernel 43,448 -> 39,389 (-9.3%); XLU ops 29,633 -> 18,370, VALU ops
88,676 -> 67,243.
- states-only kernel 40,049 -> 36,330 (-9.3%); XLU ops 26,945 -> 15,669.
The unrolled right-looking loop needs more VMEM for spills: the full kernel
would need 62.2 MiB, over the 57.6 MiB limit (0.90 x 64 MiB). Two changes make
room:
- create_allocs single-buffers the conv and recurrent state refs of a prefill
call with one sequence (the training path: one row per device). Its states
are copied in only before the first tile and out only after the last, so no
copy is in flight while a tile uses the slot. This saves 4.1 MiB, down to
58.05 MiB.
- DEFAULT_VMEM_LIMIT_FACTOR goes from 0.90 to 0.95 (60.8 MiB on TPU7x).
Verified against the previous kernels:
- v6e at the 397B per-device geometry (16,384 tokens, bf16 qkv/b/a, a packed
row of 7 documents plus padding and a one-document row, with and without
caller states). These are bit-identical: the full kernel's out, final states
and chunk_states, the states-only kernel's states, and the loss. t_inv from
both kernels differs by at most 1.5e-8 (|t_inv| <= 1). Gradients through
gdn_decoupled_conv1d, whose Pallas backward reads the saved t_inv: d_b,
d_conv_state and d_recurrent_state are bit-identical. Max |new - old| /
max |old| is 6.2e-6 for qkv, 6.5e-6 for the conv weight, 3.7e-5 for a and
4.6e-7 for A_log and dt_bias. Against an f32 HIGHEST pure-JAX reference, old
and new have the same error to 4 digits (2e-3 to 6e-3 of max).
- Pallas interpret mode on CPU with the TPU7x chip description:
- chunked_gdn inside a Pallas kernel, full and states-only modes, with and
without segments: out and state bit-identical, t_inv within 1.2e-10.
- The full and states-only kernels, 1 and 2 sequences per call, with the
interpreter patches from the bf16-qkv change (static-size DMAs, stores
through a ref reshape): final states bit-identical, t_inv within 3.7e-8,
chunk states within 8e-8 of max (f32 matmuls on CPU carry the t_inv
rounding into them), out within one bf16 ulp. The new states-only
kernel's states and t_inv equal the new full kernel's bit for bit.
- Tiny Qwen3.5 A/B on v6e (ab_model.py: gdn_cp_mode=seq, gdn_cp_state=device,
context 2, packed rows): loss and logits bit-identical. 64 of 66 gradient
leaves are bit-identical; the other two (one GDN layer's A_log and conv
kernel) differ by at most 8.3e-8 of max. Gradient norm rel diff 1.4e-14.
- 397B v7x AOT (wi_tile_dlhs 2048/2048, moe_route=device) compiles. fwd_bwd
temp is 93.31 GiB and dropless temp 67.94 GiB, both the same as the base
compiled with the same flags. Outside the 9 GDN kernel bodies and their
scoped VMEM limit (60,397,977 -> 63,753,420 bytes), the optimized HLO is
identical to the base's.
- New test test_invert_triangular_matrix_matches_float64_inverse checks the
inverse against float64 at chunk 8 to 64 (atol 1e-6; dropping one column's
update gives 0.2). gdn_bwd_pallas_test.py: 39 passed and 4 skipped on CPU;
the 2 tpu_only tests pass on v6e.
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.
…educe-scatter scatters a middle dim wo is [experts, mlp, embed]. Sharding its embed dimension over FSDP made XLA transpose each wo gradient to [embed, experts*mlp] and reduce-scatter its major dimension on the SparseCore, about 3x slower than wi's middle-dimension scatter for the same bytes (12.9 vs 4.2 ms per call on v7x), and that chain of expert-weight-gradient reduce-scatters ends every backward layer iteration. Mesh rules that define mlp_moe_out/embed_moe_out (cp-as-ep) now shard wo's mlp dimension over FSDP; rule sets without those names keep the old layout. The forward all-gather becomes [64,32,4096]->[64,1024,4096].
… buffer's valid rows Only rows [0, limit) of the ring's packed [buffer_size, ...] expert buffer hold tokens (about 56% on the bench's hot shard at ragged_buffer_factor 2.75). The grouped matmuls already skip the rest, but four XLA fusion families per layer chunk still streamed the whole buffer on every device: the unsort backward (f32[225280] + bf16[225280,4096]), the GLU forward and its backward recompute (bf16[225280,1024]), and the GLU backward, a select fusion writing both half gradients plus a separate pad/maximum pass that concatenated them (bf16[225280,2048]). - kernels/ragged/unsort_bwd_rows.py: a Pallas TensorCore kernel for the buffered branch of _ring_ragged_unsort_bwd. It computes bf16(f32(g) * w) and the f32 routing-weight row dot in 128x128 tiles over (1024, 2048) blocks, with `limit` as a prefetched scalar; blocks past the one holding row limit - 1 map to that block, so they are neither fetched nor written. The weights and the dot are lane-dense [rows / 128, 128] (free reshapes), turned into per-row columns with one 128x128 transpose each. The scatter of the dot is unchanged; off a TPU the XLA code runs. - kernels/ragged/glu_rows.py: forward and backward kernels for silu(gate) * up over the same rows, reading gate and up in place from the fused buffer. The backward writes [d_gate | d_up] in one pass, with the -0 -> +0 fix. They follow XLA's bf16 fusions exactly: the same op sequence (negate, exp, + 1, divide, multiply), in f32, rounding only the outputs (XLA keeps excess precision inside the fusion; op-by-op bf16 rounding differs from it in ~3% of elements). - moe.py: fused_glu(fused, glu, limit=None). _down_projection passes valid_token_count as the limit when the buffer is the ring's packed one and apply_ffn_activation is the plain bf16 SiLU GLU; the kernel path keeps the whole buffer as its residual, under the remat names of both halves. The dropless replay program (unpacked buffer) and every other path keep the XLA code. Rows left unwritten are never read: gmm_v2 reads lhs rows up to the sublane tile of row limit - 1, tgmm_v2 masks rows outside the groups with a select, and the dot scatter drops rows >= limit. Numerics: the weighted unsort gradient and both GLU results are bit-identical; the routing-weight dot changes summation order only. Checks: - v6e, bench shape [225280, 4096], limits 0, 1, 1023, 125952, 126156, 225280: unsort gradient bitwise equal to XLA (rows up to the end of limit's block), dot max |d| 5.3e-5. Interpret mode with the TPU7x chip description and a partial last block: bitwise as well. - v6e, GLU forward and backward against fused_glu's XLA path, every bf16 bit pattern as the gate (x64 random up/d, plus +-0, +-inf, NaN, denormals), 225280-row and partial-block buffers: 0 mismatching bits. - v6e timing at 126156 of 225280 valid rows (XLA -> kernel): unsort 4.71 -> 2.61 ms, GLU forward 1.08 -> 0.61 ms, GLU backward 1.92 -> 1.02 ms; at a full buffer the kernels are within 1-9% of XLA. - Tiny Qwen3.5 on v6e (ab_model.py, use_ragged_sort=true ragged_buffer_factor=2.75): loss, logits and all 66 gradient leaves bit-identical to the unmodified tree. - 397B v7x AOT (wi dlhs tiles 2048, moe_route=device, recompute limit 100): fwd_bwd temp 94.95 -> 92.25 GiB, dropless 56.69 GiB unchanged. module_0000 has 8 unsort_bwd_rows, 16 glu_fwd_rows and 8 glu_bwd_rows calls in place of the 16 activation, 8 unsort, 8 select and 8 concatenate fusions, with no new 225280-row copies.
In training, layout assignment gives the forward out_proj's LHS the gated RMSNorm's head-tiled layout, so the projection re-tiles it on every MXU fetch, while the backward's recomputed projection reads it token-minor. Constrain the forward value to the token-minor layout with the existing _forward_layout_constraint custom_vjp, which leaves the cotangent's layout to XLA.
… to the expert-weight gradients A custom_vjp that is the identity in the forward and an optimization barrier on the cotangents in the backward, applied to the layer input and the expert weights. The input gradient is handed on only once the expert-weight gradients exist, so each layer's weight-gradient grouped matmuls run under that layer's token reduce-scatters instead of piling up at the end of the backward iteration, where every deferred layer's sorted-token buffers would still be live.
The router kernel (embed x num_experts, f32 under float32_gate_logits) is FSDP-sharded, so every layer all-gathered it on the critical path and the backward gathered it again; on v7x those all-gathers run synchronously on the TensorCore (~22 ms each on the 397B bench). Replicate the stacked kernels before the scan (one all-gather), and give each layer's slice a backward that constrains its gradient to the kernel's sharding, so gradients stay per-layer reduce-scatters into a stacked shard rather than full-size all-reduces. (cherry picked from commit 4f5093f)
On the 397B bench (v7x) the query kernel's per-layer all-gathers are synchronous and 256-way, and the backward gathers it twice per layer (rematerialized projection and input gradient): ~0.54 s of exposed wait per micro-batch. (cherry picked from commit 035ccaf)
…, list keys in param paths
- _identity_with_grad_sharding is a module-level custom_vjp with the sharding
as a nondiff argument, as _forward_layout_constraint, instead of a new
custom_vjp per call.
- Gradient shardings use maxtext.utils.sharding.logical_to_mesh_axes, which
drops size-one mesh axes; the flax partitioning import goes.
- Param paths keep SequenceKey indices as well as DictKey keys (shared
_param_path); a leaf's trailing GetAttrKey('value') is still dropped, since
nnx.iter_modules paths do not have it.
(cherry picked from commit a0eaad4)
(cherry picked from commit ee1f10a)
…qkv) Under fused_qkv the Attention module creates qkv_proj and no query, so the gather-once selection must not touch module.query there. (cherry picked from commit dcd4a5c)
khatwanimohit
requested review from
RissyRan,
darisoy,
gagika,
igorts-git,
jiangjy1982,
michelle-yooh,
parambole,
richjames0,
shralex,
shuningjin and
zxhe-sean
as code owners
October 15, 2026 17:11
khatwanimohit
requested review from
A9isha,
NuojCheng,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
dipannita08,
gobbleturk,
hengtaoguo,
huytransformer,
vipannalla and
xibinliu
as code owners
October 15, 2026 17:11
![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 optimizations for Mixture of Experts (MoE) and Gated Delta Networks (GDN) on TPU, including Pallas scans for sequence-sharded context parallelism, SiLU-GLU row-wise kernels, and prefill-only optimizations. Key changes include the addition of dh_scan.py and glu_rows.py to fold chunk transitions and compute activations over valid rows, as well as optimizations to gather scanned kernels once before sequential layer execution. Feedback on the PR highlights a critical TPU hardware hazard in unsort_bwd_rows.py where asynchronous HBM writes to dot_ref across grid iterations can lead to stale reads and silent data corruption. To resolve this, it is recommended to write partial dot products to a separate dimension in HBM and perform the final accumulation using jnp.sum after the Pallas call.
Comment on lines
+81
to
+98
| # Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes. | ||
| w_col = jnp.transpose(jnp.broadcast_to(w_ref[pl.ds(r, 1), :], (_LANES, _LANES))) | ||
| if masked: | ||
| row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0) | ||
| valid = row_ids < limit | ||
| acc = jnp.zeros((_LANES, _LANES), jnp.float32) | ||
| for c in range(num_col_tiles): | ||
| cols = pl.ds(c * _LANES, _LANES) | ||
| g = g_ref[rows, cols].astype(jnp.float32) | ||
| if masked: | ||
| g = jnp.where(valid, g, 0.0) | ||
| grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype) | ||
| acc = acc + g * y_ref[rows, cols].astype(jnp.float32) | ||
| # Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns. | ||
| part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True) | ||
| if num_col_blocks > 1: | ||
| part = part + jnp.where(j == 0, 0.0, dot_ref[pl.ds(r, 1), :]) | ||
| dot_ref[pl.ds(r, 1), :] = part |
There was a problem hiding this comment.
Accumulating the partial dot products across grid iterations by reading and writing to the same HBM location (dot_ref) introduces a critical hardware hazard on TPUs. Since HBM writes are asynchronous and have high latency, successive grid iterations of j will read stale values of dot_ref before the previous write commits, leading to silent data corruption (incorrect accumulation). Additionally, frequent HBM read-modify-write operations inside the loop degrade performance.
To resolve this, we can write the partial dots of each column block to a separate location in HBM by adding a num_col_blocks dimension to dot_ref. This completely avoids overlapping writes and stale reads. The final sum can then be computed efficiently using jnp.sum in JAX/XLA after the Pallas call.
# Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes.
w_col = jnp.transpose(jnp.broadcast_to(w_ref[0, pl.ds(r, 1), :], (_LANES, _LANES)))
if masked:
row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0)
valid = row_ids < limit
acc = jnp.zeros((_LANES, _LANES), jnp.float32)
for c in range(num_col_tiles):
cols = pl.ds(c * _LANES, _LANES)
g = g_ref[rows, cols].astype(jnp.float32)
if masked:
g = jnp.where(valid, g, 0.0)
grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype)
acc = acc + g * y_ref[rows, cols].astype(jnp.float32)
# Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns.
part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True)
dot_ref[0, pl.ds(r, 1), :] = part
Comment on lines
+123
to
+158
| def row_map(i, j, limit_ref): | ||
| del j | ||
| return jnp.minimum(i, _last_block(limit_ref[0])), 0 | ||
|
|
||
| data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map) | ||
| lane_dense_spec = pl.BlockSpec((_BLOCK_ROWS // _LANES, _LANES), row_map) | ||
| grad, dot = pl.pallas_call( | ||
| functools.partial(_unsort_bwd_rows_kernel, num_col_blocks=num_col_blocks), | ||
| grid_spec=pltpu.PrefetchScalarGridSpec( | ||
| num_scalar_prefetch=1, | ||
| grid=(num_row_blocks, num_col_blocks), | ||
| in_specs=[data_spec, data_spec, lane_dense_spec], | ||
| out_specs=[data_spec, lane_dense_spec], | ||
| ), | ||
| out_shape=[ | ||
| jax.ShapeDtypeStruct((rows, hidden), gathered.dtype), | ||
| jax.ShapeDtypeStruct((rows // _LANES, _LANES), jnp.float32), | ||
| ], | ||
| compiler_params=pltpu.CompilerParams( | ||
| dimension_semantics=("arbitrary", "arbitrary"), | ||
| vmem_limit_bytes=48 * 1024 * 1024, | ||
| ), | ||
| cost_estimate=pl.CostEstimate( | ||
| flops=4 * rows * hidden, | ||
| transcendentals=0, | ||
| bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows, | ||
| ), | ||
| interpret=interpret, | ||
| name="unsort_bwd_rows", | ||
| )( | ||
| jnp.reshape(limit.astype(jnp.int32), (1,)), | ||
| gathered, | ||
| sorted_tokens, | ||
| jnp.reshape(weights.astype(jnp.float32), (rows // _LANES, _LANES)), | ||
| ) | ||
| return grad, jnp.reshape(dot, (rows,)) |
There was a problem hiding this comment.
Update the Pallas call setup to support the hazard-free column-wise accumulation. We define separate block specs for the input weights (which does not depend on j) and the output dot (which writes to a unique slice j for each column block). After the Pallas call, we sum the partial results along the column block dimension using jnp.sum.
def row_map_weights(i, j, limit_ref):
del j
return 0, jnp.minimum(i, _last_block(limit_ref[0])), 0
def row_map_dot(i, j, limit_ref):
return j, jnp.minimum(i, _last_block(limit_ref[0])), 0
data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map)
weights_spec = pl.BlockSpec((1, _BLOCK_ROWS // _LANES, _LANES), row_map_weights)
dot_spec = pl.BlockSpec((1, _BLOCK_ROWS // _LANES, _LANES), row_map_dot)
grad, dot = pl.pallas_call(
functools.partial(_unsort_bwd_rows_kernel, num_col_blocks=num_col_blocks),
grid_spec=pltpu.PrefetchScalarGridSpec(
num_scalar_prefetch=1,
grid=(num_row_blocks, num_col_blocks),
in_specs=[data_spec, data_spec, weights_spec],
out_specs=[data_spec, dot_spec],
),
out_shape=[
jax.ShapeDtypeStruct((rows, hidden), gathered.dtype),
jax.ShapeDtypeStruct((num_col_blocks, rows // _LANES, _LANES), jnp.float32),
],
compiler_params=pltpu.CompilerParams(
dimension_semantics=("arbitrary", "arbitrary"),
vmem_limit_bytes=48 * 1024 * 1024,
),
cost_estimate=pl.CostEstimate(
flops=4 * rows * hidden,
transcendentals=0,
bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows,
),
interpret=interpret,
name="unsort_bwd_rows",
)(
jnp.reshape(limit.astype(jnp.int32), (1,)),
gathered,
sorted_tokens,
jnp.reshape(weights.astype(jnp.float32), (1, rows // _LANES, _LANES)),
)
return grad, jnp.sum(jnp.reshape(dot, (num_col_blocks, rows)), axis=0)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
1 participant
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.
All of our Qwen3.5-397B v7x RL-trainer changes in one branch: the measured set from #5759, the follow-ups from #5760 and #5761 that have not been benched yet, and the router/query gather-once commits from #5741. All are fixed behaviour (no new config flags). Flags used with them are in google/tunix#2896.
Changes, in order of expected performance impact
Savings are per micro-batch on the 128-chip v7x trainer bench. Measured means a bench A/B; otherwise it is an estimate from profile/HLO analysis.
gdn_cp_dh_scan, forwardgdn_cp_m_scan; plus right-looking 64x64 triangular inverse in the forward kernelsstop_gradienton frozen parameters (no gate weight-gradient matmuls or reduce-scatters)The measured set (rows 1, 2, 5, 6, 8, 10-13) with the tunix#2896 flags ran at 9.796 s per micro-batch, 89.77 s/step, 3,285 tok/s/chip, against 10.616 s, 97.15 s/step, 3,036 tok/s/chip for atwigg/mlperf with #5741. Run log: https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.cluster_name%3D%22bodaborg-tpu7x-gsc%22%0Aresource.labels.namespace_name%3D%22priority-dev%22%0Alabels.%22k8s-pod%2Fjobset_sigs_k8s_io%2Fjobset-name%22%3D%22mkb-plg1%22%0Alabels.%22k8s-pod%2Fjobset_sigs_k8s_io%2Freplicatedjob-name%22%3D%22proc%22%0Aresource.labels.container_name%3D%22main%22;timeRange=2026-10-15T13:00:00Z%2F2026-10-15T14:00:00Z?project=cloud-tpu-shared-capacity
Related changes outside MaxText, also not benched yet: a bit-exact splash-attention fast path for tiles inside one segment (tokamax patch), and left-aligned packing in Tunix.
Numerics
Compile (397B, v7x, tpu7x-256, bench config with the tunix#2896 flags)
Compiles; fwd_bwd temp 82.42 GiB (94.29 GiB for the measured set alone; 89.26 GiB for atwigg/mlperf with #5741), dropless 56.79 GiB.