FP8 GEMM, FP8 QKV and FP8 Sol-Attn with Examples of Dense&Sol-sparse Wan2.1

[Uxito-Ada]

Summary

Add (1) FP8-Dense and (2) FP8-Sol sparse Wan 2.1, built on (a) FP8 GEMM, (b) FP8 QKV and (c) FP8 fused Sol-attn kernels.

Implementation

• All 300 Wan layers to dynamic W8A8 E4M3 tf-kernel GEMMs.
• Self-attention Q/K/V projections reuse one activation quantization across three GEMMs, while text encoding and VAE execution remain BF16.
• Add fused Q/K/V quantizer with RoPE. QK use one E4M3 scale per block (64 tokens), and V uses per-channel.
• Add an SM90 CuTe Sol-attn: E4M3 QK and PV with FP32 accumulation. Dequantization are fused inside, and it consumes FP8 inputs directly.
• Correctness for Wan with masking: tail padding, two-way KV splitting, and route-length restoration for FP8 PV.
• Expose BF16 Dense, BF16 Sol, FP8 Dense, and FP8 Sol through one example. FP8 Dense uses the same E4M3 Q/K/V CuTe QK/PV as FP8 Sol.

Benchmark

Environment and parameters:

• GPU: one NVIDIA H100 80GB
• Model: Wan2.1 T2V-1.3B
• Resolution: 832x480
• Frames: 81
• Inference steps: 50
• Solver: UniPC
• CFG scale: 5.0
• Sigma shift: 5.0
• Seed: 42
• SOL dense timesteps: 10
• SOL dense layers: 1
• SOL tau: 1.0
• Each configuration runs in a separate process
• Timing starts after pipeline loading and includes first-execution kernel/JIT
  costs

The bars show peak allocated GPU memory. The line shows generation throughput. @lzx1413

It is seen that :

• FP8 Sol-Attn is the best overall configuration: 36.2% faster than BF16 Dense, while using 0.417 GiB less peak memory (-2.6%).
• FP8 also improves the dense: FP8 Dense is 2.9% faster than BF16 Dense and reduces peak allocation by 0.417 GiB (-2.6%).
• FP8 removes the Sol memory overhead and adds incremental speed: FP8 Sol is 4.3% faster and uses 1.293 GiB less peak memory (-7.6%) than BF16 Sol.
• Different from native Sol, FP8 Sol shows equal peak allocation to FP8 dense in the help of our sol CuTe kernel, and FP8 Sol is 32.3% faster than FP8 Dense.

Video Generation Comparsion

Prompt:

Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.

	Dense	Sol
BF16	https://github.com/user-attachments/assets/2a338e71-3d3d-4c0c-ab01-851a0b3684b0	https://github.com/user-attachments/assets/9dcd44f7-6d49-42c4-916c-d683dcee0051
FP8	https://github.com/user-attachments/assets/c40b3114-8c99-4f4c-91fa-2f091c030ca7	https://github.com/user-attachments/assets/31068c1c-62e0-42ba-9588-4bb174e97f9b

Testing

• Unit tests pass
• Manual H100 video generation performed for all four ablations
• Throughput, peak-memory, PSNR, and SSIM benchmarks added

Commands and results:

ruff check telefuser tests examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py
# All checks passed

ruff format --check telefuser tests examples/wan_video/wan21_1_3b_text_to_video_optimized_h100.py
# 526 files already formatted

pytest tests/unit -q
# 1631 passed, 11 skipped

Checklist

• Code follows the project coding standards and passes PR-scoped Ruff lint and format checks
• Repository-wide pre-commit run --all-files passes:
• Unit tests pass locally
• New tests cover the added functionality
• English and Chinese attention documentation and the Wan README are updated
• Commit messages are clear and descriptive
• Proposed PR title follows the [TYPE] Brief description convention

Related Issues

N/A

GPU Architecture Support

• SM80 (Ampere/Ada): native FP8 Sol-Attn is not enabled
• SM90 (Hopper H100): native E4M3 CuTe QK/PV mainloop
• SM100+ (Blackwell): existing BF16 Sol-Attn remains unchanged; native FP8 path is not enabled

Performance Impact

See the benchmark section.