Skip to content

Qwen3-Next: sequential GDN paths (lax.scan recurrence, solve_triangular) bound TPU step time #4347

Description

@bzantium

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.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions