Skip to content

[TRTLLM-13767][feat] integrate FP4 Conv3d into parallel Wan VAE - #17262

Draft
taianz-nv wants to merge 3 commits into
NVIDIA:mainfrom
taianz-nv:user/taianz/TRTLLM-13767-fp4-parallel-wan-vae
Draft

[TRTLLM-13767][feat] integrate FP4 Conv3d into parallel Wan VAE#17262
taianz-nv wants to merge 3 commits into
NVIDIA:mainfrom
taianz-nv:user/taianz/TRTLLM-13767-fp4-parallel-wan-vae

Conversation

@taianz-nv

@taianz-nv taianz-nv commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

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:

  • vendors the block-scaled CuTeDSL implicit-GEMM Conv3d kernel and its dynamic preferred-cluster scheduler;
  • detects ModelOpt NVFP4 VAE checkpoints and converts eligible native Wan Conv3d modules automatically;
  • prequantizes weights once and supports calibrated static or dynamic activation quantization;
  • automatically fuses RMSNorm + SiLU + FP4 quantization when the calibrated static path and module structure permit it;
  • always folds Conv3d bias into the kernel and uses the residual epilogue when the caller provides a shape-compatible residual;
  • uses the 256x256, 2CTA, 2x1-cluster tactic selected for the parallel Wan shapes; and
  • adds structural, composition, quantization-parity, and real B200 kernel tests.

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:

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

  • Full changed-file pre-commit suite: passed.
  • NSC B200 job 1619867:
    • 21 FP4 structural/composition tests passed.
    • 6 real GPU kernel tests passed, including automatic bias + residual fusion and asymmetric spatial padding through the product wrapper.
  • The kernel reference checks cover output parity and a nonzero-convolution quality threshold.

Draft follow-ups

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
taianz-nv force-pushed the user/taianz/TRTLLM-13767-fp4-parallel-wan-vae branch from 8b7be32 to d3d7d12 Compare August 5, 2026 02:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant