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

@@ -44,6 +44,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.
 
@@ -230,34 +239,45 @@ 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``).
+
+        One ``all_to_all_single`` over a packed buffer, 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.
         """
         if self.world_size == 1:
             return input_
         world_size = self.world_size
+        # positive dims: unflatten below adds a dim, which would shift a negative one
+        scatter_dim %= input_.dim()
+        gather_dim %= input_.dim()
+        even_scatter = scatter_sizes is None or len(set(scatter_sizes)) == 1
+        even_gather = gather_sizes is None or len(set(gather_sizes)) == 1
+        if even_scatter and even_gather:
+            # [.., N, ..] -> [P, .., N/P, ..]: piece r goes to rank r, and the
+            # piece received from rank r has the same shape
+            send = input_.unflatten(scatter_dim, (world_size, -1))
+            send = send.movedim(scatter_dim, 0).contiguous()
+            recv = _all_to_all_single(self.device_group, send, torch.empty_like(send))
+            return recv.movedim(0, gather_dim).flatten(gather_dim, gather_dim + 1)
         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"
-            )
-            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)
-        ]
+            scatter_sizes = [input_.size(scatter_dim) // world_size] * world_size
+        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]] * 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,27 @@ 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])
+    to_heads = dict(scatter_dim=head_dim, gather_dim=seq_dim, gather_sizes=seq_sizes)
+    q = sp_group.all_to_all(q, **to_heads)
+    k = sp_group.all_to_all(k, **to_heads)
+    v = sp_group.all_to_all(v, **to_heads)
+    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):

mstar/model/cosmos3/submodules.py

@@ -2187,11 +2187,6 @@ def get_accelerator_graph_configs(self, device, tp_world_size: int = 1):
         else:
             capture_batch_sizes = list(self.gen_capture_batch_sizes)
         dtype = self.transformer.proj_in.weight.dtype
-        # The 2D (TP x SP) captured denoise graph hits a CUDA illegal access at
-        # 480p+ in the combined Ulysses all-gather + TP all-reduce path (pure-TP and
-        # pure-SP capture cleanly; only the combination at this scale). For 2D,
-        # capture only the smallest tier; 480p+ runs eager (correct, just no graph).
-        two_d = self.transformer.sp_group.world_size > 1 and self.transformer.comm_group.world_size > 1
         self._capture_layout: dict[tuple, dict] = {}
         configs = []
         # The capture is bit-faithful at every resolution (the rotary uses a
@@ -2201,8 +2196,6 @@ def get_accelerator_graph_configs(self, device, tp_world_size: int = 1):
         # Capture only below COSMOS3_GRAPH_MAX_LATENT_AREA (latent H*W).
         max_area = int(os.environ.get(
             "COSMOS3_GRAPH_MAX_LATENT_AREA", self.config.graph_max_latent_area))
-        if two_d:
-            max_area = min(max_area, 1000)  # 256p latent 240 captures; 480p 1560 does not
         # Video tiers (height, width, frames): a plain clip length and/or the
         # windowed rollout's window; one graph per latent shape serves t2v,
         # i2v and chained windows (see ``_capture_key``).
@@ -2389,15 +2382,12 @@ def forward_captured(
         layout = self._capture_layout[tuple(latents.shape[1:])]
         rids = engine_inputs.request_ids
         if latents.shape[0] == 1:
-            # Captured into the denoise CUDA graph. Under SP the Ulysses
-            # exchange must be all-gather (all-to-all is grouped p2p send/recv —
-            # not graph-replayable); the flag holds across warmup/capture/replay
-            # so those kernels compile during eager warmup. No-op without SP.
+            # Captured into the denoise CUDA graph.
             cond_v, uncond_v = self.transformer.denoise_step_batched_cfg(
                 latents[0], vision_timesteps[0], position_ids_cond[0], position_ids_uncond[0],
                 layout["vision_token_shapes"], layout["vision_noisy_frame_indexes"],
                 layout["mse_gen_indexes"], CFG_BATCHED_LABEL, attn,
-                prefer_all_gather=True, noisy_token_mask=noisy_token_mask[0],
+                noisy_token_mask=noisy_token_mask[0],
             )
             return {rids[0]: {"cond_v": [cond_v], "uncond_v": [uncond_v]}}
         reqs = [