[NSagan271]

On main, which now has #897 (the shared Ulysses all-to-all) and #375.

Developed with the AGENTS.md / skills harness from #370.

TP and SP

The dit is TP-enabled; its linears and across-heads qk-norms already shard by head, and the ragged resources shard with the engine's group. configs/ltx2_5_tp2.yaml: the DiT over both GPUs, text encoder on rank 1, decoders and upsampler on rank 0.

It is also SP-enabled (Ulysses): configs/ltx2_5_sp2.yaml gives each rank the whole DiT and half the video tokens, same placement otherwise. Only the video self-attention needs every video token, so it is the only exchange: it goes through the shared ulysses_attention (#897), one all_to_all_single each way, captured in the CUDA graph. The cross-attentions from the video run on the rank's own tokens against the replicated text/audio, the audio stream's attentions run whole on each rank, and the audio queries over the video attend the rank's shard through the same cross-attention resource, asking it for the log-sum-exp (below), and attention_over_sharded_keys (shared, beside ulysses_attention) merges the ranks' results. The ragged resources those use keep all their heads over SP (RaggedAttentionConfig.shard_over_sp=False) and get the rank's video span. Per denoise step this is 38 ms of collectives against 122 ms when every attention exchanges (2×H100). Moving from this PR's own exchange to the shared one changed nothing measurable: on 2×H200, 19.6 / 21.9 / 22.3 videos/min at c = 1 / 2 / 4 either way, and byte-identical 544×960 and two-stage 1088×1920 video and audio.

544×960×121, 2×H200	1 request	2 concurrent	4 concurrent
diffusers	6.54 s	9.2/min	9.2/min
SGLang 0.5.22, TP2	5.15 s	12.0/min	11.9/min
SGLang 0.5.22, Ulysses 2	4.86 s	12.5/min	12.4/min
vllm-omni 0.30.0, TP2	4.06 s	16.8/min	16.8/min
vllm-omni 0.30.0, Ulysses 2	3.42 s	20.5/min	20.5/min
M* TP2	3.37 s	19.6/min	20.4/min
M* SP2	3.04 s	21.9/min	22.3/min

How these were measured (this table and the two-stage one below): this PR at 2c40491, the distilled checkpoint in bf16 decoding with the conv VAE, benchmark/ltx2_5/bench.py (--backend mstar|vllm-omni|sglang) closed loop with the same 8 prompts and seeds, 20 measured requests per level (10 for two-stage) after warmup. Each row is the mean of two server processes, which agree within 3%; no request failed.

• vllm-omni 0.30.0 with vllm 0.30.0: --usp 2 or --tensor-parallel-size 2, with ltx2_use_conv_vae (its default decoder is the diffusion decoder).
• SGLang 0.5.22 diffusion: --ulysses-degree 2 or --tp-size 2, --enable-torch-compile, CPU offload off (by default it offloads the DiT and text encoder for video, even on an H200). sglang-omni 0.1.9 does not serve LTX-2.5.
• diffusers: sequential calls with device_map="balanced", so its throughput is 1/latency.

Served PSNR against the oracle, recorded on 2×H200 (test/ltx2_5/compare_media.py): SP2 31.5 dB, TP2 31.0 dB; the reference against itself under the other SDPA kernel 32.1 dB. On 2×H100 with that box's recording it was SP2 32.3 dB, TP2 32.1 dB, reference 32.1 dB.

Against vllm-omni's Ulysses-2 serving, M* SP2 has 11% lower single-request latency and 7% / 9% higher throughput at 2 / 4 concurrent; SGLang's throughput stays flat as concurrency rises. M* runs bf16 throughout; the DiT step is ~70% GEMM, and a sizing experiment with rowwise FP8 linears cut the SP2 step by about a quarter.

Five more changes in this PR:

• The mp4 is encoded at CRF 18 with x264's ultrafast preset on all threads, as Cosmos3 and vllm-omni do, from one contiguous copy of the frames: 0.23 s instead of 1.2 s per 544×960×121 request on an H100, and less lossy (oracle frames through the codec: 39.9 dB vs 39.0 dB).
• vae_compile (on in the configs, off by default and in the parity harness) compiles the video VAE's decoder, warmed at load for vae_compile_shapes, which is separate from the DiT's capture shapes (default: the capture shapes; the SP2 config adds the two-stage 1088×1920, which its DiT doesn't capture); other shapes decode eager. The decoder rather than vae.decode is compiled, so a tiled decode compiles four tile shapes, and the compiled path runs a channels_last_3d copy of the weights (777 MB). On an H100: 544×960×121 371 → 240 ms, tiled 1088×1920 1988 → 1315 ms, ~55 dB from eager; served PSNR is unchanged. Load takes ~1.5 min longer.
• RaggedAttentionConfig.workspace_mb lets a ragged resource size its FlashInfer workspaces. LTX's six labels × two slots held 6 GB at the 512 MB default; it declares 128 MB (1.5 GB), which serves 544×960 and the two-stage 1088×1920. This is what makes SP2 fit next to the text encoder.
• Ragged attention returns the log-sum-exp on request: run(..., return_lse=True) on the wrappers, the managers and RaggedAttentionCallable also gives [total_tokens, H] fp32 in natural log (FlashInfer's lse_base="ln") from the same plan (FlashInfer backend; Add an opt-in ragged FA3 varlen backend for cacheless attention. #863's DENSE backend raises if asked). SP's audio-over-video merge uses it, so no layer picks a kernel, and gathers each rank's output and log-sum-exp in one collective. SP2 on 2×H200, alternated with the version that called torch's flash kernel directly (two servers each): 22.3 videos/min at c=4 either way, single-request 3.05 s vs 3.02 s, two-stage 10.3 s vs 10.2 s.
• flashinfer-python is pinned to 0.7.2 in every extra (as Host overhead 5/N: cheaper resource planning, admit and commit #852 does): return_lse passes lse_base, which needs 0.7.1 or later.

Placement sweep around the TP2 DiT (text encoder / decoders on rank 0 or 1): within noise.

Two-stage recipe

recipe: "two_stage": stage 1 at half size (8 sigmas) → the x2 latent upsampler (new latent_upsampler node) → 3 sigmas at full size from the upsampled latents blended with noise at σ0. The reference threads one generator through both stages, so stage 2 replays stage 1's draws first and continues the stream. Uses #373's initial_loop_back and per-walk schedule. Frames above 720p decode in VAE tiles, as the card does for this decode. The refine walks replay the single-stage capture of their size, since the DiT step is the same forward; the TP2 config captures 544×960 and 1088×1920.

The request's geometry lives under kwargs["request"], not under the step-metadata keys: the conductor merges each pass's step_metadata into metadata.kwargs, so a request fact stored under a step-metadata key is overwritten by the previous walk's value. adding_models.rst now says so.

1088×1920×121, 2×H200	1 request	2 concurrent
diffusers	20.5 s	2.9/min
SGLang 0.5.22, Ulysses 2	15.0 s	4.0/min
vllm-omni 0.30.0, Ulysses 2	10.6 s	6.3/min
M* TP2, both stages captured, compiled VAE	10.6 s	6.2/min
M* SP2, stage 2 uncaptured, compiled VAE	9.7 s	6.75/min

Correctness (record_oracle_two_stage.py, parity_harness.py two_stage_seed): the upsampler and both of stage 2's starting tensors (generator continuation included) are bit-exact with the oracle. Served video against an H200 oracle, through H.264 (codec floor 40.6 dB): 29.8 dB on TP2 and SP2. On H100 earlier: 30.5 dB against 30.2 dB reference-vs-itself.

Tests

On the rebased head (2×H200, torch 2.12.1+cu132, FlashInfer 0.7.2), on main at b47e35b; after a last rebase over #917 (tests and CI only), ruff and the GPU files below were rerun on the pushed 8a26112 with the same result:

• ruff check . clean.
• CPU Core with the GPUs hidden (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): 7782 passed, 905 skipped, 50 failed; main (b47e35b) run alongside fails the same 50 (7767 passed), all in seven files unrelated to this PR (test_loop_output_ownership.py, SHM shutdown, send registration, eval compare, serving metrics, client error text).
• On GPU, the files touching ragged or diffusion attention, sequence parallelism, LTX-2.5 and Cosmos3 (test_sequence_parallel.py, test_ltx2_5_sequence_parallel.py, test_ltx2_5_model.py, test_ragged_attention.py, test_ragged_dense_attention.py, test_attn_backend_auto.py, test_ragged_attention_cpu.py, test_context_only_attention.py, test_openai_router.py, mstar/model/cosmos3/tests, test_cosmos3_edge_model.py): 467 passed, 31 skipped.
• Served TP2 and SP2, alternated with the head before the rebase onto main, two restarts each: level (544×960 at 1 and 4 concurrent, two-stage at 2 concurrent), no failed requests, PSNR unchanged (31.0 / 31.5 dB at 544×960 for TP2 / SP2 against 32.1 dB reference-vs-itself; 29.8 dB two-stage).

New tests: stage-1 and refine walk structure, the two-stage schedule (replaying the conductor's step-metadata merge between walks), two-stage size validation; for SP, a toy DiT on two gloo ranks against one rank (test_ltx2_5_sequence_parallel.py) and the SP step declaration; the compiled VAE decoder's shape gate and vae_compile_shapes; the ragged resource's log-sum-exp (eager, through a captured graph re-planned before replay, and the callable) against torch.logsumexp; attention_over_sharded_keys on 1–3 gloo ranks against one rank. The two-stage parity check (parity_harness.py two_stage_seed) runs the upsampler node and is bit-exact against the oracle.

🤖 Generated with Claude Code

[dghosef]

Code looks fine.

Shows a ~50%-~150% performance improvement over SGLang but loses to vllm-omni with Ulysses SP.

config	GPUs	c=1 latency (s)	videos/min c=1	videos/min c=2	videos/min c=4
vllm-omni --usp 2 (Ulysses SP)	2	3.54	16.87 ± 0.03	19.74 ± 0.01	19.87 ± 0.00
vllm-omni --tensor-parallel-size 2 (tune run)	2	—	14.23	16.31	—
mstar TP2 (ltx2_5_tp2.yaml)	2	4.65	12.91 ± 0.05	17.76 ± 0.39	18.82 ± 0.31
mstar default (ltx2_5.yaml)	2	6.04	9.93 ± 0.02	13.95 ± 0.03	14.82 ± 0.00
SGLang --tp-size 2 --enable-torch-compile	2	7.82	7.67 ± 0.10	7.61 ± 0.05	7.77 ± 0.05
diffusers (compile, native attention)	1	5.35	11.22 ± 0.05	—	—

Here is the breakdown of a single request:

stage	where	time	share of 4.65 s
DiT, 8 denoising steps	both GPUs (TP2, lockstep)	≈ 2.7 s (8 × ~336 ms)	~58%
VAE + audio decode	GPU 0 only	≈ 1.7–2.0 s (one ~2 s step per request)	~37–40%
Gemma text encoder	GPU 1	≈ 0.3–0.6 s (partly overlapped)	~5–10%
host overhead (scheduling, planning, prepare_inputs)	CPU	< 10 ms	negligible

My hypothesis parallelizing the decode will help

[Gaurav-Shah05]

I ran ltx2_5.yaml, TP2 and SP2 on this branch, and your PSNR numbers against diffusers reproduce, two-stage included. Left couple of inline comments below.

[Gaurav-Shah05]

CommGroup.all_to_all, which ulysses_attention uses, already does this exchange, but it splits the tensor into pieces and concatenates what it receives, so it copies, while this sends one buffer with all_to_all_single. I'd move this version into CommGroup.all_to_all, so there is one all-to-all and ulysses_attention drops the copies too.

[NSagan271]

Done in #897, which this PR is now stacked on: CommGroup.all_to_all sends one packed buffer with all_to_all_single (even and uneven splits), ulysses_attention takes seq_dim/head_dim, and LTX's video self-attention calls it, so this file no longer has its own exchange. On LTX SP2 the swap measured the same throughput with byte-identical outputs.

[Gaurav-Shah05]

thanks, I will review #897 next

[Gaurav-Shah05]

At batch 1 this .contiguous returns z unchanged, so its batch stride differs from the warm-up tensor's, and torch.compile recompiles the decoder on the first batch-1 request, about 10 s in my TP2 runs. z.clone(memory_format=torch.channels_last_3d) gives it the warm-up's strides.

[NSagan271]

Fixed in 5edba42 with z.clone(memory_format=torch.channels_last_3d). Confirmed with TORCH_LOGS=recompiles on TP2 (2×H200): before, the first batch-1 request recompiled the decoder on stride mismatch at index 0. expected 1044480, actual 128; after, no recompiles once the server is up, over single-stage and two-stage requests, with byte-identical video and audio.

[Gaurav-Shah05]

Nit: This check runs diffusers' upsampler, and the re-normalize after it is a copy of LTXLatentUpsamplerSubmodule.run_batch, so it compares diffusers with itself and never runs the node's code. I'd have record_oracle_two_stage.py save stage 1's final latents, and call run_batch on them here.

[NSagan271]

Done in c8d8a80. record_oracle_two_stage.py now also saves stage 1's final latents, packed and still normalized (the node's input), and stage 2's starting latents before the noise blend (its output). The harness builds the served upsampler node through LTX25Model.get_submodule, calls run_batch on the former, compares with the latter, and seeds stage 2 from the node's output. On a freshly recorded oracle (2×H200) all four are bit-exact: max_abs 0 for video and audio, before and after the noise blend.

[Gaurav-Shah05]

shard_over_sp and workspace_mb aren't in the RaggedAttentionSpec row of docs/adding_models.rst. Can you add them, with a note that by default the resource splits heads across TP times SP ranks, and shard_over_sp=False splits them across TP only?

[NSagan271]

Added in 38decd3, in the RaggedAttentionSpec row: by default heads split across TP × SP, shard_over_sp=False across TP only, and workspace_mb with its default. It isn't a YAML override key, so it's not in the deployment-overrides table.

[dghosef]

Ran it on 2xH200 at 71fa555 - SP2/TP2 match your numbers within 1-4%, PSNR 32.8 dB vs the reference's own 32.1.
544x960, c=1 latency / c=4 videos/min:

• M* SP2: 3.10 s / 22.6
• vllm-omni Ulysses-2: 3.39 s / 20.7
• M* TP2: 3.42 s / 20.7
• SGLang Ulysses-2: 5.0 s / 12.2
• diffusers (1 GPU): 6.41 s / 9.4
  Two-stage 1088x1920: SP2 9.81 s vs vllm-omni 10.09 s at c=1, tie at c=4.