Fuse the Gated Delta Net inter-chunk scan into a Pallas TPU kernel

[bzantium]

Description

Fuses the sequential parts of the Qwen3-Next Gated Delta Net into Pallas TPU kernels: the inter-chunk recurrence, the UT-transform triangular inversion, and head/tile batching to keep the MXU pipelined through both. Together: 3.373s → 2.571s/step (+31% throughput) at the reference config.

FIXES: #4347

Changes

• Inter-chunk scan kernel — the recurrence runs as one kernel per (batch, heads) with the state in VMEM, instead of a lax.scan whose state round-trips HBM every chunk. Backward walks the chunks in reverse (FLA-style checkpointing); initial/final state are first-class, so the custom_vjp covers the full state chain. The chunk-parallel WY math stays in XLA on purpose — TPU grid cells run sequentially per core, and fusing batched matmuls into the walk measured slower.
• Triangular inversion kernel — replaces solve_triangular (row-sequential, no MXU use) for A = (I+S)^-1. Blockwise ladder: with X = (I+S_b)^-1 at block size b, the size-2b inverse is exactly X - X @ S_l @ X, so doubling from b=1 needs 2·log2(C) full-width matmuls. It never forms powers of S (those overflow f32 once |S| > 1), so intermediates stay at the true inverse's scale — stable wherever the answer is representable. Products use manual bf16x3 operand splitting (Mosaic's default f32 dot truncates to bf16, which breaks (I+S)A = I enough to destabilize the recurrence; Precision.HIGH is unsupported). Backward is analytic: dS = -(A^T dA A^T), strict-lower masked.
• Head/tile batching — the walks are sequential, but heads (scan) and chunk tiles (inversion) are independent. Batching them per grid cell (scan fwd 8 / bwd 4 heads, inversion 8 tiles — bounded by the 16MB scoped-VMEM limit) turns each dot into a batched dot_general the MXU can pipeline instead of paying fill/drain per short matmul.

Peak memory is unchanged — the eliminated HBM traffic was latency-hidden, which is also why the gains come specifically from the sequential paths that cannot overlap.

Performance

v5e-256, Qwen3-Next dense-8B (emb 4096, 36 layers, GDN 16/32 heads, head dim 128), seq 4096, bs 1, bf16, FSDP, gdn_chunk_size=128, remat full, synthetic data:

build	step time	tokens/s/device
main	3.373s	1,214
+ scan kernel	3.119s	1,313
+ head batching	3.076s	1,331
+ inversion ladder	2.571s	1,593

Loss trajectories match to bf16 rounding level at every step.

Tests

python3 -m pytest tests/unit/gated_delta_network_kernel_test.py — 9 tests on CPU (interpret mode): scan kernel vs a lax.scan reference (forward, final state, all gradients, non-zero initial state and final-state cotangent, bf16 dtypes); inversion vs solve_triangular (random, all-ones, unnormalized gaussian; forward + gradients). End-to-end training verified on v5e-256 (table above). pyink / pylint clean.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov]

Codecov Report

❌ Patch coverage is 98.63014% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/qwen3.py	85.71%	1 Missing and 1 partial ⚠️
...c/maxtext/kernels/attention/gated_delta_network.py	99.51%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

[github-actions]

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.

[WandLZhang]

Independent reproduction, on a dense GatedDeltaNet config: d_model 4096, 32 layers, one full-attention layer in four, vocab_size=512, bf16, FSDP, gdn_chunk_size=128, sequence 8192, batch 1 per device, synthetic data. Both columns are the same checkout, with the control built by removing these two commits.

	main	this PR	gain
v6e-8	3,216 tok/s/chip, 137.1 TFLOP/s	3,838, 163.6	+19.4%
v5p, 4 chips, batch 1	3,121	3,883	+24.4%
v5p, 4 chips, batch 2	3,711	4,470	+20.4%

Loss matches the control to three decimals in every pair. It rebases onto current main with one conflict, an import block in models/qwen3.py where both sides are wanted.

Two things I hit that aren't in the description.

Peak memory falls by about a third. The description says peak memory is unchanged. On this config sequence 16384 runs on a v6e-8 with the kernel and OOMs without it. Where both OOM, the reported requirement falls 31% at 32768 (86.39G to 59.59G) and 35% at 65536 (259.68G to 169.83G). Worth claiming, since it buys a sequence length main can't reach.

The if use_pallas: branch returns from jax_chunk_gated_delta_rule above everything below it, so any work later in that function is skipped. We carry a patch that makes the inter-chunk recurrence context-parallel by composing each device's affine map at that point, and with the kernel on it never runs: every device does an isolated recurrence from a zero state. Training loss doesn't show it, since 60 steps at ctx=2 match to three decimals. Comparing the output tensors under a two-shard context mesh does. The shard that should inherit state is off by 7.9e-1 max, against 5.7e-2 for the shard that has nothing upstream to inherit, and the final recurrent state by 3.78. This only reaches callers who run context parallelism on this path, so it may be out of scope, but it's worth a line in the kernel docstring.

Would you consider reopening? The gain is large and it's a rebase rather than new work. Happy to carry that.

[WandLZhang]

cc @mmcsa

[bzantium]

@WandLZhang Reopened. Thanks for taking the time to reproduce this, and for offering to carry the rebase.

