Qwen3-Next GDN hardcodes `Precision.HIGHEST`, ignoring `config.matmul_precision` — is this intentional?

[Noman654]

Summary

The Qwen3-Next Gated DeltaNet path in src/maxtext/models/qwen3.py hardcodes jax.lax.Precision.HIGHEST in
9 places and never reads config.matmul_precision. Every other attention layer in the repo honors that config
(43 usages in src/maxtext/layers/). On TPU this means the GDN matmuls always run 6 bf16 MXU passes
instead of the 1 implied by the default config.

I'd like to ask whether this is deliberate before sending a patch, because the git history suggests it was
a conscious choice at one point — and there is existing evidence in this repo that bf16 can destabilize the
delta-rule recurrence. Details below, including the evidence against changing it.

Verified at e190ed60 (current main at time of writing).

The inconsistency

src/maxtext/configs/base.yml:128:

matmul_precision: "default"

Honored across the codebase — 43 usages in src/maxtext/layers/, e.g. layers/attention_mla.py:348:

logits = jnp.einsum("bthd, bsd -> btsh", q, k, precision=self.config.matmul_precision)

Not honored by GDN — src/maxtext/models/qwen3.py, 9 sites across all 3 entry points:

106  prec = jax.lax.Precision.HIGHEST                     # naive_jax_chunk_gated_delta_rule (@58)
149  prec = jax.lax.Precision.HIGHEST                     #   └─ scan_body (@146)
256  S        = jnp.matmul(k_beta, k_c.swapaxes(-1,-2), precision=jax.lax.Precision.HIGHEST)
                                                          # jax_chunk_gated_delta_rule (@183)
275  u_chunks = jnp.matmul(A, v_beta.astype(jnp.float32), precision=jax.lax.Precision.HIGHEST)
279  w_chunks = jnp.matmul(A, k_beta_g,                   precision=jax.lax.Precision.HIGHEST)
303  prec = jax.lax.Precision.HIGHEST                     #   └─ scan_body (@301), inter-chunk scan
407  v_prime    = jnp.matmul(k_cumdecay, state, precision=jax.lax.Precision.HIGHEST)
                                                          # jax_ar_gated_delta_rule (@363), decode path
413  attn_inter = jnp.matmul(q_g, state,        precision=jax.lax.Precision.HIGHEST)
420  jnp.matmul(k.astype(jnp.float32)[...,None], v_new[...,None,:], precision=jax.lax.Precision.HIGHEST)

Notably, the call sites already forward other config values into these same functions — matmul_precision is
simply absent from the list (qwen3.py:878):

        return jax_chunk_gated_delta_rule(
            query=q, key=k, value=v, g=g_val, beta=beta_val,
            chunk_size=cfg.gdn_chunk_size,
            initial_state=init_h,
            use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
            compute_dtype=cfg.dtype,
        )

(All 3 call sites: 839 AR/decode, 878 sharded, 892 non-sharded.)

What the history suggests (why I think this may be intentional)

This looks less like an oversight than like a rationale that got separated from its code:

• 0e03921 (2025-10-20) introduced the reference impl with an explicit comment:

  # NOTE: Precision set to HIGHEST for numerical accuracy.
  prec = jax.lax.Precision.HIGHEST

• 060fecf (2026-02-13, the optimized GDN impl) carried Precision.HIGHEST into the new chunked fast path
  and removed that NOTE comment. At main today, no rationale for HIGHEST survives anywhere in the file.

So the precision choice appears deliberate in origin, but it was never reconciled with matmul_precision, and
the reasoning is no longer discoverable in the source. That's really what I'm asking about.

Measurements

All numbers are from one configuration and should not be read as general.
TPU v7x (Ironwood), JAX 0.10.2 / libtpu 0.0.42.1, isolated GDN kernel, fwd+bwd,
B=4, S=4096, H=32, K=V=128, gdn_chunk_size=64, use_qk_norm_in_gdn=True:

