Merge pull request #4364 from huytransformer:htn/load-balanced-cp-chunk-masks

PiperOrigin-RevId: 944548246

Author: Google-ML-Automation

src/maxtext/configs/pyconfig_deprecated.py

@@ -249,12 +249,11 @@ def validate_keys(keys):
         keys["global_rampup_samples"],
     )
 
-  context_parallel_size = get_context_parallel_size(keys)
   # TODO remove after b/435512699 resolved
-  if context_parallel_size > 1 and keys["context_parallel_load_balance"] and keys["attention_type"] == "chunk":
-    raise ValueError("Currently load-balanced context parallelism is not supported for chunk attention.")
 
-  validate_context_parallel_strategy_ring(context_parallel_size, keys["context_parallel_strategy"], keys["hardware"])
+  validate_context_parallel_strategy_ring(
+      keys["context_parallel_size"], keys["context_parallel_strategy"], keys["hardware"]
+  )
 
   if keys["mtp_eval_target_module"] < 0:
     raise ValueError("mtp_eval_target_module cannot be negative. Set to 0 to disable evaluation.")
@@ -816,6 +815,7 @@ def user_init(raw_keys):
 
     raw_keys["num_slices"] = max_utils.get_num_slices(raw_keys)
     raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys)
+    raw_keys["context_parallel_size"] = get_context_parallel_size(raw_keys)
     raw_keys = create_parallelisms_list(raw_keys)
     raw_keys = set_and_validate_pipeline_config(raw_keys)
 

src/maxtext/layers/attention_op.py

