Reject KV head counts that cannot be sharded across the tensor-parallel mesh.

src/maxtext/layers/attentions.py

@@ -572,6 +572,75 @@ def _validate_kv_heads(self) -> None:
     if self.num_query_heads % self.num_kv_heads != 0:
       raise ValueError("Invalid num_kv_heads for GQA.")
 
+  def _validate_kv_head_sharding(
+      self,
+      kernel_axes: Tuple[Optional[str], ...],
+  ) -> None:
+    """Validates that the key/value head dimension can be sharded evenly.
+
+    Attention heads are atomic under tensor parallelism.  The `kv_heads`
+    logical axis of the key/value projection is split across whichever mesh
+    axes `logical_axis_rules` maps it to -- by default `tensor`,
+    `tensor_sequence` and `autoregressive` -- so the KV head count has to be
+    divisible by their combined size.  Note that `num_kv_heads` is the
+    per-layer count, which for models such as Gemma 4 is `global_num_kv_heads`
+    on global attention layers and `num_kv_heads` elsewhere.
+
+    Ulysses context parallelism shards the KV heads a second time, over the
+    context axis, using a runtime all-to-all rather than the weight's logical
+    axes, so that axis is folded in explicitly.  `types.py` separately requires
+    `num_kv_heads` to be divisible by the context-parallel size on its own;
+    checking the product here is what catches a head count that clears each
+    factor individually but not both together.
+
+    Without this check an over-sharded mesh fails much later, either with an
+    opaque XLA divisibility error or by silently leaving the projection
+    unsharded until `assert_params_sufficiently_sharded` trips.
+
+    Args:
+      kernel_axes: Logical axis names of the key/value projection kernel.
+
+    Raises:
+      ValueError: If the KV heads cannot be split evenly across the mesh.
+    """
+    if "kv_heads" not in kernel_axes:
+      # The projection is replicated, so the head count is unconstrained.
+      return
+
+    # Size-one mesh axes are already dropped, so this is the exact shard count.
+    kv_heads_index = kernel_axes.index("kv_heads")
+    kv_head_axes = self._logical_to_mesh_axes(kernel_axes)[kv_heads_index]
+    if kv_head_axes is None:
+      return
+    if isinstance(kv_head_axes, str):
+      kv_head_axes = (kv_head_axes,)
+    kv_head_axes = list(kv_head_axes)
+
+    # Ulysses exchanges sequence ownership for head ownership through an
+    # all-to-all, so the context axis shards KV heads too even though no
+    # logical rule says so.
+    if self.config.context_parallel_strategy.lower() == "ulysses":
+      ulysses_axis = self.config.context_sharding
+      if ulysses_axis not in kv_head_axes:
+        kv_head_axes.append(ulysses_axis)
+
+    kv_parallelism = 1
+    for axis in kv_head_axes:
+      kv_parallelism *= self.mesh.shape.get(axis, 1)
+
+    if self.num_kv_heads % kv_parallelism != 0:
+      raise ValueError(
+          f"num_kv_heads ({self.num_kv_heads}) for {self.attention_type}"
+          f" attention layers must be divisible by {kv_parallelism}, the"
+          f" combined size of the mesh axes {kv_head_axes} that shard KV heads."
+          " Attention heads are atomic under tensor parallelism and cannot be"
+          " split across more shards than there are heads. Either reduce the"
+          " parallelism on those axes, raise the KV head count"
+          " (`base_num_kv_heads`, or `global_num_kv_heads` for global attention"
+          " layers), or move the parallelism onto an axis that does not shard"
+          " KV heads (e.g. fsdp)."
+      )
+
   def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> None:
     """Initializes the query, key, value, and output projections."""
     if self.config.fused_qkv:
@@ -648,6 +717,7 @@ def init_kv_w(self, inputs_kv_shape: Tuple) -> nnx.Module:
         if self.config.ici_context_autoregressive_parallelism > 1
         else ("embed", "kv_heads", "kv_head_dim")
     )
