Ulysses attention: repeat KV heads when there are fewer than context ranks

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.
 

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)

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)

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)

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(

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)."""