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

src/maxtext/configs/base.yml

@@ -1275,6 +1275,13 @@ gdn_num_key_heads: 16
 gdn_num_value_heads: 32
 # Chunk size for the parallel scan algorithm in the Gated Delta Net.
 gdn_chunk_size: 64
+# Precision for the matmuls inside the Gated Delta Rule kernels. Choose one of default, high, and highest.
+# These matmuls are deliberately computed on float32 operands, and "highest" is what makes that upcast
+# meaningful on TPU: it emulates float32 with 6 bf16 MXU passes, where "default" would truncate the
+# operands back to bf16. Lowering it is therefore a numerics change, not just a speed knob, and the
+# delta rule's (I + S)^-1 is sensitive to it. Kept separate from matmul_precision, which stays free to
+# describe the ordinary projections, and defaulted to highest to preserve the kernel's original behavior.
+gdn_matmul_precision: "highest"
 # Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.
 use_qk_norm_in_gdn: true
 # The ratio of dimension to apply ROPE on

src/maxtext/configs/types.py

@@ -1063,6 +1063,14 @@ class Qwen3Next(BaseModel):
       64,
       description="Chunk size for the parallel scan algorithm in the Gated Delta Net.",
   )
+  gdn_matmul_precision: MatmulPrecision = Field(
+      MatmulPrecision.HIGHEST,
+      description=(
+          "Precision for the matmuls inside the Gated Delta Rule kernels. These matmuls are deliberately "
+          "computed on float32 operands, and 'highest' is what makes that upcast meaningful on TPU, so this "
+          "defaults to 'highest' to preserve the kernel's original behavior."
+      ),
+  )
   use_qk_norm_in_gdn: bool = Field(
       True,
       description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.",

src/maxtext/models/qwen3.py

@@ -56,9 +56,23 @@
 
 
 def naive_jax_chunk_gated_delta_rule(
-    query, key, value, g, beta, chunk_size=64, initial_state=None, use_qk_norm_in_gdn=False
+    query,
+    key,
+    value,
+    g,
+    beta,
+    chunk_size=64,
+    initial_state=None,
+    use_qk_norm_in_gdn=False,
+    precision=jax.lax.Precision.HIGHEST,
 ):
-  """Naive implementation of the Gated Delta Rule in jax."""
+  """Naive implementation of the Gated Delta Rule in jax.
+
+  This is the numerical reference used by the unit tests, so `precision` defaults
+  to `jax.lax.Precision.HIGHEST` and should be left at that default when the
+  function is used as the fp32 baseline. It is exposed only so tests can drive
+  the kernel at other precisions.
+  """
   initial_dtype = query.dtype
   if use_qk_norm_in_gdn:
     query = l2norm(query, dim=-1, eps=1e-6)
@@ -103,7 +117,7 @@ def naive_jax_chunk_gated_delta_rule(
   g_diff_exp = jnp.exp(g_diff_tril).astype(jnp.float32)
   decay_mask = g_diff_exp
 
-  prec = jax.lax.Precision.HIGHEST
+  prec = precision
   attn = -jnp.matmul(k_beta_c, jnp.swapaxes(key_c, -1, -2), precision=prec) * decay_mask
   attn = jnp.where(mask, 0.0, attn)
 
@@ -146,7 +160,7 @@ def inner_attn_body(i, attn_val):
   def scan_body(prev_state, x):
     q_i, k_i, v_i, k_cumdecay_i, g_i, decay_mask_i = x
     last_recurrent_state = prev_state
-    prec = jax.lax.Precision.HIGHEST
+    prec = precision
 
     attn_i = jnp.matmul(q_i, jnp.swapaxes(k_i, -1, -2), precision=prec) * decay_mask_i
     attn_i = jnp.where(mask_inter, 0.0, attn_i)
@@ -190,8 +204,17 @@ def jax_chunk_gated_delta_rule(
     initial_state: None | Array = None,
     use_qk_norm_in_gdn: bool = False,
     compute_dtype: jnp.dtype = jnp.bfloat16,
+    matmul_precision: str = "highest",
 ) -> tuple[Array, None | Array]:
-  """Optimized JAX implementation of Gated Delta Rule."""
+  """Optimized JAX implementation of Gated Delta Rule.
+
+  Args:
+    matmul_precision: Precision for the delta-rule matmuls, as accepted by
+      `jax.lax.Precision`. Most of these matmuls run on operands this function
+      deliberately upcasts to float32, so "highest" (the default, and the value
+      this path used before it was configurable) is what makes that upcast
+      meaningful on TPU. Lowering it truncates those operands back to bf16.
+  """
   # =========================================================================
   # STAGE 1: PREPARATION & PADDING
   # =========================================================================
@@ -253,7 +276,7 @@ def to_chunk_scalar(x):
   k_beta = k_c * beta_c[..., None]
 
   # S Matrix Calculation
-  S = jnp.matmul(k_beta, k_c.swapaxes(-1, -2), precision=jax.lax.Precision.HIGHEST)
+  S = jnp.matmul(k_beta, k_c.swapaxes(-1, -2), precision=matmul_precision)
   S = S.astype(jnp.float32)
 
   # Apply mask BEFORE exp to prevent 'inf' gradients
@@ -272,11 +295,11 @@ def to_chunk_scalar(x):
 
   # 5. WY Factors
   v_beta = v_c * beta_c[..., None]
-  u_chunks = jnp.matmul(A, v_beta.astype(jnp.float32), precision=jax.lax.Precision.HIGHEST)
+  u_chunks = jnp.matmul(A, v_beta.astype(jnp.float32), precision=matmul_precision)
   u_chunks = u_chunks.astype(compute_dtype)
 
   k_beta_g = k_beta.astype(jnp.float32) * jnp.exp(g_cumsum)[..., None]
-  w_chunks = jnp.matmul(A, k_beta_g, precision=jax.lax.Precision.HIGHEST)
+  w_chunks = jnp.matmul(A, k_beta_g, precision=matmul_precision)
   w_chunks = w_chunks.astype(compute_dtype)
 
   # =========================================================================
@@ -300,7 +323,7 @@ def to_chunk_scalar(x):
 
   def scan_body(h, args):
     w, u, q, k, g = args
-    prec = jax.lax.Precision.HIGHEST
+    prec = matmul_precision
 
     # --- Output Computation ---
     # 1. Inter-chunk: q(dtype) * exp(g)(f32) -> f32
@@ -369,8 +392,15 @@ def jax_ar_gated_delta_rule(
     initial_state: Array,
     use_qk_norm_in_gdn: bool = False,
     compute_dtype: jnp.dtype = jnp.bfloat16,
+    matmul_precision: str = "highest",
 ) -> tuple[Array, Array]:
-  """Highly optimized step for Autoregressive Decoding (seq_len == 1)."""
+  """Highly optimized step for Autoregressive Decoding (seq_len == 1).
+
+  Args:
+    matmul_precision: Precision for the delta-rule matmuls, as accepted by
+      `jax.lax.Precision`. Defaults to "highest", preserving the precision this
+      path used before the value was configurable.
+  """
   # Shapes: q, k (B, 1, H, K_dim) | v (B, 1, H, V_dim) | g, beta (B, 1, H)
   initial_dtype = query.dtype
 
@@ -404,20 +434,20 @@ def jax_ar_gated_delta_rule(
 
   # v_prime = state @ (k_beta * exp(g))
   k_cumdecay = (k_beta.astype(jnp.float32) * g_exp)[..., None, :]  # (B, H, 1, K)
-  v_prime = jnp.matmul(k_cumdecay, state, precision=jax.lax.Precision.HIGHEST).squeeze(-2)
+  v_prime = jnp.matmul(k_cumdecay, state, precision=matmul_precision).squeeze(-2)
 
   v_new = v_beta.astype(jnp.float32) - v_prime
 
   # Core Output
   q_g = (q.astype(jnp.float32) * g_exp)[..., None, :]  # (B, H, 1, K)
-  attn_inter = jnp.matmul(q_g, state, precision=jax.lax.Precision.HIGHEST).squeeze(-2)
+  attn_inter = jnp.matmul(q_g, state, precision=matmul_precision).squeeze(-2)
 
   attn_intra = jnp.sum(q.astype(jnp.float32) * k.astype(jnp.float32), axis=-1, keepdims=True)
   core_attn_out = attn_inter + attn_intra * v_new
 
   # State Update: new_state = state * exp(g) + k^T @ v_new
   new_state = state * g_exp[..., None] + jnp.matmul(
-      k.astype(jnp.float32)[..., None], v_new[..., None, :], precision=jax.lax.Precision.HIGHEST
+      k.astype(jnp.float32)[..., None], v_new[..., None, :], precision=matmul_precision
   )
 
   # Restore sequence dimension
@@ -838,6 +868,7 @@ def extract_state(c_in, v_len):
           initial_state=recurrent_state,  # pyrefly: ignore[bad-argument-type]
           use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
           compute_dtype=cfg.dtype,
+          matmul_precision=cfg.gdn_matmul_precision,
       )
     elif self.mesh is not None:
       logical_rules = get_logical_axis_rules()
@@ -878,6 +909,7 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
             initial_state=init_h,
             use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
             compute_dtype=cfg.dtype,
+            matmul_precision=cfg.gdn_matmul_precision,
         )
 
       core_attn_out, next_recurrent_state = shard_mapped_delta_rule(query, key, value, g, beta, recurrent_state_arg)
@@ -892,6 +924,7 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
           initial_state=recurrent_state,
           use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
           compute_dtype=cfg.dtype,
+          matmul_precision=cfg.gdn_matmul_precision,
       )
 
     if model_mode != MODEL_MODE_TRAIN and active_cache is not None:

