Enable contiguous context-parallel attention for softmax attention layers

Signed-off-by: Haoran Zhang <haoranz@nvidia.com>

Author: Mellonta

megatron/core/context_parallel/sequence_roll.py

@@ -5,18 +5,16 @@
 import torch.distributed as dist
 
 
-def _continuation_mask(tensor, packed_seq_params, cp_group, dims):
+def _continuation_mask(tensor, packed_seq_params, cp_group, dims, sp_group=None):
     """Return the local rows whose next token belongs to the same real document."""
-    size = cp_group.size() if cp_group is not None else 1
     rank = cp_group.rank() if cp_group is not None else 0
+    if sp_group is not None:
+        rank = rank * sp_group.size() + sp_group.rank()
     length = tensor.shape[dims]
-    rows = torch.arange(rank * length, (rank + 1) * length, device=tensor.device)
-    if packed_seq_params is None:
-        ends = torch.full_like(rows, length * size)
-        return rows + 1 < ends
     real = packed_seq_params.cu_seqlens_q
     physical = packed_seq_params.cu_seqlens_q_padded
     physical = real if physical is None else physical
+    rows = torch.arange(rank * length, (rank + 1) * length, device=tensor.device)
     seq = torch.bucketize(rows, physical[1:], right=True)
     in_pack = seq < physical.numel() - 1
     seq = seq.clamp_max(physical.numel() - 2)
@@ -25,75 +23,106 @@ def _continuation_mask(tensor, packed_seq_params, cp_group, dims):
     return valid
 
 
+def exchange_neighbor_tokens(tokens, group, directions, *, wrap=False):
+    """Batch boundary exchanges, preserving each token's shape and dtype.
+
+    Each direction is -1 (send left, receive right) or +1 (send right, receive left).
+    Missing neighbors return None unless wrap is enabled. Callers own boundary fills
+    and autograd; this helper only performs communication on the supplied group.
+    """
+    rank = group.rank() if group is not None else 0
+    size = group.size() if group is not None else 1
+    received, ops = [], []
+    for token, direction in zip(tokens, directions):
+        send_rank, recv_rank = rank + direction, rank - direction
+        if wrap:
+            send_rank, recv_rank = send_rank % size, recv_rank % size
+        if 0 <= send_rank < size:
+            ops.append(
+                dist.P2POp(
+                    dist.isend, token.contiguous(), dist.get_global_rank(group, send_rank), group
+                )
+            )
+        recv = None
+        if 0 <= recv_rank < size:
+            recv = torch.empty_like(token, memory_format=torch.contiguous_format)
+            ops.append(dist.P2POp(dist.irecv, recv, dist.get_global_rank(group, recv_rank), group))
+        received.append(recv)
+    # P2POp retains each contiguous send buffer until the batch completes.
+    for work in dist.batch_isend_irecv(ops) if ops else []:
+        work.wait()
+    return received
+
+
 class _RollLeft(torch.autograd.Function):
     """One differentiable left roll, including the reverse neighbor exchange."""
 
     @staticmethod
-    def forward(ctx, tensor, group, dims):
+    def forward(ctx, tensor, group, dims, sp_group):
         """Shift local values and receive the next rank's first row."""
-        ctx.group, ctx.dims = group, dims
+        ctx.group, ctx.dims, ctx.sp_group = group, dims, sp_group
         result = torch.roll(tensor, -1, dims)
-        result.select(dims, -1).zero_()
-        if group is not None and group.size() > 1:
-            rank, size = group.rank(), group.size()
-            send = tensor.select(dims, 0).contiguous()
-            recv = torch.empty_like(send)
-            ops = []
-            if rank > 0:
-                ops.append(
-                    dist.P2POp(dist.isend, send, dist.get_global_rank(group, rank - 1), group)
-                )
-            if rank + 1 < size:
-                ops.append(
-                    dist.P2POp(dist.irecv, recv, dist.get_global_rank(group, rank + 1), group)
-                )
-            for work in dist.batch_isend_irecv(ops):
-                work.wait()
-            if rank + 1 < size:
-                result.select(dims, -1).copy_(recv)
+        boundary = tensor.select(dims, 0)
+        if sp_group is None or sp_group.rank() == 0:
+            (boundary,) = exchange_neighbor_tokens((boundary,), group, (-1,))
+            if boundary is None:
+                boundary = torch.zeros_like(tensor.select(dims, 0))
+        # Owners are ordered CP rank first, then SP rank. Relay the next CP
+        # shard's first token from SP rank zero to the last SP rank.
+        if sp_group is not None and sp_group.size() > 1:
+            (boundary,) = exchange_neighbor_tokens((boundary,), sp_group, (-1,), wrap=True)
+        result.select(dims, -1).copy_(boundary)
         return result
 
     @staticmethod
     def backward(ctx, grad):
         """Return the boundary gradient to its original token owner."""