+    self._validate_kv_head_sharding(kernel_axes)
 
     return DenseGeneral(
         in_features_shape=self.convert_dense_general_inputs_shape(inputs_kv_shape),

tests/unit/attention_test.py

@@ -3901,5 +3901,127 @@ def test_generate_attention_mask_compressed_all_modes(self):
     self.assertEqual(mask_none.shape[-1], kv_len)
 
 
+class KVHeadShardingTest(parameterized.TestCase):
+  """Tests that KV heads must divide the mesh axes that shard `kv_heads`.
+
+  Attention heads are atomic under tensor parallelism, so a mesh that shards
+  `kv_heads` more ways than there are heads has to be rejected. The mesh is
+  faked here to keep the test hermetic on a single-device host; only
+  `mesh.shape` is consulted when resolving logical axes onto mesh axes.
+  """
+
+  _KV_KERNEL_AXES = ("embed", "kv_heads", "kv_head_dim")
+  _NUM_KV_HEADS = 2
+  _EMBED_DIM = 16
+  _INDIVISIBLE = r"num_kv_heads \(2\).*must be divisible"
+  _INDIVISIBLE_BY_4 = r"num_kv_heads \(2\).*must be divisible by 4"
+
+  def setUp(self):
+    super().setUp()
+    self.cfg = pyconfig.initialize(
+        [sys.argv[0], get_test_config_path()],
+        per_device_batch_size=1.0,
+        run_name="test",
+        enable_checkpointing=False,
+        max_target_length=128,
+    )
+    self.inputs_kv_shape = (1, self.cfg.max_target_length, self._EMBED_DIM)
+    mesh = Mesh(maxtext_utils.create_device_mesh(self.cfg), self.cfg.mesh_axes)
+    # A single-device mesh shards nothing, so construction always succeeds; each
+    # test then swaps in a fake mesh to exercise the sharding check.
+    self.attention = Attention(
+        config=self.cfg,
+        num_query_heads=self._NUM_KV_HEADS * 2,
+        num_kv_heads=self._NUM_KV_HEADS,
+        head_dim=self.cfg.head_dim,
+        max_target_length=self.cfg.max_target_length,
+        max_prefill_predict_length=self.cfg.max_prefill_predict_length,
+        inputs_q_shape=self.inputs_kv_shape,
+        inputs_kv_shape=self.inputs_kv_shape,
+        mesh=mesh,
+        attention_kernel="dot_product",
+        dtype=self.cfg.dtype,
+        dropout_rate=self.cfg.dropout_rate,
+        attention_type=self.cfg.attention_type,
+        model_mode=MODEL_MODE_PREFILL,
+        rngs=nnx.Rngs(params=0, dropout=jax.random.PRNGKey(42)),
+    )
+
+  def _set_mesh_shape(self, **mesh_shape):
+    """Replaces the attention mesh with one reporting `mesh_shape`."""
+    self.attention.mesh = types.SimpleNamespace(shape=mesh_shape)
+
+  def _use_ulysses(self):
+    """Switches the layer onto Ulysses context parallelism.
+
+    Set directly rather than through `pyconfig` because a genuine Ulysses
+    config additionally demands TPU hardware, flash attention with Tokamax
+    Splash and several other options that are irrelevant here.
+    """
+    self.attention.config.context_parallel_strategy = "ulysses"
+
+  @parameterized.named_parameters(
+      # `kv_heads` maps to tensor x tensor_sequence x autoregressive, so each of
+      # those axes, and their product, constrains the KV head count.
+      ("tensor", {"tensor": 4}),
+      ("tensor_sequence", {"tensor_sequence": 4}),
+      ("autoregressive", {"autoregressive": 4}),
+      (
+          "product_of_axes",
+          {"tensor": 2, "tensor_sequence": 2, "autoregressive": 2},
+      ),
+  )
+  def test_indivisible_kv_heads_rejected(self, mesh_shape):
+    self._set_mesh_shape(**mesh_shape)
+    with self.assertRaisesRegex(ValueError, self._INDIVISIBLE):
+      self.attention.init_kv_w(inputs_kv_shape=self.inputs_kv_shape)
+
+  @parameterized.named_parameters(
+      ("exactly_divisible", {"tensor": 2}),
+      ("size_one_axes_ignored", {"tensor": 2, "tensor_sequence": 1}),
+      # fsdp shards `embed`, not `kv_heads`, so it places no constraint.
+      ("axis_that_does_not_shard_kv_heads", {"fsdp": 4}),
+      ("unsharded", {}),
+  )
+  def test_divisible_kv_heads_accepted(self, mesh_shape):
+    # The validator is called directly rather than through `init_kv_w`, which
+    # would go on to initialize parameters against the fake mesh.
+    self._set_mesh_shape(**mesh_shape)
+    self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES)  # pylint: disable=protected-access
+
+  def test_replicated_kernel_axes_skip_validation(self):
+    """A replicated KV projection is unconstrained even on an over-sharded mesh."""
+    self._set_mesh_shape(tensor=4)
+    self.attention._validate_kv_head_sharding((None, None, None))  # pylint: disable=protected-access
+
+  def test_context_axis_ignored_without_ulysses(self):
+    """Only Ulysses shards KV heads over the context axis."""
+    self._set_mesh_shape(context=4)
+    self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES)  # pylint: disable=protected-access
+
+  def test_ulysses_context_axis_rejected(self):
+    """Ulysses shards KV heads over context via all-to-all, not a logical rule."""
+    self._use_ulysses()
+    self._set_mesh_shape(context=4)
+    with self.assertRaisesRegex(ValueError, self._INDIVISIBLE_BY_4):
+      self.attention.init_kv_w(inputs_kv_shape=self.inputs_kv_shape)
+
+  def test_ulysses_multiplies_with_tensor_parallelism(self):
+    """The binding constraint is the product of the context and tensor axes.
+
+    Two KV heads clear `context`=2 and `tensor`=1 individually, and `types.py`
+    only checks the context factor, so the combined degree of 4 is caught here.
+    """
+    self._use_ulysses()
+    self._set_mesh_shape(context=2, tensor=2)
+    with self.assertRaisesRegex(ValueError, self._INDIVISIBLE_BY_4):
+      self.attention.init_kv_w(inputs_kv_shape=self.inputs_kv_shape)
+
+  def test_ulysses_divisible_accepted(self):
+    self._use_ulysses()
+    self._set_mesh_shape(context=2)
+    self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES)  # pylint: disable=protected-access
+
+
 if __name__ == "__main__":
   unittest.main()