Add gdn_matmul_precision opt-in knob for Qwen3-Next Gated DeltaNet

[Noman654]

What

Adds an opt-in gdn_matmul_precision config for the Qwen3-Next Gated DeltaNet (GDN) delta-rule matmuls, as requested by @parambole in #4510 (Option 2). Default is "highest", so out-of-the-box numerics are unchanged; "default"/"high" are a pure opt-in for users who want the speedup after validating their setup.

Why

The GDN kernels in models/qwen3.py pinned jax.lax.Precision.HIGHEST in the delta-rule matmuls. Per #4510 that pin is deliberate: the delta rule solves an intra-chunk triangular system (I+S)A=I and runs an inter-chunk recurrent scan on operands the kernel upcasts to float32, and HIGHEST is what makes that upcast meaningful on the TPU MXU. This keeps that behavior by default and only exposes a knob to lower it.

What changed

• base.yml: new gdn_matmul_precision: "highest".
• jax_chunk_gated_delta_rule / jax_ar_gated_delta_rule: new matmul_precision param (default "highest") threaded to the 7 delta-rule matmuls.
• Three call sites in Qwen3NextGatedDeltaNet forward cfg.gdn_matmul_precision.
• naive_jax_chunk_gated_delta_rule: gains an optional precision= (default HIGHEST, unchanged) so it stays the numerical reference in qwen3_next_vs_reference_test.py, per @parambole's request.

Tests

New tests/unit/qwen3_next_gdn_precision_test.py (6 tests): the default arg is bit-identical to explicit HIGHEST for all three kernels (zero behavior change), the configured precision reaches the lowered HLO, and default/high/highest all run. pytest green locally.

Open question

The two production kernels take a string matmul_precision="highest" (following the attention_mla.py precedent of passing the config string to jnp.matmul), while naive takes precision=jax.lax.Precision.HIGHEST as written in the request. Happy to make them consistent either way — let me know your preference.

Validation of the "default" path on real data (chunk 64 and 128) is summarized in a comment on #4510.

Addresses #4510.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces a new configuration parameter gdn_matmul_precision to control the precision of matmuls inside the Gated Delta Rule kernels in Qwen3, defaulting to "highest" to preserve original behavior. The relevant JAX-based delta rule functions have been updated to accept and apply this precision setting, and a new unit test suite has been added to verify that the configuration is correctly respected. There are no review comments, and I have no feedback to provide.