-        group, dims = ctx.group, ctx.dims
+        group, dims, sp_group = ctx.group, ctx.dims, ctx.sp_group
         result = torch.roll(grad, 1, dims)
-        result.select(dims, 0).zero_()
-        if group is not None and group.size() > 1:
-            rank, size = group.rank(), group.size()
-            send = grad.select(dims, -1).contiguous()
-            recv = torch.empty_like(send)
-            ops = []
-            if rank + 1 < size:
-                ops.append(
-                    dist.P2POp(dist.isend, send, dist.get_global_rank(group, rank + 1), group)
-                )
-            if rank > 0:
-                ops.append(
-                    dist.P2POp(dist.irecv, recv, dist.get_global_rank(group, rank - 1), group)
-                )
-            for work in dist.batch_isend_irecv(ops):
-                work.wait()
-            if rank > 0:
-                result.select(dims, 0).copy_(recv)
-        return result, None, None
+        boundary = grad.select(dims, -1)
+        if sp_group is not None and sp_group.size() > 1:
+            (boundary,) = exchange_neighbor_tokens((boundary,), sp_group, (1,), wrap=True)
+        if sp_group is None or sp_group.rank() == 0:
+            (boundary,) = exchange_neighbor_tokens((boundary,), group, (1,))
+            if boundary is None:
+                boundary = torch.zeros_like(grad.select(dims, -1))
+        result.select(dims, 0).copy_(boundary)
+        return result, None, None, None
 
 
-def roll_contiguous(tensor, dims, cp_group, packed_seq_params=None, fill_value=0):
-    """Roll a dense or packed contiguous shard by one, preserving real document boundaries."""
+def roll_contiguous(tensor, dims, cp_group, packed_seq_params=None, fill_value=0, *, sp_group=None):
+    """Roll a contiguous shard by one, preserving real document boundaries.
+
+    With sequence parallelism, each CP shard is further split in SP rank order.
+    Both boundary exchanges are differentiable and use the supplied groups.
+    """
     if tensor.shape[dims] == 0:
         return tensor.clone()
-    valid = _continuation_mask(tensor, packed_seq_params, cp_group, dims)
+    cp_size = cp_group.size() if cp_group is not None else 1
+    sp_size = sp_group.size() if sp_group is not None else 1
+    if cp_size == sp_size == 1:
+        result = torch.roll(tensor, -1, dims)
+        result.select(dims, -1).fill_(fill_value)
+    else:
+        result = _RollLeft.apply(tensor, cp_group, dims, sp_group)
+        if fill_value != 0:
+            last_cp = cp_group is None or cp_group.rank() == cp_size - 1
+            last_sp = sp_group is None or sp_group.rank() == sp_size - 1
+            if last_cp and last_sp:
+                result.select(dims, -1).fill_(fill_value)
+    if packed_seq_params is None:
+        return result
+    valid = _continuation_mask(tensor, packed_seq_params, cp_group, dims, sp_group)
     shape = [1] * tensor.ndim
     shape[dims] = tensor.shape[dims]
-    return _RollLeft.apply(tensor, cp_group, dims).masked_fill(~valid.view(shape), fill_value)
+    return result.masked_fill_(~valid.view(shape), fill_value)
 
 
 def roll_contiguous_fields(tensors, cp_group, packed_seq_params=None, fill_values=None):
     """Roll non-differentiable token fields together, preserving each dtype and fill value.
 
     This batches one depth's IDs/positions/masks; activations continue to use the
-    differentiable roll above. No cross-layer cache or model-wide state is introduced.
+    differentiable roll above.
     """
     present = [tensor for tensor in tensors if tensor is not None]
     if not present:
