[Attention] Add trtllm diffusion attention backend with Skip-Softmax (#5283)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

docs/user_guide/diffusion/attention_backends.md

@@ -14,6 +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`. |
 | `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`. |
@@ -68,11 +69,12 @@ vllm-omni serve <model> \
     --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
 ```
 
-Backends may also accept backend-specific parameters via `extra`:
+A backend that needs configuration exposes it as a typed field on the spec. Today the only
+one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)):
 
 ```bash
---diffusion-attention-config.per_role.self.backend SPARSE_BLOCK \
---diffusion-attention-config.per_role.self.extra.block_size 128
+--diffusion-attention-config.default.backend TRTLLM_ATTN \
+--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
 ```
 
 ### Programmatic API
@@ -101,10 +103,15 @@ A plain dict is also accepted and normalized to `AttentionConfig`.
 
 Auto-route preference, in order:
 
-1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
-2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
-3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
-4. `TORCH_SDPA` — last resort
+1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model is mask-free** (see below)
+2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
+3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
+4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
+5. `TORCH_SDPA` — last resort
+
+`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
+
+`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell because it is the only backend that can enable Skip-Softmax, not because its dense kernel is faster — dense, the two are comparable. It cannot honor an attention mask, so the auto-default only applies to pipelines verified mask-free (the Wan family); mask-using pipelines keep `CUDNN_ATTN`. `TRTLLM_ATTN` stays available everywhere via explicit `--diffusion-attention-backend TRTLLM_ATTN` (which raises clearly if a mask is received).
 
 The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
 
@@ -119,6 +126,47 @@ Auto-route preference:
 
 `CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
 
+## TRTLLM_ATTN Backend and Skip-Softmax
+
+`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
+that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
+(algorithm: [design doc](../../design/feature/skip_softmax.md)).
+
+Enable it through the typed `skip_softmax` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
+| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
+| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
+```
+
+Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
+    ),
+)
+```
+
+**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
+as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
+sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
+`skip_softmax` block.
+
+Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
+`TRTLLM_ATTN` raises rather than silently degrading.
+
 ## 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`.