Context parallel all-gather

GitOrigin-RevId: 08b9fea

Author: babych-d

axlearn/common/attention.py

@@ -97,6 +97,7 @@
 import jax
 from absl import logging
 from jax import numpy as jnp
+from jax._src.mesh import thread_resources
 
 from axlearn.common import param_init
 from axlearn.common.attention_bias import (
@@ -166,6 +167,7 @@
     save_and_offload_only_these_names_regex,
     shapes,
     split_prng_key,
+    with_sharding_constraint,
 )
 
 
@@ -1064,6 +1066,20 @@ def forward(
         q_proj, k_proj, v_proj = jnp.split(
             proj, [cfg.num_heads, cfg.num_heads + cfg.num_kv_heads], axis=-2
         )
+        # This sharding hint is needed since compiler sometimes will generate large allgather
+        # before the split and then slice, which is not the ideal compilation. Ensure sharding
+        # after the split to ensure allgather is inserted after the split.
+        axis_names = thread_resources.env.physical_mesh.axis_names
+        batch_axes = tuple(x for x in axis_names if x in ("data", "fsdp")) or None
+        spec = PartitionSpec(
+            batch_axes,
+            "seq" if "seq" in axis_names else None,
+            "model" if "model" in axis_names else None,
+            None,
+        )
+        q_proj = with_sharding_constraint(q_proj, spec)
+        k_proj = with_sharding_constraint(k_proj, spec)
+        v_proj = with_sharding_constraint(v_proj, spec)
         return self.Output(query=q_proj, key=k_proj, value=v_proj)
 
 

axlearn/common/attention_bias.py

@@ -440,10 +440,17 @@ def partition_spec(
         self, mha_dim_to_partition_spec: dict[str, PartitionSpec]
     ) -> Union[BaseAttentionBias, PartitionSpec]:
         # Segment IDs: [batch_size, seq_len].
-        q_spec = mha_dim_to_partition_spec["btnh"]
-        if q_spec == PartitionSpec(None):
+        # We use the partition spec of KV (which are not sequence sharded) for segment ids. This is
+        # because Splash requires two seg ids, q_seg and kv_seg. Therefore, we pass a not seq
+        # sharded seg ids into the shard map, and manually shard it inside for q_seg and not
+        # shard it for kv_seg.
+        kv_spec = mha_dim_to_partition_spec["bsnh"]
+        if kv_spec == PartitionSpec(None):
             return PartitionSpec(None)
-        return PartitionSpec(q_spec[0], q_spec[1])
+
+        if kv_spec[1] is not None:
+            raise ValueError("The partition spec of `s` in `bsnh` should be None.")
+        return PartitionSpec(kv_spec[0], kv_spec[1])
 
 
 class MaskFn(Protocol):

axlearn/common/flash_attention/common.py

@@ -19,6 +19,7 @@
 from axlearn.common.attention import compute_gqa_context, compute_gqa_logits, softmax_with_biases
 from axlearn.common.attention_bias import BaseAttentionBias, MaskFn, SegmentIdAttentionBias
 from axlearn.common.config import Configurable, config_class
+from axlearn.common.flash_attention.types import FlashAttentionWithShardMapSpecs
 from axlearn.common.kv_cache.base_kv_cache import BaseKVCache
 from axlearn.common.kv_cache.kv_cache import KVCache
 from axlearn.common.kv_cache.paged_kv_cache import PagedKVCache, reconstruct_kv
@@ -292,6 +293,27 @@ def _check_block_size(
             return False
         return True
 
+    def build(
+        self, input_batch: Nested[Tensor | BaseAttentionBias]  # pylint: disable=unused-argument
+    ) -> FlashAttentionWithShardMapSpecs:
+        """Builds sharding specifications for the flash attention operation.
+
+        This method is called to build the sharding map specifications that will be used
+        to distribute the attention computation across devices. The base implementation simply
+        wraps `self` in a FlashAttentionShardWithMapSpecs, but subclasses may override this to
+        create more sophisticated specs based on the input batch characteristics.
+
+        Args:
+            input_batch: A dict containing the input tensors and attention biases. While unused
+                in the base implementation, this parameter allows subclasses to inspect input
+                properties when building custom sharding specifications.
+
+        Returns:
+            FlashAttentionShardWithMapSpecs containing the attention function and any
+                additional sharding metadata required for distributed execution.
+        """
+        return FlashAttentionWithShardMapSpecs(fn=self)
+
 
 class BaseSingleStepDecoding(BaseFlashAttention):
     """Wraps the common checks for single step decoding kernels."""

axlearn/common/flash_attention/layer.py

@@ -220,7 +220,7 @@ def _compute_attention(
         # Get logit sink parameter if configured.
         logit_sink = self.parameters.get("sink", None)
 
-        jit_attn = flash_attention_implementation(
+        flash_specs = flash_attention_implementation(
             backend=backend,
             query=q_proj,
             key=k_proj,
@@ -236,7 +236,7 @@ def _compute_attention(
             page_tables=page_indices,
             backend_overrides=cfg.backend_overrides,
         )
-        if jit_attn is None:
+        if flash_specs is None:
             # Fall back to standard attention if no backend kernels are supported.
             return super()._compute_attention(
                 mode=mode,
@@ -265,7 +265,6 @@ def _compute_attention(
         v_proj = with_sharding_constraint(v_proj, cfg.mha_dim_to_partition_spec[kv_partition])
 
         # We need to manually partition pallas | jax-triton calls.
-        # Note: shard_map doesn't support kwargs.
         input_batch_specs = {
             # Q [batch_size, seq_len, num_heads, per_head_dim].
             "query": cfg.mha_dim_to_partition_spec["btnh"],
@@ -285,9 +284,10 @@ def _compute_attention(
             ),
             # PagedKVCache's page indices.
             "page_tables": cfg.mha_dim_to_partition_spec.get("bs", PartitionSpec(None)),
+            **flash_specs.additional_in_specs,
         }
         partitioned_mha = shard_map(
-            jit_attn,
+            flash_specs.fn,
             mesh=thread_resources.env.physical_mesh,
             in_specs=(input_batch_specs,),
             # O [batch_size, seq_len, num_heads, per_head_dim].
@@ -307,6 +307,7 @@ def _compute_attention(
             "bias": attention_logit_biases,
             "logit_sink": logit_sink,
             "page_tables": page_indices,
+            **flash_specs.additional_kwargs,
         }
         outputs = with_sharding_constraint(
             partitioned_mha(
@@ -345,7 +346,7 @@ def default_mha_dim_to_partition_spec(
     Returns:
         A dictionary keyed by MHA tensor dims with partition spec values.
     """
-    batch_axis_names = tuple(el for el in mesh_axis_names if el != "model")
+    batch_axis_names = tuple(el for el in mesh_axis_names if el in ["data", "fsdp"])
     tp_axis_name = "model" if "model" in mesh_axis_names else None
     return {
         "btnh": PartitionSpec(batch_axis_names, None, tp_axis_name, None),
@@ -369,7 +370,7 @@ def default_output_dim_to_partition_spec(
     Returns:
         A dictionary keyed by FlashAttention output tensor dims with partition spec values.
     """
-    batch_axis_names = tuple(el for el in mesh_axis_names if el not in ["seq", "model"])
+    batch_axis_names = tuple(el for el in mesh_axis_names if el in ["data", "fsdp"])
     tp_axis_name = "model" if "model" in mesh_axis_names else None
     sp_axis_name = "seq" if "seq" in mesh_axis_names else None
     return {