@@ -110,34 +139,27 @@ def roll_contiguous_fields(tensors, cp_group, packed_seq_params=None, fill_value
         raise ValueError("Each token field needs one fill value.")
     if length == 0:
         return tuple(None if t is None else t.clone() for t in tensors)
-    valid = _continuation_mask(reference, packed_seq_params, cp_group, -1)
-    rank = cp_group.rank() if cp_group is not None else 0
-    size = cp_group.size() if cp_group is not None else 1
-    outputs, received, sends, ops = [], [], [], []
+    invalid = (
+        ~_continuation_mask(reference, packed_seq_params, cp_group, -1)
+        if packed_seq_params is not None
+        else None
+    )
+    boundaries = exchange_neighbor_tokens(
+        [tensor[..., 0] for tensor in present], cp_group, [-1] * len(present)
+    )
+    received = iter(boundaries)
+    outputs = []
     for tensor, fill in zip(tensors, fill_values):
         if tensor is None:
             outputs.append(None)
             continue
         result = torch.roll(tensor, -1, -1)
-        result[..., -1].fill_(fill)
+        boundary = next(received)
+        if boundary is not None:
+            result[..., -1].copy_(boundary)
+        else:
+            result[..., -1].fill_(fill)
+        if invalid is not None:
+            result.masked_fill_(invalid, fill)
         outputs.append(result)
-        if rank > 0:
-            send = tensor[..., 0].contiguous()
-            sends.append(send)  # Retain temporary send storage until all work completes.
-            ops.append(
-                dist.P2POp(dist.isend, send, dist.get_global_rank(cp_group, rank - 1), cp_group)
-            )
-        if rank + 1 < size:
-            recv = torch.empty_like(tensor[..., 0])
-            received.append((result, recv))
-            ops.append(
-                dist.P2POp(dist.irecv, recv, dist.get_global_rank(cp_group, rank + 1), cp_group)
-            )
-    for work in dist.batch_isend_irecv(ops) if ops else []:
-        work.wait()
-    for result, recv in received:
-        result[..., -1].copy_(recv)
-    for result, fill in zip(outputs, fill_values):
-        if result is not None:
-            result.masked_fill_(~valid, fill)
     return tuple(outputs)

megatron/core/context_parallel/utils.py

@@ -178,6 +178,12 @@ def _build_packed_seq_params(
         if (layout == "contiguous" and cp_size > 1) or use_logical_qkv_seqlens
         else physical_cu_seqlens
     )
+    if layout == "contiguous" and cp_size > 1 and pad_between_seqs is None:
+        # TE's sync-free default uses object identity. Independently materialized
+        # but equal boundaries are still unpadded and must remain eligible for FA.
+        pad_between_seqs = cu_seqlens_padded is not None and not torch.equal(
+            cu_seqlens, cu_seqlens_padded
+        )
 
     max_seqlen = int(batch['max_seqlen'].item())
     local_cp_size = batch.get('local_cp_size')

megatron/core/extensions/transformer_engine.py

@@ -2231,6 +2231,11 @@ def forward(
 _te_dpa_supports_softcap = (
     "softcap" in inspect.signature(te.pytorch.DotProductAttention.__init__).parameters
 )
+_te_dpa_supports_no_load_balance = (
+    hasattr(getattr(te.pytorch, "CPLoadBalancingStrategy", None), "NO_LOAD_BALANCE")
+    and "load_balancing_strategy"
+    in inspect.signature(te.pytorch.DotProductAttention.set_context_parallel_group).parameters
+)
 
 
 class TEDotProductAttention(te.pytorch.DotProductAttention):
@@ -2444,6 +2449,39 @@ def __init__(
             **extra_kwargs,
         )
 
+        if self.config.attention_cp_layout == "contiguous":
+            cp_comm_type = cp_comm_type or self.cp_comm_type
+            cp_groups = self.cp_group if isinstance(self.cp_group, list) else [self.cp_group]
+            if any(get_pg_size(group) > 1 for group in cp_groups):
+                if cp_comm_type != "all_gather":
+                    raise ValueError("Contiguous attention CP requires cp_comm_type='all_gather'.")
+                if not _te_dpa_supports_no_load_balance:
+                    raise RuntimeError(
+                        "Contiguous attention CP requires a Transformer Engine build with "
+                        "CPLoadBalancingStrategy.NO_LOAD_BALANCE support."
+                    )
+                if self.attention_type != "self":
+                    raise ValueError("Contiguous attention CP only supports self-attention.")
+            self.set_context_parallel_group(
+                self.cp_group, self.cp_global_ranks, self.cp_stream, cp_comm_type
+            )
+
+    def set_context_parallel_group(
+        self, cp_group, cp_global_ranks, cp_stream, cp_comm_type="p2p", load_balancing_strategy=None
+    ):
+        """Keep TE's partition strategy consistent when a caller rebinds the CP group."""
+        cp_groups = cp_group if isinstance(cp_group, list) else [cp_group]
+        if self.config.attention_cp_layout == "contiguous" and any(
+            get_pg_size(group) > 1 for group in cp_groups
+        ):
+            load_balancing_strategy = te.pytorch.CPLoadBalancingStrategy.NO_LOAD_BALANCE
+        kwargs = {}
+        if load_balancing_strategy is not None:
+            kwargs["load_balancing_strategy"] = load_balancing_strategy
+        super().set_context_parallel_group(
+            cp_group, cp_global_ranks, cp_stream, cp_comm_type, **kwargs
+        )
+
     @contextmanager
     def _temporary_runtime_context_parallel_group(
         self, packed_seq_params: Optional[PackedSeqParams]
@@ -2462,11 +2500,11 @@ def _temporary_runtime_context_parallel_group(
             if runtime_cp_group.size() == 1:
                 # Dynamic CP metadata retains the singleton group, while TE
                 # must see CP disabled for this microbatch.
-                super().set_context_parallel_group(None, None, None, self.cp_comm_type)
+                self.set_context_parallel_group(None, None, None, self.cp_comm_type)
             else:
                 if TEDotProductAttention.cp_stream is None:
                     TEDotProductAttention.cp_stream = torch.cuda.Stream()
-                super().set_context_parallel_group(
+                self.set_context_parallel_group(
                     runtime_cp_group,
                     torch.distributed.get_process_group_ranks(runtime_cp_group),
                     TEDotProductAttention.cp_stream,
@@ -2475,9 +2513,9 @@ def _temporary_runtime_context_parallel_group(
             yield
         finally:
             if original_cp_group is None or original_cp_group.size() == 1:
-                super().set_context_parallel_group(None, None, None, self.cp_comm_type)
+                self.set_context_parallel_group(None, None, None, self.cp_comm_type)
             else:
-                super().set_context_parallel_group(
+                self.set_context_parallel_group(
                     original_cp_group,
                     original_cp_global_ranks,
                     TEDotProductAttention.cp_stream,
@@ -2542,6 +2580,50 @@ def _forward(
         )
         qkv_format = packed_seq_kwargs.get('qkv_format', self.qkv_format)
 
+        if self.config.attention_cp_layout == "contiguous" and get_pg_size(self.cp_group) > 1:
+            if attention_bias is not None:
+                raise ValueError("Contiguous attention CP does not support attention bias.")
+            if qkv_format == "sbhd":
+                if attn_mask_type != AttnMaskType.causal:
+                    raise ValueError(
+                        "Padded contiguous attention CP requires THD sequence metadata."
+                    )
+                # Adapt dense inputs to the THD format required by TE's no-load-balance CP.
+                from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+                    get_full_cu_seqlens,
+                )
+
+                seq_length = query.size(0) * self.cp_group.size()
+                cu_seqlens = get_full_cu_seqlens(1, seq_length, query.device)
+                sample_params = PackedSeqParams(
+                    qkv_format="thd",
+                    cu_seqlens_q=cu_seqlens,
+                    cu_seqlens_kv=cu_seqlens,
+                    cu_seqlens_q_padded=cu_seqlens,
+                    cu_seqlens_kv_padded=cu_seqlens,
+                    max_seqlen_q=seq_length,
+                    max_seqlen_kv=seq_length,
+                    pad_between_seqs=False,
+                )
+                # Copy each batch once; each sample is then a contiguous THD view.
+                samples = [x.transpose(0, 1).contiguous().unbind(0) for x in (query, key, value)]
+                outputs = [
+                    self._forward(
+                        q,
+                        k,
+                        v,
+                        None,
+                        attn_mask_type,
+                        packed_seq_params=sample_params,
+                        num_splits=num_splits,
+                        bf16_backward=bf16_backward,
+                    )
+                    for q, k, v in zip(*samples)
+                ]
+                return outputs[0].unsqueeze(1) if len(outputs) == 1 else torch.stack(outputs, dim=1)
+            if qkv_format != "thd":
+                raise ValueError("Contiguous attention CP requires SBHD or THD input.")
+
         attention_bias_kwargs = {}
         if attention_bias is not None:
             assert is_te_min_version("1.2.0"), (

megatron/core/transformer/attention.py

@@ -1395,6 +1395,17 @@ def forward_pre_attn_and_core_attn(
         if no_rope:
             rotary_pos_emb = None
 
+        if (
+            self.config.attention_cp_layout == "contiguous"
+            and runtime_cp_group is not None
+            and runtime_cp_group.size() > 1
+            and any(
+                value is not None
+                for value in (rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, rotary_pos_cos_sin)
+            )
+        ):
+            raise ValueError("Contiguous attention CP does not yet support rotary embeddings.")
+
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
         if inference_context and inference_context.is_dynamic_batching():