tests/unit/qwen3_next_gdn_precision_test.py

@@ -0,0 +1,152 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests that the Qwen3-Next Gated DeltaNet kernels honor gdn_matmul_precision.
+
+The delta-rule matmuls historically pinned `jax.lax.Precision.HIGHEST`. These
+tests lock in the opt-in knob: the default reproduces that pinned behavior
+exactly, and the configured value actually reaches the matmuls rather than being
+accepted and ignored.
+"""
+
+import unittest
+import jax
+import jax.numpy as jnp
+from maxtext.models import qwen3
+
+
+def _inputs(batch=1, seq_len=128, num_heads=2, head_dim=64, seed=0):
+  """Returns (query, key, value, g, beta, initial_state) for the delta rule."""
+  keys = jax.random.split(jax.random.PRNGKey(seed), 6)
+  query = jax.random.normal(keys[0], (batch, seq_len, num_heads, head_dim), jnp.float32)
+  key = jax.random.normal(keys[1], (batch, seq_len, num_heads, head_dim), jnp.float32)
+  value = jax.random.normal(keys[2], (batch, seq_len, num_heads, head_dim), jnp.float32)
+  g = -jax.nn.softplus(jax.random.normal(keys[3], (batch, seq_len, num_heads), jnp.float32))
+  beta = jax.nn.sigmoid(jax.random.normal(keys[4], (batch, seq_len, num_heads), jnp.float32))
+  initial_state = jax.random.normal(keys[5], (batch, num_heads, head_dim, head_dim), jnp.float32) * 0.1
+  return query, key, value, g, beta, initial_state
+
+
+class Qwen3NextGdnPrecisionTest(unittest.TestCase):
+  """The delta-rule matmuls must take their precision from the caller."""
+
+  def test_chunk_rule_default_arg_preserves_previous_behavior(self):
+    """Omitting matmul_precision must reproduce the previously hardcoded HIGHEST."""
+    query, key, value, g, beta, initial_state = _inputs()
+    kwargs = {"chunk_size": 64, "use_qk_norm_in_gdn": True, "initial_state": initial_state}
+
+    omitted, omitted_state = qwen3.jax_chunk_gated_delta_rule(query, key, value, g, beta, **kwargs)
+    pinned, pinned_state = qwen3.jax_chunk_gated_delta_rule(
+        query, key, value, g, beta, matmul_precision="highest", **kwargs
+    )
+
+    self.assertTrue(jnp.array_equal(omitted, pinned))
+    self.assertTrue(jnp.array_equal(omitted_state, pinned_state))
+
+  def test_ar_rule_default_arg_preserves_previous_behavior(self):
+    """The decode path must be unchanged for callers that omit the argument."""
+    query, key, value, g, beta, initial_state = _inputs(seq_len=1)
+    kwargs = {"initial_state": initial_state, "use_qk_norm_in_gdn": True}
+
+    omitted, omitted_state = qwen3.jax_ar_gated_delta_rule(query, key, value, g, beta, **kwargs)
+    pinned, pinned_state = qwen3.jax_ar_gated_delta_rule(query, key, value, g, beta, matmul_precision="highest", **kwargs)
+
+    self.assertTrue(jnp.array_equal(omitted, pinned))
+    self.assertTrue(jnp.array_equal(omitted_state, pinned_state))
+
+  def test_naive_rule_default_arg_preserves_baseline(self):
+    """The reference kernel must stay at HIGHEST by default so tests keep their baseline."""
+    query, key, value, g, beta, initial_state = _inputs()
+    kwargs = {"chunk_size": 64, "use_qk_norm_in_gdn": True, "initial_state": initial_state}
+
+    omitted, omitted_state = qwen3.naive_jax_chunk_gated_delta_rule(query, key, value, g, beta, **kwargs)
+    pinned, pinned_state = qwen3.naive_jax_chunk_gated_delta_rule(
+        query, key, value, g, beta, precision=jax.lax.Precision.HIGHEST, **kwargs
+    )
+
+    self.assertTrue(jnp.array_equal(omitted, pinned))
+    self.assertTrue(jnp.array_equal(omitted_state, pinned_state))
+
+  def test_chunk_rule_precision_reaches_the_matmuls(self):
+    """The requested precision must show up in the lowered HLO, not be dropped.
+
+    This is the property under test: the value has to reach jnp.matmul rather
+    than being accepted and ignored. It is asserted on the HLO because on CPU
+    precision does not change the emitted numerics.
+    """
+    query, key, value, g, beta, initial_state = _inputs()
+
+    def lower(precision):
+      def fn(q, k, v, gate, b):
+        return qwen3.jax_chunk_gated_delta_rule(
+            q,
+            k,
+            v,
+            gate,
+            b,
+            chunk_size=64,
+            use_qk_norm_in_gdn=True,
+            initial_state=initial_state,
+            matmul_precision=precision,
+        )
+
+      return jax.jit(fn).lower(query, key, value, g, beta).as_text()
+
+    self.assertIn("HIGHEST", lower("highest"))
+    self.assertNotIn("HIGHEST", lower("default"))
+
+  def test_ar_rule_precision_reaches_the_matmuls(self):
+    """Same property for the autoregressive decode path."""
+    query, key, value, g, beta, initial_state = _inputs(seq_len=1)
+
+    def lower(precision):
+      def fn(q, k, v, gate, b):
+        return qwen3.jax_ar_gated_delta_rule(
+            q,
+            k,
+            v,
+            gate,
+            b,
+            initial_state=initial_state,
+            use_qk_norm_in_gdn=True,
+            matmul_precision=precision,
+        )
+
+      return jax.jit(fn).lower(query, key, value, g, beta).as_text()
+
+    self.assertIn("HIGHEST", lower("highest"))
+    self.assertNotIn("HIGHEST", lower("default"))
+
+  def test_configured_precisions_are_accepted(self):
+    """Every value gdn_matmul_precision may take in base.yml must run."""
+    query, key, value, g, beta, initial_state = _inputs()
+
+    for precision in ("default", "high", "highest"):
+      with self.subTest(precision=precision):
+        out, _ = qwen3.jax_chunk_gated_delta_rule(
+            query,
+            key,
+            value,
+            g,
+            beta,
+            chunk_size=64,
+            use_qk_norm_in_gdn=True,
+            initial_state=initial_state,
+            matmul_precision=precision,
+        )
+        self.assertTrue(bool(jnp.isfinite(out).all()))
+
+
+if __name__ == "__main__":
+  unittest.main()