Merge pull request #4537 from huytransformer:htn/tokamax-ring-load-balance

PiperOrigin-RevId: 952535031

Author: Google-ML-Automation

src/maxtext/configs/types.py

@@ -3421,7 +3421,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.packing:
         raise ValueError("TPU Tokamax ring attention does not support packing yet.")
       if self.context_parallel_load_balance:
-        raise ValueError("TPU Tokamax ring attention does not support context_parallel_load_balance yet.")
+        if context_parallel_size % 2 != 0:
+          raise ValueError("TPU Tokamax ring load balancing requires an even context_parallel_size.")
+        if self.mtp_num_layers > 0:
+          raise ValueError("TPU Tokamax ring attention with context_parallel_load_balance=True does not support MTP.")
       if self.use_ragged_attention:
         raise ValueError("TPU Tokamax ring attention does not support ragged attention.")
       if self.attention_sink:

src/maxtext/kernels/attention/tokamax_ring_attention.py

@@ -25,11 +25,13 @@
 
 import jax
 from jax.experimental import pallas as pl
+import numpy as np
 
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.kernels.tokamax_splash_attention import ring_attention_kernel
 from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as tokamax_splash_kernel
 from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as tokamax_splash_mask
+from maxtext.utils import max_utils
 
 
 def is_context_parallel_ring_requested(config: Any) -> bool:
