[mickqian]

Motivation

LTX-2/2.3 two-stage runs stage 2, the distilled refiner, without CFG. Under --enable-cfg-parallel, LTX2RefinementStage was MAIN_RANK_ONLY, or MAIN_RANK_ONLY_AND_SEND_TO_OTHERS with parallel VAE decode. So on 2 GPUs one GPU sat idle for all of stage 2. At 768x512x121 on 2xH200 that is 3 steps at 0.394 s each.

TensorRT-LLM runs stage 2 as Ulysses 2 on the same ranks. In a 4xB300 benchmark round, sglang's stage 2 took 0.143 s/step x3 under breakable CUDA graphs, while TensorRT-LLM took 0.30 s for the whole of stage 2.

Sharding only pays off once stage 2 is long enough. At CI's 768x512x25 (1536 stage-2 tokens), it made LTX2RefinementStage slower, so the new path is gated on the stage-2 sequence length.

Modifications

At stage-2 entry, every CFG rank already holds the same inputs:

• stage 1 combines the guided prediction with an all-reduce, so all CFG ranks hold the same latents after every step;
• the latent upsampler is replicated;
• stage 2 draws its re-noise from a generator re-seeded from the request seed. LTX-2 resets the batch generators and LTX-2.3 builds a fresh one, so every rank draws the same noise.

So a stage without CFG can be sequence-sharded over the CFG ranks instead of leaving them idle:

• parallel_state.py: when cfg > 1 and sp == 1, the CFG ranks also form one pure-Ulysses group (get_cfg_sequence_parallel_group()). SequenceParallelGroupCoordinator.over_group builds it as a view over the CFG group's communicators. It adds no process group or NCCL communicator, and destroying the CFG group releases it.
• use_sequence_parallel_group() points the SP accessors at that group for one scope. Everything that reads get_sp_group() follows it: latent shard and gather, RoPE offsets, the Ulysses all-to-alls. The exchanges use the usual SP transports: the 2-rank CUDA-IPC all-to-all by default, otherwise NCCL.
• stages/base.py: new StageParallelismType.SEQUENCE_PARALLEL_OVER_CFG_RANKS. parallel_executor.py runs such a stage on every CFG rank inside that scope.
• The executor reads each stage's parallelism per batch through the new PipelineStage.parallelism_type_for(batch), which defaults to parallelism_type.
• LTX2RefinementStage.parallelism_type_for returns the new type only when the group exists and stage 2 has at least CFG_SP_MIN_TOKENS_PER_RANK (2304) video tokens per rank. It reads the sequence length of the packed stage-2 latents, which every rank holds with the same shape, so all ranks pick the same path. Shorter requests keep the previous main-rank path, and so do CFG ranks that have SP of their own (sp > 1). A warmup takes the same path as serving at the same shape, so breakable CUDA graphs capture what serving replays.
• On the new path, stage 2 ends with the SP gather, so every rank holds the full latents for decode. With parallel decode, the post-stage batch broadcast no longer runs.

Tests: test/unit/test_sequence_parallel_over_cfg_ranks.py runs 2 gloo ranks on CPU with cfg=2, sp=1. It checks:

• the group is a view over the CFG communicators (same device group, Ulysses 2, ring 1), and the global SP size stays 1;
• the executor runs an LTX2RefinementStage batch one token per rank short of the threshold on the main rank only, and the other rank keeps its batch;
• a batch at the threshold runs on both ranks with SP world 2, and SP 1 is restored afterwards;
• inside the scope, LTX-2's video and audio latents round-trip exactly through shard and gather, the shard's start frame follows the rank, and Ulysses attention matches single-rank SDPA for both sharded and replicated inputs.

No files are shared with #45942, #45943, #45944 or #45948, and this PR merges cleanly with all four. With parallel VAE decode, stage 2 now leaves the full latents on both ranks instead of broadcasting the batch. #45942 changes the decode itself, so the two PRs are independent. Stage 2's exchange runs on the 2-rank IPC transport, which #45948 (N-rank transport) does not touch.

Accuracy Tests

2xH200, LTX-2 two-stage 768x512x121, --enable-cfg-parallel, against main:

• PSNR 33.4–35.7 dB, high-frequency energy unchanged. The output is not bitwise equal: stage-2 attention now runs on two head shards, which reorders the sums.
• Breakable CUDA graphs: no hang, and the output is bitwise equal to eager in every arm. Main captures 3 graphs, with stage 2 on rank 0 only. This PR captures 4, with stage 2 on both ranks at (1, 3072, 128).

Speed Tests and Profiling

2xH200, LTX-2 two-stage 768x512x121, --enable-cfg-parallel. Server benchmark with a fresh server per run, arms interleaved.

	main	this PR
