[TRTLLM-14123][feat] add ParallelVAE_TrtllmWan for native Wan VAE by taianz-nv · Pull Request #16066 · NVIDIA/TensorRT-LLM

taianz-nv · 2026-07-07T12:00:40Z

Summary by CodeRabbit

New Features
- Added native TensorRT-LLM Wan VAE support for parallel encoding and decoding.
- Added automatic registration and selection of the native Wan VAE implementation.
- Diffusers VAE usage can now be enabled explicitly through the TRTLLM_USE_DIFFUSER_VAE environment setting.
Bug Fixes
- Improved VAE loading behavior by removing unnecessary configuration-based restrictions.
Tests
- Added multi-GPU validation for native Wan VAE encoding and decoding across width and height splits.

Description

Adds ParallelVAE_TrtllmWan, a tensor-parallel (parallel_vae_size > 1) wrapper for the native WanVAE (added in #15555), so multi-GPU Wan VAE decode/encode uses the native VAE instead of falling back to the diffusers AutoencoderKLWan. It mirrors the existing ParallelVAE_Wan by parameterizing the replaced conv3d / attention module classes (_conv3d_cls / _attn_cls) and overriding only those in a small subclass; vae_loader no longer forces diffusers under parallel_vae_size > 1; ParallelVAEFactory registers WanVAE -> ParallelVAE_TrtllmWan. ParallelVAE_Wan is kept because it wraps the diffusers AutoencoderKLWan still used by Cosmos3 and the TRTLLM_USE_DIFFUSER_VAE fallback.

This PR also includes a channels-last (NHWC) fix in HaloExchangeConv / HaloExchangeConv2dStride2 (conv.py): the halo exchange built row-major send/recv buffers, so the boundary torch.cat fell back to row-major and forced every halo'd conv to re-convert the full tensor. A shared _spatial_channels_last_format helper now materialises the halo slices in x's memory format before the cat, keeping it channels-last so the native VAE convs (which run in channels-last) skip the re-conversion. Output is bit-identical.

Performance & accuracy (4×GB200, BF16, Wan2.2-TI2V-5B, 720p / 81 frames)

Comparison	Result	Accuracy
channels-last halo fix (parallel_vae_size=4 decode, with vs without)	1.31× (1.243 → 0.951 s)	bit-identical
parallel VAE on vs off (parallel_vae_size 4 vs 1, VAE-only decode)	2.88× (2.735 → 0.951 s)	54.3 dB PSNR (near-lossless)
end-to-end, true 4-GPU (cfg×ulysses DiT + parallel VAE) vs parallel_vae_size=1	1.65× @10 steps / 1.18× @50 steps	50.4 / 57.1 dB PSNR

Test Coverage

tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py — TestParallelVAETrtllmDecode / TestParallelVAETrtllmEncode: native WanVAE parallel decode/encode match single-GPU (width + height split, 2-GPU). The wrapper is built via ParallelVAEFactory.from_vae, so the WanVAE -> ParallelVAE_TrtllmWan registry mapping is exercised (with an explicit isinstance assert). The small test VAE keeps a mid-block WanAttentionBlock, so _replace_attention on native attention is covered.
tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_conv.py — TestSpatialChannelsLastFormat: CPU unit tests for the channels-last halo helper (memory-format detection + cat preservation).
tests/unittest/_torch/visual_gen/test_wan_vae.py — updated _use_native_wan_vae / load_wan_vae tests for the removed visual_gen_mapping parameter (backend selection is now env-only via TRTLLM_USE_DIFFUSER_VAE).

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
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.

coderabbitai · 2026-07-10T07:03:59Z

📝 Walkthrough

Walkthrough

Wan VAE loading now uses environment-controlled backend selection. ParallelVAE_Wan supports configurable module targets, with ParallelVAE_TrtllmWan adapting native WanVAE modules. Factory registration and multi-GPU encode/decode equivalence tests were added.

Changes

Wan VAE integration

Layer / File(s)	Summary
Configurable Wan parallel adapter `tensorrt_llm/_torch/visual_gen/models/wan/parallel_vae.py`	Module replacement targets are configurable, and `ParallelVAE_TrtllmWan` targets native WanVAE convolution and attention modules.
Environment-controlled VAE loading `tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py`, `tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan*.py`	The loader API no longer accepts `visual_gen_mapping`; backend selection depends on `TRTLLM_USE_DIFFUSER_VAE`, with both pipeline call sites updated.
Factory registration and distributed validation `tensorrt_llm/_torch/visual_gen/modules/vae/parallel_vae_interface.py`, `tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py`	Native WanVAE resolves to the new parallel adapter, with 2-GPU width/height split encode and decode equivalence tests. Estimated code review effort: 3 (Moderate) \| ~20 minutes

Suggested reviewers: achartier, yechank-nvidia, Superjomn, 2ez4bz, yuxianq

Sequence Diagram(s)

sequenceDiagram
  participant Pipeline
  participant load_wan_vae
  participant ParallelVAEFactory
  participant ParallelVAE_TrtllmWan
  participant WanVAE
  Pipeline->>load_wan_vae: load checkpoint, device, and dtype
  load_wan_vae->>WanVAE: construct native backend when env flag is zero
  WanVAE->>ParallelVAEFactory: resolve native WanVAE type
  ParallelVAEFactory->>ParallelVAE_TrtllmWan: create parallel adapter
  ParallelVAE_TrtllmWan->>WanVAE: replace native convolution and attention modules

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 41.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the main change and follows the required [ticket][type] format.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections with substantive details.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 2

🧹 Nitpick comments (1)

tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py (1)
64-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Document the public loader API.

Line 64 is called across pipeline modules but has no Google-style docstring documenting its backend-selection behavior, arguments, and return value.

As per coding guidelines, externally called functions should have docstrings, and their arguments should be documented.
🤖 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 `@tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py` around lines 64 -
68, The public load_wan_vae function lacks API documentation. Add a Google-style
docstring to load_wan_vae describing its backend-selection behavior,
checkpoint_dir, device, and dtype arguments, and the returned nn.Module.
Source: Coding guidelines

🤖 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.

Inline comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py`:
- Around line 140-172: Add complete type annotations to the new helper and test
functions, including concrete parameter types and -> None return annotations for
test_* methods; cover the distributed helper functions and all additions through
the referenced range, including _create_small_trtllm_vae and
_create_parallel_trtllm_vae.
- Around line 140-172: Coverage is insufficient in
tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py: the native VAE
disables attention and directly constructs the wrapper. Add an attention-enabled
WanVAE case to exercise _replace_attention() for WanAttentionBlock, and add a
ParallelVAEFactory test asserting supports(WanVAE) and from_vae(...) return
ParallelVAE_TrtllmWan.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py`:
- Around line 64-68: The public load_wan_vae function lacks API documentation.
Add a Google-style docstring to load_wan_vae describing its backend-selection
behavior, checkpoint_dir, device, and dtype arguments, and the returned
nn.Module.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 16e51f25-3f16-4631-b0ee-5a82930318ce

📥 Commits

Reviewing files that changed from the base of the PR and between c92ac90 and 3d4d22d.

📒 Files selected for processing (6)

tensorrt_llm/_torch/visual_gen/models/wan/parallel_vae.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py
tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py
tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py
tensorrt_llm/_torch/visual_gen/modules/vae/parallel_vae_interface.py
tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py

NVShreyas

LGTM overall. Naive question - why is WANVae needed if WAN can use diffusers VAE directly?

taianz-nv · 2026-07-14T14:43:50Z

Thanks @NVShreyas ! The diffusers AutoencoderKLWan is a third-party module we can’t easily optimize, while the native WanVAE lets us directly apply TRT-LLM optimizations (channels-last kernel, FP4 quant, kernel fusion). On single-GPU VAE decode (Wan2.2-TI2V-5B, 720p), channels-last gives 1.29× vs diffusers, and FP4 conv quant + RMSNorm/SiLU→quant fusion (WIP) brings it to ~2.0×. This will matter more for few-step or distilled models (e.g., FastWan)

taianz-nv · 2026-07-15T07:07:12Z

/bot run

tensorrt-cicd · 2026-07-15T07:12:57Z

PR_Github #59386 [ run ] triggered by Bot. Commit: 668ac14 Link to invocation

tensorrt-cicd · 2026-07-15T11:30:34Z

PR_Github #59386 [ run ] completed with state SUCCESS. Commit: 668ac14
/LLM/main/L0_MergeRequest_PR pipeline #47860 completed with status: 'ABORTED'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

Add a tensor-parallel wrapper (ParallelVAE_TrtllmWan) for the native WanVAE, enabling multi-GPU (parallel_vae_size > 1) VAE decode/encode with the native VAE instead of falling back to the diffusers AutoencoderKLWan. - ParallelVAE_Wan is parameterized on the replaced conv3d/attention module classes (_conv3d_cls/_attn_cls); ParallelVAE_TrtllmWan is a small subclass that overrides only those to target the native wan_vae classes. - vae_loader drops the parallel-VAE gate so the native WanVAE is used when parallel_vae_size > 1; ParallelVAEFactory registers WanVAE -> ParallelVAE_TrtllmWan. - HaloExchangeConv / HaloExchangeConv2dStride2 preserve channels-last across the halo cat (conv.py). The native WanVAE convs run in channels-last (NHWC); the halo exchange built row-major send/recv buffers, so the boundary cat fell back to row-major and every halo'd conv re-converted the full tensor. A shared _spatial_channels_last_format helper materialises the halo slices in x's memory format before the cat, keeping it channels-last so the conv's channels-last conversion becomes a no-op. Bit-identical output; ~23% faster parallel decode (1.31x, Wan2.2-TI2V-5B 720p, 4x GB200). - Tests: the multi-GPU parallel-VAE tests cover the native adapter (decode/encode, width/height split) and route through ParallelVAEFactory so the WanVAE -> ParallelVAE_TrtllmWan registration is exercised; a CPU unit test covers the channels-last halo helper; test_wan_vae is updated for the vae_loader signature change (the visual_gen_mapping parameter was removed). ParallelVAE_Wan is kept for the diffusers AutoencoderKLWan still used by Cosmos3 and the TRTLLM_USE_DIFFUSER_VAE fallback. Verified VAE-only on 4x GB200 (BF16): native-parallel matches diffusers-parallel / single-GPU to ~58-60 dB PSNR (parallelization adds negligible error) and is faster than both the diffusers parallel VAE and the single-GPU native VAE (Wan2.2-TI2V-5B 720p: 2.29x over single-GPU; Wan2.1-T2V-1.3B 480p: 1.24x). Signed-off-by: Taian Zhang <taianz@nvidia.com>

taianz-nv · 2026-07-15T11:49:38Z

/bot run

tensorrt-cicd · 2026-07-15T11:55:59Z

PR_Github #59447 [ run ] triggered by Bot. Commit: 389bcb9 Link to invocation

tensorrt-cicd · 2026-07-15T15:11:50Z

PR_Github #59447 [ run ] completed with state SUCCESS. Commit: 389bcb9
/LLM/main/L0_MergeRequest_PR pipeline #47915 completed with status: 'SUCCESS'

CI Report

Link to invocation

github-actions Bot assigned taianz-nv Jul 7, 2026

taianz-nv force-pushed the dev-taianz-parallel-vae branch 4 times, most recently from 402a01c to 3d4d22d Compare July 8, 2026 11:42

taianz-nv marked this pull request as ready for review July 10, 2026 06:59

taianz-nv requested a review from a team as a code owner July 10, 2026 06:59

coderabbitai Bot reviewed Jul 10, 2026

View reviewed changes

Comment thread tests/unittest/_torch/visual_gen/multi_gpu/test_parallel_vae.py Outdated

zhenhuaw-me reviewed Jul 10, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/modules/vae/parallel_vae_interface.py

taianz-nv force-pushed the dev-taianz-parallel-vae branch 3 times, most recently from 2f4f1a1 to 668ac14 Compare July 13, 2026 15:15

taianz-nv changed the title ~~[TRTLLM-14123][feat] add ParallelVAE_NativeWan for native Wan VAE~~ [TRTLLM-14123][feat] add ParallelVAE_TrtllmWan for native Wan VAE Jul 13, 2026

zhenhuaw-me requested a review from NVShreyas July 14, 2026 08:46

zhenhuaw-me reviewed Jul 14, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/models/wan/vae_loader.py

zhenhuaw-me added the VisualGen label Jul 14, 2026

NVShreyas approved these changes Jul 14, 2026

View reviewed changes

taianz-nv force-pushed the dev-taianz-parallel-vae branch from 668ac14 to 389bcb9 Compare July 15, 2026 11:49

zhenhuaw-me approved these changes Jul 16, 2026

View reviewed changes

taianz-nv merged commit ef90708 into NVIDIA:main Jul 16, 2026
7 checks passed

Conversation

taianz-nv commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Performance & accuracy (4×GB200, BF16, Wan2.2-TI2V-5B, 720p / 81 frames)

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

NVShreyas left a comment

Choose a reason for hiding this comment

Uh oh!

taianz-nv commented Jul 14, 2026

Uh oh!

taianz-nv commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

taianz-nv commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

taianz-nv commented Jul 7, 2026 •

edited

Loading

coderabbitai Bot commented Jul 10, 2026 •

edited

Loading