[TRTLLM-13767][feat] integrate FP4 Conv3d into parallel Wan VAE - #17262
Draft
taianz-nv wants to merge 3 commits into
Draft
[TRTLLM-13767][feat] integrate FP4 Conv3d into parallel Wan VAE#17262taianz-nv wants to merge 3 commits into
taianz-nv wants to merge 3 commits into
Conversation
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
Signed-off-by: Taian Zhang <taianz@nvidia.com>
taianz-nv
force-pushed
the
user/taianz/TRTLLM-13767-fp4-parallel-wan-vae
branch
from
August 5, 2026 02:40
8b7be32 to
d3d7d12
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integrate the Blackwell NVFP4 Conv3d kernel into the native Wan VAE and make the FP4 path compose with the existing parallel VAE wrapper.
This change:
There are no per-fusion environment variables or new public VisualGen API. Fusion selection is an internal implementation decision derived from the checkpoint, module structure, and parallel output contract. The BF16 path does not import or initialize the FP4 kernel.
Dependency
This PR requires the CUTLASS DSL 4.6.1 upgrade in #17274. The dependency and its optional-attention compatibility changes are intentionally absent from this diff.
Parallel-VAE scope
This PR is intentionally independent of the other parallel-VAE optimizations:
mainand is not part of this diff.The current mainline halo wrapper delegates the pointwise RMSNorm/SiLU fusion capabilities, so FP4 Conv3d works in the parallel VAE. The rank-local residual remains unfused in the parallel path because mainline halo computes an expanded output and strips it after the convolution. Once #17001 lands, its local-width output contract can enable the residual epilogue without bringing the halo implementation into this PR. Bias fusion and the single-GPU residual epilogue are active here.
Validation
Draft follow-ups