[Feature] Support VAE tiling parallel encode for LTX-2

[abtonmoy]

Purpose

Add distributed tiled VAE encode to DistributedAutoencoderKLLTX2Video (encode-side counterpart of #4277's decode path, following the Wan recipe from #2368), and switch LTX2Pipeline to the distributed VAE class so LTX-2 I2V (vae.encode of the reference image) and the latent upsampler (vae.encode of full videos) can use it. LTX-2 also gains the existing decode-parallel path it was not wired to before. Behavior is unchanged unless vae_patch_parallel_size > 1.

• encode_tile_split / encode_tile_exec / encode_tile_merge + tiled_encode override, mirroring diffusers' sequential tiled_encode semantics exactly.
• Per-tile latent output shapes are precomputed so the LTX2VaeExecutor known-metadata gather (no extra all-reduce) works for encode too; tiles whose sample dims do not divide the spatial compression ratio automatically fall back to the metadata-gather executor path.
• broadcast_result=True for encode (latents are consumed by the denoiser on all ranks), unlike decode.

Updates the support matrices in docs/design/feature/vae_parallel.md and docs/user_guide/diffusion_features.md (LTX-2 and LTX-2.3 → ✅ (encode/decode)). With #4739 merged, LTX23ImageToVideoPipeline now VAE-encodes the reference image through DistributedAutoencoderKLLTX2Video, so LTX-2.3 I2V picks up the parallel encode path from this PR directly.

Known limitation: LTX2TwoStagesPipeline does not expose model.vae at the top level, so the registry's VAE-parallel wiring does not engage for it. This is pre-existing and applies to decode as well.

Test Plan

pytest tests/diffusion/distributed/test_autoencoder_kl_ltx2_encode.py -m "cpu and core_model"

12 CPU-only unit tests (same tier/markers as test_autoencoder_kl_wan_encode.py): tile split grid + shape metadata, metadata fallback guard for non-divisible tiles, causal passthrough, executor dispatch (operator wiring + broadcast flag), and exact equivalence of split→exec→merge vs a sequential tiled encode reference.

Additionally verified locally (not committed):

• Bit-exact (max diff 0.0) against diffusers' own AutoencoderKLLTX2Video.tiled_encode on a real randomly-initialized VAE, across causal=True/False/None, frame counts 1/5/9/13, batch > 1, and rectangular inputs; predicted tile shapes matched real encoder outputs for every tile.
• 2-process gloo run through the full tiled_encode → LTX2VaeExecutor.execute path: both the known-metadata fast path and the stripped-metadata fallback bit-exact against sequential on both ranks.

Multi-GPU e2e on a real checkpoint still needs CI (developed on a single-GPU machine); the CPU unit-test tier matches what #2368 shipped with.

Test Result

12 passed in 3.33s

• The purpose of the PR is clearly explained.
• The test plan describes how reviewers can verify the changes.
• The test results are attached.
• References the related issues/PRs ([Feature] Support vae tiling parallel encode #2368, [Feat] Add vae-decode-parallel for LTX-2.3 #4277).

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[abtonmoy]

Rebased onto main after #4739 landed. Since LTX23ImageToVideoPipeline now VAE-encodes the reference image via DistributedAutoencoderKLLTX2Video, LTX-2.3 I2V directly benefits from the parallel encode path added here - updated the feature matrix accordingly (LTX-2.3 -> encode/decode). All 12 unit tests still pass after the rebase.

[hsliuustc0106]

Please add docstrings for this encode-tile path (encode_tile_split, encode_tile_exec, encode_tile_merge, and tiled_encode) that define the scope, input/return tensor shapes, grid/tile metadata contract, and fallback condition for non-divisible tiles. This adds a new distributed VAE contract, so the design should be explicit rather than only recoverable from the tests.

[abtonmoy]

Done - added docstrings to all four encode-tile methods covering the input/return tensor shapes, the tile_spec metadata contract consumed by LTX2VaeExecutor (including the predicted output-shape formula), and the non-divisible-tile fallback to the metadata-gather path. Also rebased onto latest main. All 12 unit tests still pass.

[abtonmoy]

@hsliuustc0106 Gentle ping - the docstring feedback has been addressed and CI is green. Anything else needed on this one? The same encode-parallel pattern for Qwen-Image edit pipelines is also ready in #4933.

[wtomin]

To test this ‎tests/e2e/online_serving/test_ltx2_expansion.py‎ file, the label diffusion-x2v needs to be added. @hsliuustc0106 @Gaohan123 Could you please help to add it?

[mglyn]

This still appears useful: current main has distributed tiled decode, while distributed tiled encode can reduce the LTX I2V encoder cost. The branch now conflicts in the shared VAE runtime and its validation is centered on the older pipelines. Would it make sense to rebase the focused encode delta and validate/document it with a current LTX-2.5 I2V workload?

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 15 days

@abtonmoy this pull request has had no human commit, comment or review since 2026-08-24. Per repository policy it may be closed if it stays inactive.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[abtonmoy]

Thanks @mglyn — rebased onto current main and validated on LTX-2.5. Reporting the perf result plainly, including where it does not help.

Rebase / focused delta. The branch is now 4 commits on top of main and merges clean. The pipeline change the PR originally carried is gone: main's LTX component profiles already point video_vae_cls at DistributedAutoencoderKLLTX2Video, so the encode methods on the shared autoencoder are the whole delta. Remaining scope is the encode ops + docstrings, the CPU unit tests, one e2e case, and two doc rows.

e2e case moved. tests/e2e/online_serving/test_ltx2_expansion.py was removed when the LTX serving tests were split per version, so the case is now test_ltx25_i2v_vae_patch_parallel in tests/e2e/online_serving/test_ltx25.py — full one-stage LTX2Pipeline on 2xH100, --use-hsdp --hsdp-shard-size 2 --vae-patch-parallel-size 2 --vae-use-tiling, I2V at 576x576. It passes (1 passed, ~87s). The existing single-card smokes run at 256x256, which stays under the VAE's 512 tile threshold and never reaches the tiled path; at 576x576 the reference-image encode splits into a 2x2 tile grid (verified against the LTX-2.5 VAE config: spatial_compression_ratio=32, tile_sample_min=512, stride=448). The shared assertion helper decodes the returned video and checks width/height/frames/fps, so the output through the parallel path is checked, not just the exit status.

Numbers. I2V, 2xH100, HSDP shard-size 2, --enforce-eager, 9 frames, median of 3 after a warmup; the only variable is --vae-patch-parallel-size:

model	resolution	VAE patch parallel	e2e latency
LTX-2.5	576x576	off	6.30s
LTX-2.5	576x576	encode & decode = 2	6.29s
LTX-2	1024x576	off	6.31s
LTX-2	1024x576	encode & decode = 2	5.44s

So on LTX-2.5 the change is neutral at this resolution — 0.01s is inside run-to-run noise — while on LTX-2 it is about 14% faster. The difference looks structural rather than incidental: LTX-2.5's VAE compresses 32x spatially against LTX-2's 8x, so a 576x576 reference image becomes an 18x18 latent and the VAE is a very small share of a ~6s request. There is simply little there to parallelize. I tried 1280x1280 to find a point where the VAE carries more weight, but LTX-2.5 rejects that request (400), and I stopped rather than keep guessing at resolutions.

Where that leaves it. Correctness is covered: the encode tiles match diffusers' sequential tiled_encode bit-exactly on a real AutoencoderKLLTX2Video, a 2-process run is bit-exact on both ranks, and the e2e case above passes. The performance case is real on LTX-2 and not demonstrable on LTX-2.5 at the resolution I could run, so the honest framing is that this completes the encode half of a pair whose decode half already ships on the same autoencoder, and pays off on the configurations where the VAE is a meaningful share of the request.

Happy to go either way: land it on the symmetry/correctness argument, or if you would rather see an LTX-2.5 configuration where the encode cost is actually material, point me at a resolution/frame count you consider representative and I will measure that instead.

[hsliuustc0106]

@mglyn PTAL

[nagisa-kunhah]

The PR currently reports a successful 2xH100 E2E run, but does not include numerical parity or peak-memory measurements. Since this changes a numerically sensitive tiled VAE encode path and is intended for distributed execution, could we add:

1. a numerical comparison between sequential and distributed tiled_encode (for example, max_abs_diff/mean_abs_diff with the tested dtype, resolution, and tolerance)
2. peak GPU memory per rank for the 1-GPU and 2-GPU configurations, measured around the VAE encode?

[nagisa-kunhah]

Suggestion: The current unit tests cover encode_tile_split, encode_tile_exec, encode_tile_merge, and the operator wiring, but they do not exercise LTX2VaeExecutor.execute() with real.

Could we add a CPU/Gloo two-process integration test, similar to

vllm-omni/tests/diffusion/distributed/test_wan_spatial_shard.py

Lines 542 to 543 in d349338

	@pytest.mark.parametrize("split_dim", ["height", "width"])
	def test_spatial_shard_decode_matches_reference(split_dim: str):

[mglyn]

Thanks for the suggestion. I separated VAE encoding from the full pipeline so that DiT, decoder, and scheduling overhead do not affect the measurement.

Test setup: 2x NVIDIA H200, BF16, input shape (1, 3, 1, 544, 960). Each configuration was warmed up for 20 iterations and measured for 100 iterations. For PP=2, latency is the slower of the two ranks for each iteration.

Configuration	Encoder latency (median)	P10-P90	Peak allocated memory / rank	Incremental peak during encode / rank
PP=1, tiling off	9.13 ms	9.10-9.31 ms	1528.60 MiB	130.03 MiB
PP=1, tiling on	38.44 ms	38.26-40.16 ms	1480.69 MiB	82.13 MiB
PP=2, tiling on	20.85 ms	20.69-21.87 ms	1480.57 MiB	82.00 MiB

PP=2 provides a 1.84x speedup over sequential tiled encoding, reducing encoder latency by 45.8%. A direct comparison of the raw encoder moments shows that PP=1 tiled and PP=2 tiled are bit-exact.

The results also show that untiled encoding remains the fastest option at 960x544. Tiling reduces the incremental peak allocation of the encode call by about 48 MiB, but introduces substantial overhead at this resolution. Therefore, encode parallelism is best positioned as a way to recover part of the tiling overhead when users enable tiling for memory-constrained or larger inputs, rather than as a default performance optimization at this resolution.

[abtonmoy]

Thanks @mglyn — isolating the encoder is the measurement I should have made. My end-to-end runs could not separate the VAE from DiT and scheduling, which is why LTX-2.5 looked neutral, and your harness answers the question directly. I will take your conclusion as the framing for this PR: at 960x544 untiled is the fastest option, tiling is an opt-in for memory headroom, and encode parallelism recovers part of what that opt-in costs. Concretely, tiling adds 29.3 ms over untiled (9.13 -> 38.44 ms) and PP=2 gives back 17.6 ms of it (38.44 -> 20.85 ms), so roughly 60% of the tiling overhead, or the 45.8% reduction in tiled encode latency you quote. Not a default performance win at this resolution.

The memory axis is the one I had missed entirely. I compared PP=1-tiled against PP=2-tiled and never measured untiled, so I did not see that tiling itself buys about 48 MiB of incremental peak per rank (130.03 -> 82.13 MiB). That is the trade this feature sits on top of, and it is a better argument for the change than the neutral end-to-end result I reported earlier.

I had measured peak memory and parity independently before your comment landed, and they agree. LTX-2.5 VAE, BF16, 2xH100, 2-rank NCCL group, the same input encoded both ways, max_memory_allocated reset around each encode. I am not re-reporting latency: my probe timed a single iteration per configuration, so your 100-iteration medians are the number to use.

resolution	vae_patch_parallel_size	rank0 peak / activations	rank1 peak / activations	max_abs_diff
576x576	1 (off)	1690.3 / 276.8 MiB	1690.3 / 276.8 MiB	—
576x576	2 (on)	1690.3 / 276.8 MiB	1484.2 / 70.8 MiB	0.0
1024x1024	1 (off)	1727.1 / 277.7 MiB	1727.1 / 277.7 MiB	—
1024x1024	2 (on)	1726.5 / 277.1 MiB	1726.5 / 277.1 MiB	0.0

Same conclusion as yours: bit-exact, and peak memory per rank does not improve — 1690.3 vs 1690.3 and 1727.1 vs 1726.5, against your 1480.69 vs 1480.57. Tiled encode already bounds peak at one tile's activations, and at 1024x1024 the balancer gives every rank at least one full 512x512 tile, so both configurations peak on the same tile. It does not regress either: the extra allocation is latent-space and small — the packed tile buffer plus its gather is 1.5 MiB at 576x576 and 2.5 MiB at 1024x1024 — and it does not move a peak set by one tile's encoder activations. The 576x576 rank1 drop to 70.8 MiB is not a general result — that grid is 2x2 with very uneven tiles (512^2, 512x128, 128x512, 128^2), so rank0 takes the one large tile and rank1 the three small ones.

---

@nagisa-kunhah both points are addressed.

1. CPU/Gloo integration test. Added to tests/diffusion/distributed/test_autoencoder_kl_ltx2_encode.py, following the test_wan_spatial_shard.py pattern you linked (mp.spawn + Manager().dict(), gloo world group, per-case master port). Three cases, max_abs_diff = mean_abs_diff = 0.0 on every rank in all three:

• grid3x3 — 24x24 into a (3, 3) grid, 9 tiles, on a tiny real AutoencoderKLLTX2Video (randomly initialized, ~24k params, no download), driven through LTX2VaeExecutor.execute(): task balancing, the packed-tile all_gather, the rank-0 unpack/merge, and the final broadcast.
• grid1x1_idle_rank — 8x8 into a (1, 1) grid, 1 tile, two ranks: rank 1 gets no work, contributes an all-zero buffer to the gather, and still has to come out of the broadcast holding the full latent.
• test_distributed_tiled_encode_metadata_fallback_gloo — 28x26, whose 2-wide edge tiles do not divide the compression ratio, so the predicted shapes are dropped and the dynamic metadata-gather path runs instead (shape all-reduce, per-tile metadata tensor, second gather, dynamic unpack).

Each case asserts the grid it produced and which gather path it took, so the coverage is checked rather than assumed — the fallback case fails if it silently takes the fast path. Results are checked on every rank rather than rank 0 only, since tiled_encode uses broadcast_result=True: every rank's denoiser consumes the latents, unlike decode. 16 tests in the file pass.

On why that last case uses the mock encoder: a real LTX-2 encoder rejects tiles whose extents do not divide spatial_compression_ratio — its space-to-depth downsampler raises unflatten: Provided sizes [-1, 2] don't multiply up to the size of dim 4 (3). The fallback is unreachable with a real VAE and exists as a safety net, so driving the executor with the mock is the only way to cover it through real collectives.

2. Parity and peak memory. Covered by the table above and by @mglyn's independent H200 run, both bit-exact, with the honest result on memory being that peak per rank is unchanged. This is a latency optimization for the tiled path, not a memory one, and I would not want the PR to imply otherwise.

One note on tiering: these cases keep tensors and collectives on gloo, but init_distributed_environment still picks a device index (vllm_omni/diffusion/distributed/parallel_state.py:420), so like the existing CPU collective tests they need at least one visible accelerator — on a GPU-less container tests/diffusion/distributed/test_comm.py::test_4d_identity fails the same way, so this is pre-existing rather than specific to these tests.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 10 days

@abtonmoy this pull request has had no human commit, comment or review since 2026-09-13. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 21cf0c9f6b28 refers to an earlier head; the pull request now points at 5518f3209ba9.

[hsliuustc0106]

All three 09-11 comments are resolved: docstrings now cover the full encode-tile path with shapes and the tile_spec metadata contract, the (28, 26) non-divisible case is in the parametrize (exercising the edge-tile fallback the existing cases missed), and the import is at module scope. Mergeable and checks green on this head. Good to merge.

[hsliuustc0106]

Correction to my review just now: I approved saying checks were green, but the readthedocs build (34532905) is actually failing on this head. I re-checked the docs diffs — the added table row and cell updates are well-formed GFM and both files already exist in the nav on main — so this looks infra-side rather than PR-caused, but please rerun/confirm the RTD build before merging so it isn't masked by this approval.

[abtonmoy]

Thanks @hsliuustc0106 — agreed it is not PR-caused, though I do not think a rerun on its own clears it. Build 34532905 is from 09-13 and died at command 3 of 5, before mkdocs ever ran:

git fetch origin --force --prune --prune-tags --depth 50 pull/4831/head:external-4831   exit 0
git checkout --force 21cf0c9f6b283fa4880e46dadb41401430b578e2    exit 128
  fatal: reference is not a tree: 21cf0c9f6b283fa4880e46dadb41401430b578e2

RTD fetches refs/pull/4831/head, and upstream still has that ref at the pre-09-13 head while the PR head is 21cf0c9f:

$ git ls-remote https://github.com/vllm-project/vllm-omni 'refs/pull/4831/*'
95f7db2ea90e6574e7a712655c1dd5700bba67c9	refs/pull/4831/head
ac3c813cd0065a3109f9fb138d6243111aae58ec	refs/pull/4831/merge

The head ref is the stale one; merge is current and does contain 21cf0c9f.

The commit really is absent from what RTD fetched, and still is today, so I would expect a rerun to hit the same error. A fresh push should resync the head ref.

If that is fine with you I will merge current main in, which also picks up #8065: it commented out the typing-extensions, pillow and psutil inventories because RTD rate-limits them in strict builds, and this branch still enables all three under fail_on_warning: true, so a build could clear the checkout and then fail at mkdocs build. The merge is clean.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[hsliuustc0106]

@abtonmoy this PR is labeled ready + high priority, but CI is failing on the latest commit:

• docs/readthedocs.org:vllm-omni

Could you please take a look and push an update to get CI green? Once the checks pass we can proceed with review/merge. Thanks!

[abtonmoy]

Pushed the main merge (5518f320) and readthedocs is green — build 34927239, docs at https://vllm--4831.org.readthedocs.build/projects/vllm-omni/en/4831/

That confirms the cause was the ref rather than the diff: refs/pull/4831/head was stuck at 95f7db2e and the push resynced it to 5518f320, so RTD could finally check out the head it had been asked to build.

Everything reported is green now: readthedocs, pre-commit, build (3.11), build (3.12), DCO.

It still shows as blocked because no buildkite/vllm-omni* context has reported on this head. Comparing against other PRs, those runs look like they are triggered by the ready label event rather than by pushes — ready has been applied here since 08-11, so the pushes on 09-08, 09-13 and 10-04 never kicked one off. Could you remove and re-apply ready so Buildkite runs against 5518f320? I cannot change labels myself.

The merge is sync-only; the encode delta is still the same 5 files.

[vllm-omni-review-bot]

Omni ReviewBot review

1 actionable finding(s).

CI at 5518f3209ba9 (2026-10-09T23:48:15.299188+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Full review analysis

Scan:

Category	Result
Tests / verification	1 finding(s) below
Security	no finding reported
Docs / comments	no finding reported
Behavior / compatibility	no finding reported
Correctness	no finding reported

Validated:

• [resolved] (28,26) parametrize covers ragged-edge merge (322). Residual: that case does not call LTX2VaeExecutor; the gloo fallback test does (608).
• [resolved] Gloo tests drive LTX2VaeExecutor.execute (505,608). Residual: e2e is LTX-2.5 I2V smoke; LTX-2.5 decode is DiffVAE, not KL tiled_decode.
• [resolved] (28,26) added to merge-equivalence parametrize (test_autoencoder_kl_ltx2_encode.py:312-329); residual: real encoder cannot reach that fallback
• [resolved] nagisa Gloo executor request: test_distributed_tiled_encode_matches_sequential_gloo and _metadata_fallback_gloo drive LTX2VaeExecutor.execute(). Residual: no GPU numerical-parity assert in the e2e smoke.
• [claim-verified] LTX-2/2.3/2.5 user-guide cells flipped to encode/decode (diffusion_features.md:193-195); design table lists only LTX-2/2.3 (vae_parallel.md:469).
• [claim-verified] LTX23 I2V inherits via LTX23_COMPONENT_PROFILE.video_vae_cls and LTXI2VConditioningMixin; no dedicated 2/2.3 GPU e2e in this diff.

1 actionable finding(s).

Verdict: REQUEST CHANGES

Findings

• [P1] This PR adds test_ltx25_i2v_vae_patch_parallel with hardware_marks(res={"cuda":… — .buildkite/cuda/test-weekly.yml

Evidence for This PR adds test_ltx25_i2v_vae_patch_parallel with hardware_marks(res={"cuda":…

This PR adds test_ltx25_i2v_vae_patch_parallel with hardware_marks(res={"cuda": "H100"}, num_cards=2). Unchanged weekly step Diffusion · LTX · H100 · 2-GPU (.buildkite/cuda/test-weekly.yml:575-582) still lists only tests/e2e/online_serving/test_ltx.py with -m "slow and diffusion and H100 and cards_2". The 1-GPU LTX step that lists test_ltx25.py uses cards_1, which deselects VAE_PARALLEL_MARKS (cards_2 is not cards_1). Wire test_ltx25.py into that 2-GPU step so a tiled I2V encode regression cannot ship on green weekly.

Evidence: Trigger: the PR adds a 2-card serving smoke that weekly never collects. tests/e2e/online_serving/test_ltx25.py:83 VAE_PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) and :113 def test_ltx25_i2v_vae_patch_parallel( (file-level pytestmark is slow+diffusion). Adverse effect: that case is the only online-serving coverage of tiled I2V encode (576>512, --vae-patch-parallel-size 2) and it is not collected by weekly, so a serving-path encode regression can ship on a green weekly. Unchanged by this diff, present in the PR-time tree: .buildkite/cuda/test-weekly.yml:580 tests/e2e/online_serving/test_ltx.py and :581 -m "slow and diffusion and H100 and cards_2" — the 2-GPU LTX step never names test_ltx25.py. Same file, also unchanged: :424 tests/e2e/online_serving/test_ltx25.py with :427 -m "slow and diffusion and H100 and cards_1". Marker contract, unchanged: tests/helpers/mark.py:149 return getattr(pytest.mark, f"cards_{num_cards}"); tests/helpers/tests/test_mark.py:48 assert "cards_1" not in names for num_cards=2 — so the 1-GPU expression cannot select the new case.

Suggestion: tests/e2e/online_serving/test_ltx.py
tests/e2e/online_serving/test_ltx25.py

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] This PR adds test_ltx25_i2v_vae_patch_parallel with hardware_marks(res={"cuda":… — .buildkite/cuda/test-weekly.yml

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.