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

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions