[mickqian]

Motivation

With more than one decode rank (CFG parallel, TP or SP), the LTX VAE installs spatially sharded convs that exchange halo rows with the neighbouring rank. Those convs expect each rank to hold its own height split. But:

• _decode checked tiling before the sharded path;
• tiled_decode, _temporal_tiled_decode and decode_chunk called the decoder on whole tensors while the sharded convs were installed.

So each rank blended its neighbour's copy of the same tile into its own edge rows.

Decoding a real LTX-2 latent (768x512x121) on 2xH200, with 2 decode ranks compared against 1:

• main: rank 0's max error is 135/149 levels, and 36.9%/41.7% of the pixels in the bottom 32 rows are off by more than 2 levels. Rank 1 mirrors this at the top.
• this PR: both ranks are identical, at 58.2–58.5 dB against the 1-rank untiled decode (bf16 noise).

This affects the LTX-2, LTX-2.3 and LTX-2.5 conv VAE on multiple GPUs with default flags, and SANA-WM's streaming decode_chunk. The Wan and Hunyuan VAEs are not affected.

Modifications

ltx_2_vae.py: every decoder call now either runs on this rank's height split or bypasses the sharded convs (new _decode_spatial).

• With tiling on, the height split replaces tiling when each rank's rows fit in one tile. A rank therefore still decodes at most one tile's area at once.
• Otherwise every rank tiles the whole clip under disable_spatial_parallel_decode(), exactly as on one rank. Large frames keep per-rank tiling because at 1920x1088 the untiled sharded path costs 28–38 GB more per rank.
• Framewise decoding applies the same rule to each temporal tile, and the streaming decode_chunk bypasses the sharded convs.

pipeline_configs/base.py: drop the "vae_sp requires vae_tiling" check. vae_sp gates nothing at runtime, and the height split never uses tiles, so --vae-tiling false now works on its own.

New test test/unit/test_ltx_vae_parallel_decode.py: two gloo ranks on CPU decode a tiny LTX-2 VAE with sharded convs and compare against one rank on 5 paths:

• tiled, below the auto threshold;
• tiled with spatial_shard;
• height split;
• framewise;
• decode_chunk.

Follow-up, not fixed here: the diffusers 2D VAE has the same ordering bug, but only with tiling plus the opt-in --vae-config.parallel-decode-mode spatial_shard.

Accuracy Tests

• Real LTX-2 VAE on 2xH200: see Motivation. Both ranks now decode identically, at 58.2–58.5 dB against one rank.
• test_ltx_vae_parallel_decode.py: main fails on all 5 paths, with max abs error 1.00–1.24 across the two ranks. This PR matches one rank exactly (0.0).

Speed Tests and Profiling

LTX-2 on 2xH200. On main, the halo traffic on whole tiles made the tiled decode slower than on one rank.

	main	this PR
VAE only, 768x512	0.342 s	0.135 s (1-rank tiled: 0.265 s)
pipeline decode stage, 768x512	0.665 s	0.418 s
decode, 1920x1088	3.63 s, peak 82.4 GB	3.07 s, peak 79.9 GB

At 768x512, e2e is −4.7% / −2.8% over two interleaved rounds, and memory is 1.7 GB lower per rank.

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 #39889302861
Latest PR Test (Extra): ❌ Run #39889302814
Latest PR Test (AMD ROCm 10): ❌ Run #39889302907

[mickqian]

/tag-and-rerun-ci