Sequence parallel: one captured all-to-all for Ulysses; drop the all-gather exchange

configs/cosmos3_nano_sp2.yaml

@@ -12,10 +12,8 @@ resources:
 # so each rank attends the full sequence over half the heads. Weights are
 # replicated (no tensor parallelism here). The VAE decoder is small and runs
 # un-sharded on rank 0; the DiT all-gathers its final latents, so the decoder
-# reads them directly. The denoise CUDA-graph captures the image step under SP:
-# during capture the Ulysses exchange switches to all-gather (the all-to-all is
-# point-to-point send/recv, which a CUDA graph cannot replay), so serving
-# graph-on graphs t2i. Video runs eager (dense, uncaptured) on the all-to-all.
+# reads them directly. The denoise CUDA-graph captures the image step, all-to-all
+# included. Video runs eager (dense, uncaptured).
 node_groups:
   - node_names: ["dit"]
     ranks: [0, 1]

configs/cosmos3_nano_sp4.yaml

@@ -12,8 +12,7 @@ resources:
 # so each rank attends the full sequence over a quarter of the heads (sp_size 4
 # divides the 8 KV heads -> 2 heads/rank). Weights are replicated (no tensor
 # parallelism). A 4-GPU correctness cross-check. The denoise CUDA-graph captures
-# the image step via the all-gather Ulysses exchange (the all-to-all is
-# point-to-point send/recv, which a CUDA graph cannot replay).
+# the image step, all-to-all included.
 node_groups:
   - node_names: ["dit"]
     ranks: [0, 1, 2, 3]

configs/cosmos3_nano_tp2_sp2.yaml

@@ -13,9 +13,8 @@ resources:
 # attention). The 4 ranks form a [sp_size, tp_size] grid: TP groups {0,1},{2,3};
 # SP groups {0,2},{1,3}. Used to validate TP x SP composition on Nano (the small
 # model fits on one GPU, so this is a correctness cross-check, not a perf config).
-# The denoise CUDA-graph captures the image step: the SP Ulysses exchange uses
-# all-gather under capture (the all-to-all is point-to-point send/recv, which a
-# CUDA graph cannot replay); the TP row-parallel projections all-reduce as usual.
+# The denoise CUDA-graph captures the image step, the SP all-to-all and the TP
+# all-reduce included.
 node_groups:
   - node_names: ["dit"]
     ranks: [0, 1, 2, 3]

mstar/distributed/communication.py

@@ -1,4 +1,5 @@
 import socket
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import Any
@@ -44,6 +45,15 @@ def rendezvous_store(init_method: str, rank: int, world_size: int, timeout: time
     )
 
 
+@torch.compiler.disable
+def _all_to_all_single(
+    group: dist.ProcessGroup, send: torch.Tensor, recv: torch.Tensor,
+    recv_splits: list[int] | None = None, send_splits: list[int] | None = None,
+) -> torch.Tensor:
+    dist.all_to_all_single(recv, send, recv_splits, send_splits, group=group)
+    return recv
+
+
 class CommGroup:
     """A communication group over one axis of the worker device mesh.
 
@@ -210,12 +220,12 @@ def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
 
     def all_to_all(
         self,
-        input_: torch.Tensor,
+        input_: torch.Tensor | Sequence[torch.Tensor],
         scatter_dim: int,
         gather_dim: int,
         scatter_sizes: list[int] | None = None,
         gather_sizes: list[int] | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Redistribute ``input_`` across the group: split it into
         ``world_size`` pieces along ``scatter_dim`` (piece i goes to rank i)
         and concatenate the pieces received from every rank along
@@ -230,34 +240,80 @@ def all_to_all(
         uneven split / gather (e.g. a sequence length not divisible by the
         group size); when ``None`` the respective dimension is split evenly
         (and ``scatter_dim`` must then be divisible by ``world_size``).
+
+        ``input_`` may be several tensors that differ only in the extent of
+        ``scatter_dim``, split evenly (q, k and v with their own head counts):
+        they go in one exchange and come back as a tuple. Each exchange is one
+        ``all_to_all_single`` over a buffer packed in a single copy, which
+        replays from a CUDA graph. The collective is hidden from torch.compile,
+        so a compiled caller breaks around this call; tracing it as a
+        functional collective instead gave an illegal memory access in
+        Cosmos3's batched capture.
         """
