Add ModelOpt Triton attention kernels for WAN2.2 diffusion (sparse, skip-softmax, NVFP4) - #1190

Open

yeyu-nvidia wants to merge 63 commits into

mainfrom

yeyu/sage-attention-diffusion

Commits on Apr 8, 2026

Add vLLM integration for modelopt sparse attention

kaix-nv
authored and
yeyu-nvidia
committed
Replace impl to support decode

kaix-nv
authored and
yeyu-nvidia
committed
Remove the monkey-patch code and unify impl replacement

kaix-nv
authored and
yeyu-nvidia
committed
Fix per-layer sparse config and sliding_window

kaix-nv
authored and
yeyu-nvidia
committed
Add Wan2.2 SageAttention example for diffusion inference

yeyu-nvidia
and
claude
committed
Add sageattn call counters for diagnostics

yeyu-nvidia
and
claude
committed
Add SA2 kernel variants (sage2-fp16, sage2-fp8) and full benchmark

yeyu-nvidia
and
claude
committed
Handle missing SA2 kernels gracefully with version detection

yeyu-nvidia
and
claude
committed
Add GPU SM compatibility check and graceful kernel fallback

yeyu-nvidia
and
claude
committed
Add FP8 E4M3 attention and accuracy metrics to wan2 sage attention example

yeyu-nvidia
and
claude
committed
Set 5B model as default in wan2 sage attention example

yeyu-nvidia
and
claude
committed
Add CLIP score to --compare mode for semantic quality evaluation

yeyu-nvidia
and
claude
committed
Fix CLIP scoring: add --clip-model flag, HF_TOKEN support, error handling

yeyu-nvidia
and
claude
committed
Add NVFP4 E2M1 attention kernel to wan2 sage attention example

yeyu-nvidia
and
claude
committed
Fix CLIP score range comment and update output message

yeyu-nvidia
and
claude
committed
Fix nvfp4: quantize Q/K only, keep V in original precision

yeyu-nvidia
and
claude
committed
Fix nvfp4: quantize V to FP8 E4M3 instead of original precision

yeyu-nvidia
and
claude
committed
Fix nvfp4: keep V in BF16 to preserve visual quality

yeyu-nvidia
and
claude
committed
Fix nvfp4: quantize post-softmax P instead of Q/K (SA3-faithful)

yeyu-nvidia
and
claude
committed
Replace nvfp4 kernel with int4: 16-level P quantization for diffusion

yeyu-nvidia
and
claude
committed
Fix nvfp4: use per-tile NVFP4 quantization of P (64x64 tiles)

yeyu-nvidia
and
claude
committed
Fix nvfp4 OOM: chunked row processing to avoid full NxN attention matrix

yeyu-nvidia
and
claude
committed
Add --nvfp4-tile: configurable per-tile NVFP4 P quantization granularity

yeyu-nvidia
and
claude
committed
Add diffusers_triton backend for WAN2.2 sparse attention

yeyu-nvidia
and
claude
committed
Add unit tests for diffusers WAN sparse attention plugin

yeyu-nvidia
and
claude
committed
Fix TestForwardShape: use m.disable() instead of proc._enabled=False

yeyu-nvidia
and
claude
committed
Add NVFP4 P-matrix attention quantization for WAN2.2 via mtq.quantize()

yeyu-nvidia
and
claude
committed
Fix OOM in _QuantWanAttnProcessor: chunk attention rows to avoid N×N float32 matrix

yeyu-nvidia
and
claude
committed
Add NVFP4 P-matrix quantization to Triton flash-attention kernel

yeyu-nvidia
and
claude
committed
Fix WanSparseAttentionModule.forward: call processor directly

yeyu-nvidia
and
claude
committed
Fix WanAttention sparse registration order: diffusers plugin before HF generic
yeyu-nvidia
committed
Lower default skip-softmax threshold and add --skip-threshold CLI arg

yeyu-nvidia
and
claude
committed
Fix skip-softmax threshold formula: remove erroneous * sm_scale factor

yeyu-nvidia
and
claude
committed
Address PR review comments
yeyu-nvidia
committed
Address remaining PR review comments in triton_fa.py
yeyu-nvidia
committed
Revert skip-softmax threshold formula change: restore * sm_scale
yeyu-nvidia
committed
Build SageAttention as standalone quantization feature

yeyu-nvidia
and
claude
committed
Add nvfp4 kernel choice for standalone SageAttention in example script

yeyu-nvidia
and
claude
committed

Commits on Apr 15, 2026

Commits on May 26, 2026

Add systematic accuracy/speed benchmark harness for NVFP4 attention

yeyu-nvidia
and
claude
committed

Commits on May 28, 2026