precision	fwd+bwd	vs HIGHEST
HIGHEST (current)	24.12 ms	1.00×
HIGH	18.82 ms	1.28×
DEFAULT	13.92 ms	1.73×

Accuracy vs an fp32 reference (8 seeds × 4 decay strengths, worst case), at default:
output cos 0.999984, final-state cos 0.999990, min gradient cos 0.999952 (over q, k, v, g, beta).

Long horizon: recurrent state chained 30× (122,880 tokens, weak decay = worst case for compounding) —
state cos 0.999992, flat across all 30 chains, no drift.

Full-model 500-step seed-locked A/B (Qwen3-Next-80B-A3B MoE, synthetic data, HIGHEST vs DEFAULT):
mean |Δloss| 0.0012 over steps ≤100 and 0.0000 over steps ≥400; identical final loss; no NaN.
The gap shrinks with horizon rather than compounding.

Evidence against changing this

I want to flag the counter-evidence myself, since it's material:

• @bzantium's Fuse the Gated Delta Net inter-chunk scan into a Pallas TPU kernel #4348 reports that bf16 "breaks (I+S)A = I enough to destabilize the recurrence" at
  gdn_chunk_size=128, and works around it with manual bf16x3 operand splitting.
• My results are at chunk 64 with use_qk_norm_in_gdn=True, where I measured max|S| ≈ 0.2 — the
  (I+S) inverse is very well conditioned there, which likely explains why 1 bf16 pass suffices for me.
  That conditioning argument cuts both ways: it implies nothing about other chunk sizes or without qk-norm.
• Both results can be true. My data does not show bf16 is safe in general.
• My convergence A/B used synthetic data, which approaches a memorization floor early; it stresses
  horizon/compounding but is not real-data optimization dynamics.
• On my own pipeline the end-to-end payoff is small (GDN ≈16% of step, which is comms-bound). The kernel win
  should matter far more on GDN-dominated dense configs — where Fuse the Gated Delta Net inter-chunk scan into a Pallas TPU kernel #4348 measured +31%.

Taken with the removed NOTE comment above, it's entirely possible HIGHEST is load-bearing and should stay.

Question

Is the hardcoded HIGHEST deliberate (delta-rule numerical stability), or an oversight that predates
matmul_precision being plumbed through the other layers?

Depending on your answer I'm happy to send either:

1. Honor the config — plumb matmul_precision through the GDN entry points and forward
   cfg.matmul_precision at the 3 call sites, following the existing gdn_chunk_size / use_qk_norm_in_gdn /
   dtype convention. Fixes the inconsistency, but does change default numerics (HIGHEST → "default").
2. Opt-in knob — add gdn_matmul_precision defaulting to "highest". Zero behavior change; the speedup is
   available to anyone who validates it for their config.

I'd lean toward (1) as the more correct fix, but given the stability concern and the history above I don't want
to change default numerics for everyone on the strength of a single-config result. Happy to follow your
preference.

A related sub-question: should naive_jax_chunk_gated_delta_rule (@58) be plumbed too, or stay pinned at
HIGHEST deliberately as the numerical reference? tests/unit/qwen3_next_vs_reference_test.py:643,1129 already
passes matmul_precision=highest explicitly — which the GDN hardcode currently makes a no-op for this path.

This touches the same file as #4348, so I'm happy to coordinate ordering / rebase behind it.

Environment: JAX 0.10.2, libtpu 0.0.42.1, TPU v7x.

[saisharan0103]

This appears to hit the Qwen3-Next GDN path in src/maxtext/models/qwen3.py, where the delta-rule matmuls use Precision.HIGHEST directly while the surrounding attention code paths use config.matmul_precision. The strongest signal is that all three GDN entry points seem to make the same precision choice, so I would inspect that path first rather than treating one callsite as drift. For a PR, I would start by threading config.matmul_precision through GDN behind the existing default behavior. Was HIGHEST kept there specifically for delta-rule stability?