@@ -245,6 +247,11 @@ def build_splash_config(
   block_q_dkv = min(config.sa_block_q_dkv, q_seq_len_per_shard)
   block_kv_dkv = min(config.sa_block_kv_dkv, kv_seq_len_per_shard)
   block_kv_dkv_compute = min(config.sa_block_kv_dkv_compute, kv_seq_len_per_shard)
+  if config.context_parallel_load_balance and block_q_dkv % tokamax_splash_kernel.NUM_LANES != 0:
+    raise ValueError(
+        "TPU Tokamax ring attention with context_parallel_load_balance=True requires "
+        f"sa_block_q_dkv ({block_q_dkv}) to be a multiple of {tokamax_splash_kernel.NUM_LANES} after clamping."
+    )
   return tokamax_splash_kernel.SplashConfig(
       block_q=block_q,
       block_kv=block_kv,
@@ -270,11 +277,18 @@ def build_splash_config(
   )
 
 
-def _make_causal_mask(shape: tuple[int, int], context_parallel_size: int):
+def _make_causal_mask(shape: tuple[int, int], context_parallel_size: int, *, load_balanced: bool = False):
   """Builds a lazy causal mask for ring attention."""
   if context_parallel_size <= 1:
     raise ValueError("context_parallel_size must be > 1 for ring attention.")
-  return tokamax_splash_mask.CausalMask(shape=shape, shard_count=context_parallel_size)
+  mask = tokamax_splash_mask.CausalMask(shape=shape, shard_count=context_parallel_size)
+  if load_balanced:
+    sequence_indices = max_utils.reorder_mask_load_balancing(
+        np.arange(shape[0], dtype=np.int32), context_parallel_size, 0
+    )
+    mask.q_sequence = sequence_indices
+    mask.kv_sequence = sequence_indices
+  return mask
 
 
 def make_sharded_ring_attention_kernel(
@@ -301,6 +315,7 @@ def make_sharded_ring_attention_kernel(
   mask = _make_causal_mask(
       (query.shape[2], key.shape[2]),
       context_parallel_size,
+      load_balanced=config.context_parallel_load_balance,
   )
 
   @functools.partial(jax.jit, static_argnames=["single_head_mask"])

src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py

@@ -604,6 +604,7 @@ def mask_info_spec(mask_info):
           if mask_info.partial_mask_blocks is not None
           else None,
           q_sequence=_resolve_spec(mask_info.q_sequence),
+          kv_sequence=jax.sharding.PartitionSpec() if mask_info.kv_sequence is not None else None,
       )
 
     return RingSplashAttentionKernel(

src/maxtext/kernels/tokamax_splash_attention/ring_attention_utils.py

@@ -41,13 +41,16 @@ def slice_if_exists(arr: jax.Array | None):
       block_mask=slice_if_exists(mask_info.block_mask),
       partial_mask_blocks=mask_info.partial_mask_blocks,  # partial mask blocks are global
       q_sequence=mask_info.q_sequence,  # Q sequence stays stationary
+      kv_sequence=slice_if_exists(mask_info.kv_sequence),
   )
 
 
 def offset_q_sequence_for_kv_shard(mask_info: MaskInfo, kv_shard_idx: jax.Array, kv_seq_len: int) -> MaskInfo:
   """Converts lazy mask Q ids to the current local KV coordinate frame."""
   if mask_info.q_sequence is None:
     return mask_info
+  if mask_info.kv_sequence is not None:
+    return mask_info
 
   kv_shard_offset = jnp.asarray(kv_shard_idx, dtype=mask_info.q_sequence.dtype) * kv_seq_len
   return mask_info._replace(q_sequence=mask_info.q_sequence - kv_shard_offset)

src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py

@@ -194,6 +194,7 @@ def _apply_mask_and_soft_cap(
     mask_value: float,
     mask_ref,
     q_sequence_ref,
+    kv_sequence_ref,
     q_segment_ids_ref,
     kv_segment_ids_ref,
     *,
@@ -207,30 +208,38 @@ def _apply_mask_and_soft_cap(
 ) -> tuple[jax.Array, jax.Array | None]:
   assert mask_ref is None or q_sequence_ref is None
   assert (q_sequence_ref is None) == (mask_function is None)
+  assert kv_sequence_ref is None or mask_function is not None
 
   masks = []
   if has_partial_mask:
     if mask_ref is not None:
       mask = mask_ref[:, k_slice] if k_in_lanes else mask_ref[k_slice, :]
       masks.append(mask)
     elif mask_function is not None:
-      # Compute the mask using the given q_sequence indices.
-      # KV indices are computed on the fly. This works because we only support Q
-      # sequence sharding. If we wanted to compute Q indices too, then we would
-      # need to keep into account the current shard along Q sequence.
+      # Compute the mask using original Q positions. K/V positions are computed
+      # on the fly unless explicit original K/V positions are provided.
 
       if k_in_lanes:
         assert q_sequence_ref.shape == (bq, NUM_LANES)
 
-        k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (bq, k_slice.size), 1)
+        if kv_sequence_ref is None:
+          k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (bq, k_slice.size), 1)
+        else:
+          k_sequence = jnp.broadcast_to(kv_sequence_ref[:1, k_slice], (bq, k_slice.size))
 
         repeats, rem = divmod(k_slice.size, NUM_LANES)
         assert rem == 0
         q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats))  # [bq, k_slice.size]
       else:
         assert q_sequence_ref.shape == (NUM_SUBLANES, bq)
 
-        k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (k_slice.size, bq), 0)
+        if kv_sequence_ref is None:
+          k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (k_slice.size, bq), 0)
+        else:
+          repeats, rem = divmod(bq, NUM_LANES)
+          if rem:
+            raise NotImplementedError(f"block_q must be a multiple of {NUM_LANES}")
+          k_sequence = jnp.tile(kv_sequence_ref[k_slice, :], (1, repeats))
         q_sequence = q_sequence_ref[:1, :]  # [1, bq]
         q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))
 
@@ -290,6 +299,7 @@ def flash_attention_kernel(
     sinks_ref,
     mask_ref,
     q_sequence_ref,
+    kv_sequence_ref,
     max_logit_value_ref,
     # Outputs
     o_ref,
@@ -394,6 +404,7 @@ def body(kv_compute_index, _, has_partial_mask=False):
         mask_value,
         mask_ref,
         q_sequence_ref,
+        kv_sequence_ref,
         q_segment_ids_ref,
         kv_segment_ids_ref,
         attn_logits_soft_cap=attn_logits_soft_cap,
@@ -669,6 +680,13 @@ def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_):
     q_sequence = None
     in_specs.append(None)
 