@@ -615,6 +615,7 @@ def generate_attention_mask(
       previous_chunk: Any = None,
       bidirectional_mask: Any = None,
       compressed_mask: Optional[Array] = None,
+      segment_positions: Array | None = None,
   ) -> Array | None:
     """Generates a combined attention mask for Transformer models.
 
@@ -675,6 +676,9 @@ def generate_attention_mask(
       compressed_mask: Optional `Array`. A pre-computed attention mask for
         compressed kv blocks (e.g., DeepSeek-V4 compressed attention). If provided, it is
         concatenated with the dynamically generated uncompressed mask.
+      segment_positions: Optional `Array` of shape `[batch_size,
+        q_sequence_length]`. Identifies original positions for load-balanced
+        context-parallel inputs.
 
     Returns:
       An `Array` representing the attention mask, with shape
@@ -710,20 +714,33 @@ def generate_attention_mask(
     elif model_mode == MODEL_MODE_AUTOREGRESSIVE and q_seq_len == 1:
       # In autoregression, the query position is the last position in the KV sequence.
       next_pos = kv_seq_len - 1
+    use_segment_positions = (
+        segment_positions is not None
+        and self.config.context_parallel_load_balance
+        and self.mesh.shape.get(self.config.context_sharding, 1) > 1
+        and previous_chunk is None
+        and model_mode != MODEL_MODE_AUTOREGRESSIVE
+    )
+    if use_segment_positions:
+      position_row_ids = segment_positions[:, :, None]
+      position_col_ids = segment_positions[:, None, :]
 
     causal_mask = None
     if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in (
         AttentionType.FULL,
         AttentionType.COMPRESSED,
     ):
-      mask_shape = (q_seq_len, kv_seq_len)
-      # row_ids indicates the position of query
-      # col_ids indicates the position of kv
-      row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
-      col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
-      # Attention mask for chunked prefill is generated in the same way
-      # as mentioned in SARATHI - https://arxiv.org/abs/2308.16369
-      causal_mask = (col_ids <= row_ids + next_pos)[None, None, None, :, :]
+      if use_segment_positions:
+        causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :]
+      else:
+        mask_shape = (q_seq_len, kv_seq_len)
+        # row_ids indicates the position of query
+        # col_ids indicates the position of kv
+        row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
+        col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
+        # Attention mask for chunked prefill is generated in the same way
+        # as mentioned in SARATHI - https://arxiv.org/abs/2308.16369
+        causal_mask = (col_ids <= row_ids + next_pos)[None, None, None, :, :]
 
     output_mask = None
     if (mask is not None) and (causal_mask is not None):
@@ -737,11 +754,17 @@ def generate_attention_mask(
       if self.sliding_window_size is None:
         raise ValueError("Sliding_window_size must be set if Local Sliding attention type")
 
-      row_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (q_seq_len, 1), 0) + next_pos
-      col_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (1, kv_seq_len), 1)
+      if use_segment_positions:
+        row_ids_sliding = position_row_ids
+        col_ids_sliding = position_col_ids
+      else:
+        row_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (q_seq_len, 1), 0) + next_pos
+        col_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (1, kv_seq_len), 1)
       sliding_mask = (col_ids_sliding > (row_ids_sliding - self.sliding_window_size)) & (
           col_ids_sliding <= row_ids_sliding
       )
+      if use_segment_positions:
+        sliding_mask = sliding_mask[:, None, None, :, :]
       output_mask = sliding_mask * output_mask
     elif self.attention_type == AttentionType.COMPRESSED:
       if compressed_mask is None:
@@ -772,12 +795,17 @@ def generate_attention_mask(
       return jnp.concatenate([uncompressed_mask, compressed_mask], axis=-1)
 
     elif self.attention_type == AttentionType.CHUNK and output_mask is not None:
-      mask_shape = (q_seq_len, kv_seq_len)
-      chunk_mask = _generate_chunk_attention_mask(
-          mask_shape=(q_seq_len, kv_seq_len),
-          chunk_size=self.chunk_attn_window_size,
-          q_offset=next_pos,
-      )
+      if use_segment_positions:
+        same_chunk = (position_row_ids // self.chunk_attn_window_size) == (
+            position_col_ids // self.chunk_attn_window_size
+        )
+        chunk_mask = (same_chunk & (position_row_ids >= position_col_ids))[:, None, None, :, :]
+      else:
+        chunk_mask = _generate_chunk_attention_mask(
+            mask_shape=(q_seq_len, kv_seq_len),
+            chunk_size=self.chunk_attn_window_size,
+            q_offset=next_pos,
+        )
       output_mask = chunk_mask * output_mask
 
     if bidirectional_mask is not None:
@@ -976,6 +1004,7 @@ def apply_attention(
           decoder_segment_ids,
           model_mode,
           previous_chunk,
+          segment_positions=segment_positions,
           bidirectional_mask=bidirectional_mask,
           sinks=sinks,
           indexer_mask=indexer_mask,
@@ -1301,10 +1330,17 @@ def create_sa_config(config, query, key, attn_logits_soft_cap):
       if self.chunk_attn_window_size is None:
         raise ValueError("chunk_attn_window_size must be set for chunk attention type")
 
-      mask &= ChunkedCausalMask(
-          shape=(query.shape[2], key.shape[2]),
-          chunk_size=self.chunk_attn_window_size,
-      )
+      if use_load_balanced_cp:
+        mask &= LoadBalancedChunkedCausalMask(
+            shape=(query.shape[2], key.shape[2]),
+            chunk_size=self.chunk_attn_window_size,
+            cp_size=cp_size,
+        )
+      else:
+        mask &= ChunkedCausalMask(
+            shape=(query.shape[2], key.shape[2]),
+            chunk_size=self.chunk_attn_window_size,
+        )
 
     max_logit_value = None
     if self.config.use_tokamax_splash:
@@ -1588,6 +1624,8 @@ def cudnn_flash_attention(
         and self.config.context_parallel_strategy == "ring"
         and self.config.context_parallel_load_balance
     )
+    if using_context_parallelism and self.attention_type == AttentionType.CHUNK:
+      raise ValueError("Chunk attention is not supported with CUDNN context parallelism.")
 
     # Initialize default attention configuration
     sliding_window_size = None
@@ -1815,6 +1853,7 @@ def apply_attention_dot(
       decoder_segment_ids: Array | None,
       model_mode: str = MODEL_MODE_TRAIN,
       previous_chunk: Any = None,
+      segment_positions: Array | None = None,
       bidirectional_mask: Any = None,
       sinks: Array | None = None,
       indexer_mask: Array | None = None,
@@ -1880,6 +1919,7 @@ def apply_attention_dot(
         previous_chunk,
         bidirectional_mask,
         compressed_mask=compressed_mask,
+        segment_positions=segment_positions,
     )
 
     if self.config.moba:
@@ -2254,3 +2294,21 @@ def __init__(
     # LocalMask uses shard_count for mask-shard validation. cp_size is the
     # context-parallel size used for the load-balanced query order.
     self.q_sequence = _load_balanced_q_sequence(shape, cp_size)
+
+
+class LoadBalancedChunkedCausalMask(ChunkedCausalMask):
+  """Lazy chunked mask with load-balanced query positions."""
+
+  def __init__(
+      self,
+      shape: tuple[int, int],
+      chunk_size: int,
+      cp_size: int,
+      shard_count: int = 1,
+  ):
+    super().__init__(
+        shape=shape,
+        chunk_size=chunk_size,
+        shard_count=shard_count,
+    )
+    self.q_sequence = _load_balanced_q_sequence(shape, cp_size)