[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/.nav.yml

@@ -116,6 +116,7 @@ nav:
       - design/feature/hsdp.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
+      - design/feature/skip_softmax.md
       - design/feature/async_chunk.md
       - design/feature/vae_parallel.md
       - design/feature/diffusion_step_execution.md

docs/design/feature/skip_softmax.md

@@ -0,0 +1,88 @@
+# Skip-Softmax
+
+Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
+and how to pick an operating point — is in
+[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
+the algorithm.
+
+## Motivation
+
+In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
+concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
+output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
+detects, per block of keys, when a block cannot matter and skips its softmax and its value
+multiply.
+
+It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
+opt-in and off by default.
+
+## The online-softmax pass
+
+Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
+kernel maintains three running values:
+
+- `m` — the largest score seen so far,
+- `l` — the running denominator `Σ exp(sⱼ − m)`,
+- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,
+
+and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
+updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
+`m` grows keeps the single pass numerically exact.
+
+## The skip test
+
+Once a tile's scores are known, its largest score `tile_max` is compared against the running
+maximum:
+
+```text
+if exp(tile_max − running_max) < threshold:
+    skip this tile          # do not compute its softmax weights or its P·V accumulation
+```
+
+`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
+receive: if even the tile's best key is far below the current maximum, every key in the tile is
+negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
+dropped. The tile's contribution to `l` and `O` is simply skipped.
+
+## What this bounds
+
+Two properties of the test shape the achievable speedup:
+
+- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
+  score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
+  and the value matmul are comparable in cost, so skipping every eligible tile removes at most
+  roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
+
+- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
+  collectively negligible; a single important key keeps the whole tile. How many tiles actually
+  qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
+  operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
+  skipped.
+
+## The threshold
+
+The per-tile threshold is normalized by sequence length:
+
+```text
+threshold = factor / seqlen
+```
+
+`factor` comes from one of two sources:
+
+- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
+  per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
+  skipped tiles on the calibration data.
+- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
+  the threshold themselves.
+
+## Timestep gating
+
+`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
+turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
+schedule). The early steps set the global structure of the output and their errors propagate through
+every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
+
+`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
+`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
+because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
+stays dense when `disabled_until_timestep` is set.

docs/design/module/dit_module.md

@@ -498,7 +498,7 @@ class Attention(nn.Module):
             role_category=role_category,
         )
         self.attn_impl_cls = attn_backend_cls.get_impl_cls()
-        self.attention = self.attn_impl_cls(..., backend_kwargs=spec.extra if spec else None)
+        self.attention = self.attn_impl_cls(..., backend_kwargs=spec.backend_kwargs() if spec else None)
 ```
 
 **Available Backends**:
@@ -540,7 +540,7 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ
 
 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
 
-`AttentionSpec.extra` is forwarded to the backend constructor as `backend_kwargs`, so backend-specific parameters (e.g. SparseBlock `block_size`) can be set without changing model code.
+A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code.
 
 For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
 

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`.