On the numbers, our original run was v5e-256 at around +8%, so it's good to see +19-24% on v6e/v5p. Nice that it holds up on the newer chips.

On peak memory, you're right that "unchanged" was too conservative. The recurrent state stays in VMEM across the inter-chunk walk instead of getting written back to HBM every scan step, so the drop makes sense. A 31-35% reduction that lets you fit a sequence length main OOMs on is worth calling out, I'll update the description.

On the context-parallel issue, good catch. The use_pallas branch returns before the code below it, so the fused path only threads a single initial_state/final_h and never composes state across sequence-sharded shards. Each shard just starts from its own h0. It doesn't surface in this PR's configs because main keeps the sequence replicated (batch/head sharding only), so there's no upstream state for a shard to inherit, which is why ctx=1 loss matches to three decimals. It only shows up once you shard the sequence, like your ctx=2 comparison. I added a note on the branch (dd5f9b0) so anyone running context parallelism knows to compose per-shard state around the call. If you want to upstream your affine-composition patch on top of this, happy to coordinate.

Branch is rebased and pushed. I'll take the import conflict in models/qwen3.py.

[WandLZhang]

Thanks for reopening, and for the NOTE.

Re-verified against your rebase. Applying our patch to the new base b18ceac97 reproduces the branch exactly in both files, and the kernel blob is unchanged, so the +19% to +24% above still stands as measured. We carry the two commits plus your note locally until this lands.

On the context-parallel side, the note covers it for callers. The composition itself is in #4932. The inter-chunk step is

h_new = exp(g_last) * h + k_g^T (u - w h) = (exp(g_last) I - k_g^T w) h + k_g^T u = A h + B

A and B don't depend on h, so affine maps compose associatively and each device can fold its local chunks into one (A, B) pair, all-gather the D pairs, take the exclusive prefix for its incoming state, and replay locally. A and B are 128x128 for every published Qwen 3.5 size, so the gathered volume stays small. It needs jax.checkpoint on the local scan body; without it the autodiff residuals reach 103 GB at sequence 262,144.

I can send that as a PR on top once yours merges. Let me know.

[bzantium]

@WandLZhang thanks for re-verifying against the rebase, glad the numbers held.

The affine-composition approach makes sense to me. Folding each device's local chunks into a single (A, B) pair and taking the exclusive prefix over the gathered pairs is a clean way to do it, and with A, B at 128x128 the all-gather stays cheap. Good catch on needing jax.checkpoint around the local scan body too, that residual blowup at long sequence would be easy to miss.

Please do send the context-parallel PR on top once this lands, that's a natural follow-up and I'd be happy to review it. I'll ping the maintainers to take another look here so it can merge.

[WandLZhang]

Thanks. I'll send the context-parallel PR once this lands.

One design question:

gdn_inter_chunk_scan takes h0 and returns (o, h_final), so it evaluates A·h0 + B for a given h0. The exclusive prefix needs A on its own: each device has to publish its affine map before it knows its inbound state, otherwise the devices serialise.

Two options:

Return it. _fwd_pallas already forms the pieces, gamma = exp(g_last) and the k_gᵀw term. An optional third output of [B, H, D_k, D_k], off by default so nothing changes for existing callers.

Or form A_loc in XLA next to the kernel. It's a product of (exp(g_last)·I − k_gᵀw) over the device's chunks, 128x128 each. At 1M over 64 devices with gdn_chunk_size=64 that's about 244 small matmuls per device, negligible against the attention in the same layer. The sequential walk stays in your kernel.

I lean to the second: your signature stays untouched and the composition stays where the collective already is. If returning A is close to free, the first is tidier.

[WandLZhang]

Rethinking the ordering.

At 1M the two can't both be live. The composition needs each device's A, and the fused path returns before it, so our patch sets use_pallas false whenever a context axis is set. The 1M run I mentioned used the XLA scan, not your kernel.

So they're disjoint rather than stacked. I can send the context-parallel PR against main and you rebase #4348 on top, which is the models/qwen3.py conflict you already absorbed once. Let me know.

[bzantium]

@WandLZhang makes sense, disjoint is the right call. Gating use_pallas off whenever a context axis is set keeps the fused path for the common case and hands the sequence-sharded case to the XLA scan where A is available for the prefix. That also lines up with the note on the branch, the caller is expected to handle composition, and falling back to the scan is exactly that.

I checked the fallback is loss-preserving before agreeing: the kernel path and the XLA scan path match to about 6e-6 relative on the output with no initial state, and 7e-6 output / 2e-5 final-state with an h0 passed in (float32, interpret mode). So switching to the scan under a context axis changes speed, not math, and state threading through h0 behaves the same on both paths, which is what the replay relies on.

On ordering, since they're disjoint either way works and it's just the one models/qwen3.py import block. #4348 is already rebased, green, and mergeable, so landing it first and rebasing your PR on top is probably the least work overall, but if the maintainers would rather take the context-parallel PR first I'm happy to rebase #4348 onto it instead. Either way, ping me when yours is up.

[WandLZhang]

@bzantium it's up: #4968, against main as we discussed. Small — models/gdn_cp.py plus the two pspec fixes, 121 lines. No dependency on #4348, so rebase whenever suits.