From 4dfe6cc44420fd06ad1819415972fa3d344dfded Mon Sep 17 00:00:00 2001 From: Param Bole Date: Thu, 15 Oct 2026 13:01:36 -0700 Subject: [PATCH] Ulysses attention: repeat KV heads when there are fewer than context ranks Ulysses context parallelism hands each context rank a contiguous block of num_query_heads / cp query heads, so it required num_kv_heads to be divisible by the context parallel size: each rank had to receive whole KV heads. Qwen3.5 has 2 KV heads and trains at context 4 (8 at 256 chips), so it could not use Ulysses. When the context size is a multiple of the KV heads left on each device, repeat each KV head cp / num_kv_heads times in place ([kv0, kv0, kv1, kv1]) before the all-to-all. Rank r then receives a copy of KV head r // replication, which is the head its query block uses under Splash's grouped pairing (query head h uses KV head h // (H_q / H_kv)). Tiling the heads instead ([kv0, kv1, kv0, kv1]) would pair half of the query heads with the wrong KV head. The repeat is linear, so autodiff sums the dK/dV of the copies back onto each KV head. Each rank receives exactly the one KV head it needs, so the K/V all-to-all moves no more bytes than gathering that head. The config, layer KV-head sharding and attention-op validators accept the case. USP keeps the old requirement because its exchange does not repeat heads. Other strategies trace exactly as before. GatedDeltaNet layers keep their own sequence-parallel context parallelism, which does not read context_parallel_strategy. Tests: helper unit tests (replication factor, grouped order, query/KV pairing at Ulysses 4 and 8, validator flag); a forced 4-device CPU test of the collectives and of AttentionOp.tpu_flash_attention (with a dense stand-in for the Splash kernel), Ulysses vs all-gather vs dense on output and dQ/dK/dV, with tiled-heads and missing-group-sum positive controls that must fail; config and layer-validator acceptance; TPU-only Ulysses cases with 2 KV heads. --- docs/guides/optimization/sharding.md | 2 +- src/maxtext/configs/types.py | 7 +- .../kernels/attention/ulysses_attention.py | 46 ++- src/maxtext/layers/attention_op.py | 6 + src/maxtext/layers/attentions.py | 28 +- tests/unit/attention_test.py | 53 ++- tests/unit/configs_value_test.py | 27 ++ tests/unit/ulysses_attention_test.py | 67 ++++ .../unit/ulysses_kv_head_replication_test.py | 327 ++++++++++++++++++ 9 files changed, 536 insertions(+), 27 deletions(-) create mode 100644 tests/unit/ulysses_kv_head_replication_test.py diff --git a/docs/guides/optimization/sharding.md b/docs/guides/optimization/sharding.md index b0618d293a..d955c3bf52 100644 --- a/docs/guides/optimization/sharding.md +++ b/docs/guides/optimization/sharding.md @@ -260,7 +260,7 @@ Note in general there are many flavors of CP such as ring attention, which in th MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips. -MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention. +MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, a query head count divisible by the context parallel size and a KV head count that is divisible by it or divides it, including after tensor-parallel head sharding (with fewer KV heads than context ranks, each KV head is repeated to the ranks whose query heads use it, so 2 KV heads run at context 4 or 8), matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention. MaxText also supports `context_parallel_strategy=usp` ([USP](https://arxiv.org/abs/2405.07719), Ulysses over ring) on the same TPU Tokamax Splash path for training. To use it, set `context_parallel_strategy=usp`, `ici_context_parallelism` for the ring size, and `ici_context_usp_ulysses_parallelism` for the Ulysses size (total context parallelism is their product). USP factors the context parallelism into a ring dimension on the `context` mesh axis and a Ulysses dimension on the `context_usp_ulysses` mesh axis: the Ulysses all-to-all exchanges sequence ownership for head ownership over the Ulysses axis at each fixed ring position, and the ring kernel then rotates K and V across the ring axis inside each head subset. The strategy is hybrid-only: both dimensions must be greater than one, and the single-dimension endpoints are the existing `ring` and `ulysses` strategies. It shares the Ulysses restrictions (explicit positive ICI-only sizes, `attention=flash` with Tokamax Splash, global causal attention, head counts divisible by the Ulysses size, no MQA, no dropout, no multi-token prediction, no dKV megacore) and additionally requires `max_target_length` divisible by the total context parallelism and by the ring size squared; load balancing requires an even ring size. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 8f9afcd803..7f7171de76 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -6609,10 +6609,13 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) if self.num_kv_heads == 1: raise ValueError("TPU Ulysses attention does not support MQA with context_parallel_size > 1.") - if self.num_kv_heads % context_parallel_size != 0: + # Fewer KV heads than context ranks is fine when the ranks are a multiple of them: each KV head is then + # repeated across the ranks whose query heads use it (ulysses_attention.repeat_kv_heads_for_ulysses). + if self.num_kv_heads % context_parallel_size != 0 and context_parallel_size % self.num_kv_heads != 0: raise ValueError( "TPU Ulysses attention requires num_kv_heads " - f"({self.num_kv_heads}) to be divisible by context_parallel_size ({context_parallel_size})." + f"({self.num_kv_heads}) to be divisible by context_parallel_size ({context_parallel_size}) or to " + "divide it." ) self._validate_usp_context_parallelism() self._validate_context_parallel_attention_load_balance(context_parallel_size) diff --git a/src/maxtext/kernels/attention/ulysses_attention.py b/src/maxtext/kernels/attention/ulysses_attention.py index bcf7221b52..aff3a2b1a9 100644 --- a/src/maxtext/kernels/attention/ulysses_attention.py +++ b/src/maxtext/kernels/attention/ulysses_attention.py @@ -18,6 +18,7 @@ from typing import Any import jax +import jax.numpy as jnp from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.utils import sharding @@ -158,8 +159,13 @@ def validate_head_sharding( head_dim_kv: int, ulysses_size: int, attention_label: str, + allow_kv_head_replication: bool = False, ) -> None: - """Validates local head counts before the Ulysses head/sequence exchange.""" + """Validates local head counts before the Ulysses head/sequence exchange. + + With `allow_kv_head_replication`, fewer local KV heads than Ulysses ranks are accepted when the ranks are a + multiple of them: the caller then repeats each KV head with `repeat_kv_heads_for_ulysses` before the exchange. + """ q_head_axes = sharding.mesh_axes_for_dim(axis_names_q[head_dim_q]) kv_head_axes = sharding.mesh_axes_for_dim(axis_names_kv[head_dim_kv]) q_head_shards = sharding.mesh_axes_size(mesh, q_head_axes, label=attention_label) @@ -195,13 +201,47 @@ def validate_head_sharding( f"{attention_label} requires local query heads " f"({local_query_heads}) to be divisible by the Ulysses exchange size ({ulysses_size})." ) - if local_kv_heads % ulysses_size != 0: + if local_kv_heads % ulysses_size != 0 and not (allow_kv_head_replication and ulysses_size % local_kv_heads == 0): raise ValueError( f"{attention_label} requires local KV heads " - f"({local_kv_heads}) to be divisible by the Ulysses exchange size ({ulysses_size})." + f"({local_kv_heads}) to be divisible by the Ulysses exchange size ({ulysses_size})" + + (" or to divide it." if allow_kv_head_replication else ".") ) +def kv_head_replication(local_kv_heads: int, ulysses_size: int) -> int: + """How many copies of each local KV head the Ulysses exchange needs. + + The all-to-all hands rank r the r-th contiguous block of heads. When the local KV heads are a multiple of the + ranks, every rank already receives whole KV heads (1). When the ranks are a multiple of the KV heads, each KV + head is repeated `ulysses_size // local_kv_heads` times so that every rank receives one. + """ + if local_kv_heads % ulysses_size == 0: + return 1 + if ulysses_size % local_kv_heads == 0: + return ulysses_size // local_kv_heads + raise ValueError( + f"TPU Ulysses attention requires local KV heads ({local_kv_heads}) and the Ulysses exchange size " + f"({ulysses_size}) to divide one another." + ) + + +def repeat_kv_heads_for_ulysses(tensor: Any, replication: int): + """Repeats each KV head of `[B, H_kv, S/U, D]` in place, `[kv0, kv0, kv1, kv1]` for 2 heads and 2 copies. + + The copies of a head must be adjacent. Splash pairs query head h with KV head h // (H_q / H_kv), so a query + head's KV head is the one at the same relative position in the head dimension. The all-to-all keeps that + position: rank r receives query heads [r * H_q / U, (r + 1) * H_q / U) and KV head r of the repeated tensor, + which is a copy of KV head r // replication, the one those query heads use. Tiling the heads instead + (`[kv0, kv1, kv0, kv1]`) would give rank 1 the second KV head for query heads of the first group. + + `jnp.repeat` is linear, so its transpose sums the gradients of the copies back onto each KV head. + """ + if replication == 1: + return tensor + return jnp.repeat(tensor, replication, axis=1) + + def ulysses_all_to_all(tensor: Any, ulysses_axis: str): """Moves `[B, H, S/U, D]` to `[B, H/U, S, D]`.""" return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=1, concat_axis=2, tiled=True) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 3b9f319aa3..38f93b23fb 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -811,6 +811,7 @@ def __init__( head_dim_kv=1, ulysses_size=self.mesh.shape[context_axis], attention_label="TPU Ulysses attention", + allow_kv_head_replication=True, ) ulysses_attention.validate_dkv_sharding( axis_names_q=axis_names_q, @@ -2460,6 +2461,11 @@ def wrap_flash_attention( return attention_output, None if use_ulysses: + # With fewer local KV heads than context ranks (e.g. 2 KV heads at context 4), repeat each KV head so + # that every rank receives the one its block of query heads uses. + kv_replication = ulysses_attention.kv_head_replication(key.shape[1], self.mesh.shape[context_axis]) + key = ulysses_attention.repeat_kv_heads_for_ulysses(key, kv_replication) + value = ulysses_attention.repeat_kv_heads_for_ulysses(value, kv_replication) query = ulysses_attention.ulysses_all_to_all(query, context_axis) key = ulysses_attention.ulysses_all_to_all(key, context_axis) value = ulysses_attention.ulysses_all_to_all(value, context_axis) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index bc273ca1a9..4b78d71078 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -16,6 +16,7 @@ import dataclasses import functools +import math import os from typing import Any, Iterable, Optional, Tuple, Union, cast @@ -458,10 +459,14 @@ def _validate_kv_head_sharding( 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. + axes, so that axis is folded in explicitly. `types.py` separately checks + `num_kv_heads` against 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. The one exception is a context axis + that is a multiple of the KV heads left on each device by the other axes: + Ulysses then repeats each KV head across the context ranks whose query heads + use it (`ulysses_attention.repeat_kv_heads_for_ulysses`) instead of splitting + the KV heads, so the context axis places no constraint on them. Without this check an over-sharded mesh fails much later, either with an opaque XLA divisibility error or by silently leaving the projection @@ -489,17 +494,22 @@ def _validate_kv_head_sharding( kv_head_axes = (kv_head_axes,) kv_head_axes = list(kv_head_axes) + def shards(axes): + return math.prod(self.mesh.shape.get(axis, 1) for axis in 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. + # logical rule says so -- unless it is a multiple of the KV heads left per + # device, in which case those heads are repeated rather than split. 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) + local_kv_heads, remainder = divmod(self.num_kv_heads, shards(kv_head_axes)) + replicated = remainder == 0 and self.mesh.shape.get(ulysses_axis, 1) % local_kv_heads == 0 + if not replicated: + kv_head_axes.append(ulysses_axis) - kv_parallelism = 1 - for axis in kv_head_axes: - kv_parallelism *= self.mesh.shape.get(axis, 1) + kv_parallelism = shards(kv_head_axes) if self.num_kv_heads % kv_parallelism != 0: raise ValueError( diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index c43dcc1a3b..2e274001ce 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -3050,7 +3050,7 @@ def ring_loss(lnx): f"dq_reduction_steps={dq_reduction_steps}, ring_scan_unroll={ring_scan_unroll}, packing={packing}.", ) - def _ulysses_test_config(self, ici_context_parallelism, packing=False): + def _ulysses_test_config(self, ici_context_parallelism, packing=False, **heads): return pyconfig.initialize( [sys.argv[0], get_test_config_path()], **self.config_arguments, @@ -3062,6 +3062,7 @@ def _ulysses_test_config(self, ici_context_parallelism, packing=False): use_jax_splash=False, packing=packing, dtype="float32", + **heads, ) def _ulysses_test_modules(self, cfg_cp, mesh_cp, lnx): @@ -3104,12 +3105,20 @@ def _ulysses_test_modules(self, cfg_cp, mesh_cp, lnx): {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2, "packing": False}, {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4, "packing": False}, {"testcase_name": "ulysses_size_4_packed", "ici_context_parallelism": 4, "packing": True}, + # 2 KV heads at Ulysses size 4: each KV head is repeated to two context ranks. + { + "testcase_name": "ulysses_size_4_packed_kv_repeated", + "ici_context_parallelism": 4, + "packing": True, + "kv_heads": 2, + }, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallelism, packing): + def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallelism, packing, kv_heads=None): """Test equivalence between dot_product and flash attention + Ulysses context parallelism""" - cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing) + heads = {} if kv_heads is None else {"base_num_query_heads": 8, "base_num_kv_heads": kv_heads} + cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing, **heads) devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) if packing: @@ -3149,12 +3158,14 @@ def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallel {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2, "packing": False}, {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4, "packing": False}, {"testcase_name": "ulysses_size_4_packed", "ici_context_parallelism": 4, "packing": True}, + {"testcase_name": "ulysses_size_4_packed_kv_repeated", "ici_context_parallelism": 4, "packing": True, "kv_heads": 2}, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ulysses_context_parallel_grad(self, ici_context_parallelism, packing): + def test_tpu_flash_attention_ulysses_context_parallel_grad(self, ici_context_parallelism, packing, kv_heads=None): """Test input-gradient equivalence between dot_product and flash attention + Ulysses context parallelism""" - cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing) + heads = {} if kv_heads is None else {"base_num_query_heads": 8, "base_num_kv_heads": kv_heads} + cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing, **heads) devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) if packing: @@ -6930,7 +6941,6 @@ class KVHeadShardingTest(parameterized.TestCase): _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): """Builds an attention layer with two KV heads on a single-device mesh.""" @@ -7022,32 +7032,51 @@ def test_context_axis_ignored_without_ulysses(self): 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._set_mesh_shape(context=4) + """Ulysses shards KV heads over context via all-to-all, not a logical rule. + + Two KV heads and a context axis of 3 neither split evenly nor repeat + evenly, so the context axis counts against the KV heads. + """ + self._set_mesh_shape(context=3) with ( self._use_ulysses(), - self.assertRaisesRegex(ValueError, self._INDIVISIBLE_BY_4), + self.assertRaisesRegex(ValueError, r"num_kv_heads \(2\).*must be divisible by 3"), ): 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` + Six KV heads clear `context`=2 and `tensor`=2 individually, and `types.py` only checks the context factor, so the combined degree of 4 is caught here. + The 3 KV heads left per device by `tensor` cannot be repeated evenly over + `context` either. """ self._set_mesh_shape(context=2, tensor=2) + self.attention.num_kv_heads = 6 with ( self._use_ulysses(), - self.assertRaisesRegex(ValueError, self._INDIVISIBLE_BY_4), + self.assertRaisesRegex(ValueError, r"num_kv_heads \(6\).*must be divisible by 4"), ): - self.attention.init_kv_w(inputs_kv_shape=self.inputs_kv_shape) + self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES) # pylint: disable=protected-access def test_ulysses_divisible_accepted(self): self._set_mesh_shape(context=2) with self._use_ulysses(): self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES) # pylint: disable=protected-access + @parameterized.named_parameters( + ("two_heads_context_4", {"context": 4}), + ("two_heads_context_8", {"context": 8}), + # `tensor` leaves one KV head per device, which Ulysses repeats to both context ranks. + ("one_local_head_context_2", {"context": 2, "tensor": 2}), + ) + def test_ulysses_repeated_kv_heads_accepted(self, mesh_shape): + """A context axis that is a multiple of the local KV heads repeats them instead of splitting them.""" + self._set_mesh_shape(**mesh_shape) + with self._use_ulysses(): + self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES) # pylint: disable=protected-access + class MLAAbsorbedMQATest(attention_test_util.MLATestBase): """Comprehensive unit test suite for Absorbed-Query Latent MQA (PR 1 / Milestone 2).""" diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 685a66fe34..3188d57c3e 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -578,6 +578,32 @@ def test_tpu_ulysses_config_validation_accepts_initial_config(self): self.assertEqual(config.ici_context_parallelism, 4) self.assertFalse(config.context_parallel_load_balance) + def test_tpu_ulysses_config_validation_accepts_fewer_kv_heads_than_context_ranks(self): + """2 KV heads at context 4 or 8 (Qwen3.5) are repeated across the context ranks.""" + for ici_context_parallelism in (4, 8): + with self.subTest(ici_context_parallelism=ici_context_parallelism): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ulysses", + "context_parallel_load_balance=False", + f"ici_context_parallelism={ici_context_parallelism}", + "base_num_query_heads=32", + "base_num_kv_heads=2", + "hardware=tpu", + "packing=True", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + self.assertEqual(config.num_kv_heads, 2) + self.assertEqual(config.context_parallel_strategy, "ulysses") + def test_tpu_ulysses_config_validation_accepts_packing(self): argv = [ "", @@ -642,6 +668,7 @@ def test_tpu_ulysses_config_validation_rejects_unsupported_configs(self): (["base_num_kv_heads=1"], [], "MQA"), (["base_num_query_heads=18"], [], "requires num_query_heads"), (["base_num_kv_heads=10"], [], "requires num_kv_heads"), + (["base_num_query_heads=24", "base_num_kv_heads=6"], [], r"requires num_kv_heads \(6\).*or to divide it"), (["attention_type=mla"], [], "global causal attention"), (["attention_type=local_sliding", "sliding_window_size=128"], [], "global causal attention"), (["attention_type=chunk", "chunk_attn_window_size=128"], [], "global causal attention"), diff --git a/tests/unit/ulysses_attention_test.py b/tests/unit/ulysses_attention_test.py index 963cff31d1..686989a313 100644 --- a/tests/unit/ulysses_attention_test.py +++ b/tests/unit/ulysses_attention_test.py @@ -21,6 +21,7 @@ from absl.testing import absltest import jax import jax.numpy as jnp +import numpy as np from maxtext.common.common_types import MODEL_MODE_PREFILL from maxtext.common.common_types import MODEL_MODE_TRAIN @@ -251,6 +252,72 @@ def test_validate_head_sharding_requires_q_and_kv_head_axes_to_match(self): attention_label="TPU Ulysses attention", ) + def test_validate_head_sharding_accepts_kv_head_replication_only_when_allowed(self): + """Two KV heads at Ulysses size 4 need replication, which only the Ulysses path (not USP) performs.""" + mesh = types.SimpleNamespace(shape={"context": 4}) + kwargs = { + "axis_names_q": (None, None, "context", None), + "axis_names_kv": (None, None, "context", None), + "mesh": mesh, + "num_query_heads": 32, + "num_kv_heads": 2, + "head_dim_q": 1, + "head_dim_kv": 1, + "ulysses_size": 4, + "attention_label": "TPU Ulysses attention", + } + with self.assertRaisesRegex( + ValueError, r"local KV heads \(2\) to be divisible by the Ulysses exchange size \(4\)\.$" + ): + ulysses_attention.validate_head_sharding(**kwargs) + ulysses_attention.validate_head_sharding(**kwargs, allow_kv_head_replication=True) + # One local KV head left by tensor sharding is repeated to every rank too. + tensor_sharded = (None, "tensor", "context", None) + ulysses_attention.validate_head_sharding( + **{ + **kwargs, + "axis_names_q": tensor_sharded, + "axis_names_kv": tensor_sharded, + "mesh": types.SimpleNamespace(shape={"context": 4, "tensor": 2}), + }, + allow_kv_head_replication=True, + ) + with self.assertRaisesRegex(ValueError, "or to divide it"): + ulysses_attention.validate_head_sharding( + **{**kwargs, "num_query_heads": 36, "num_kv_heads": 6}, allow_kv_head_replication=True + ) + + def test_kv_head_replication(self): + cases = {(2, 4): 2, (2, 8): 4, (1, 4): 4, (4, 4): 1, (8, 4): 1, (16, 2): 1} + for (local_kv_heads, ulysses_size), expected in cases.items(): + with self.subTest(local_kv_heads=local_kv_heads, ulysses_size=ulysses_size): + self.assertEqual(ulysses_attention.kv_head_replication(local_kv_heads, ulysses_size), expected) + for local_kv_heads, ulysses_size in ((3, 4), (6, 4), (4, 6)): + with self.subTest(local_kv_heads=local_kv_heads, ulysses_size=ulysses_size): + with self.assertRaisesRegex(ValueError, "divide one another"): + ulysses_attention.kv_head_replication(local_kv_heads, ulysses_size) + + def test_repeat_kv_heads_pairs_every_query_head_with_its_kv_head(self): + """After the exchange, every rank's query heads find their own KV head (32 query / 2 KV heads, Ulysses 4 and 8). + + Splash pairs local query head i with local KV head i // (local query heads / local KV heads). The exchange + hands rank r the r-th contiguous block of heads, simulated here by splitting the head dimension. + """ + num_query_heads, num_kv_heads = 32, 2 + key = jnp.arange(num_kv_heads, dtype=jnp.int32).reshape(1, num_kv_heads, 1, 1) + for ulysses_size in (4, 8): + with self.subTest(ulysses_size=ulysses_size): + replication = ulysses_attention.kv_head_replication(num_kv_heads, ulysses_size) + repeated = ulysses_attention.repeat_kv_heads_for_ulysses(key, replication) + np.testing.assert_array_equal(repeated[0, :, 0, 0], np.repeat(np.arange(num_kv_heads), replication)) + query_heads = np.arange(num_query_heads) + for rank, kv_block in enumerate(np.split(np.asarray(repeated[0, :, 0, 0]), ulysses_size)): + rank_query_heads = np.split(query_heads, ulysses_size)[rank] + local_group = len(rank_query_heads) // len(kv_block) + paired = kv_block[np.arange(len(rank_query_heads)) // local_group] + np.testing.assert_array_equal(paired, rank_query_heads // (num_query_heads // num_kv_heads)) + self.assertIs(ulysses_attention.repeat_kv_heads_for_ulysses(key, 1), key) + def test_ulysses_all_to_all_moves_heads_to_sequence(self): tensor = jnp.ones((1, 16, 8, 2)) diff --git a/tests/unit/ulysses_kv_head_replication_test.py b/tests/unit/ulysses_kv_head_replication_test.py new file mode 100644 index 0000000000..36d0e7df99 --- /dev/null +++ b/tests/unit/ulysses_kv_head_replication_test.py @@ -0,0 +1,327 @@ +# 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. +"""Checks Ulysses attention with fewer KV heads than context ranks on a forced 4-device CPU mesh. + +With 2 KV heads at context 4, each KV head is repeated twice before the all-to-all so that every rank receives +the KV head its contiguous block of query heads uses. The checks run in a subprocess so the forced device count +takes effect before JAX initializes: + + - Collectives. The real repeat and all-to-all helpers around a dense GQA kernel match a dense single-device + reference in the output and in dQ/dK/dV, on 8 query heads and on 32 (the Qwen3.5 head count). + - Attention op. `AttentionOp.tpu_flash_attention`, on a packed row with a long segment, short segments and a + padded tail: the Ulysses path matches the all-gather path and a dense reference, output and dQ/dK/dV. + +Each check has positive controls that must fail the same gate: tiling the KV heads (`[kv0, kv1, kv0, kv1]`), which +pairs half of the query heads with the wrong KV head, and a repeat whose backward keeps only the first copy's +gradient instead of summing the copies. + +The attention op runs with a dense stand-in for the Tokamax Splash kernel (same call, mask and grouped query/KV +head pairing), not with the kernel in Pallas TPU interpret mode: every interpreted kernel call waits on a barrier +across all devices, and XLA:CPU may run a collective that does not depend on the kernel before it on one device and +after it on another, so the devices can wait on each other forever. The kernel itself is covered on TPU by the +Ulysses cases in attention_test.py. +""" + +import functools +import os +import subprocess +import sys +from unittest import mock + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +import numpy as np +import pytest + +from maxtext.kernels.attention import ulysses_attention + +_CHILD_MARKER = "ULYSSES_KV_HEAD_REPLICATION_CHECKS_PASSED" +# Normalized max error max|a - b| / max|b| in float32; the paths differ only in reduction order. +_TOLERANCE = 1e-5 +_NUM_DEVICES = 4 + + +@pytest.mark.cpu_only +def test_ulysses_kv_head_replication_on_cpu_mesh(): + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" + env["JAX_PLATFORMS"] = "cpu" + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert _CHILD_MARKER in result.stdout + + +def _normalized_max_error(actual, expected) -> float: + actual = np.asarray(actual, np.float64) + expected = np.asarray(expected, np.float64) + return float(np.max(np.abs(actual - expected)) / max(np.max(np.abs(expected)), 1e-30)) + + +def _tile_kv_heads(tensor, replication): + """Positive control: tiled copies, which give rank 1 of 4 the second KV head for first-group query heads.""" + return tensor if replication == 1 else jnp.tile(tensor, (1, replication, 1, 1)) + + +def _repeat_without_group_sum(tensor, replication): + """Positive control: correct copies, but a backward that drops every copy's gradient except the first.""" + if replication == 1: + return tensor + + @jax.custom_vjp + def repeat(x): + return jnp.repeat(x, replication, axis=1) + + def repeat_fwd(x): + return repeat(x), None + + def repeat_bwd(_, cotangent): + return (cotangent[:, ::replication],) + + repeat.defvjp(repeat_fwd, repeat_bwd) + return repeat(tensor) + + +@jax.tree_util.register_pytree_node_class +class _DenseSplashKernel: + """Dense stand-in for a Tokamax Splash kernel: the same call, mask and grouped query/KV head pairing. + + Splash pairs query head h with KV head h // (H_q / H_kv). The mask is materialized, so it is sharded over the + query sequence exactly as the real kernel's mask metadata is, and each shard holds the rows of its queries. + """ + + def __init__(self, mask): + self.mask = mask + + def tree_flatten(self): + return (self.mask,), None + + @classmethod + def tree_unflatten(cls, _, children): + return cls(*children) + + def __call__(self, query, key, value, segment_ids, sinks=None, max_logit_value=None): + del sinks, max_logit_value + group = query.shape[0] // key.shape[0] + key, value = jnp.repeat(key, group, axis=0), jnp.repeat(value, group, axis=0) + logits = jnp.einsum("hqd,hkd->hqk", query, key, precision="highest") + allowed = self.mask & (segment_ids.q[:, None] == segment_ids.kv[None, :]) + probs = jax.nn.softmax(jnp.where(allowed[None], logits, -1e30), axis=-1) + return jnp.einsum("hqk,hkd->hqd", probs, value, precision="highest").astype(query.dtype) + + +def _make_dense_splash_kernel(mask, config, q_seq_shards): + del config, q_seq_shards + return _DenseSplashKernel(jnp.asarray(mask[:, :])) + + +def _segment_ids(seq_len: int, row: int = 0) -> np.ndarray: + """A packed row. Row 0: a long segment over several context shards, five short segments, then padding (id 0). + + Row 1: sixteen short segments whose boundaries are not aligned to the context shards. + """ + positions = np.arange(seq_len) + if row == 1: + boundaries = np.linspace(0, seq_len, 17)[1:-1].astype(int) + 3 + return (np.searchsorted(boundaries, positions, side="right") + 1).astype(np.int32) + long_end = (seq_len * 9) // 16 + pad_start = (seq_len * 15) // 16 + short = np.searchsorted(np.linspace(long_end, pad_start, 6)[1:-1], positions, side="right") + 2 + return np.where(positions < long_end, 1, np.where(positions < pad_start, short, 0)).astype(np.int32) + + +def _dense_attention(query, key, value, segment_ids): + """Causal, same-segment GQA attention over `[B, H, S, D]` (padding attends to padding, as in Splash).""" + seq_len = query.shape[2] + group = query.shape[1] // key.shape[1] + key, value = jnp.repeat(key, group, axis=1), jnp.repeat(value, group, axis=1) + logits = jnp.einsum("bhqd,bhkd->bhqk", query, key, precision="highest") + allowed = jnp.tril(jnp.ones((seq_len, seq_len), bool))[None, None] & ( + segment_ids[:, None, :, None] == segment_ids[:, None, None, :] + ) + probs = jax.nn.softmax(jnp.where(allowed, logits, -1e30), axis=-1) + return jnp.einsum("bhqk,bhkd->bhqd", probs, value, precision="highest") + + +def _value_and_grads(fn, args, cotangent, mesh, with_grads=True): + def loss(*xs): + return jnp.sum(fn(*xs).astype(jnp.float32) * cotangent) + + with jax.set_mesh(mesh): + out = jax.jit(fn)(*args) + grads = jax.jit(jax.grad(loss, argnums=(0, 1, 2)))(*args) if with_grads else () + return [np.asarray(out, np.float32)] + [np.asarray(g, np.float32) for g in grads] + + +def _check(name, result, reference, controls): + """`result` must match `reference`; every `controls` entry (label, values, field index) must miss it.""" + errors = [_normalized_max_error(r, e) for r, e in zip(result, reference)] + print(f"{name}: normalized max error out/dq/dk/dv = {errors}", flush=True) + assert max(errors) <= _TOLERANCE, f"{name}: {errors}" + for label, values, index in controls: + error = _normalized_max_error(values[index], reference[index]) + print(f"{name} positive control {label}: {error}", flush=True) + assert error > 1000 * _TOLERANCE, f"{name}: positive control {label} passed the gate ({error})" + + +def _run_collective_checks(num_query_heads: int): + """The repeat and all-to-all helpers around a dense kernel vs a dense single-device reference.""" + mesh = Mesh(np.array(jax.devices()), ("context",)) + ulysses_axis = "context" + batch, num_kv_heads, seq_len, head_dim = 1, 2, 64, 8 + rng = np.random.default_rng(0) + query = jnp.asarray(rng.standard_normal((batch, num_query_heads, seq_len, head_dim)), jnp.float32) + key = jnp.asarray(rng.standard_normal((batch, num_kv_heads, seq_len, head_dim)), jnp.float32) + value = jnp.asarray(rng.standard_normal((batch, num_kv_heads, seq_len, head_dim)), jnp.float32) + cotangent = jnp.asarray(rng.standard_normal((batch, num_query_heads, seq_len, head_dim)), jnp.float32) + segment_ids = jnp.asarray(_segment_ids(seq_len)[None]) + data_spec, segment_spec = P(None, None, "context", None), P(None, "context") + + def ulysses(repeat_fn): + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(data_spec, data_spec, data_spec, segment_spec), + out_specs=data_spec, + check_vma=False, + ) + def attention(query, key, value, segment_ids): + replication = ulysses_attention.kv_head_replication(key.shape[1], mesh.shape[ulysses_axis]) + key, value = repeat_fn(key, replication), repeat_fn(value, replication) + query, key, value = (ulysses_attention.ulysses_all_to_all(x, ulysses_axis) for x in (query, key, value)) + segment_ids = jax.lax.all_gather(segment_ids, ulysses_axis, axis=1, tiled=True) + output = _dense_attention(query, key, value, segment_ids) + return ulysses_attention.inverse_ulysses_all_to_all(output, ulysses_axis) + + return lambda q, k, v: attention(q, k, v, segment_ids) + + args = (query, key, value) + reference = _value_and_grads(lambda q, k, v: _dense_attention(q, k, v, segment_ids), args, cotangent, mesh) + result = _value_and_grads(ulysses(ulysses_attention.repeat_kv_heads_for_ulysses), args, cotangent, mesh) + tiled = _value_and_grads(ulysses(_tile_kv_heads), args, cotangent, mesh, with_grads=False) + no_sum = _value_and_grads(ulysses(_repeat_without_group_sum), args, cotangent, mesh) + _check( + f"collectives {num_query_heads}q/{num_kv_heads}kv", + result, + reference, + [("tiled KV heads (out)", tiled, 0), ("no group sum (dk)", no_sum, 2), ("no group sum (dv)", no_sum, 3)], + ) + + +def _run_attention_op_checks(): + """`AttentionOp.tpu_flash_attention`, Ulysses vs all-gather vs dense, with the dense stand-in for the kernel.""" + # pylint: disable=import-outside-toplevel + from flax.linen import partitioning as nn_partitioning + + from maxtext.common.common_types import MODEL_MODE_TRAIN + from maxtext.configs import pyconfig + from maxtext.layers import attention_op + from maxtext.utils import globals as maxtext_globals + from maxtext.utils import maxtext_utils + + seq_len, num_query_heads, num_kv_heads, head_dim, block = 512, 8, 2, 128, 128 + cfg = pyconfig.initialize( + [ + None, + os.path.join(maxtext_globals.MAXTEXT_CONFIGS_DIR, "base.yml"), + "run_name=ulysses_kv_head_replication_test", + "hardware=cpu", + "custom_mesh_and_rule=cp-as-ep", + "ici_fsdp_parallelism=1", + f"ici_context_parallelism={_NUM_DEVICES}", + "ici_expert_parallelism=1", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + # The layer is built for all-gather; the Ulysses case switches the strategy at call time (below). + "context_parallel_strategy=all_gather", + "context_parallel_load_balance=False", + f"max_target_length={seq_len}", + f"base_num_query_heads={num_query_heads}", + f"base_num_kv_heads={num_kv_heads}", + f"head_dim={head_dim}", + *(f"sa_{name}={block}" for name in ("block_q", "block_kv", "block_kv_compute", "block_q_dkv")), + *(f"sa_{name}={block}" for name in ("block_kv_dkv", "block_kv_dkv_compute")), + "dq_reduction_steps=3", + f"per_device_batch_size={2 / _NUM_DEVICES}", + "packing=True", + "dataset_type=synthetic", + "enable_checkpointing=false", + "skip_jax_distributed_system=True", + "enable_tensorboard=False", + ] + ) + mesh = maxtext_utils.get_mesh_from_config(cfg) + with nn_partitioning.axis_rules(cfg.logical_axis_rules), jax.set_mesh(mesh): + op = attention_op.AttentionOp( + config=cfg, + mesh=mesh, + attention_kernel="flash", + max_target_length=seq_len, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + dtype=jnp.float32, + ) + + rng = np.random.default_rng(1) + batch = cfg.micro_batch_size_to_train_on + sharding = NamedSharding(mesh, P(None, "context")) + query = rng.standard_normal((batch, seq_len, num_query_heads, head_dim)).astype(np.float32) / np.sqrt(head_dim) + key = rng.standard_normal((batch, seq_len, num_kv_heads, head_dim)).astype(np.float32) + value = rng.standard_normal((batch, seq_len, num_kv_heads, head_dim)).astype(np.float32) + cotangent = rng.standard_normal((batch, seq_len, num_query_heads, head_dim)).astype(np.float32) + segment_ids = np.stack([_segment_ids(seq_len, row % 2) for row in range(batch)]) + query, key, value, cotangent, segment_ids = ( + jax.device_put(x, sharding) for x in (query, key, value, cotangent, segment_ids) + ) + + def flash(strategy, repeat_fn=ulysses_attention.repeat_kv_heads_for_ulysses): + def attention(q, k, v): + with ( + mock.patch.dict(cfg.get_keys(), {"context_parallel_strategy": strategy}), + mock.patch.object(attention_op.tokamax_splash_kernel, "make_splash_mha", _make_dense_splash_kernel), + mock.patch.object(ulysses_attention, "repeat_kv_heads_for_ulysses", repeat_fn), + nn_partitioning.axis_rules(cfg.logical_axis_rules), + ): + return op.tpu_flash_attention(q, k, v, segment_ids, model_mode=MODEL_MODE_TRAIN)[0] + + return attention + + def dense(q, k, v): + q, k, v = (jnp.transpose(x, (0, 2, 1, 3)) for x in (q, k, v)) + return jnp.transpose(_dense_attention(q, k, v, segment_ids), (0, 2, 1, 3)) + + args = (query, key, value) + reference = _value_and_grads(dense, args, cotangent, mesh) + all_gather = _value_and_grads(flash("all_gather"), args, cotangent, mesh) + result = _value_and_grads(flash("ulysses"), args, cotangent, mesh) + tiled = _value_and_grads(flash("ulysses", _tile_kv_heads), args, cotangent, mesh, with_grads=False) + no_sum = _value_and_grads(flash("ulysses", _repeat_without_group_sum), args, cotangent, mesh) + controls = [("tiled KV heads (out)", tiled, 0), ("no group sum (dk)", no_sum, 2), ("no group sum (dv)", no_sum, 3)] + _check("attention op, ulysses vs all_gather", result, all_gather, controls) + _check("attention op, ulysses vs dense", result, reference, controls) + + with nn_partitioning.axis_rules(cfg.logical_axis_rules), jax.set_mesh(mesh): + hlo = jax.jit(flash("ulysses")).lower(*args).compile().as_text() + # q, k, v in and the output back; no K/V all-gather (the one all-gather is the int32 segment IDs). + assert hlo.count(" all-to-all(") == 4, hlo.count(" all-to-all(") + assert hlo.count(" all-gather(") + hlo.count(" all-gather-start(") == 1 + + +if __name__ == "__main__": + assert len(jax.devices()) == _NUM_DEVICES, jax.devices() + _run_collective_checks(num_query_heads=8) + _run_collective_checks(num_query_heads=32) + _run_attention_op_checks() + print(_CHILD_MARKER)