From 1573b1a992312d4721204ace8c4352b5787fdf43 Mon Sep 17 00:00:00 2001 From: Jackson Zhao Date: Fri, 7 Aug 2026 20:10:12 -0700 Subject: [PATCH] Reject KV head counts that cannot be sharded across the tensor-parallel mesh. Attention heads are atomic under tensor parallelism, so the per-layer KV head count has to be divisible by the combined size of the mesh axes that `logical_axis_rules` maps the `kv_heads` logical axis onto (by default tensor x tensor_sequence x autoregressive). Previously an over-sharded mesh -- for example gemma4-26b, whose global attention layers use `global_num_kv_heads: 2`, with `ici_tensor_parallelism=4` -- failed much later with an opaque XLA divisibility error, or silently left the key/value projection unsharded. `Attention.init_kv_w` now raises a `ValueError` naming the head count, the shard count, the mesh axes responsible, and how to fix it. PiperOrigin-RevId: 961244739 --- src/maxtext/layers/attentions.py | 70 ++++++++++++++++++ tests/unit/attention_test.py | 122 +++++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 3825819a29..c0f396656a 100644 --- a/src/maxtext/layers/attentions.py +++ b/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), diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index dc7ea03f43..9a067b50e8 100644 --- a/tests/unit/attention_test.py +++ b/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()