stage-2 step, eager	0.394 s	0.280 s (−29%)
stage-2 step, breakable CUDA graphs	0.365 s	0.2345 s (−36%)
e2e, default profile (eager)	7.336 s	6.877 s (−6.3%)
e2e, breakable CUDA graphs	5.966 s	5.386 s (−9.7%)
peak memory GPU0 / GPU1, breakable CUDA graphs	85.2 / 78.1 GB	80.8 / 80.6 GB

• The compile profile went from 10.508 s to 9.757 s, but that was measured with this change stacked on other pending changes.
• Eager peak memory is 0.8–1.8 GB higher.
• VAE decode is slightly slower on this branch under breakable CUDA graphs: 0.665 s to 0.715 s. The e2e numbers above include it.
• 121 frames is above the threshold below, so the gate keeps these numbers. LTX-2.3 is measured in the next section.

Sequence-length gate

2xH200, two-stage 768x512, --enable-cfg-parallel --dit-layerwise-offload false --performance-mode speed, eager. Server benchmark, one warm request per shape, then medians of the measured requests. LTX-2 25/57/89/121 frames ran in 2 interleaved rounds (n≥4); the other shapes ran on one server per arm (n=2–4). The server rounds the frame count so the latent frames split evenly over 2 GPUs, so the table lists the served frames. "Step" is [LTX2RefinementStage] average time per step, and "e2e" is the server's request time.

LTX-2:

frames (served)	stage-2 tokens	per rank	main step	ungated step	gated step	main e2e	ungated e2e	gated path
25	1536	768	203 ms	277 ms	200 ms	6.21 s	6.40 s	main rank
57	3072	1536	210 ms	280 ms	209 ms	6.48 s	6.68 s	main rank
73	3840	1920	254 ms	280 ms	254 ms	6.74 s	6.65 s	main rank
89	4608	2304	304 ms	282 ms	274 ms	6.93 s	6.75 s	SP over CFG
105	5376	2688	363 ms	282 ms	–	7.04 s	6.80 s	SP over CFG
121	6144	3072	397 ms	280 ms	277 ms	7.38 s	6.91 s	SP over CFG

• The ungated step stays at about 280 ms at every length. Eager stage 2 is host-bound at these sizes, and Ulysses adds about 75 ms of host work per step. So the new path wins only once the main rank's step passes that floor, between 1920 and 2304 tokens per rank (about 2.1k by interpolation). The threshold is 2304, the first measured win (−7% per step).
• End to end, the crossover is lower, at about 1.6–1.9k tokens per rank. Between stage 2 and decode, the main-rank path also spends 80–160 ms broadcasting the batch, which the new path skips. At 1920 tokens per rank, the gate gives up about 30 ms e2e to keep the faster stage 2.
• Stage 1 varies by about ±100 ms between server launches. That is why the e2e columns are noisier than the step columns.
• CI's ltx_2_two_stage_t2v settings (default serve args, 768x512x25): LTX2RefinementStage takes 590 ms on main, 841 ms ungated, and 609 ms gated.

LTX-2.3 (Lightricks/LTX-2.3, same flags):

frames (served)	per rank	main step	ungated step	gated step	main e2e	ungated e2e	gated path
25	768	147 ms	221 ms	145 ms	5.61 s	5.76 s	main rank
73	1920	226 ms	224 ms	–	6.16 s	6.18 s	main rank
89	2304	265 ms	282 ms	287 ms	6.57 s	6.49 s	SP over CFG
105	2688	320 ms	307 ms	–	7.67 s	7.50 s	SP over CFG
121	3072	352 ms	229 ms	226 ms	8.16 s	7.66 s	SP over CFG

• On LTX-2.3, the ungated step does not grow monotonically with length: 224, 282, 307 and 229 ms at 1920 to 3072 tokens per rank. At 2304 tokens per rank, the new path makes its stage 2 6% slower. The request is still about 74 ms faster, because the batch broadcast is skipped.

Breakable CUDA graphs, with warmup at the serving shape (LTX-2, n=2):

frames	main step / graphs	ungated step / graphs	gated step / graphs	main e2e	ungated e2e	gated e2e
25	87.1 ms / 3	91.5 ms / 4	87.4 ms / 3	2.53 s	2.39 s	2.52 s
121	367.4 ms / 3	236.6 ms / 4	237.2 ms / 4	6.02 s	5.44 s	5.43 s

• The gated server captures the graphs of the path it serves: 3 graphs at 25 frames (stage 2 on rank 0 only), 4 at 121 frames.
• Graphs remove the host floor, so at 25 frames the new path is only 5% slower per step. The request is 141 ms faster because the batch broadcast is skipped. The eager-measured threshold gives up that win for short requests under graphs.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #39951066058
Latest PR Test (Extra): ✅ Run #39951066122
Latest PR Test (AMD ROCm 10): ❌ Run #39951066039

[mickqian]

/tag-and-rerun-ci

[mickqian]

/tag-and-rerun-ci extra