[TRTLLM-14029][feat] LTX-2 tile-parallel VAE decode#15753
Conversation
📝 WalkthroughWalkthroughAdds a tile-parallel VAE decoding path for LTX-2 distributed inference. A new ChangesLTX-2 Tile-Parallel VAE Decode
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py (1)
177-179: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd coverage for the pipeline routing layer.
This test sufficiently covers
tile_parallel_decodeparity, but coverage is still insufficient for the pipeline integration added inpipeline_ltx2.py. Add a follow-up test in this file, or a newtests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_pipeline_parallel_vae.py, that verifiessetup_parallel_vae()enables the flag fromvisual_gen_mappingand that decode routes throughtile_parallel_decodewithvgm.vae_group.As per path instructions, tests under
tests/**should be reviewed for TensorRT-LLM coverage sufficiency and actionable follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py` around lines 177 - 179, The current parity test in TestLTX2ParallelVAEDecode only validates tile_parallel_decode behavior, but it does not cover the new pipeline routing in pipeline_ltx2.py. Add a follow-up test in this file or a new test_ltx2_pipeline_parallel_vae.py that exercises setup_parallel_vae() and verifies it enables the visual_gen_mapping flag, then confirm decode is routed through tile_parallel_decode using vgm.vae_group. Use the existing symbols setup_parallel_vae, visual_gen_mapping, tile_parallel_decode, and vgm.vae_group so the coverage targets the pipeline integration directly.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py`:
- Around line 177-179: The current parity test in TestLTX2ParallelVAEDecode only
validates tile_parallel_decode behavior, but it does not cover the new pipeline
routing in pipeline_ltx2.py. Add a follow-up test in this file or a new
test_ltx2_pipeline_parallel_vae.py that exercises setup_parallel_vae() and
verifies it enables the visual_gen_mapping flag, then confirm decode is routed
through tile_parallel_decode using vgm.vae_group. Use the existing symbols
setup_parallel_vae, visual_gen_mapping, tile_parallel_decode, and vgm.vae_group
so the coverage targets the pipeline integration directly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 482adfa0-78b8-4f89-8d98-db2ce38b56e1
📒 Files selected for processing (3)
tensorrt_llm/_torch/visual_gen/models/ltx2/parallel_vae.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.pytests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
3 times, most recently
from
8ea67aa to
c921add
Compare
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
c921add to
f71fcec
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57110 [ run ] triggered by Bot. Commit: |
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
f71fcec to
6a3a492
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57148 [ run ] triggered by Bot. Commit: |
|
PR_Github #57110 [ run ] completed with state |
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
6a3a492 to
df96232
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57152 [ run ] triggered by Bot. Commit: |
|
PR_Github #57148 [ run ] completed with state |
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
df96232 to
47502a3
Compare
|
/bot run --disable-fail-fast |
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
47502a3 to
43e5191
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57156 [ run ] triggered by Bot. Commit: |
|
PR_Github #57152 [ run ] completed with state |
|
PR_Github #57158 [ run ] triggered by Bot. Commit: |
|
PR_Github #57156 [ run ] completed with state |
|
PR_Github #57158 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57227 [ run ] triggered by Bot. Commit: |
|
PR_Github #57227 [ run ] completed with state |
Distribute tiled_decode's overlapping tiles across vae_ranks (LPT by input volume), blend each tile into a full-size buffer, all_reduce the buffer + weights, and divide -- output matches serial tiled_decode within bf16 re-association. setup_parallel_vae sets the global _parallel_vae_enabled flag directly (LTX-2 decodes via video_decoder, not vae, so the base factory never engages); decode_video_fn routes to tile_parallel_decode on vae_ranks. Per-rank peak memory stays flat (one tile in flight per rank, same as serial). Tile.blend_mask assembles the per-tile blend mask on-device from its 1-D components, used by both the serial and tile-parallel decode paths. Building the full tile-sized mask on the CPU with a per-tile H2D copy left the decode host-bound -- the host mask work does not overlap the GPU forward, leaving the GPU ~50% idle in a busy multi-process pipeline. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
Stage 2 (upsample + refinement denoise) runs rank-0 only; rank 0 then broadcasts the refined latents to every vae_rank and all vae_ranks decode collectively via tile_parallel_decode. Fixes the regression where enabling parallel VAE let all vae_ranks fall through into the rank-0-only Stage 2 (decode_latents' rank set widened from {0} to vae_ranks). Non-parallel two-stage keeps the serial rank-0 decode. Validated on B200: bf16 + nvfp4 two-stage, parallel (pvs2) and serial (pvs1).
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
get_nvfp4_input_scale decided whether the fused AdaLN norm emits a pre-quantized Fp4QuantizedTensor from the input_scale attribute alone. The two-stage distilled-LoRA merge dequantizes FP4 weights to bf16 and swaps quant_method to UnquantizedLinearMethod but leaves input_scale on the module, so the fused norm still emitted FP4 for a now-unquantized F.linear -> 'linear(): input must be Tensor, not Fp4QuantizedTensor' at Stage 2. Gate on isinstance(quant_method, NVFP4LinearMethod) so the fusion follows the live quant_method. Prerequisite for nvfp4 two-stage (crashes at Stage 2 warmup without it, independent of CUDA graph / parallel VAE). Single-stage nvfp4 unchanged (quant_method stays NVFP4). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
luyiyun1021
force-pushed
the
dev/ltx2-vae-tile-parallel
branch
from
43e5191 to
aedf74e
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57406 [ run ] triggered by Bot. Commit: |
luyiyun1021
changed the title
… idle vae_ranks assign_tiles_lpt logs the per-rank tile-count/volume distribution and max/min imbalance ratio at debug (rank 0), and warns once when parallel_vae_size exceeds the tile count so some vae_ranks receive zero tiles (idle at the all_reduce; suggests lowering parallel_vae_size). The load array is computed deterministically on every rank, so no extra communication is needed. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #57417 [ run ] triggered by Bot. Commit: |
|
PR_Github #57406 [ run ] completed with state |
|
PR_Github #57417 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57535 [ run ] triggered by Bot. Commit: |
|
PR_Github #57535 [ run ] completed with state |
Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
4 participants
Description
Tile-parallel VAE decode for LTX-2. The VAE decode (
VideoDecoder.tiled_decode) is a large, embarrassingly-parallel pass over ~27 overlapping spatial/temporal tiles that, served serially, dominates the non-transformer wall time at high resolution. This PR distributes those tiles across the VAE ranks and reassembles the result, so decode latency scales with the number of VAE ranks while per-rank peak memory stays flat.Approach
tile_parallel_decode(newparallel_vae.py): assigntiled_decode's tiles acrossvae_ranksby LPT on input volume (deterministic, identical partition on every rank), decode each assigned tile, blend it into a full-size output buffer, thenall_reducethe buffer + the overlap-weight buffer overvae_groupand divide. Output matches serialtiled_decodeup to bf16 summation re-association in the tile-overlap regions (a few ulp).setup_parallel_vae(override inpipeline_ltx2.py): LTX-2 decodes viaself.video_decoder(notself.vae), so the base factory never engages — set the global_parallel_vae_enabledflag directly from config + mapping.decode_video_fnthen routes totile_parallel_decodeonvae_ranks, else the original serialtiled_decode.Tile.blend_maskassembles the per-tile blend mask on-device from its 1-D components, used by both the serial and tile-parallel paths. Assembling the full tile-sized mask on the CPU and copying it to the GPU per tile makes the decode loop host-bound: in a busy multi-process pipeline that host work does not overlap the GPU forward and leaves the GPU ~50% idle. On-device assembly keeps the decode GPU-bound (see E2E below).Performance (B200)
VAE module bench (real VAE size, bf16
VideoDecoder)tile_parallel_decodevs serialtiled_decode, latent →(1, 3, 121, 768, 1280)(27 tiles), both on the shipped on-device blend mask:Per-rank peak memory is constant ~4.0 GB across world sizes (no memory blow-up from parallelism).
End-to-end (FA4 backend, 8-GPU cfg=2×uly=4, 768×1280×121, 40-step, guidance=4.0, CUDA-graph + torch.compile)
Tile-parallel VAE (pvs=8) vs serial VAE (pvs=1), same seed:
The clean, same-basis decode comparison is the module bench above (serial 873 ms → tile-parallel 129 ms, 6.74×). In the FA4 pipeline the parallel VAE decode is 0.14 s (GPU-timed via CUDA events), matching that module-bench floor; that decode saving is what surfaces as the ~10% (0.80 s) E2E wall reduction, with the percentage diluted by the fixed, VAE-orthogonal remainder (text-encode + the 4.4 s denoise + result IPC) that this change does not touch. Denoise is unchanged, confirming the change is VAE-only. The on-device blend mask is what lets the in-pipeline decode reach the module-bench floor — with the mask built on the CPU the decode is ~0.25 s here (GPU ~50% idle waiting on the host mask).
Numerical accuracy
tile_parallel_decodeis numerically equivalent to serialtiled_decode: identical tiles and blend masks, so the only difference is bf16 summation re-association in the tile-overlap regions — per-rank partial sums combined byall_reducevs sequential accumulation on one GPU. Direct parity on the real VAE, world=8, full 768×1280×121 output (same latent,tile_parallel_decodevs single-GPUtiled_decode):So parallelizing the decode adds no algorithmic error — only floating-point re-association at the bf16 ULP level. The multi-GPU parity unit test (below) asserts rel max-abs < 0.02 and passes.
Perceptually this is invisible. On a deterministic attention backend (same seed → identical denoise, so only the VAE decode differs), parallel (pvs=8) vs serial (pvs=1) decoded video is mean LPIPS 0.017 — at/below the run-to-run noise floor (0.021 for two separate serial runs), i.e. indistinguishable (left = serial, right = tile-parallel; per-frame LPIPS baked in):
Tests
tests/unittest/_torch/visual_gen/multi_gpu/test_ltx2_parallel_vae.py— multi-GPU (world=2) parity: a small bf16VideoDecoder(random weights, broadcast so every rank is identical),tile_parallel_decodevs single-GPU serialtiled_decode, asserting relative max-abs diff< 0.02(bf16 ulp from overlap-region re-association). Passes on B200 (2-GPU).PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.