Skip to content

feat(thor): add FA4 and stride-based split-N NVFP4 support for Wan2.2 - #1344

Open
Michael20070814 wants to merge 25 commits into
ModelTC:mainfrom
Michael20070814:extreme-compression-thor
Open

feat(thor): add FA4 and stride-based split-N NVFP4 support for Wan2.2#1344
Michael20070814 wants to merge 25 commits into
ModelTC:mainfrom
Michael20070814:extreme-compression-thor

Conversation

@Michael20070814

Copy link
Copy Markdown
Contributor

Summary

This PR adds the Wan2.2 acceleration path for NVIDIA Thor:

  • enables NVFP4 kernels on Thor (sm_110a);
  • adds a stride-based, batched split-N NVFP4 GEMM for both Wan FFN projections;
  • integrates FlashAttention-4 variable-length attention and FA4-backed dynamic sparse attention;
  • adds Thor configurations for NVFP4, split-N, and split-N with torch.compile.

Performance

For Wan2.2 MoE I2V on Thor (4 denoising steps, 81 frames, 480x832), the complete optimized configuration reduced end-to-end inference latency from approximately 45 s to 16 s, corresponding to a 2.81x speedup and a 64.4% latency reduction.

This result reflects the complete configuration and should not be interpreted as the isolated gain from split-N, NVFP4, or FA4.

Before the optimization:
加速前

After the optimization:
加速后

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant