[Gemma4] Enable use_qk_norm in Gemma 4 configs and fix attention weight init depth_scaling

# Description

This change resolves an issue with Gemma 4 attention weight initialization and QK normalization:

1. In `src/maxtext/layers/attentions.py`, `init_query_w` now checks `self.use_qk_norm` in addition to `self.config.use_qk_norm`. When `use_qk_norm` is enabled (or `query_pre_attn_scalar != 1.0`), `depth_scaling` is set to 1.0 instead of `sqrt(head_dim)`. This prevents dividing the initial query projection weights by `sqrt(head_dim)`, ensuring unscaled canonical weight initialization for Gemma 4.
2. In Gemma 4 model configs (`gemma4-26b.yml`, `gemma4-31b.yml`, `gemma4-e2b.yml`, `gemma4-e4b.yml`), explicitly enabled `use_qk_norm: true`.
3. In `gemma4-e2b.yml` and `gemma4-e4b.yml`, set `scan_layers: false` since per-layer embeddings and variable layer configs are not compatible with layer scanning.
4. Added `tests/unit/gemma4_canonical_test.py` to verify that canonical attention scaling (`query_pre_attn_scalar = 1.0`), `use_qk_norm`, and unscaled weight initialization hold for both dense/MoE Gemma 4 and small Gemma 4 models.

# Tests

Added unit test:
- `tests/unit/gemma4_canonical_test.py`

TAG=agy
CONV=f731b620-dff8-4a9f-b791-657890c4604f

Author: Dr-Left

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/attentions.py

@@ -665,8 +665,10 @@ 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 getattr(self.config, "use_qk_norm", False)
+        or (self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar != 1.0)
     ):
       depth_scaling = 1.0
     else:

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()