[Pallas] Fused forward+backward kernel for Gated DeltaNet-2 on TPU: 411× over jax.lax.associative_scan, and open questions on head-sharding

[Akseleu-J]

Lineage
Baseline → jax.lax.associative_scan (GDN-2) → my pure-JAX port of the NVlabs Triton kernels (WY reference) → atomic_ops, a Pallas port for TPU v5e → atomic_gdn2, a TPU-native fused forward+backward kernel (the contribution). Numbers below are vs. each of the first three stages, respectively.

Context
I ported NVlabs' Gated DeltaNet-2 Triton kernels to jax.experimental.pallas for TPU v5e. That base port is atomic_ops. I then built atomic_gdn2 on top of it — a TPU-native kernel where both forward and backward are fused. Compared to the base port, it adds:

A unified mega-backward (B2+B1+B3+B4+B5) in a single Pallas launch.

Reuse of forward residuals instead of recomputing them.

A fused custom VJP.

Performance
Setup: single v5e-1, B=4, D=128, bs2=64. Measured forward+backward against jax.lax.associative_scan (GDN-2), and against a pure-JAX WY reference.

411.75× vs associative_scan (peak gradient, H=6, L=2048).

34.50× vs the pure-JAX WY reference.

117.76× end-to-end synthetic 34.87M LM step (B=8, L=2048): 204.6 ms vs 24101.3 ms for associative_scan.

These are single-chip v5e-1 numbers at the stated shapes; I'm not claiming they transfer to v5e-8 or other configurations.

Numerical audit
Because the kernel relies on bf16x3 for matmuls, I ran an independent 80-digit Decimal CPU experiment and exact TPU integer-bin contractions to check for gradient cancellation.

Relative L2 vs JAX reference is ~2.5e-06.

bf16x3 is not numerically identical to true FP32.

I found specific FP32 sign disagreements in stress-test gradients (FP64 resolves them), and some NEW/OLD parameter-update gates fail at specific seeds.

I don't claim universal numerical equivalence. Both passing checks and observed failures are retained in the repo to map the exact numerical contract.

Tile pack limits (H=6 vs H=8)
The speedup peaks at H=6. With bs2=64, the blocking fits cleanly into VMEM/MXU. At H=8 it does not fit cleanly, so the marginal gain per head drops. To be clear: the kernel doesn't "magically get faster at H=6" — associative_scan keeps materializing O(H·L·d²) intermediates, while this kernel is O(L·d²) per chip with state in VMEM.

Distributed execution.
I trained two models on a TPU v5e-8 using standard data parallelism (model replicated, batch split across 8 chips): a 334M pure GDN-2 model and a 307M hybrid model. Both are published. So DP on v5e-8 works end-to-end for this kernel family. The questions below are about tensor parallelism (head-sharding) — questions 1–2 — plus one separate single-chip MXU question (3), to scale beyond single-chip limits.

Open questions

1. Sharded H-scaling on v5e-8.
   GDN-2 state is per-head independent, so the scan shouldn't need cross-chip sync. If the per-chip sweet spot is ~6 heads, can we shard the H axis to hit a global peak? How does jax.lax.associative_scan behave under shard_map — does it all-gather cross-chip state, and how much ICI latency does the sync add?

2. 2D sharding (H × B).
   At B=4 and H-only sharding, per-chip batch drops below 2, which breaks my bs2=64 granularity. What mesh shapes are typically used here? Is a 2D mesh with batch DP + head TP the right approach?

3. MXU factorization for pairwise decay.
   My current broadcast-reduce for pairwise decay is VPU-bound. Is there a clean way to write the [B,H,L,L] pairwise decay as an MXU matmul in Pallas?

Links: repo · scaling data · numerical audit details · model repo