Problem
In the Qwen3-Next Gated Delta Net (jax_chunk_gated_delta_rule), the two sequential computations dominate TPU step time:
- the inter-chunk recurrence runs as
lax.scan, so the recurrent state round-trips HBM on every chunk step, and
- the UT-transform inverse
A = (I+S)^-1 uses solve_triangular, which substitutes row by row and barely uses the MXU.
Neither overlaps with other work, so they bound the step directly.
Measurement
v5e-256, Qwen3-Next dense-8B config, seq 4096, bs 1, remat full (details in #4348):
| build |
step time |
tokens/s/device |
| main |
3.373s |
1,214 |
| #4348 (Pallas kernels) |
2.571s |
1,593 (+31%) |
Loss trajectories match to bf16 rounding level.
Problem
In the Qwen3-Next Gated Delta Net (
jax_chunk_gated_delta_rule), the two sequential computations dominate TPU step time:lax.scan, so the recurrent state round-trips HBM on every chunk step, andA = (I+S)^-1usessolve_triangular, which substitutes row by row and barely uses the MXU.Neither overlaps with other work, so they bound the step directly.
Measurement
v5e-256, Qwen3-Next dense-8B config, seq 4096, bs 1, remat full (details in #4348):
Loss trajectories match to bf16 rounding level.