[dghosef]

Draft. Stacked on #948 (perf/cfg-axis-core, the cfg_size axis), up to date with its head by a merge.

Wan2.2's dit had no sequence parallelism: its sp_size slot was used for CFG parallelism instead. With CFG now on its own cfg_size axis (#948), sp_size can mean what it means for every other model, Ulysses sequence parallelism, and with both axes a 4-GPU deployment can split the guidance branches and each branch's tokens. cfg2 × sp2 on 4 GPUs makes a 1280×704×81 video in 8.76 s, against 13.03 s for cfg2 on 2 (−33%), and cfg2's own output is byte-identical to the base.

Wan22DiT takes the node's sp_group. After patchify, each rank keeps a contiguous slice of the tokens, with their RoPE tables and per-token timesteps. Self-attention all-to-alls q/k/v around the ragged attention kernel (ulysses_attention; heads shard over SP through head_shards), while cross-attention, the FFN and the output head run on each rank's slice; the slices are gathered after proj_out, before unpatchify. The binding is LTX-2.5's SeqParallel, moved to mstar/model/components/distributed/sequence_parallel.py and given scatter / gather: Wan builds one per forward and passes sp.attend to its blocks, and LTX-2.5 imports it from there. The (q, k, v) -> out type that LTX-2.5 and LingBot-World each defined is one Attend beside RaggedAttentionFn, and it types Wan's attend parameters. The dit is declared SP-enabled. On one rank the scatter and gather do nothing and ulysses_attention calls the kernel directly, and Wan22DiT refuses an sp_size that doesn't divide its 24 heads, as Cosmos3 does. configs/wan22_cfg2_sp2_4gpu.yaml (cfg_size: 2, sp_size: 2) runs one guidance branch per pair of ranks, and Ulysses over each pair. The mesh is row-major over (cfg, sp, tp), so on ranks [0, 1, 2, 3] the SP pairs are [0, 1] and [2, 3], each inside one guidance branch, and the all-to-alls never cross branches. docs/models.rst lists it (split the guidance branches first with cfg_size, then add sp_size on top; sp_size must divide 24 heads; wan22_dp2.yaml stays the throughput layout), and docs/serving.rst's parallelism example shows the two axes composed. An old-style Wan config that set sp_size: 2 for the guidance split now runs Ulysses. That's correct, but slower than cfg_size: 2 on two GPUs (14.20 vs 13.03 s). Only the unmerged stack ever had that shape, so no released config changes. With Wan's dit SP-enabled, #948's placement hint ("split guidance branches with cfg_size instead") could no longer fire for any model, so this PR deletes it.

GPU rows were measured at be35b7409 against its base, #948 at 1ffb5c019; the head a641f029c differs from be35b7409 only by main merges. Eras 4× H200 (NVLink), torch 2.12.1+cu132, Wan2.2 TI2V-5B, 1280×704×81, 20 steps, guidance 5, a fixed seed, a fresh Inductor cache on every boot, 0 failed requests in every cell. Latency is c=1, 2 boots per layout in rotated order, each 1 warmup and 3 videos; peak VRAM was 42,749–43,567 MiB across layouts.

