From bb1a2e88876dfc631240124fb8ea7cd7b3bd0850 Mon Sep 17 00:00:00 2001 From: Noman654 Date: Sat, 8 Aug 2026 19:25:29 +0000 Subject: [PATCH] Add gdn_matmul_precision opt-in knob for Qwen3-Next Gated DeltaNet (#4510) The Gated DeltaNet delta-rule matmuls pinned jax.lax.Precision.HIGHEST. Per issue #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 upcast to float32, and HIGHEST is what makes that upcast meaningful on the TPU MXU. This adds an opt-in gdn_matmul_precision config (default "highest", so no behavior change) that plumbs through jax_chunk_gated_delta_rule and jax_ar_gated_delta_rule. naive_jax_chunk_gated_delta_rule keeps HIGHEST as its default and gains an optional precision= argument so it stays the numerical reference for the unit tests. The config is declared as MatmulPrecision on the Qwen3Next model in configs/types.py alongside the other gdn_* keys, so invalid values are rejected at config-validation time. --- src/maxtext/configs/base.yml | 7 + src/maxtext/configs/types.py | 8 ++ src/maxtext/models/qwen3.py | 59 ++++++-- tests/unit/qwen3_next_gdn_precision_test.py | 152 ++++++++++++++++++++ 4 files changed, 213 insertions(+), 13 deletions(-) create mode 100644 tests/unit/qwen3_next_gdn_precision_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 7a725fc4ca..0cbb8f980b 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a548422..9f244726b6 100644 --- a/src/maxtext/configs/types.py +++ b/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.", diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 7cb710bf1e..85d989c310 100644 --- a/src/maxtext/models/qwen3.py +++ b/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: diff --git a/tests/unit/qwen3_next_gdn_precision_test.py b/tests/unit/qwen3_next_gdn_precision_test.py new file mode 100644 index 0000000000..0e8d543853 --- /dev/null +++ b/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()