[NSagan271]

Stacked under #376 (LTX-2.5 TP/SP), which is rebased on this and uses it.

What changes

• CommGroup.all_to_all (mstar/distributed/communication.py) sends one packed buffer with all_to_all_single, instead of a list of contiguous pieces whose received counterparts it then concatenated. Uneven splits pass split sizes; negative dims are normalized first, as all_gather does. The collective stays behind torch.compiler.disable, so a compiled caller breaks around it and the packing and unpacking compile as their own frame. Tracing it as a functional collective (_functional_collectives.all_to_all_single) keeps one graph and is bit-exact in isolation, but served Cosmos3 SP2 then hit an illegal memory access in the batched capture's warm-up (and in review, the plain collective without the disable, or a custom_op wrapper, gave wrong batched images), so the disable stays.
• ulysses_attention takes seq_dim / head_dim (default 0, 1, the [S, H, D] layout it always assumed), so batched layouts such as LTX's [B, S, H, D] use it too. It sends q, k and v in one all-to-all: CommGroup.all_to_all takes several tensors that differ only in the scattered extent, writes rank r's piece of each one after another into rank r's block of the send buffer (one copy per tensor), and unpacks each on its own, so they come back contiguous and the copies are those of separate exchanges. That is two exchanges per attention instead of four, which matters because at small sizes each one is latency-bound (below). Different query and key/value head counts work, as long as each divides by the group size, as Ulysses already requires.
• prefer_all_gather and _ulysses_attention_via_all_gather are removed. They existed because the all-to-all was believed not to replay from a CUDA graph, so captured paths used an all-gather that moves P× the bytes. The all-to-all does replay (below). Cosmos3's captured denoise drops the flag and its plumbing, and the SP config comments say the capture includes the all-to-all.
• Cosmos3 under TP×SP captures every tier again. The cap to the smallest tier was for an illegal access seen with the all-gather exchange composed with the TP all-reduce; with the all-to-all, cosmos3_nano_tp2_sp2 on 4× H100 (with [Distributed][Bug fix] Make the members of a gloo subgroup agree on its name, so a TP x SP mesh starts #705, without which TP×SP doesn't start) captures 192×320 and 480×832, and its images match SP2's at 38 and 41 dB PSNR (run in review).

This is the exchange #376 wrote for LTX as its own sequence_to_heads / heads_to_sequence; review there asked for one all-to-all rather than two. #376 now deletes its copy and calls ulysses_attention.

Evidence that the all-to-all captures

On 2× H200, torch 2.12.1+cu132, NCCL 2.29.7:

• Raw NCCL: both dist.all_to_all (list) and all_to_all_single capture in torch.cuda.graph and replay bit-exact, including with new inputs copied into the static buffers.
• ulysses_attention under torch.compile(mode="max-autotune-no-cudagraphs"), captured: bit-exact against the compiled eager call and against one rank attending the whole sequence, for even and uneven splits, before and after new inputs.
• Served Cosmos3 SP2 t2i, against this PR's base (0b1a723):
  • Single requests: byte-identical at 320×192 and 832×480 (captured) and 1280×720 (eager), across six server boots.
  • Batched, the same 8 seeds sent 4 and 8 at a time: 18 of 32 images byte-identical; the rest at 21–45 dB, which is the base's own spread between its batched and single-request images (21–42 dB), since which requests share a batch varies. No image is corrupted. On the base, the batched captured denoise already used the list all-to-all.

Why the earlier conclusion was reached is not known; an older NCCL/torch, or autotuning during capture (which its docstring also mentions), are guesses.

Performance

Cosmos3-Nano, configs/cosmos3_nano_sp2.yaml, 2× H200, SHM arena, COSMOS3_GEN_CAPTURE_RES=192x320,480x832,720x1280 COSMOS3_GEN_CAPTURE_BS=1,4,8, so 1280×720 (over graph_max_latent_area) runs eager on both builds.

Single request, 50 steps, median of 8 after 2 warmups per server; base and this PR alternated, three server boots each, the mean of the three medians:

t2i, 1 request	base	this PR	Δ	2× H100 (in review, one boot per GPU pair)
320×192 (captured)	0.592 s	0.593 s	+0.3%	0.69 → 0.82 s, +19%
832×480 (captured)	0.991 s	0.911 s	−8.1%	1.07 → 1.14 s, +6%
1280×720 (eager)	3.43 s	3.28 s	−4.3%	3.88 → 3.62 s, −7%

That run was before q, k and v shared an exchange. Sharing it, against the version that sent them separately, alternated, three boots each:

t2i, 1 request	separate (538cbd7)	one exchange	Δ
320×192 (captured)	0.633 s	0.548 s	−13.4%
832×480 (captured)	0.927 s	0.871 s	−6.1%
1280×720 (eager)	3.44 s	2.78 s	−19.2%

Images byte-identical across the six boots, and 600 more requests (single and 8 at a time) without a failure. Against the base from the first run that is about −7%, −12% and −19% (across runs, so indicative). That run's version laid q, k and v interleaved by head and packed them twice; the head (6eb3f3c) packs once and returns each contiguous, and measured level with it on Cosmos3 (320×192 0.552 vs 0.544 s, 832×480 0.872 vs 0.866 s, 1280×720 2.760 vs 2.764 s, two boots each, images identical). On LTX-2.5 SP2 the interleaved version was ~2% slower than separate exchanges; the head is level with them: 19.52 / 21.72 / 22.30 videos/min at c=1/2/4 against 19.55 / 21.79 / 22.24, two-stage 5.78 / 6.31 against 5.77 / 6.35, video and audio byte-identical.

Where the time goes, from an nsys trace of the captured step (320×192, per request): the exchange's point-to-point ncclDevKernel_SendRecv took a median 13.9 µs a call against 8.4 µs for the base's AllGather_RING_LL, over ~7,200 calls, about 40 ms; GEMM, attention and fused-kernel time were unchanged, and both builds ran the same 50,419 kernels. Raising NCCL_P2P_LL_THRESHOLD to 1 MiB made it 2–3% slower, so the fix is fewer calls. Separately, ~240 ms of the ~650 ms request is the GPU idle between the 50 denoising steps (host work): MSTAR_TP_ASYNC_SCHED=1 cut 9% at 320×192 and 6% at 832×480 with identical images, which is outside this PR.

In the first table, the 832×480 gain is where the base captured the all-gather. 1280×720 runs eager, where the base already used the list all-to-all, so its gain is from sending one packed buffer instead. On H100 the captured sizes got slower; that is not reproduced on H200, and its cause is open: the exchange alone in a CUDA graph is faster than the base's (3.1 against 5.1 ms per forward at 320×192, measured in review), and on H200 the extra graph breaks cost nothing measurable in a captured replay (100 against 104 µs with and without).

From an earlier A/B on the same box (two boots each): t2i throughput at 832×480 c=1 1.07 against 0.98 req/s (+9.5%), at c=4 / c=8 −0.2% to +2.8%; t2v 189 frames at 320×192 / 832×480 / 1280×720 within 0.6%.

LTX-2.5 SP2 (configs/ltx2_5_sp2.yaml, 2× H200), #376 rebased on this, its own exchange vs ulysses_attention, alternated with two servers each: 19.6 / 21.9 / 22.3 videos/min at c = 1 / 2 / 4 on both (within ±1%), and the 544×960 video, the two-stage 1088×1920 video and both audio tracks byte-identical across all four servers.

Tests

• ruff check .: clean.
• CPU Core (python -m pytest test/modular test/distributed test/chatterbox test/kokoro test/omnivoice test/sampling_test mstar/model/*/tests benchmark/worker_phases, OMP_NUM_THREADS=2, on a 2× H200 box so GPU-guarded tests ran) at 26955cf: 7218 passed, 385 skipped, 21 failed; main (0b1a723) gives 7206 passed, 385 skipped, 21 failed, and the 21 failing tests are the same set (eval-compare job control, SHM shutdown cleanup, and a few capture tests). test/modular/test_waypoint_gpu.py was run separately (14 passed on both) because running it from several suites at once on one GPU hangs, on main too. At the head (538cbd7), test_sequence_parallel.py and the Cosmos3 tests: 218 passed, 28 skipped.
• New test/modular/test_sequence_parallel.py (gloo, CPU): CommGroup.all_to_all against the layout it promises on 2 and 3 ranks with even and uneven splits, for [S, H, D] and [B, S, H, D], with positive and negative dims; ulysses_attention against one rank attending the whole sequence, with equal and with fewer key/value heads than query heads. CommGroup.all_to_all of several tensors at once against one at a time. At the head (6eb3f3c), with the Cosmos3 tests: 226 passed, 28 skipped.
• GPU and served checks: above.

Not done here

• Whether the H100 slowdown remains with one exchange for q, k and v: not run on H100. The H200 trace above suggests its cause (the per-call latency of the point-to-point exchange), which the merged exchange halves.
• Under nsys with CUDA-graph node tracing, this branch failed twice with unspecified launch failure (after 4 and 9 traced requests) where the base ran ten; 1,000 requests without the profiler on this branch have not failed. Whether it is the profiler with point-to-point NCCL inside graphs or a rarer bug is open.
• The 4-GPU SP configs (cosmos3_nano_sp4, cosmos3_super_tp2_sp2) were not run; cosmos3_nano_tp2_sp2 was, in review (above).
• LTX's attend_shards (the log-sum-exp merge for a replicated query over sharded keys) stays in LTX-2.5: tensor- and sequence-parallel DiT and the two-stage distilled recipe #376.

Developed with AI assistance (Claude Code).

🤖 Generated with Claude Code

[mstar-review]

Automated review

🟡 No findings on 26955cf, light review

Description

• The benchmark section names the config, client directory, capture env vars and hardware, but not the commands that were run, the number of requests per measurement or the warmup, so the c=1 t2i numbers (−8.5% and −7.1%) can't be reproduced as written.
• The description says the 4-GPU SP configs were not run, yet this PR rewrites their comments (cosmos3_nano_sp4.yaml, cosmos3_nano_tp2_sp2.yaml) to say the captured graph includes the all-to-all. That is true of the code path but hasn't been exercised on those configs; one sentence in the description saying so would help.

Questions

• Under TP×SP (cosmos3_nano_tp2_sp2, cosmos3_super_tp2_sp2), the smallest-tier denoise graph is still captured. Before this PR it captured the all-gather exchange; now it captures the all-to-all next to the TP all-reduce, and neither 4-GPU config was run. Has any TP×SP capture of the all-to-all been run at the smallest tier? If it hits the same kind of illegal access that ruled out 480p+, those configs fail at startup rather than falling back to eager.

What the review checked

• Even path in CommGroup.all_to_all: unflatten(scatter_dim, (P, -1)) then movedim(scatter_dim, 0) makes piece r go to rank r. On receipt, movedim(0, gather_dim) then flatten(gather_dim, gather_dim+1) concatenates the pieces along gather_dim in rank order. This holds whether scatter_dim is before or after gather_dim.
• Uneven path: the send splits are the chunk numels in rank order. The shape received from rank r is this rank's own chunk shape with gather_dim set to gather_sizes[r]. The recv splits and the views match that, both for the to-heads exchange (even heads, uneven seq gather) and for the back exchange (uneven seq scatter, gather_sizes=None).
• When the divisibility assert was removed, the even path stayed safe: unflatten(..., -1) still raises on an indivisible scatter dim.
• ulysses_attention with the defaults seq_dim=0, head_dim=1 is equivalent to the old [S, H, D] calls. Cosmos3 is the only in-tree caller, and it passes positionally only up to seq_sizes.
• All prefer_all_gather plumbing (attention, layer, _sp_run_gen_layers, denoise_step_batched_cfg, forward_captured) is removed, with no remaining references in pr-head/. Warmup, capture and replay now all trace the same exchange, so the old concern about autotuning mid-capture doesn't arise from a mode switch.
• No compiled caller uses fullgraph=True around ulysses_attention (Cosmos3 compiles the denoise with fullgraph=False), so the torch.compiler.disable graph break around _all_to_all_single is allowed.
• The new test/modular/test_sequence_parallel.py is CPU and gloo only, lives under a CPU_TEST_PATHS directory, and has a unique basename.
• No stale comment still claims that the all-to-all can't be captured; the remaining all-gather mention in cosmos3/submodules.py is correctly in the past tense.

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[Gaurav-Shah05]

Thanks for the follow-up, the shape is right, but the captured single-request t2i got slower on H100, the opposite of the description's numbers.

On 2 H100s with configs/cosmos3_nano_sp2.yaml and your capture settings, each build booted once on each GPU pair, median of 8 requests after 2 warmups, with images identical to main:

	main	this PR	Δ
t2i 320x192 (captured)	0.69 s	0.82 s	1.19x ↑
t2i 832x480 (captured)	1.07 s	1.14 s	1.06x ↑
t2i 1280x720 (eager)	3.88 s	3.62 s	1.07x ↓
c=1 throughput, 832x480	0.90 req/s	0.86 req/s	1.04x ↓

• The slowdown isn't the exchange. Alone in a CUDA graph, the packed all-to-all beats main's all-gather path (3.1 against 5.1 ms per forward at 320x192), so the extra time is elsewhere in the captured step.
• The 1280x720 gain. The description credits the gains to dropping the captured all-gather, but this size is over graph_max_latent_area on SP2, so both builds run it eager, where main already used the all-to-all. Its gain comes from sending one packed buffer instead.

[Gaurav-Shah05]

With the torch.compiler.disable on _all_to_all_single, dynamo compiles all_to_all as a frame of its own, so the packing and unpacking end up in separate graphs, not in the caller's.

In my runs both ways of keeping the collective in the caller's graph, deleting the disable or wrapping it in a torch.library.custom_op as ar_fusion.py does, made the served SP2 images wrong.

So the disable is needed.

[NSagan271]

Agreed, and confirmed: compiling a caller on gloo gives four graphs (producer, packing, unpacking, consumer), so the docstring's claim was wrong. I also tried the one variant you hadn't, _functional_collectives.all_to_all_single with no disable: one graph, and bit-exact against full attention in isolation under compile + capture on 2×H200 (even and uneven), but served Cosmos3 SP2 hit an illegal memory access in FlashInfer's prefill during the bs=8 capture warm-up on all three boots. So the disable stays; 3110916 corrects the docstring and the description says what the boundary costs.

[Gaurav-Shah05]

This is a bug for negative dims. On two ranks, scatter_dim=-2 or gather_dim=-3 returns a wrong tensor with no error, where main returned the right one. unflatten adds a dim, so movedim(scatter_dim, 0) moves the piece's own dim instead of the rank dim, and movedim(0, gather_dim) goes wrong the same way. It should be fixed if all_to_all normalize both dims first, as all_gather does.

[NSagan271]

Thanks, reproduced on two gloo ranks: all four negative combinations wrong on this branch, all right on main. Fixed in 3110916 by normalizing both dims first, as all_gather does; the exchange tests now run with negative dims too, and four of them fail without the fix. No in-tree caller passed negative dims, so served output was unaffected.

[Gaurav-Shah05]

With #705 under it and this cap removed, configs/cosmos3_nano_tp2_sp2.yaml captures both 192x320 and 480x832 on 4 H100s. The images match the SP2 ones at 38 and 41 dB PSNR, so the two_d cap and this comment can go.

[NSagan271]

Removed in 538cbd7, citing your 4×H100 run (and #705 for startup) in the commit and description. I can't run 4 GPUs here, so that run is the evidence.

[dghosef]

pushed some minor test / doc changes, lgtm once you resolve with Gaurav

[NSagan271]

I profiled the captured step on H200 (nsys with CUDA-graph node tracing), base 0b1a723 against 538cbd7, per 320×192 request:

• Same kernels: both builds run 50,419 kernels, and GEMM, attention and fused-kernel time are the same. So the extra graph breaks don't add work.
• The exchange is the difference. The base's ncclDevKernel_AllGather_RING_LL takes a median 8.4 µs a call. This branch's point-to-point ncclDevKernel_SendRecv takes 13.9 µs. Over ~7,200 calls a request that's about 40 ms. On H200 it was offset by less idle between steps; on H100 it may not be, which would fit your +19% at 320×192. That's a guess, not measured.
• Raising NCCL_P2P_LL_THRESHOLD to 1 MiB didn't help. It was 2–3% slower, so the fix is fewer calls.

So q, k and v now go in one exchange (6eb3f3c): CommGroup.all_to_all takes several tensors and packs them into one buffer, each coming back contiguous. That makes two exchanges per attention instead of four. Results on 2× H200, alternated:

	separate exchanges	one exchange
Cosmos3 SP2 t2i, 1 request, 320×192	0.633 s	0.548 s (−13%)
832×480	0.927 s	0.871 s (−6%)
1280×720 (eager)	3.44 s	2.78 s (−19%)
LTX-2.5 SP2 videos/min, c=1/2/4	19.55 / 21.79 / 22.24	19.52 / 21.72 / 22.30 (level)

The Cosmos3 rows come from three boots each, measured with the earlier one-exchange layout (6f298da). The final layout measured level with it on Cosmos3; it replaced it because it is level with separate exchanges on LTX, where the earlier layout was ~2% slower. The LTX row is the final layout, two boots each. Images and videos are byte-identical throughout, and 1,000 more Cosmos3 requests (single and 8 at a time) completed without a failure.

If you can rerun your H100 comparison on 6eb3f3c, that would show whether the slowdown is gone there. I'll also run it on laude.

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: #897 on 2× H100 (Gaurav's H100 slowdown)

Cosmos3-Nano, configs/cosmos3_nano_sp2.yaml, 2× H100 80GB, SHM arena, COSMOS3_GEN_CAPTURE_RES=192x320,480x832,720x1280 COSMOS3_GEN_CAPTURE_BS=1,4,8, Rust runtime. t2i, 1 request, 50 steps, median of 8 after 2 warmups. Builds alternated, 3 server boots each, quiet host (load 17–46):

t2i, 1 request	base 0b1a723e5	26955cf11 (Gaurav's run)	head 6eb3f3c6f
320×192 (captured)	0.723 / 0.726 / 0.734 s	0.867 / 0.879 / 0.848 s (+19%)	0.731 / 0.734 / 0.728 s (+0.4%)
832×480 (captured)	1.101 / 1.102 / 1.111 s	1.165 / 1.194 / 1.172 s (+6.5%)	1.027 / 1.020 / 1.015 s (−7.6%)
1280×720 (eager)	4.055 / 4.607 / 4.053 s	3.874 / 3.813 / 3.800 s (−6%)	3.285 / 3.237 / 3.263 s (−19.5%)

• Gaurav's slowdown reproduces on 26955cf11: +19% / +6.5% / −6%, against his +19% / +6% / −7%.
• It's gone at the head: level at 320×192, faster at the other two sizes. Every head boot beats every base boot at 832×480 and 1280×720.
• It was GPU time, not host time. nsys, graph-level trace, 8 requests per size, one boot per build. Graph replay per request (rank 0):

	base	26955cf11	head
320×192	470 ms	502 ms (+6.6%)	450 ms (−4.4%)
832×480	793 ms	730 ms (−8%)	687 ms (−13%)

Idle time between graphs is about the same on all three (286–301 ms per 320×192 request), so the extra graph breaks cost nothing visible.

• Images are byte-identical across all 21 boots at every size, and no server logged errors.
• An earlier set of 3 boots per build ran while the host was heavily loaded (load average 150–1,900). Its ranking was the same, but the deltas were smaller and much noisier, so the table uses the quiet set only.

[Gaurav-Shah05]

@NSagan271 Thanks for the profile and the live test. I ran them as well and the slowdown is gone at the head. You were also right that it was the exchange, not something else in the step as I said before.

I also ran the 4-GPU configs, and SP4 with batched captures gives wrong images.

With configs/cosmos3_nano_sp4.yaml and COSMOS3_GEN_CAPTURE_BS=1,4,8, every 320x192 request came back as a damaged picture, a different one each time for the same seed.

I tried to pin down the cause and It isn't your exchange, though.

• NCCL registers the buffers that CUDA graphs capture (NCCL_GRAPH_REGISTER, on by default) whenever the memory comes from expandable_segments:True, which benchmark/cosmos3/reproduce.sh sets.

• With registration on, whether a graph comes out right depends on where its exchange buffers land in the memory pool the captures share.

• Your first version breaks the same way once its buffers are kept alive a little longer, and with NCCL_GRAPH_REGISTER=0 every variant matches main.

I'm opening a PR that turns it off by default on CUDA workers. With it, this head's SP4 images match main's exactly and the SP2 speedups hold, so this can go in right after it. I will post the numbers in the description there. Let me know your thoughts.

[NSagan271]

I'm opening a PR that turns it off by default on CUDA workers. With it, this head's SP4 images match main's exactly and the SP2 speedups hold, so this can go in right after it.

@Gaurav-Shah05 thanks for hunting down the bug! Yes, lmk when you make a PR and I can review it. I guess the long-term solution would be to make sure the exchange buffers actually land in the right memory locations, but setting NCCL_GRAPH_REGISTER=0 by default sounds like a good solution for now.

[NSagan271]

@Gaurav-Shah05 merged the PR!

[Gaurav-Shah05]

For the long term, do you think it's better to look for what breaks inside NCCL's registration rather than where the buffers live? the allocator can put them anywhere, and some placements just trigger it.

[Gaurav-Shah05]

good for merge from my end

[NSagan271]

For the long term, do you think it's better to look for what breaks inside NCCL's registration

Yeah, I think so. Lmk if you want me to take a look?