Repository navigation
397B v7x RL trainer: GDN kernel, MoE routing and engine speedups (10.62 -> 9.80 s per micro-batch on the 128-chip bench) - #5759
Closed
khatwanimohit wants to merge 11 commits into
Closed
khatwanimohit wants to merge 11 commits into
khatwanimohit wants to merge 11 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.
khatwanimohit
requested review from
A9isha,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
gagika,
gobbleturk,
hengtaoguo,
igorts-git,
jiangjy1982,
michelle-yooh,
richjames0,
shralex,
shuningjin,
vipannalla and
zxhe-sean
as code owners
October 15, 2026 16:15
khatwanimohit
requested review from
dipannita08,
huytransformer and
xibinliu
as code owners
October 15, 2026 16:15
![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 performance optimizations for GDN (Gated Decurrent Network) and MoE (Mixture of Experts) layers on TPU. Key improvements include implementing Pallas scans (gdn_cp_m_scan and gdn_cp_dh_scan) to fold sequence-sharded context parallel GDN chunk transitions in VMEM, adding a states-only PER_SEQ_STATES mode and prefill-only optimizations to skip output computations, and optimizing MoE routing by reusing sort permutations in the backward pass. Additionally, a custom VJP for GLU (fused_glu) is introduced to optimize the backward pass of prefused gate/up GEMMs. I have no feedback to provide as there are no review comments.
khatwanimohit
added a commit
to google/tunix
that referenced
this pull request
Oct 15, 2026
…ext with #5759 TRAINER_TUNED_FLAGS (on in mlperf_gbs1024_submission.sh) now also sets, each measured on the 128-chip v7x RL trainer bench (AI-Hypercomputer/maxtext#5701 harness): - --xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true: 97.24 -> 97.15 s/step. - --xla_jf_rematerialization_percent_shared_memory_limit=99: 10.226 -> 10.197 s per micro-batch. - moe_route=device: keeps the ragged sort's routing integers instead of replaying both argsorts in the backward (now tagged at the pin), 10.226 -> 10.184 s per micro-batch; +4.3 GiB compile temp, runs within the trainer's HBM. - wi_tile_dlhs_mlp_dim=2048 wi_tile_dlhs_embed_dim=2048: the wi input-gradient GEMM tiles the full contraction once. - REMAT_DECODER_LAYER_INPUT=device in the submission's tuned cell: the layer inputs stay in HBM, as in every bench run. MaxText pin: atwigg/mlperf plus AI-Hypercomputer/maxtext#5759 (GDN kernel, MoE routing and engine speedups) and #5741 (router and query kernels gathered once). With these flags the bench runs 9.796 s per micro-batch (89.77 s/step, 3,285 tok/s/chip) against 10.616 s (97.15 s/step) for atwigg/mlperf with #5741. Tests: mlperf_submission_launcher_test expects the new tuned keys and libtpu flags.
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.
Performance changes for the Qwen3.5-397B RL trainer (MaxText training engine under Tunix), measured on the 128-chip v7x trainer bench (#5701 harness, staged step boundary). All are fixed behaviour (no new config flags), and numerics stay bit-exact or change only in summation order.
Changes
stop_gradienton parameterstrainable_parameters_maskfreezes, so the gate's weight-gradient matmuls and 256-way reduce-scatters are no longer computed and then discarded (update and grad_norm unchanged).fused_glucustom_vjp returns one concatenated gate/up gradient instead of XLA's pad-and-add over the whole ragged buffer (bit-identical).lax.condfallback; bit-identical).gdn_cp_m_scan,gdn_cp_dh_scan) instead of XLA's balanced trees of HBM-bound f32 fusions; closer to HIGHEST than the old HIGH tree.Measured (128 v7x chips, 256 devices; regular micro-batch median, steps >= 2)
moe_route=device,wi_tile_dlhs_mlp_dim=2048 wi_tile_dlhs_embed_dim=2048,--xla_jf_rematerialization_percent_shared_memory_limit=99(tunix#2896)Numerics and memory
Flags used in the measurements are in google/tunix#2896.