Add Wan FP8 quantization with Sol-Attn - #30
Open
Uxito-Ada wants to merge 3 commits into
Open
Conversation
Add tf-kernel and TorchAO online FP8 choices for Wan2.1 transformer blocks, with dense and Sol-Attn runnable examples. Keep active Sol-Attn Q/K/V in BF16 while preserving Wan residual dtypes, and cover the quantization and attention interaction with focused tests. Document the H100 ablation setup and results, including a peak-memory bar and throughput line benchmark chart. Verified with 18 focused pytest cases (including the H100 Sol-Attn kernel), both example --help entry points, pre-commit hooks, and git diff --check.
Mock the FP8 wrapping entry points so the Wan quantization unit test validates filter and option wiring without requiring the optional tf-kernel package. Remove the benchmark PNG asset and its README reference while retaining the benchmark table and reproduction details. Verified the previously failing test, the 18-test focused suite, Ruff, pre-commit, and staged diff checks.
Replace the separate dense FP8 and Sol-Attn FP8 scripts with one Wan2.1 example that independently selects attention and quantization from CLI configuration. Document Dense/SOL and BF16/tf-kernel FP8/TorchAO FP8/NF4 combinations, and rename the focused tests around the consolidated interface. Verified the direct --help entry point, 18 focused tests including the H100 SOL kernel, pre-commit hooks, PR-scope diff checks, and stale-reference audit.
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
1 participant
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
Add online FP8 quantization support for Wan2.1 and make it compatible with Sol-Attn.
This PR provides:
H100 Benchmark
Environment and parameters:
costs
Prompt:
The bars show peak allocated GPU memory. The line shows generation throughput.
Compared with BF16 dense:
It is seen that Sol-Attn introduces more memory consumption because it saves a huge sparse mask, therefore, quantization is useful to shrink the memory trace when combined at the same time, and also raises the throughput. @lzx1413
Video Generation Comparision
Testing
18 passed:tests/unit/models/test_wan_video_sol_attention.pytests/unit/pipelines/wan_video/test_sol_quant_example.py--helpentry points passedgit diff --checkpassedChecklist