+  if mask_info.kv_sequence is not None:
+    kv_sequence = jax.lax.broadcast_in_dim(mask_info.kv_sequence, (NUM_SUBLANES, kv_seq_len), (1,))
+    in_specs.append(pl.BlockSpec((NUM_SUBLANES, bkv), kv_segment_ids_index_map))
+  else:
+    kv_sequence = None
+    in_specs.append(None)
+
   if max_logit_value is not None:
     # reshape to allow sublane selection for vmap-ping and shard_map-ping
     max_logit_value = jnp.broadcast_to(
@@ -819,6 +837,7 @@ def _fwd_cost_estimate(
         sinks,
         mask_info.partial_mask_blocks,
         q_sequence,
+        kv_sequence,
         max_logit_value,
     )
   out, logsumexp, l_linear, max_logits = all_out
@@ -996,6 +1015,7 @@ def _flash_attention_dq_kernel(
     di_ref,
     mask_ref,
     q_sequence_ref,
+    kv_sequence_ref,
     # Outputs
     dq_scratch_ref,
     dq_ref,
@@ -1049,6 +1069,7 @@ def body(has_partial_mask: bool = False):
         mask_value,
         mask_ref,
         q_sequence_ref,
+        kv_sequence_ref,
         q_segment_ids_ref,
         kv_segment_ids_ref,
         attn_logits_soft_cap=attn_logits_soft_cap,
@@ -1115,6 +1136,7 @@ def _flash_attention_dkv_kernel(
     di_ref,
     mask_ref,
     q_sequence_ref,
+    kv_sequence_ref,
     # aliases
     dq_alias,
     dk_alias,
@@ -1212,6 +1234,7 @@ def _load_kv(ref, layout):
         mask_value,
         mask_ref,
         q_sequence_ref,
+        kv_sequence_ref,
         q_segment_ids_ref,
         kv_segment_ids_ref,
         attn_logits_soft_cap=attn_logits_soft_cap,
@@ -1436,9 +1459,8 @@ def create_dkv_index_map(h, i, j, *_):
   mask_spec = pl.BlockSpec((None, bkv, bq), mask_index_map)
 
   q_segment_ids_index_map = unravel(lambda h, i, j: (0, i))
+  kv_segment_ids_index_map = unravel(lambda h, i, j: (j, 0))
   if segment_ids is not None:
-    kv_segment_ids_index_map = unravel(lambda h, i, j: (j, 0))
-
     q_segment_spec = pl.BlockSpec((NUM_SUBLANES, bq), q_segment_ids_index_map)
     kv_segment_spec = pl.BlockSpec((bkv, NUM_LANES), kv_segment_ids_index_map)
     q_segment_ids = jax.lax.broadcast_in_dim(segment_ids.q, (NUM_SUBLANES, q_seq_len), (1,))
@@ -1485,6 +1507,13 @@ def create_dkv_index_map(h, i, j, *_):
     q_sequence = None
     in_specs.append(None)
 
+  if mask_info.kv_sequence is not None:
+    in_specs.append(pl.BlockSpec((bkv, NUM_LANES), kv_segment_ids_index_map))
+    kv_sequence = jax.lax.broadcast_in_dim(mask_info.kv_sequence, (kv_seq_len, NUM_LANES), (0,))
+  else:
+    kv_sequence = None
+    in_specs.append(None)
+
   dq_reduction_steps = config.dq_reduction_steps
   if not dynamic_grid and kv_steps <= 3 and dq_reduction_steps == 3:
     dq_reduction_steps = None
@@ -1581,6 +1610,7 @@ def create_dkv_index_map(h, i, j, *_):
       di,
       mask_info.partial_mask_blocks,
       q_sequence,
+      kv_sequence,
   ]
   num_args = sum(1 for x in args if x is not None)
   input_output_aliases = {}
@@ -1609,6 +1639,7 @@ def _bwd_cost_estimate(
       di: jax.Array,
       partial_mask_blocks: jax.Array | None,
       q_sequence: jax.Array | None,
+      kv_sequence: jax.Array | None,
       out_shapes: list[jax.ShapeDtypeStruct],
       mask_sparsity_factor: float,
   ) -> pl.CostEstimate:
@@ -1643,6 +1674,7 @@ def _bwd_cost_estimate(
         di,
         partial_mask_blocks,
         q_sequence,
+        kv_sequence,
     ]
     input_bytes = sum(map(_bytes, inputs_))
     output_bytes = sum(map(_bytes, out_shapes))
@@ -1666,6 +1698,7 @@ def _bwd_cost_estimate(
       di,
       mask_info.partial_mask_blocks,
       q_sequence,
+      kv_sequence,
       out_shapes,
       dkv_mask_sparsity,
   )
@@ -1880,6 +1913,7 @@ def mask_info_spec(mask_info):
           if mask_info.partial_mask_blocks is not None
           else None,
           q_sequence=_resolve_spec(mask_info.q_sequence),
+          kv_sequence=(jax.sharding.PartitionSpec() if mask_info.kv_sequence is not None else None),
       )
 
     return SplashAttentionKernel(
@@ -2029,6 +2063,7 @@ def process_mask_shard(mask):
       block_mask=mask_spec,
       partial_mask_blocks=mask_spec,
       q_sequence=None,
+      kv_sequence=None,
   )
   out_specs = (
       mask_info_specs,

src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask.py

@@ -183,12 +183,15 @@ class _ComputableMask(Mask):
       to undefined softmax.
     q_sequence: Indices of Q sequence. q_sequence is reused across __getitem__
       calls which is important for compile-time performance.
+    kv_sequence: Optional indices of KV sequence. When unset, KV positions are
+      the physical column indices.
     mask_function: Function used by the SplashAttention kernel to compute the
       mask rather than loading it.
   """
 
   _shape: tuple[int, int]
   q_sequence: np.ndarray
+  kv_sequence: np.ndarray | None
   mask_function: Callable[..., Any]
 
   def __init__(
@@ -207,6 +210,7 @@ def __init__(
       )
 
     self.q_sequence = np.arange(q_seq_len, dtype=np.int32)
+    self.kv_sequence = None
 
   @property
   def shape(self) -> tuple[int, ...]:
@@ -224,7 +228,10 @@ def __getitem__(self, idx) -> np.ndarray:
     kv_slice = _fill_slice(kv_slice, self.shape[1])
 
     rows = self.q_sequence[q_slice]
-    cols = np.arange(kv_slice.start, kv_slice.stop)
+    if self.kv_sequence is None:
+      cols = np.arange(kv_slice.start, kv_slice.stop)
+    else:
+      cols = self.kv_sequence[kv_slice]
 
     return self.mask_function(rows[:, None], cols[None, :])
 
@@ -276,7 +283,12 @@ def __eq__(self, other: object):
     if not isinstance(other, type(self)):
       return NotImplemented
 
-    return self.shape == other.shape and self.offset == other.offset and np.array_equal(self.q_sequence, other.q_sequence)
+    return (
+        self.shape == other.shape
+        and self.offset == other.offset
+        and np.array_equal(self.q_sequence, other.q_sequence)
+        and np.array_equal(self.kv_sequence, other.kv_sequence)
+    )
 
   def __hash__(self):
     return hash(
@@ -285,6 +297,7 @@ def __hash__(self):
             self.shape,
             self.offset,
             self.q_sequence.tobytes() if self.q_sequence is not None else None,
+            self.kv_sequence.tobytes() if self.kv_sequence is not None else None,
         )
     )