[Kernel] Add typed SAGE for TRTLLM_ATTN (#5509)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
This commit is contained in:
SYLAR authored and GitHub committed
1 parent
d33c905d80
commit
896f05e7e6
5 files changed
+301
-94
No files matched your search
@@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
|
||||
@@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T
|
||||
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
|
||||
`TRTLLM_ATTN` raises rather than silently degrading.
|
||||
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
|
||||
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
|
||||
in low precision. Enable it through the typed `quant` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
|
||||
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
|
||||
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
|
||||
|
||||
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
|
||||
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
|
||||
|
||||
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
|
||||
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
|
||||
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
|
||||
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
|
||||
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
|
||||
(B300); `int8` QK kernels are compiled for **SM100 only**.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
@@ -180,22 +222,6 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
### FA4 not yet integrated
|
||||
|
||||
FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
@@ -353,64 +379,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
|
||||
### Mixed backends across roles
|
||||
|
||||
Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--diffusion-attention-config.per_role.self.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_mixed.mp4
|
||||
```
|
||||
|
||||
### Compare against FlashAttention
|
||||
|
||||
Unset the backend override, or explicitly use `FLASH_ATTN`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_fa3.mp4
|
||||
```
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
|
||||
- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
- model
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- frame count / step count
|
||||
- parallel config (TP / CFG-parallel / Ulysses degrees)
|
||||
|
||||
## Reproducing the Benchmark Table
|
||||
|
||||
The end-to-end numbers above were collected by running `text_to_video.py` /
|
||||
`text_to_image.py` with the same prompt and seed while varying
|
||||
`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
|
||||
backends without loading a model:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
```
|
||||
|
||||
It runs all three BF16 backends on representative DiT attention shapes and
|
||||
prints a ranking table at the end.
|
||||
Reference in new issue
Block a user