Check	Setup	Result
Latency, cfg2	2 GPUs	13.03 / 13.04 s per boot, against the base's 13.07 / 13.06 s (−0.2%, noise)
Latency, sp2 (Ulysses)	2 GPUs	14.20 / 14.25 s. With the dit uncompiled 18.03 / 18.01 s, so compiling saves 21%.
Latency, cfg2 × sp2	4 GPUs	8.76 / 8.78 s
Latency, sp4 (Ulysses)	4 GPUs	8.84 / 8.86 s
Output, cfg2	Against the base	Byte-identical video, and byte-identical across boots.
Output, SP layouts	PSNR against cfg2	Not bit-exact, though the exchange itself is exact (the CPU test matches the one-rank dit). sp2 26.1–26.5 dB, cfg2 × sp2 25.9–26.1 dB, sp4 25.6 dB. The same-code floor (two boots of the same layout: 1-GPU, sp2, cfg2 × sp2) is 28.0–28.6 dB, and the pass bar of floor − 3 is 25.04 dB, so every layout passes. Per-frame and row-band PSNR show no dip at the shard boundaries, and a contact sheet of five frames per layout is clean. Each layout is deterministic within a boot (two identical requests give identical bytes); the SP layouts differ only across boots.
Compiled vs eager dit	sp2, against 1 GPU	sp2 compiled vs eager 26.14 / 26.48 dB, against 25.76 / 27.28 dB for the same pair on 1 GPU: no compiled-path mis-lowering under SP.
Graph breaks and recompiles	TORCH_LOGS=graph_breaks,recompiles	1 GPU and cfg2 have no break in the dit. SP adds one break site, at the torch.compiler.disabled all-to-all. SP's extra recompiles all land in the first forward; a second request has 0 recompiles and 0 graph breaks on every layout.
Profile	p50 engine.forward per denoise step; nsys per rank	cfg2 500 ms (base 498), sp2 557, cfg2 × sp2 284, sp4 286. GPU-bound: kernel time per step is close to the step's wall time. sp2 − cfg2 ≈ 55 ms per step: the all-to-all is ~38 ms (against cfg2's 8.5 ms all-gather) and packing plus unfused elementwise work ~27 ms; the graph breaks cost ~1 ms of host time.
Throughput	4 GPUs, c=2 and c=4	cfg2 × sp2 0.1205 req/s at c=2 and 0.1197 at c=4, against 0.1438 req/s at c=4 for four data-parallel replicas, so data parallelism stays the throughput layout.
Rerun on 4× H100 (NVLink), head dacb760a9	cfg2 × sp2, sp4 and DP4, one boot each (cfg2 × sp2 two), NCCL_NVLS_ENABLE=0	cfg2 × sp2 8.81 s and sp4 8.88 s (within 1% of the H200 rows); PSNR 25.6–28.5 dB against 1 GPU and cfg2, inside the same-code floor (26.05 dB boot to boot, 27.9 dB between DP4 replicas); no seam at the shard boundaries; 0 recompiles and 0 graph breaks after warmup; DP4 0.183 req/s at c=4 against 0.123 for cfg2 × sp2; the 1-GPU output is byte-identical to the base. 0 failed.
LTX-2.5 SP2 (SeqParallel moved)	ltx2_5_sp2.yaml, one t2av request, 4 compiled boots per arm	Neither arm is bit-reproducible across boots. Cross-build median 34.5 dB, against a same-code minimum of 30.07 dB.
CPU tests	Linux CPU	A two-rank gloo test runs the dit under SP against the same dit on one rank in-process (the shipped sp_group=None default), on an even and an uneven token split. test_startup_graph_validation loads every shipped config, wan22_cfg2_sp2_4gpu.yaml included, through placement, which refuses it without the dit's sp_enabled_nodes. A three-rank gloo case of test_sequence_parallel (test_ulysses_attention_matches_one_rank) aborts at process teardown on some runs; its code (ulysses_attention, CommGroup.all_to_all) is the same as #948's head, where it fails intermittently too.
ruff; CPU Core vs main 9f84f18cd	Eras CPU box, OMP_NUM_THREADS=2, head a641f029c	ruff clean. CPU Core against main 9f84f18cd: 9094 passed, 1066 skipped, 2 failed, the same 2 as main (the test_eval_compare_multigpu dump_stacks pair); none fail only on this branch. GitHub CI green at a641f029c.

Not yet run on GPU: the full videos watched by eye (the contact sheets above are five frames per layout).

Follow-ups: a compile-traceable all-to-all. CommGroup.all_to_all runs through a torch.compiler.disabled _all_to_all_single, so every Ulysses exchange breaks the compiled graph, in Cosmos3, LTX-2.5 and Wan alike; making it traceable is a shared CommGroup change, not Wan work, and its gain is small: the breaks cost ~1 ms of host time per step, where the exchange and its packing cost ~55 ms. Cosmos3 keeps its own SP binding: it threads seq_sizes through each layer over an unbatched [S, H, D] stream (sequence dim 0), and its attention mixes the Ulysses exchange with a replicated prefix sliced by heads (sp_head_slice / sp_head_gather). SeqParallel is fixed to [B, L, H, D], so moving Cosmos3 onto it needs a layout field, which isn't a mechanical swap. Ulysses in the attention resource: the ragged resource already knows the node's SP group and sizes its plan at heads / (tp × sp), so doing the two all-to-alls in its run would leave each DiT only the boundary scatter and gather. Outside this PR: at c=4, two of the first four queued requests wait out the whole backlog. That also happens on one GPU and on #948's base, so it is the scheduler's service order, not the CFG or SP axis.

AI assistance (Claude) was used for this PR.

🤖 Generated with Claude Code