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

[taianz-nv]

Summary

Integrate the Blackwell NVFP4 Conv3d kernel into the native Wan VAE and make it compose with the existing parallel VAE wrapper.

• Detect ModelOpt NVFP4 VAE checkpoints and replace eligible native Wan Conv3d modules automatically.
• Prequantize weights once into the CuTeDSL kernel layout.
• Use calibrated static activation scales when present; otherwise compute a rank-local dynamic activation amax.
• Fuse RMSNorm + SiLU + FP4 quantization on the calibrated static path when the module structure permits it.
• Fold bias into the Conv3d epilogue and use the residual epilogue when the output contract is shape-compatible.\n- Issue residual TMA loads from a dedicated warp when beta != 0, allowing overlap with A/B/SFB traffic while leaving non-residual launch geometry unchanged.\n- Use the 256x256, 2CTA, 2x1-cluster tactic selected for parallel Wan shapes.
• Keep the BF16 path free of FP4 kernel imports and initialization.

No public VisualGen API or per-fusion environment variables are introduced.

Checkpoint selection and quantization behavior

The native loader selects this path when vae/config.json contains:

{"quantization_config": {"quant_algo": "NVFP4"}}

The checkpoint already contains packed FP4 weights and ModelOpt weight scales. TensorRT-LLM dequantizes those weights into the native state dict during loading, then prequantizes eligible 3x3x3 Conv3d weights once into the layout consumed by the integrated kernel.

• A scalar checkpoint input_scale selects calibrated static activation quantization and enables the fused input-preparation path.
• If a quantized Conv3d has no calibrated input_scale, the operator falls back to rank-local dynamic activation quantization using the current sharded activation amax.
• A normal BF16 checkpoint is not dynamically weight-quantized into FP4, and there is no new user-facing runtime quantization mode in this PR.

The FastWan production-best measurements use checkpoint-calibrated static activation scales. When a ModelOpt checkpoint omits input_scale for some quantized convolutions, only those convolutions use the dynamic fallback. Dynamic activation quantization remains slower because it recomputes amax and cannot use the same fused static input preparation.

Related PRs

• Required dependency: [TRTLLM-13767][chore] upgrade CUTLASS DSL stack to 4.6.1 #17274 upgrades NVIDIA CUTLASS DSL to 4.6.1.
• Parallel halo/local-output contract: [TRTLLM-14555][perf] Optimize parallel VAE halo exchange path #17001 is independent of FP4, but enables the optimized parallel residual epilogue after it lands.
• Temporal decode batching: [TRTLLM-14554][perf] Batch Wan VAE decode along the temporal axis #17002 is precision-independent and already merged.
• Fused DupUp3D: [TRTLLM-14798][perf] Fuse Wan DupUp3D output mapping #17003 is precision-independent and shared by BF16 and FP4.

This diff intentionally does not duplicate the implementations from #17001 or #17003. The performance tables below show the validated selected stack and therefore include the applicable precision-independent optimizations; they are project-level results, not an isolated attribution of every millisecond to this PR.

Performance

FastWan 2.2 TI2V 5B

NSC B200, 121 frames at 704x1280, three DMD steps. VAE-only endpoints pool 16 critical-rank samples from two independent runs; end-to-end values are medians of three measured generations after one warmup.

GPUs	Diffusers BF16 VAE (ms)	Native BF16 best (ms)	Static FP4 best (ms)	Diffusers / FP4	Native / FP4
1	4023.593	2996.796	1375.493	2.925x	2.179x
4	1377.372	851.553	453.881	3.035x	1.876x
8	845.447	437.580	238.314	3.548x	1.836x

GPUs	Diffusers BF16 E2E (s)	Native BF16 best (s)	Static FP4 best (s)	Diffusers / FP4	Native / FP4
1	6.038855	4.949657	3.357354	1.799x	1.474x
4	1.971329	1.441643	1.034641	1.905x	1.393x
8	1.191434	0.784804	0.582953	2.044x	1.346x

Wan2.2 TI2V 5B cross-check

NSC B200, 81 frames at 704x1280. End-to-end uses ten denoise steps, guidance 5, sequence length 512, and seed 1234.

The Wan2.2 checkpoint contains 48 quantized 3x3x3 convolutions: 28 use calibrated static input scales and 20 use the dynamic activation fallback. All 48 execute the FP4 kernel.

GPUs	VAE-only: Diffusers / native best / FP4 best (ms)	FP4 vs native	E2E: Diffusers / native best / FP4 best (s)	FP4 vs native
1	2691.2 / 2024.9 / 909.6	2.226x	10.693 / 10.037 / 8.924	1.125x
4	921.8 / 571.1 / 321.7	1.775x	3.118 / 2.837 / 2.545	1.115x

The denoiser is unchanged, so VAE-only gains are intentionally diluted in end-to-end latency.

Accuracy

FastWan end-to-end comparisons use Diffusers BF16 at the same GPU count as reference:

Candidate	Mean LPIPS	SSIM
Native BF16 best, 1 GPU	0.000065	0.999623
Static FP4 best, 1 GPU	0.005424	0.989049
Native BF16 best, 4 GPU	0.000069	0.999605
Static FP4 best, 4 GPU	0.005454	0.989058

The four-GPU static-FP4 shard-seam error ratio is 1.159 versus 1.207 for native BF16, so the parallel FP4 path shows no quantization-specific seam amplification. On the Wan2.2 ten-step workload, four-GPU FP4 best measures LPIPS 0.013590 versus native BF16 best and 0.013604 versus Diffusers BF16 across all 81 encoded frames.

Full interactive reports:

• FastWan performance and quality report
• Wan2.2 FP4 parallel VAE report

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.
• Kernel reference checks cover output parity and a nonzero-convolution quality threshold.\n- Provider residual-load refinement (dynamic-kernel-generator!20906 commit a28df6ded13) was revalidated on exact PR sources: NSC job 1637299 passed both focused fused bias+residual/product-path tests; direct ABBA job 1637290 passed all reference checks; full-VAE ABBA job 1637284 measured 1103.017 -> 1095.787 ms on 1 GPU and an effectively neutral 371.379 -> 371.987 ms on 4 GPUs.\n- FastWan 1/4/8-GPU VAE-only, matched end-to-end, videos, and 1/4-GPU LPIPS validation completed on NSC B200.

Draft follow-ups

• Rebase after [TRTLLM-14555][perf] Optimize parallel VAE halo exchange path #17001 and enable the parallel residual epilogue when its local-width output contract is available.
• Run the target VisualGen CI stages before marking ready.