[Gemma4] Fix the bug of use_qk_norm accidently set to False in Gemma 4 configs

src/maxtext/configs/models/gemma4-26b.yml

@@ -39,6 +39,7 @@ global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 v_norm_with_scale: false
 final_logits_soft_cap: 30.0
+use_qk_norm: true
 
 # MoE configuration
 num_experts: 128

src/maxtext/configs/models/gemma4-31b.yml

@@ -42,6 +42,7 @@ rope_max_timescale: 1000000
 global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
+use_qk_norm: true
 
 # Multimodal flags (need to set use_multimodal=true)
 vision_encoder_block: "gemma4"

src/maxtext/configs/models/gemma4-e2b.yml

@@ -43,6 +43,8 @@ rope_max_timescale: 1000000
 global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
+use_qk_norm: true
+scan_layers: false
 
 # Vision encoder flags — multimodal not yet supported for E2B / E4B.
 vision_encoder_block: "gemma4"

src/maxtext/configs/models/gemma4-e4b.yml

@@ -44,6 +44,8 @@ rope_max_timescale: 1000000
 global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
+use_qk_norm: true
+scan_layers: false
 
 # Vision encoder flags — multimodal not yet supported for E2B / E4B.
 vision_encoder_block: "gemma4"

src/maxtext/layers/attention_compressed.py

@@ -1110,7 +1110,7 @@ def __init__(
       sliding_window_size: int | None = None,
       use_ragged_attention: bool = False,
       ragged_block_size: int = 256,
-      use_qk_norm: bool = False,
+      use_qk_norm: bool | None = None,
       query_pre_attn_scalar: float | None = None,
       use_bias_in_projections: bool = False,
       # Compression Specific Parameters:
@@ -1606,7 +1606,7 @@ def compressed_attention(
     sliding_window_size: int | None = None,
     use_ragged_attention: bool = False,
     ragged_block_size: int = 256,
-    use_qk_norm: bool = False,
+    use_qk_norm: bool | None = None,
     query_pre_attn_scalar: float | None = None,
     use_bias_in_projections: bool = False,
     q_lora_rank: int = 1536,

src/maxtext/layers/attention_mla.py

@@ -477,7 +477,7 @@ def mla_as_linen(
     sliding_window_size: int | None = None,
     use_ragged_attention: bool = False,
     ragged_block_size: int = 256,
-    use_qk_norm: bool = False,
+    use_qk_norm: bool | None = None,
     query_pre_attn_scalar: float | None = None,
     use_bias_in_projections: bool = False,  # Set to True will enable bias in q, k, v, o projections
     # Temperature tuning parameters used for Llama4
@@ -615,7 +615,7 @@ def __init__(
       sliding_window_size: int | None = None,
       use_ragged_attention: bool = False,
       ragged_block_size: int = 256,
-      use_qk_norm: bool = False,
+      use_qk_norm: bool | None = None,
       query_pre_attn_scalar: float | None = None,
       use_bias_in_projections: bool = False,  # Set to True will enable bias in q, k, v, o projections
       # Temperature tuning parameters used for Llama4

src/maxtext/layers/attentions.py

@@ -126,7 +126,7 @@ def attention_as_linen(
     sliding_window_size: int | None = None,
     use_ragged_attention: bool = False,
     ragged_block_size: int = 256,
-    use_qk_norm: bool = False,
+    use_qk_norm: bool | None = None,
     query_pre_attn_scalar: float | None = None,
     use_bias_in_projections: bool = False,  # Set to True will enable bias in q, k, v, o projections
     share_kv_projections: bool = False,  # If true, Key and Value use the same projection
@@ -283,7 +283,7 @@ def __init__(
       sliding_window_size: int | None = None,
       use_ragged_attention: bool = False,
       ragged_block_size: int = 256,
-      use_qk_norm: bool = False,
+      use_qk_norm: bool | None = None,
       query_pre_attn_scalar: float | None = None,
       use_bias_in_projections: bool = False,  # Set to True will enable bias in q, k, v, o projections
       share_kv_projections: bool = False,  # If true, Key and Value use the same projection
@@ -394,7 +394,7 @@ def __init__(
     self.sliding_window_size = sliding_window_size
     self.use_ragged_attention = use_ragged_attention
     self.ragged_block_size = ragged_block_size
-    self.use_qk_norm = use_qk_norm
+    self.use_qk_norm = getattr(self.config, "use_qk_norm", False) if use_qk_norm is None else use_qk_norm
     self.query_pre_attn_scalar = query_pre_attn_scalar
     self.use_bias_in_projections = use_bias_in_projections
     self.share_kv_projections = share_kv_projections
@@ -665,9 +665,7 @@ def init_query_w(self, inputs_q_shape: Tuple) -> nnx.Module:
     #       linear transformations, which is equivalent under Adafactor.
     # We disable depth_scaling when using qk_norm or a query_pre_attn_scalar
     # to avoid applying scaling twice.
-    if getattr(self.config, "use_qk_norm", False) or (
-        self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar != 1.0
-    ):
+    if self.use_qk_norm or (self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar != 1.0):
       depth_scaling = 1.0
     else:
       depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype)

tests/unit/gemma4_canonical_test.py

@@ -0,0 +1,91 @@
+# Copyright 2026 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.
+
+"""Unit tests for canonical Gemma 4 attention scaling and weight initialization."""
+
+import os
+import unittest
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+from maxtext.configs import pyconfig
+from maxtext.common import common_types
+from maxtext.models import gemma4, gemma4_small
+from maxtext.utils.globals import MAXTEXT_REPO_ROOT
+
+
+class Gemma4CanonicalAttentionTest(unittest.TestCase):
+  """Tests that Gemma 4 models follow canonical attention scaling and unscaled weight init."""
+
+  def setUp(self):
+    super().setUp()
+    self.base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml")
+
+  def test_gemma4_26b_attention_config(self):
+    config = pyconfig.initialize(
+        ["", self.base_config_path],
+        model_name="gemma4-26b",
+        enable_dropout=False,
+    )
+    self.assertTrue(config.use_qk_norm, "gemma4-26b should enable use_qk_norm in config")
+
+    mesh = jax.sharding.Mesh(jax.devices()[:1], ("data",))
+    rngs = nnx.Rngs(0)
+    layer = gemma4.Gemma4DecoderLayer(
+        config=config,
+        mesh=mesh,
+        model_mode=common_types.MODEL_MODE_PREFILL,
+        rngs=rngs,
+        attention_type=gemma4.AttentionType.LOCAL_SLIDING,
+        layer_idx=0,
+    )
+    # Canonical Gemma 4 attention uses query_pre_attn_scalar = 1.0 (unscaled logits)
+    self.assertEqual(layer.self_attention.query_pre_attn_scalar, 1.0)
+    self.assertTrue(layer.self_attention.use_qk_norm)
+
+    # Initial query weights should NOT be divided by sqrt(head_dim) = 16.0
+    # Expected standard deviation for fan_in=2816 is 1/sqrt(2816) ~= 0.0188
+    q_kernel = layer.self_attention.query.kernel[...]
+    std_q = float(jnp.std(q_kernel))
+    self.assertGreater(std_q, 0.01, f"Query kernel std ({std_q}) should not be divided by depth_scaling")
+
+  def test_gemma4_small_attention_config(self):
+    for model_name in ["gemma4-e2b", "gemma4-e4b"]:
+      config = pyconfig.initialize(
+          ["", self.base_config_path],
+          model_name=model_name,
+          enable_dropout=False,
+      )
+      self.assertTrue(config.use_qk_norm, f"{model_name} should enable use_qk_norm in config")
+
+      mesh = jax.sharding.Mesh(jax.devices()[:1], ("data",))
+      rngs = nnx.Rngs(0)
+      layer = gemma4_small.Gemma4SmallDecoderLayer(
+          config=config,
+          mesh=mesh,
+          model_mode=common_types.MODEL_MODE_PREFILL,
+          layer_idx=0,
+          rngs=rngs,
+      )
+      self.assertEqual(layer.self_attention.query_pre_attn_scalar, 1.0)
+      self.assertTrue(layer.self_attention.use_qk_norm)
+
+      q_kernel = layer.self_attention.query.kernel[...]
+      std_q = float(jnp.std(q_kernel))
+      self.assertGreater(std_q, 0.01, f"Query kernel std ({std_q}) should not be divided by depth_scaling")
+
+
+if __name__ == "__main__":
+  unittest.main()