+        tensors = (input_,) if isinstance(input_, torch.Tensor) else tuple(input_)
         if self.world_size == 1:
-            return input_
+            return input_ if isinstance(input_, torch.Tensor) else tensors
         world_size = self.world_size
-        if scatter_sizes is None:
-            assert input_.size(scatter_dim) % world_size == 0, (
-                f"all_to_all: scatter_dim {scatter_dim} size "
-                f"{input_.size(scatter_dim)} not divisible by world_size "
-                f"{world_size}; pass scatter_sizes for an uneven split"
+        # positive dims: unflatten below adds a dim, which would shift a negative one
+        scatter_dim %= tensors[0].dim()
+        gather_dim %= tensors[0].dim()
+        if scatter_sizes is not None and len(set(scatter_sizes)) > 1:
+            assert len(tensors) == 1, "all_to_all: several tensors need an even scatter"
+            out = self._all_to_all_uneven_scatter(tensors[0], scatter_dim, gather_dim, scatter_sizes, gather_sizes)
+            return out if isinstance(input_, torch.Tensor) else (out,)
+        # per tensor [.., N, ..] -> [P, .., N/P, ..]; rank r's block of the send buffer holds
+        # its piece of each tensor one after another, so each tensor unpacks contiguous
+        blocks = [x.unflatten(scatter_dim, (world_size, -1)).movedim(scatter_dim, 0) for x in tensors]
+        numels = [blk[0].numel() for blk in blocks]
+        send = blocks[0].new_empty(world_size, sum(numels))
+        for blk, piece in zip(blocks, send.split(numels, dim=1), strict=True):
+            piece.view(blk.shape).copy_(blk)
+        if gather_sizes is None or len(set(gather_sizes)) == 1:
+            # the block received from rank r has this rank's own block's layout
+            recv = _all_to_all_single(self.device_group, send, torch.empty_like(send))
+            outs = [
+                piece.view(blk.shape).movedim(0, gather_dim).flatten(gather_dim, gather_dim + 1)
+                for blk, piece in zip(blocks, recv.split(numels, dim=1), strict=True)
+            ]
+        else:
+            # rank r's block: each tensor's piece with the gathered dim at rank r's extent
+            shapes = [
+                [blk.shape[1:gather_dim + 1] + (n,) + blk.shape[gather_dim + 2:] for blk in blocks]
+                for n in gather_sizes
+            ]
+            rank_numels = [[int(torch.Size(shape).numel()) for shape in row] for row in shapes]
+            recv = _all_to_all_single(
+                self.device_group, send.flatten(), send.new_empty(sum(map(sum, rank_numels))),
+                [sum(row) for row in rank_numels], [sum(numels)] * world_size,
             )
-            chunk = input_.size(scatter_dim) // world_size
-            scatter_sizes = [chunk] * world_size
-        send = [
-            t.contiguous()
-            for t in torch.split(input_, scatter_sizes, dim=scatter_dim)
-        ]
+            per_rank = [
+                [piece.view(shape) for piece, shape in zip(block.split(row), shape_row, strict=True)]
+                for block, row, shape_row in zip(recv.split([sum(row) for row in rank_numels]), rank_numels, shapes,
+                                                 strict=True)
+            ]
+            outs = [torch.cat([pieces[i] for pieces in per_rank], dim=gather_dim) for i in range(len(tensors))]
+        return outs[0] if isinstance(input_, torch.Tensor) else tuple(outs)
+
+    def _all_to_all_uneven_scatter(
+        self, input_: torch.Tensor, scatter_dim: int, gather_dim: int,
+        scatter_sizes: list[int], gather_sizes: list[int] | None,
+    ) -> torch.Tensor:
+        chunks = input_.split(scatter_sizes, dim=scatter_dim)
+        send = torch.cat([chunk.contiguous().flatten() for chunk in chunks])
+        mine = list(chunks[self.rank].shape)
         if gather_sizes is None:
-            gather_sizes = [send[self.rank].size(gather_dim)] * world_size
-        base_shape = list(send[self.rank].shape)
-        recv = []
-        for r in range(world_size):
-            shape = list(base_shape)
-            shape[gather_dim] = gather_sizes[r]
-            recv.append(
-                torch.empty(shape, dtype=input_.dtype, device=input_.device)
-            )
-        dist.all_to_all(recv, send, group=self.device_group)
-        return torch.cat(recv, dim=gather_dim).contiguous()
+            gather_sizes = [mine[gather_dim]] * self.world_size
+        shapes = [mine[:gather_dim] + [n] + mine[gather_dim + 1:] for n in gather_sizes]
+        numels = [int(torch.Size(shape).numel()) for shape in shapes]
+        recv = _all_to_all_single(
+            self.device_group, send, input_.new_empty(sum(numels)),
+            numels, [chunk.numel() for chunk in chunks],
+        )
+        pieces = [
+            piece.view(shape)
+            for piece, shape in zip(recv.split(numels), shapes, strict=True)
+        ]
+        return torch.cat(pieces, dim=gather_dim)
 
 
 @dataclass

mstar/model/components/distributed/sequence_parallel.py

@@ -14,10 +14,12 @@
 These helpers take an ``sp_group`` (a :class:`CommGroup` over the SP axis of the
 mesh; trivial when ``sp_size == 1``) and are model-independent: any attention that
 projects to ``[tokens, heads, head_dim]`` and calls a ``run_attention(q, k, v)``
-can use :func:`ulysses_attention`. ``seq_sizes`` is the per-rank token count along
-the sequence (need not be equal — sequences indivisible by ``sp`` are handled
-without padding); the scatter/gather of the residual stream at the model boundary
-uses the same split.
+can use :func:`ulysses_attention`, which takes ``seq_dim`` / ``head_dim`` for other
+layouts such as ``[batch, tokens, heads, head_dim]``; :func:`sp_head_slice` and
+:func:`sp_head_gather` assume ``[tokens, heads, head_dim]``. ``seq_sizes`` is the
+per-rank token count along the sequence (need not be equal — sequences indivisible
+by ``sp`` are handled without padding); the scatter/gather of the residual stream
+at the model boundary uses the same split.
 """
 from __future__ import annotations
 
@@ -88,69 +90,25 @@ def ulysses_attention(
     v: torch.Tensor,
     run_attention: Callable[..., torch.Tensor],
     seq_sizes: list[int],
-    prefer_all_gather: bool = False,
+    seq_dim: int = 0,
+    head_dim: int = 1,
 ) -> torch.Tensor:
     """Run attention under Ulysses SP.
 
     ``q``: ``[seq/P, Hq, D]``; ``k``/``v``: ``[seq/P, Hkv, D]`` (this rank's
-    tokens, TP-local heads). Returns ``[seq/P, Hq, D]``. The head counts must be
-    divisible by ``sp`` (Ulysses shards heads). When the group is trivial this is
-    a passthrough — byte-identical to the non-SP path.
-
-    ``prefer_all_gather`` selects the all-gather collective instead of the
-    all-to-all (see :func:`_ulysses_attention_via_all_gather`). The caller sets
-    it on the denoise forward that the CUDA graph captures: the all-to-all is
-    grouped point-to-point send/recv and does not replay from a captured graph,
-    whereas all-gather (a true collective, like the TP all-reduce) does. It must
-    be set consistently across warmup, capture and replay so the all-gather
-    kernels are compiled and autotuned during eager warmup, not mid-capture
-    (autotuning synchronizes, which is illegal while a graph is recording).
-    Eager paths (video, uncaptured resolutions) leave it off for the lighter
-    all-to-all. Both produce identical results."""
+    tokens, TP-local heads). Returns ``[seq/P, Hq, D]``. Any other layout names
+    its sequence and head dims (``[B, seq/P, H, D]`` is ``seq_dim=1,
+    head_dim=2``). The head counts must be divisible by ``sp`` (Ulysses shards
+    heads). When the group is trivial this is a passthrough — byte-identical to
+    the non-SP path."""
     if sp_group.world_size == 1:
         return run_attention(q=q, k=k, v=v)
-    if prefer_all_gather:
-        return _ulysses_attention_via_all_gather(sp_group, q, k, v, run_attention)
-    # scatter heads, gather sequence: [seq/P, H, D] -> [seq, H/P, D]
-    q = sp_group.all_to_all(q, scatter_dim=1, gather_dim=0, gather_sizes=seq_sizes)
-    k = sp_group.all_to_all(k, scatter_dim=1, gather_dim=0, gather_sizes=seq_sizes)
-    v = sp_group.all_to_all(v, scatter_dim=1, gather_dim=0, gather_sizes=seq_sizes)
-    out = run_attention(q=q, k=k, v=v)  # [seq, Hq/P, D] (attends [UND-prefix | GEN])
+    # scatter heads, gather sequence: [seq/P, H, D] -> [seq, H/P, D], q, k and v in one
+    # exchange (the exchanges are latency-bound at small sizes)
+    q, k, v = sp_group.all_to_all((q, k, v), scatter_dim=head_dim, gather_dim=seq_dim, gather_sizes=seq_sizes)
+    out = run_attention(q=q, k=k, v=v)
     # scatter sequence, gather heads: [seq, H/P, D] -> [seq/P, H, D]
-    out = sp_group.all_to_all(out, scatter_dim=0, gather_dim=1, scatter_sizes=seq_sizes)
-    return out
-
-
-def _ulysses_attention_via_all_gather(
-    sp_group: CommGroup,
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    run_attention: Callable[..., torch.Tensor],
-) -> torch.Tensor:
-    """CUDA-graph-capturable Ulysses attention built from all-gather.
-
-    Same result as :func:`ulysses_attention`'s all-to-all, different collective:
-    each rank all-gathers the full sequence and attends over its own head-group,
-    then all-gathers the full heads back and keeps its own sequence shard. The
-    all-to-all would move fewer bytes, but it is grouped send/recv and does not
-    replay from a CUDA graph; all-gather is a true collective and does. Assumes
-    an even sequence split across the group (the captured resolutions guarantee
-    it). Head counts are divisible by the group size (the Ulysses constraint)."""
-    world_size, rank = sp_group.world_size, sp_group.rank
-    # [seq/P, H, D] -> all-gather sequence -> [seq, H, D] -> keep head-group rank
-    q_full = sp_group.all_gather(q, dim=0)
-    k_full = sp_group.all_gather(k, dim=0)
-    v_full = sp_group.all_gather(v, dim=0)
-    hq, hkv = q_full.size(1) // world_size, k_full.size(1) // world_size
-    q = q_full[:, rank * hq:(rank + 1) * hq, :].contiguous()
-    k = k_full[:, rank * hkv:(rank + 1) * hkv, :].contiguous()
-    v = v_full[:, rank * hkv:(rank + 1) * hkv, :].contiguous()
-    out = run_attention(q=q, k=k, v=v)  # [seq, Hq/P, D]
-    # [seq, Hq/P, D] -> all-gather heads -> [seq, Hq, D] -> keep sequence shard rank
-    out = sp_group.all_gather(out, dim=1)
-    sl = out.size(0) // world_size
-    return out[rank * sl:(rank + 1) * sl, :, :].contiguous()
+    return sp_group.all_to_all(out, scatter_dim=seq_dim, gather_dim=head_dim, scatter_sizes=seq_sizes)
 
 
 def sp_head_slice(sp_group: CommGroup, x: torch.Tensor) -> torch.Tensor:

mstar/model/cosmos3/components/transformer.py

@@ -403,7 +403,6 @@ def forward_gen(
         self, gen_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
         attend: AttentionCallable,
         seq_sizes: list[int] | None = None,
-        prefer_all_gather: bool = False,
     ) -> torch.Tensor:
         H, Hkv, D = self.num_attention_heads, self.num_key_value_heads, self.head_dim
         q = self.norm_added_q(self.add_q_proj(gen_seq).view(-1, H, D))
@@ -414,12 +413,9 @@ def forward_gen(
         # Ulysses all-to-all wraps the attend: gen_seq is sequence-sharded
         # across the SP group, so attention runs over the full sequence at
         # tp*sp head-degree, then the result is re-sharded back. Trivial SP
-        # group -> passthrough (byte-identical to the non-SP path). The captured
-        # denoise forward sets prefer_all_gather (the all-to-all does not replay
-        # from a CUDA graph; all-gather does).
+        # group -> passthrough (byte-identical to the non-SP path).
         out = ulysses_attention(
             self.sp_group, q, k, v, attend, seq_sizes,
-            prefer_all_gather=prefer_all_gather,
         ).reshape(-1, H * D)
         return self.to_add_out(out)
 
@@ -508,11 +504,10 @@ def forward_gen(
         self, gen_seq: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
         attend: AttentionCallable,
         seq_sizes: list[int] | None = None,
-        prefer_all_gather: bool = False,
     ) -> torch.Tensor:
         gen_norm = self.input_layernorm_moe_gen(gen_seq)
         attn_out = self.self_attn.forward_gen(
-            gen_norm, cos, sin, attend, seq_sizes, prefer_all_gather
+            gen_norm, cos, sin, attend, seq_sizes
         )
         residual = gen_seq + attn_out
         return residual + self.mlp_moe_gen(self.post_attention_layernorm_moe_gen(residual))
@@ -1051,7 +1046,7 @@ def commit_window(
         sin = torch.cat(sin_parts, dim=0) if len(sin_parts) > 1 else sin_parts[0]
         self._sp_run_gen_layers(gen_seq, cos, sin, label, attn)
 
-    def _sp_run_gen_layers(self, gen_seq, cos, sin, label, attn, prefer_all_gather=False):
+    def _sp_run_gen_layers(self, gen_seq, cos, sin, label, attn):
         """Run the generation layer stack, sequence-parallel-sharded across the
         SP group when active. ``gen_seq``/``cos``/``sin`` are the FULL sequence
         (identical on every SP rank); returns the FULL post-layer sequence.
@@ -1083,9 +1078,7 @@ def _sp_run_gen_layers(self, gen_seq, cos, sin, label, attn, prefer_all_gather=F
         sin = scatter_sequence(sp, sin, seq_sizes)
         for i, layer in enumerate(self.layers):
             attend.set_layer_idx(i)
-            gen_seq = layer.forward_gen(
-                gen_seq, cos, sin, attend, seq_sizes, prefer_all_gather
-            )
+            gen_seq = layer.forward_gen(gen_seq, cos, sin, attend, seq_sizes)
         return gather_sequence(sp, gen_seq, seq_sizes)
 
     def denoise_step(
@@ -1197,7 +1190,6 @@ def denoise_step_batched_cfg(
         sound_noisy_frame_indexes: list[torch.Tensor] | None = None,
         sound_mse_gen_indexes: torch.Tensor | None = None,
         sound_timesteps: torch.Tensor | None = None,
-        prefer_all_gather: bool = False,
         noisy_token_mask: torch.Tensor | None = None,
     ):
         """Conditional and unconditional generation in one batched pass.
@@ -1256,10 +1248,7 @@ def denoise_step_batched_cfg(
         cos = torch.cat([cos_c, cos_u], dim=0)
         sin = torch.cat([sin_c, sin_u], dim=0)
 
-        gen_seq = self._sp_run_gen_layers(
-            gen_seq, cos, sin, label, attn,
-            prefer_all_gather=prefer_all_gather,
-        )
+        gen_seq = self._sp_run_gen_layers(gen_seq, cos, sin, label, attn)
         gen_out = self.norm_moe_gen(gen_seq)
 
         def _decode(out):