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

recipes/README.md

@@ -66,6 +66,7 @@ recipes/
 | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
 | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
+| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |

recipes/Wan-AI/Wan2.2-T2V.md

@@ -0,0 +1,186 @@
+# Wan2.2 Text-to-Video
+
+> Text-to-video serving (Wan2.2 14B), with optional Skip-Softmax sparse attention on Blackwell
+
+## Summary
+
+- Vendor: Wan-AI
+- Model: `Wan-AI/Wan2.2-T2V-A14B-Diffusers`
+- Task: Text-to-video generation
+- Mode: Online serving with the OpenAI-compatible API (offline `Omni` also supported)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy the Wan2.2 14B text-to-video model with vLLM-Omni. The
+standard plan is multi-card serving with sequence/CFG parallelism (below). On
+**datacenter Blackwell** (sm_100 / sm_103) you can additionally enable
+**Skip-Softmax** through the `TRTLLM_ATTN` backend for extra speed — see the
+Blackwell section.
+
+## References
+
+- Upstream model card: <https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers>
+- Attention backends guide: [`docs/user_guide/diffusion/attention_backends.md`](../../docs/user_guide/diffusion/attention_backends.md)
+- Skip-Softmax design: [`docs/design/feature/skip_softmax.md`](../../docs/design/feature/skip_softmax.md)
+
+## Hardware Support
+
+## GPU
+
+### 8x NVIDIA H20 / H100 / A100 (standard multi-card serving)
+
+The recommended parallel strategy depends on whether the checkpoint needs
+classifier-free guidance:
+
+1. **Distilled model (no CFG)** — higher throughput; use when the checkpoint does
+   not require negative-prompt computation.
+2. **Official open-source model (with CFG)** — uses CFG parallelism to run the
+   negative and positive samples for the original released weights.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA driver with a CUDA runtime supported by your PyTorch build
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+**Distilled model (no CFG, recommended for distilled checkpoints):**
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --omni \
+  --use-hsdp \
+  --usp 8 \
+  --vae-patch-parallel-size 8 \
+  --vae-use-tiling
+```
+
+**Official open-source model (with CFG):**
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --omni \
+  --use-hsdp \
+  --cfg-parallel-size 2 \
+  --usp 4 \
+  --vae-patch-parallel-size 8 \
+  --vae-use-tiling
+```
+
+The official model splits the positive and negative CFG branches with
+`--cfg-parallel-size 2`, keeping Ulysses sequence parallelism at `--usp 4` so the
+two branches cover all 8 GPUs (`usp * cfg = 8`).
+
+#### Verification
+
+For online serving client examples and request formats, see
+[`examples/online_serving/text_to_video`](../../examples/online_serving/text_to_video).
+
+#### Notes
+
+- **Key flags:**
+  - `--omni` — enables vLLM-Omni diffusion serving.
+  - `--use-hsdp` — Hybrid Sharded Data Parallelism for the 14B DiT weights.
+  - `--usp <N>` — Unified (Ulysses) Sequence Parallelism degree.
+  - `--cfg-parallel-size <N>` — CFG parallelism; set to 2 for the official model,
+    omit for distilled checkpoints.
+  - `--vae-patch-parallel-size 8` / `--vae-use-tiling` — parallel + tiled VAE
+    decoding; disabling patch parallelism can significantly increase VAE latency.
+
+## Accelerating with TRTLLM_ATTN + Skip-Softmax (datacenter Blackwell)
+
+On **datacenter Blackwell** (sm_100 / sm_103), the mask-free Wan pipeline
+auto-routes to the `TRTLLM_ATTN` backend, which can enable **Skip-Softmax** — a
+sparse-attention mode that trades a little fidelity for speed. This composes with
+the multi-card plan above (Ulysses SP is compatible; ring SP is not).
+
+### Environment
+
+- Same as above, plus: datacenter Blackwell (sm_100 / sm_103), `head_dim == 128`,
+  and `flashinfer` (already a hard dependency of vLLM — no extra install).
+
+### Command
+
+Dense BF16 (Skip-Softmax off — `TRTLLM_ATTN` is already the Blackwell auto-route
+for mask-free Wan, so no flags are needed):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni
+```
+
+With Skip-Softmax (calibrated curve — the `(a, b)` curve is read from the
+checkpoint's `sparse_attention_config`; you only pick the operating point):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --omni \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
+```
+
+Offline, the same config is a typed object (or the equivalent dict) — values are
+validated at construction:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
+
+llm = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    diffusion_attention_config=AttentionConfig(
+        default=AttentionSpec(
+            backend="TRTLLM_ATTN",
+            skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
+        ),
+    ),
+)
+```
+
+Calibration-free path (no checkpoint curve needed — set an absolute threshold
+instead of `target_sparsity`):
+
+```bash
+--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+    "skip_softmax": {"threshold": 0.02, "disabled_until_timestep": 0.86}}}'
+```
+
+### Skip-Softmax controls
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `target_sparsity` | finite, `[0, 1]` | Operating point on the checkpoint's calibrated curve (`a·exp(b·s)`). |
+| `threshold` | finite, `≥ 0` | Direct skip threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
+| `disabled_until_timestep` | finite, `[0, 1]` | Keeps early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
+
+**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** (near-lossless)
+and raise `target_sparsity` only as far as your quality bar allows — the output
+diverges from dense as it climbs. Gate the noisy early steps first (`D`), then
+raise sparsity; the timestep gate buys more speed per unit of divergence than
+sparsity does. The gain grows with sequence length, since Skip-Softmax only
+accelerates attention.
+
+Measured on B300 / SM103, 1280×720 / 81f / 50 steps, torch.compile, official
+ModelOpt calibration (speedup and divergence relative to dense):
+
+| config | speedup | LPIPS (divergence vs dense) |
+|---|---|---|
+| dense | 1.000x | - |
+| `s=0.65, D=0.86` | 1.062x | 0.041 |
+| `s=0.75, D=1.00` | 1.181x | 0.377 |
+| `s=1.00` | 1.202x | 0.597 |
+
+### Notes
+
+- **Requirements:** datacenter Blackwell (sm_100 / sm_103), `head_dim == 128`, and
+  `flashinfer`. Elsewhere, or without FlashInfer, selecting `TRTLLM_ATTN` raises
+  rather than silently degrading.
+- **Known limitations:**
+  - Skip-Softmax is not supported under ring sequence parallelism; use Ulysses SP
+    (`--usp`), or run Skip-Softmax without ring.
+  - `TRTLLM_ATTN` is BF16-only (no attention-level quantization).

tests/diffusion/attention/test_attention_config.py

@@ -31,21 +31,27 @@ from vllm_omni.diffusion.data import (
     parse_attention_config,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 class TestAttentionSpec:
-    def test_construct_no_extra(self):
+    def test_construct_no_skip_softmax(self):
         spec = AttentionSpec(backend="FLASH_ATTN")
-        assert spec.extra == {}
+        assert spec.skip_softmax is None
 
-    def test_mapping_extra_normalized(self):
-        spec = AttentionSpec(backend="SAGE_ATTN", extra={"quant": "int8"})
-        assert spec.backend == "SAGE_ATTN"
-        assert spec.extra == {"quant": "int8"}
+    def test_skip_softmax_mapping_coerced(self):
+        spec = AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
+        assert spec.backend == "TRTLLM_ATTN"
+        assert spec.skip_softmax.target_sparsity == 0.5
 
     def test_invalid_backend_type(self):
         with pytest.raises(TypeError):
             AttentionSpec(backend=123)  # type: ignore[arg-type]
 
+    def test_skip_softmax_rejected_on_non_trtllm(self):
+        with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"):
+            AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5})
+
 
 class TestAttentionConfig:
     def test_empty_config(self):
@@ -57,13 +63,13 @@ class TestAttentionConfig:
         config = AttentionConfig(
             default={"backend": "FLASH_ATTN"},
             per_role={
-                "self": {"backend": "SPARSE_BLOCK", "extra": {"block_size": 128}},
+                "self": {"backend": "TRTLLM_ATTN", "skip_softmax": {"target_sparsity": 0.5}},
                 "cross": "SAGE_ATTN",
             },
         )
         assert config.default.backend == "FLASH_ATTN"
-        assert config.per_role["self"].backend == "SPARSE_BLOCK"
-        assert config.per_role["self"].extra == {"block_size": 128}
+        assert config.per_role["self"].backend == "TRTLLM_ATTN"
+        assert config.per_role["self"].skip_softmax.target_sparsity == 0.5
         assert config.per_role["cross"].backend == "SAGE_ATTN"
 
     def test_constructor_flattens_nested_per_role_tree(self):
@@ -178,10 +184,10 @@ class TestAttentionConfig:
         config = AttentionConfig(
             default=AttentionSpec(backend="FLASH_ATTN"),
             per_role={
-                "self": AttentionSpec(backend="SPARSE_BLOCK", extra={"block_size": 128}),
+                "self": AttentionSpec(backend="SPARSE_BLOCK"),
                 "cross": AttentionSpec(backend="SAGE_ATTN"),
                 "ltx2.audio_self": AttentionSpec(backend="FLASH_ATTN"),
-                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN", extra={"causal_window": 64}),
+                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN"),
             },
         )
 
@@ -200,7 +206,6 @@ class TestAttentionConfig:
         # audio-to-video → exact match
         spec, _ = config.resolve_with_source("ltx2.audio_to_video", "cross")
         assert spec.backend == "FLASH_ATTN"
-        assert spec.extra == {"causal_window": 64}
 
         # video-to-audio → category fallback to "cross"
         assert config.resolve_with_source("ltx2.video_to_audio", "cross")[0].backend == "SAGE_ATTN"
@@ -364,12 +369,13 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
             head_size,
             attention_config=None,
             role_category=None,
+            allow_trtllm_default=False,
         ):
             captured["role"] = role
             captured["head_size"] = head_size
             captured["role_category"] = role_category
             captured["attention_config"] = attention_config
-            return _FakeBackend, AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})
+            return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
 
         class _FakeRingParallelAttention:
             def __init__(self, sp_group, attn_backend_pref=None):
@@ -391,7 +397,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         od_config = SimpleNamespace(
             diffusion_attention_config=AttentionConfig(
                 default=AttentionSpec(backend="FLASH_ATTN"),
-                per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
+                per_role={"cross": AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})},
             ),
             parallel_config=SimpleNamespace(ring_degree=2),
             diffusion_kv_cache_dtype=None,
@@ -414,12 +420,12 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         assert captured["role_category"] == "cross"
         assert captured["head_size"] == 64
         assert captured["attention_config"] is od_config.diffusion_attention_config
-        assert attn.backend_pref == "TORCH_SDPA"
-        assert attn.attention.kwargs["backend_kwargs"] == {"block_size": 128}
+        assert attn.backend_pref == "TRTLLM_ATTN"
+        assert attn.attention.kwargs["backend_kwargs"] == {"target_sparsity": 0.5}
         assert attn.attention.kwargs["qkv_layout"] == "BSND"
         assert attn.use_ring is True
         assert attn.ring_runner is not None
-        assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"
+        assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN"
 
 
 class TestDiffusionKvCacheQuantization:
@@ -453,7 +459,10 @@ class TestDiffusionKvCacheQuantization:
         monkeypatch.setattr(
             layer_mod,
             "get_attn_backend_for_role",
-            lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None),
+            lambda role, head_size, attention_config=None, role_category=None, allow_trtllm_default=False: (
+                _FakeBackend,
+                None,
+            ),
         )
         monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
         monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())

tests/diffusion/attention/test_trtllm_attn.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
+from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
+
+
+def _has_trtllm_attn() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    if torch.cuda.get_device_capability()[0] < 10:
+        return False
+    try:
+        import flashinfer  # noqa: F401
+    except Exception:
+        return False
+    return tg.HAS_FLASHINFER
+
+
+requires_trtllm_attn = pytest.mark.skipif(
+    not _has_trtllm_attn(), reason="requires Blackwell SM100+ GPU with flashinfer"
+)
+
+
+def _impl(**bk):
+    return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk)
+
+
+def test_skip_config_pure_resolution():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    assert not SkipSoftmaxConfig().enabled
+    assert SkipSoftmaxConfig().resolve_factor(4096, timestep=0.3) is None
+
+    cfg = SkipSoftmaxConfig.from_backend_kwargs({"skip_softmax_threshold": 0.01, "disabled_until_timestep": 0.6})
+    assert cfg.enabled and cfg.gated
+    assert cfg.resolve_factor(4096, timestep=0.9) is None
+    assert cfg.resolve_factor(4096, timestep=0.3) == pytest.approx(0.01 * 4096)
+    assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096)
+
+
+def test_skip_factor_none_without_curve():
+    assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
+
+
+def _calibrated(a, b, **bk):
+    impl = _impl(**bk)
+    impl.set_layer_calibration(a, b)
+    return impl
+
+
+def test_skip_factor_from_curve():
+    assert _calibrated(100.0, 0.0, target_sparsity=0.5)._resolve_skip_factor(4096) == pytest.approx(100.0)
+    f2 = _calibrated(2.0, 1.0, target_sparsity=0.5)._resolve_skip_factor(4096)
+    assert f2 == pytest.approx(2.0 * math.exp(0.5))
+
+
+def test_skip_factor_none_until_stamped():
+    assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
+
+
+def test_skip_factor_from_threshold_no_calibration():
+    impl = _impl(skip_softmax_threshold=0.01)
+    assert impl._resolve_skip_factor(4096) == pytest.approx(0.01 * 4096)
+    assert impl._resolve_skip_factor(8192) == pytest.approx(0.01 * 8192)
+
+    impl2 = _calibrated(100.0, 0.0, skip_softmax_threshold=0.02, target_sparsity=0.5)
+    assert impl2._resolve_skip_factor(4096) == pytest.approx(0.02 * 4096)
+
+
+def test_skip_timestep_gate(monkeypatch):
+    from vllm_omni.diffusion import forward_context as fc
+
+    impl = _calibrated(100.0, 0.0, target_sparsity=0.5, disabled_until_timestep=0.6)
+    ctx = fc.ForwardContext(denoise_timestep=0.9)
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    assert impl._resolve_skip_factor(4096) is None
+    ctx.denoise_timestep = 0.3
+    assert impl._resolve_skip_factor(4096) == pytest.approx(100.0)
+
+
+def test_denoise_progress_mixin(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion import forward_context as fc
+
+    class _Pipe(fc.DenoiseProgressMixin):
+        scheduler = types.SimpleNamespace(config=types.SimpleNamespace(num_train_timesteps=1000))
+
+    ctx = fc.ForwardContext()
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    _Pipe().record_denoise_step(3, 250)
+    assert ctx.denoise_step_idx == 3
+    assert ctx.denoise_timestep == pytest.approx(0.25)
+
+    _Pipe().record_denoise_step(4)
+    assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25)
+
+
+@pytest.mark.parametrize(
+    "bk, field",
+    [
+        ({"target_sparsity": 1.5}, "target_sparsity"),
+        ({"target_sparsity": -0.1}, "target_sparsity"),
+        ({"disabled_until_timestep": 1.2}, "disabled_until_timestep"),
+        ({"skip_softmax_threshold": -1.0}, "skip_softmax_threshold"),
+        ({"target_sparsity": float("nan")}, "target_sparsity"),
+        ({"skip_softmax_threshold": float("inf")}, "skip_softmax_threshold"),
+    ],
+)
+def test_skip_control_validation_rejects_out_of_range(bk, field):
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    with pytest.raises(ValueError, match=field):
+        SkipSoftmaxConfig.from_backend_kwargs(bk)
+
+
+def test_skip_control_validation_accepts_valid():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
+
+    cfg = SkipSoftmaxConfig.from_backend_kwargs(
+        {"target_sparsity": 0.65, "disabled_until_timestep": 0.86, "skip_softmax_threshold": 0.01}
+    )
+    assert cfg.target_sparsity == pytest.approx(0.65)
+    assert cfg.disabled_until_timestep == pytest.approx(0.86)
+    assert cfg.threshold == pytest.approx(0.01)
+
+
+def test_ring_sp_with_skip_softmax_rejected():
+    import types
+
+    from vllm_omni.diffusion.attention.layer import Attention
+
+    fake = types.SimpleNamespace(attention=_impl(target_sparsity=0.5), ring_runner=None)
+    with pytest.raises(NotImplementedError, match="ring sequence parallelism"):
+        Attention._run_ring_attention(fake, None, None, None, None)
+
+
+def test_mask_free_allowlist_gates_trtllm_default():
+    from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata as meta
+
+    assert meta("WanPipeline").attention_mask_free is True
+    assert meta("WanVACEPipeline").attention_mask_free is True
+    assert meta("FluxPipeline").attention_mask_free is False
+    assert meta(None).attention_mask_free is False
+
+
+def test_masked_layer_raises():
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    mask = torch.ones(b, s, dtype=torch.bool)
+    mask[:, s // 2 :] = False
+
+    class _Meta:
+        attn_mask = mask
+
+    with pytest.raises(ValueError, match="does not support attn_mask"):
+        _impl().forward_cuda(q, k, v, _Meta())
+
+
+def test_propagate_calibration_045_schema(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    cfg = AttentionConfig(default=AttentionSpec(backend="TRTLLM_ATTN"))
+    tf = TransformerConfig.from_dict(
+        {
+            "sparse_attention_config": {
+                "config_groups": {
+                    "group_0": {
+                        "algorithm": "skip_softmax",
+                        "targets": ["WanAttention"],
+                        "ignore": ["blocks.0.attn2", "blocks.1.attn2"],
+                        "threshold_scale_factor": {
+                            "formula": "a * exp(b * target_sparsity)",
+                            "coefficients": {"a": 2142.7, "b": 4.28},
+                        },
+                        "target_sparsity": 0.5,
+                    }
+                }
+            }
+        }
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.trtllm_calibration.parse_secondary_expert_calibration",
+        lambda *a, **k: None,
+    )
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(
+        types.SimpleNamespace(
+            diffusion_attention_config=cfg,
+            model="/x",
+        ),
+        tf,
+    )
+    spec, _ = cfg.resolve_with_source("self")
+    cal = spec.skip_calibration["by_expert"]["transformer"]
+    assert cal["a"] == pytest.approx(2142.7)
+    assert cal["b"] == pytest.approx(4.28)
+    assert cal["ignore"] == ["blocks.0.attn2", "blocks.1.attn2"]
+    assert "transformer_2" not in spec.skip_calibration["by_expert"]
+
+
+def test_target_sparsity_without_calibration_raises():
+    import types
+
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    def _cfg(skip_softmax):
+        return types.SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(
+                default=AttentionSpec(backend="TRTLLM_ATTN", skip_softmax=skip_softmax)
+            ),
+            model="/models/no-such-calibration",
+        )
+
+    with pytest.raises(ValueError, match="no skip-softmax calibration"):
+        OmniDiffusionConfig._propagate_skip_softmax_calibration(
+            _cfg({"target_sparsity": 0.5}), TransformerConfig.from_dict({})
+        )
+
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(_cfg({"threshold": 0.02}), TransformerConfig.from_dict({}))
+
+
+def _sdpa_ref(q, k, v, scale):
+    qp, kp, vp = (x.permute(0, 2, 1, 3) for x in (q, k, v))
+    return F.scaled_dot_product_attention(qp, kp, vp, is_causal=False, scale=scale).permute(0, 2, 1, 3).float()
+
+
+@requires_trtllm_attn
+def test_bf16_dense_matches_sdpa():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 256, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    out = _impl().forward_cuda(q, k, v, None).float()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
+
+
+@requires_trtllm_attn
+def test_skip_tiny_threshold_matches_dense():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 512, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+
+    out = _impl(skip_softmax_threshold=1e-9).forward_cuda(q, k, v, None).float()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.02, f"skip(tiny) rel err {rel:.4f} should be ~dense"
+
+
+@requires_trtllm_attn
+def test_skip_threshold_engages():
+    torch.manual_seed(0)
+    B, S, H, D = 2, 512, 8, 128
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    dense = _impl().forward_cuda(q, k, v, None).float()
+    skipped = _impl(skip_softmax_threshold=2.0).forward_cuda(q, k, v, None).float()
+    assert torch.isfinite(skipped).all()
+    assert (skipped - dense).abs().mean() / dense.abs().mean() > 0.02
+
+
+@requires_trtllm_attn
+def test_gqa_matches_sdpa():
+    torch.manual_seed(0)
+    B, S, Hq, Hkv, D = 1, 256, 64, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q = torch.randn(B, S, Hq, D, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(B, S, Hkv, D, device="cuda", dtype=torch.bfloat16)
+    out = (
+        TrtllmAttentionImpl(Hq, D, scale, causal=False, num_kv_heads=Hkv, backend_kwargs={})
+        .forward_cuda(q, k, v, None)
+        .float()
+    )
+    qp, kp, vp = (x.permute(0, 2, 1, 3) for x in (q, k, v))
+    ref = (
+        F.scaled_dot_product_attention(qp, kp, vp, is_causal=False, scale=scale, enable_gqa=True)
+        .permute(0, 2, 1, 3)
+        .float()
+    )
+    rel = (out - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.01, f"GQA rel err {rel:.4f} vs SDPA(GQA) too high"
+
+
+@requires_trtllm_attn
+def test_unsupported_head_dim_raises():
+    q = k = v = torch.randn(1, 32, 8, 64, device="cuda", dtype=torch.bfloat16)
+    with pytest.raises(Exception, match="[Uu]nsupported head dim"):
+        _impl().forward_cuda(q, k, v, None)
+
+
+def test_skip_end_to_end_config_path(monkeypatch):
+    import types
+
+    from vllm_omni.diffusion import forward_context as fc
+    from vllm_omni.diffusion.data import (
+        AttentionConfig,
+        AttentionSpec,
+        OmniDiffusionConfig,
+        TransformerConfig,
+    )
+
+    cfg = AttentionConfig(
+        default=AttentionSpec(
+            backend="TRTLLM_ATTN",
+            skip_softmax={"target_sparsity": 0.65, "disabled_until_timestep": 0.86},
+        )
+    )
+
+    tf = TransformerConfig.from_dict(
+        {"sparse_attention_config": {"threshold_scale_factor": {"prefill": {"a": 7.93, "b": 8.61}}}}
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.trtllm_calibration.parse_secondary_expert_calibration",
+        lambda *a, **k: None,
+    )
+    OmniDiffusionConfig._propagate_skip_softmax_calibration(
+        types.SimpleNamespace(
+            diffusion_attention_config=cfg,
+            model="/x",
+        ),
+        tf,
+    )
+    spec, _ = cfg.resolve_with_source("self")
+    cal = spec.skip_calibration["by_expert"]["transformer"]
+    assert cal["a"] == 7.93 and cal["b"] == 8.61
+
+    impl = TrtllmAttentionImpl(
+        8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=spec.backend_kwargs()
+    )
+    impl.set_layer_calibration(cal["a"], cal["b"])
+    expected = 7.93 * math.exp(8.61 * 0.65)
+
+    ctx = fc.ForwardContext(denoise_timestep=0.3)
+    monkeypatch.setattr(fc, "_forward_context", ctx)
+    assert impl._resolve_skip_factor(4096) == pytest.approx(expected)

tests/diffusion/attention/test_trtllm_calibration.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
+    is_ignored,
+    layer_match_names,
+    parse_sparse_attention_config,
+    resolve_layer_calibration,
+    select_expert,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_layer_match_names_strips_expert_prefix():
+    names = layer_match_names("transformer_2.blocks.5.attn1")
+    assert "transformer_2.blocks.5.attn1" in names
+    assert "blocks.5.attn1" in names
+
+
+def test_layer_match_names_strips_orig_mod():
+    names = layer_match_names("transformer.blocks.5._orig_mod.attn1")
+    assert "transformer.blocks.5.attn1" in names or "blocks.5.attn1" in names
+
+
+def test_is_ignored_matches_relative_pattern():
+    assert is_ignored("transformer.blocks.5.attn2", ["blocks.*.attn2"])
+    assert not is_ignored("transformer.blocks.5.attn1", ["blocks.*.attn2"])
+
+
+def test_is_ignored_empty_patterns():
+    assert not is_ignored("transformer.blocks.5.attn1", [])
+    assert not is_ignored("transformer.blocks.5.attn1", None)
+
+
+def test_is_ignored_ancestor_match_with_attn_suffix():
+    assert is_ignored("transformer.blocks.0.attn2.attn", ["blocks.0.attn2"])
+
+    assert not is_ignored("transformer.blocks.11.attn2.attn", ["blocks.1.attn2"])
+
+    assert not is_ignored("transformer.blocks.5.attn1.attn", ["blocks.0.attn2", "blocks.1.attn2"])
+
+
+def test_select_expert_longest_prefix_wins():
+    keys = ("transformer", "transformer_2")
+    assert select_expert("transformer_2.blocks.5.attn1", keys) == "transformer_2"
+    assert select_expert("transformer.blocks.5.attn1", keys) == "transformer"
+
+
+def test_select_expert_no_match():
+    assert select_expert("blocks.5.attn1", ("transformer", "transformer_2")) is None
+
+
+def _two_expert_calib():
+    return {
+        "by_expert": {
+            "transformer": {
+                "a": 2142.7,
+                "b": 4.28,
+                "target_sparsity": 0.45,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2"],
+            },
+            "transformer_2": {
+                "a": 314.68,
+                "b": 6.17,
+                "target_sparsity": 0.45,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2"],
+            },
+        }
+    }
+
+
+def test_resolve_routes_to_correct_expert():
+    calib = _two_expert_calib()
+    hi = resolve_layer_calibration("transformer.blocks.5.attn1", calib)
+    lo = resolve_layer_calibration("transformer_2.blocks.5.attn1", calib)
+    assert hi["a"] == pytest.approx(2142.7)
+    assert lo["a"] == pytest.approx(314.68)
+
+
+def test_resolve_ignored_layer_is_dense():
+    assert resolve_layer_calibration("transformer.blocks.5.attn2", _two_expert_calib()) is None
+
+
+def test_resolve_single_transformer():
+    calib = {
+        "a": 104.76,
+        "b": 7.81,
+        "target_sparsity": 0.45,
+        "formula": "a*exp(b*target_sparsity)",
+        "ignore": ["refiner_blocks.*"],
+    }
+    assert resolve_layer_calibration("transformer_blocks.5.attn", calib)["a"] == pytest.approx(104.76)
+    assert resolve_layer_calibration("refiner_blocks.0.attn", calib) is None
+
+
+def test_resolve_none_calibration():
+    assert resolve_layer_calibration("transformer.blocks.5.attn1", None) is None
+
+
+def test_resolve_unmatched_component_is_uncalibrated():
+    calib = {"by_expert": {"transformer": {"a": 1.0, "b": 2.0, "ignore": []}}}
+    assert resolve_layer_calibration("transformer.blocks.5.attn1", calib)["a"] == pytest.approx(1.0)
+    assert resolve_layer_calibration("transformer_2.blocks.5.attn1", calib) is None
+    assert resolve_layer_calibration("blocks.5.attn1", calib) is None
+
+
+def test_parse_045_config_groups_schema():
+    sac = {
+        "config_groups": {
+            "group_0": {
+                "algorithm": "skip_softmax",
+                "ignore": ["blocks.0.attn1"],
+                "target_sparsity": 0.45,
+                "threshold_scale_factor": {
+                    "formula": "a * exp(b * target_sparsity)",
+                    "coefficients": {"a": 2142.7, "b": 4.28},
+                },
+            }
+        }
+    }
+    parsed = parse_sparse_attention_config(sac)
+    assert parsed["a"] == pytest.approx(2142.7)
+    assert parsed["b"] == pytest.approx(4.28)
+    assert parsed["ignore"] == ["blocks.0.attn1"]
+
+
+def test_parse_unsupported_formula_raises():
+    sac = {
+        "config_groups": {
+            "group_0": {
+                "algorithm": "skip_softmax",
+                "threshold_scale_factor": {"formula": "a * b ** target_sparsity", "coefficients": {"a": 1.0, "b": 2.0}},
+            }
+        }
+    }
+    with pytest.raises(ValueError, match="unsupported skip-softmax formula"):
+        parse_sparse_attention_config(sac)
+
+
+def test_parse_empty_returns_none():
+    assert parse_sparse_attention_config(None) is None
+    assert parse_sparse_attention_config({}) is None
+
+
+class _FakeImpl:
+    def __init__(self):
+        self.stamped = None
+
+    def set_layer_calibration(self, a, b):
+        self.stamped = (a, b)
+
+
+def _fake_pipeline():
+    import torch.nn as nn
+
+    def _attn_layer():
+        layer = nn.Module()
+        layer.attention = _FakeImpl()
+        return layer
+
+    def _wan_attn():
+        wa = nn.Module()
+        wa.attn = _attn_layer()
+        return wa
+
+    def _transformer():
+        tf = nn.Module()
+        blocks = nn.ModuleList()
+        for _ in range(3):
+            blk = nn.Module()
+            blk.attn1 = _wan_attn()
+            blk.attn2 = _wan_attn()
+            blocks.append(blk)
+        tf.blocks = blocks
+        return tf
+
+    pipe = nn.Module()
+    pipe.transformer = _transformer()
+    pipe.transformer_2 = _transformer()
+    return pipe
+
+
+def test_apply_to_pipeline_routes_and_ignores():
+    from vllm_omni.diffusion.attention.backends.trtllm_calibration import apply_to_pipeline
+
+    pipe = _fake_pipeline()
+
+    calib = {
+        "by_expert": {
+            "transformer": {
+                "a": 2142.7,
+                "b": 4.28,
+                "target_sparsity": 0.5,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2", "blocks.0.attn1"],
+            },
+            "transformer_2": {
+                "a": 314.68,
+                "b": 6.17,
+                "target_sparsity": 0.5,
+                "formula": "a*exp(b*target_sparsity)",
+                "ignore": ["blocks.*.attn2", "blocks.0.attn1"],
+            },
+        }
+    }
+    stamped = apply_to_pipeline(pipe, calib)
+
+    assert stamped == 4
+
+    def impl(expert, blk, attn):
+        return getattr(getattr(pipe, expert).blocks[blk], attn).attn.attention
+
+    assert impl("transformer", 1, "attn1").stamped == pytest.approx((2142.7, 4.28))
+
+    assert impl("transformer_2", 1, "attn1").stamped == pytest.approx((314.68, 6.17))
+
+    assert impl("transformer", 1, "attn2").stamped is None
+
+    assert impl("transformer", 0, "attn1").stamped is None

vllm_omni/diffusion/attention/backends/registry.py

@@ -64,6 +64,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
     FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
     FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
+    TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+from dataclasses import dataclass, replace
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+
+def _validate_control(value, name: str, lo: float, hi: float | None) -> float | None:
+    if value is None:
+        return None
+    v = float(value)
+    if not math.isfinite(v) or v < lo or (hi is not None and v > hi):
+        rng = f"in [{lo}, {hi}]" if hi is not None else f">= {lo}"
+        raise ValueError(f"{name} must be finite and {rng}; got {value!r}.")
+    return v
+
+
+@dataclass(frozen=True)
+class SkipSoftmaxConfig:
+    threshold: float | None = None
+    target_sparsity: float | None = None
+    disabled_until_timestep: float = 0.0
+    a: float | None = None
+    b: float | None = None
+
+    @classmethod
+    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> "SkipSoftmaxConfig":
+        bk = backend_kwargs or {}
+        return cls(
+            threshold=_validate_control(bk.get("skip_softmax_threshold"), "skip_softmax_threshold", 0.0, None),
+            target_sparsity=_validate_control(bk.get("target_sparsity"), "target_sparsity", 0.0, 1.0),
+            disabled_until_timestep=_validate_control(
+                bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0
+            ),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return self.threshold is not None or (
+            self.target_sparsity is not None and self.a is not None and self.b is not None
+        )
+
+    @property
+    def configured(self) -> bool:
+        return self.threshold is not None or self.target_sparsity is not None
+
+    @property
+    def gated(self) -> bool:
+        return self.disabled_until_timestep > 0.0
+
+    def resolve_factor(self, seqlen: int, timestep: float | None) -> float | None:
+        if self.threshold is not None:
+            factor = self.threshold * seqlen
+        elif self.target_sparsity is not None and self.a is not None and self.b is not None:
+            factor = self.a * math.exp(self.b * self.target_sparsity)
+        else:
+            return None
+        if self.gated and timestep is not None and timestep > self.disabled_until_timestep:
+            return None
+        return factor
+
+
+try:
+    from flashinfer.prefill import trtllm_ragged_attention_deepseek
+
+    HAS_FLASHINFER = True
+except Exception as e:  # pragma: no cover - import guard
+    HAS_FLASHINFER = False
+    logger.warning(
+        "FlashInfer is unavailable; TRTLLM_ATTN backend will not work. Reason: %s",
+        e,
+    )
+
+
+def _workspace_bytes() -> int:
+    import vllm.envs as envs
+
+    return getattr(envs, "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", 394 * 1024 * 1024)
+
+
+class TrtllmAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [128]
+
+    @staticmethod
+    def get_name() -> str:
+        return "TRTLLM_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["TrtllmAttentionImpl"]:
+        return TrtllmAttentionImpl
+
+
+class TrtllmAttentionImpl(AttentionImpl):
+    _workspace: torch.Tensor | None = None
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.softmax_scale = softmax_scale
+        self.causal = causal
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+
+        self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs)
+        self._warned_missing_timestep = False
+
+    def set_layer_calibration(self, a: float, b: float) -> None:
+        self.skip = replace(self.skip, a=a, b=b)
+
+    def _resolve_skip_factor(self, seqlen: int) -> float | None:
+        if not self.skip.enabled:
+            return None
+
+        timestep = None
+        if self.skip.gated:
+            from vllm_omni.diffusion.forward_context import get_forward_context
+
+            timestep = getattr(get_forward_context(), "denoise_timestep", None)
+            if timestep is None:
+                if not self._warned_missing_timestep:
+                    logger.warning(
+                        "TRTLLM skip: disabled_until_timestep=%s set but this pipeline does not "
+                        "publish denoise_timestep; staying dense. Have the pipeline call "
+                        "DenoiseProgressMixin.record_denoise_step to enable timestep gating.",
+                        self.skip.disabled_until_timestep,
+                    )
+                    self._warned_missing_timestep = True
+                return None
+        return self.skip.resolve_factor(seqlen, timestep)
+
+    @classmethod
+    def _get_workspace(cls, device: torch.device) -> torch.Tensor:
+        nbytes = _workspace_bytes()
+        ws = cls._workspace
+        if ws is None or ws.device != device or ws.numel() < nbytes:
+            ws = torch.zeros(nbytes, dtype=torch.uint8, device=device)
+            cls._workspace = ws
+        return ws
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
+        if attn_mask is not None:
+            raise ValueError(
+                "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with "
+                "mask_sp_padding=True and a non-divisible seq len). Either set "
+                "parallel_config.mask_sp_padding=False, or select a mask-capable backend "
+                "(e.g. CUDNN_ATTN / TORCH_SDPA)."
+            )
+
+        if not HAS_FLASHINFER:
+            raise ImportError(
+                "TRTLLM_ATTN backend requires flashinfer. Install it or select "
+                "another backend via --diffusion-attention-backend."
+            )
+
+        batch, q_len, num_q_heads, head_dim = query.shape
+        kv_len, num_kv_heads = key.shape[1], key.shape[2]
+
+        device = query.device
+
+        q = query.reshape(batch * q_len, num_q_heads, head_dim).contiguous()
+        k = key.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous()
+        v = value.reshape(batch * kv_len, num_kv_heads, head_dim).contiguous()
+
+        seq_lens = torch.full((batch,), kv_len, dtype=torch.int32, device=device)
+        cu_seq_lens_q = torch.arange(0, (batch + 1) * q_len, step=q_len, dtype=torch.int32, device=device)
+        cu_seq_lens_kv = torch.arange(0, (batch + 1) * kv_len, step=kv_len, dtype=torch.int32, device=device)
+        workspace = self._get_workspace(device)
+
+        bmm1_scale = self.softmax_scale
+        bmm2_scale = 1.0
+
+        _skip_factor = self._resolve_skip_factor(kv_len)
+
+        out = trtllm_ragged_attention_deepseek(
+            query=q,
+            key=k,
+            value=v,
+            workspace_buffer=workspace,
+            seq_lens=seq_lens,
+            max_q_len=q_len,
+            max_kv_len=kv_len,
+            bmm1_scale=bmm1_scale,
+            bmm2_scale=bmm2_scale,
+            o_sf_scale=-1.0,
+            batch_size=batch,
+            window_left=-1,
+            cum_seq_lens_q=cu_seq_lens_q,
+            cum_seq_lens_kv=cu_seq_lens_kv,
+            enable_pdl=False,
+            is_causal=self.causal,
+            return_lse=False,
+            skip_softmax_threshold_scale_factor=_skip_factor,
+        )
+        return out.reshape(batch, q_len, num_q_heads, head_dim)

vllm_omni/diffusion/attention/backends/trtllm_calibration.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import fnmatch
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_FORMULA = "a*exp(b*target_sparsity)"
+
+
+def parse_sparse_attention_config(sac: dict | None) -> dict | None:
+    if not sac:
+        return None
+    group: dict = {}
+    for g in (sac.get("config_groups") or {}).values():
+        if isinstance(g, dict) and (g.get("algorithm") == "skip_softmax" or g.get("threshold_scale_factor")):
+            group = g
+            break
+    tsf = group.get("threshold_scale_factor") or sac.get("threshold_scale_factor") or {}
+    ab = tsf.get("coefficients") or tsf.get("prefill") or {}
+    if "a" not in ab or "b" not in ab:
+        return None
+
+    formula = (tsf.get("formula") or _SUPPORTED_FORMULA).replace(" ", "")
+    if formula != _SUPPORTED_FORMULA:
+        raise ValueError(
+            f"unsupported skip-softmax formula {tsf.get('formula')!r}; this backend implements '{_SUPPORTED_FORMULA}'."
+        )
+
+    return {"a": float(ab["a"]), "b": float(ab["b"]), "ignore": list(group.get("ignore") or [])}
+
+
+def parse_secondary_expert_calibration(model: str | None) -> dict | None:
+    from vllm.transformers_utils.config import get_hf_file_to_dict
+
+    if not model:
+        return None
+    try:
+        cfg = get_hf_file_to_dict("transformer_2/config.json", model)
+        return parse_sparse_attention_config(cfg.get("sparse_attention_config")) if cfg else None
+    except Exception as exc:  # noqa: BLE001
+        logger.warning(
+            "Could not read transformer_2 skip-softmax calibration (%s); transformer_2 will stay uncalibrated (dense).",
+            exc,
+        )
+        return None
+
+
+def propagate_skip_softmax_calibration(specs: list, model: str | None, tf_config: Any) -> None:
+    calibration = parse_sparse_attention_config(
+        tf_config.get("sparse_attention_config") if tf_config is not None else None
+    )
+    if calibration is not None:
+        by_expert = {"transformer": calibration}
+        ckpt2 = parse_secondary_expert_calibration(model)
+        if ckpt2 is not None:
+            by_expert["transformer_2"] = ckpt2
+        calibration = {"by_expert": by_expert}
+
+    if calibration is None:
+        for spec in specs:
+            ss = spec.skip_softmax
+            if ss is not None and ss.target_sparsity is not None and ss.threshold is None:
+                raise ValueError(
+                    f"target_sparsity was requested but the checkpoint for model '{model}' carries no "
+                    f"skip-softmax calibration. Either set skip_softmax.threshold for the "
+                    f"calibration-free path or load a calibrated checkpoint."
+                )
+        return
+
+    for spec in specs:
+        if spec.skip_calibration is None:
+            spec.skip_calibration = calibration
+
+    experts = list(calibration["by_expert"])
+    logger.info("Loaded skip-softmax calibration from checkpoint (%d curve(s): %s).", len(experts), ", ".join(experts))
+
+
+def apply_skip_softmax_calibration(attention_config: Any, model: Any) -> None:
+    cfg = attention_config
+    specs = getattr(cfg, "per_role", None)
+    calibration = None
+    for spec in (getattr(cfg, "default", None), *(specs.values() if specs else ())):
+        calibration = calibration or getattr(spec, "skip_calibration", None)
+    if not calibration:
+        return
+
+    stamped = apply_to_pipeline(model, calibration)
+    logger.info("Skip-softmax: stamped calibration onto %d attention layer(s).", stamped)
+
+
+_EXPERT_PREFIXES = ("transformer.", "transformer_2.")
+
+
+def layer_match_names(module_name: str) -> tuple[str, ...]:
+    names = {module_name, module_name.replace("._orig_mod.", ".")}
+    for name in tuple(names):
+        for pfx in _EXPERT_PREFIXES:
+            if name.startswith(pfx):
+                names.add(name[len(pfx) :])
+    return tuple(names)
+
+
+def is_ignored(module_name: str, ignore_patterns) -> bool:
+    names = layer_match_names(module_name)
+    for p in ignore_patterns or ():
+        for n in names:
+            if fnmatch.fnmatch(n, p) or fnmatch.fnmatch(n, p + ".*"):
+                return True
+    return False
+
+
+def select_expert(module_name: str, expert_keys) -> str | None:
+    for key in sorted(expert_keys or (), key=len, reverse=True):
+        if module_name == key or module_name.startswith(key + "."):
+            return key
+    return None
+
+
+def resolve_layer_calibration(module_name: str, calibration: dict) -> dict | None:
+    if not calibration:
+        return None
+    by_expert = calibration.get("by_expert")
+    if by_expert:
+        key = select_expert(module_name, by_expert.keys())
+        if key is None:
+            return None
+        entry = by_expert[key]
+    else:
+        entry = calibration
+    if is_ignored(module_name, entry.get("ignore")):
+        return None
+    return {"a": entry.get("a"), "b": entry.get("b")}
+
+
+def apply_to_pipeline(pipeline, calibration: dict) -> int:
+    if not calibration:
+        return 0
+    stamped = 0
+    for name, module in pipeline.named_modules():
+        impl = getattr(module, "attention", None)
+        set_layer = getattr(impl, "set_layer_calibration", None)
+        if set_layer is None:
+            continue
+        per = resolve_layer_calibration(name, calibration)
+        if per and per.get("a") is not None and per.get("b") is not None:
+            set_layer(per["a"], per["b"])
+            stamped += 1
+    return stamped

vllm_omni/diffusion/attention/layer.py

@@ -77,14 +77,20 @@ class Attention(nn.Module):
         config = get_current_diffusion_config_or_none()
         attention_config = config.diffusion_attention_config if config is not None else None
 
+        from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
+
+        model_class_name = getattr(config, "model_class_name", None) if config is not None else None
+        allow_trtllm_default = get_diffusion_model_metadata(model_class_name).attention_mask_free
+
         attn_backend_cls, spec = get_attn_backend_for_role(
             role=role,
             head_size=head_size,
             attention_config=attention_config,
             role_category=role_category,
+            allow_trtllm_default=allow_trtllm_default,
         )
         if spec is not None:
-            backend_kwargs = spec.extra or None
+            backend_kwargs = spec.backend_kwargs()
             self.backend_pref = spec.backend
             logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
         else:
@@ -99,6 +105,7 @@ class Attention(nn.Module):
             causal=causal,
             num_kv_heads=num_kv_heads,
             qkv_layout=qkv_layout,
+            prefix=prefix,
             backend_kwargs=backend_kwargs,
         )
         # Instantiate fallback backend for float32 support
@@ -310,6 +317,13 @@ class Attention(nn.Module):
             )
 
     def _run_ring_attention(self, query, key, value, attn_metadata):
+        skip = getattr(self.attention, "skip", None)
+        if skip is not None and getattr(skip, "configured", False):
+            raise NotImplementedError(
+                "Skip-Softmax (TRTLLM_ATTN) is not supported with ring sequence parallelism: "
+                "the ring path bypasses the backend, so the skip config would be silently ignored. "
+                "Use Ulysses SP instead, or remove the skip_softmax config."
+            )
         # Delegate to RingParallelAttention strategy if available
         if self.ring_runner is not None:
             return self.ring_runner.run_attention(

vllm_omni/diffusion/attention/selector.py

@@ -51,18 +51,20 @@ def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
 def _cached_get_backend_cls(
     backend_name: str | None,
     head_size: int,
+    allow_trtllm_default: bool = True,
 ) -> type[AttentionBackend]:
-    """Cache backend class resolution by (backend_name, head_size).
+    """Cache backend class resolution by (backend_name, head_size, allow_trtllm_default).
 
     This ensures platform validation (compute capability checks, package
-    availability, etc.) runs only once per unique (backend_name, head_size)
-    combination, avoiding repeated log messages.
+    availability, etc.) runs only once per unique combination, avoiding
+    repeated log messages.
     """
     from vllm_omni.platforms import current_omni_platform
 
     backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
         selected_backend=backend_name,
         head_size=head_size,
+        allow_trtllm_default=allow_trtllm_default,
     )
     return _load_backend_cls(backend_cls_path)
 
@@ -97,6 +99,7 @@ def get_attn_backend_for_role(
     head_size: int,
     attention_config: AttentionConfig | None = None,
     role_category: str | None = None,
+    allow_trtllm_default: bool = True,
 ) -> tuple[type[AttentionBackend], AttentionSpec | None]:
     """
     Get attention backend for a specific attention role.
@@ -140,7 +143,7 @@ def get_attn_backend_for_role(
         return backend_cls, spec
 
     # 2. Platform default
-    backend_cls = _cached_get_backend_cls(None, head_size)
+    backend_cls = _cached_get_backend_cls(None, head_size, allow_trtllm_default)
     _log_backend_resolution(
         role=role,
         role_category=role_category,

vllm_omni/diffusion/data.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
+import math
 import os
 import random
 from collections.abc import Callable, Mapping
@@ -1070,6 +1071,19 @@ class OmniDiffusionConfig:
         """
         self.tf_model_config = tf_config
         self._propagate_quantization_from_tf_config(tf_config)
+        self._propagate_skip_softmax_calibration(tf_config)
+
+    def _propagate_skip_softmax_calibration(self, tf_config: "TransformerConfig") -> None:
+        cfg = getattr(self, "diffusion_attention_config", None)
+        if not isinstance(cfg, AttentionConfig):
+            return
+        specs = [s for s in (cfg.default, *cfg.per_role.values()) if s is not None]
+
+        from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
+            propagate_skip_softmax_calibration,
+        )
+
+        propagate_skip_softmax_calibration(specs, self.model, tf_config)
 
     def update_multimodal_support(self) -> None:
         # Resolve serving-visible multimodal behavior from shared metadata
@@ -1323,23 +1337,78 @@ class DiffusionRequestAbortedError(RuntimeError):
     """Raised when a diffusion request ends via user-visible abort."""
 
 
+def _in_range(value: Any, name: str, lo: float, hi: float | None) -> float | None:
+    """Validate an optional numeric control at config-parse time."""
+    if value is None:
+        return None
+    x = float(value)
+    if not math.isfinite(x) or x < lo or (hi is not None and x > hi):
+        rng = f"in [{lo}, {hi}]" if hi is not None else f">= {lo}"
+        raise ValueError(f"{name} must be finite and {rng}; got {value!r}.")
+    return x
+
+
+@dataclass
+class SkipSoftmaxSpec:
+    """User-facing skip-softmax controls for the TRTLLM_ATTN backend.
+
+    ``target_sparsity`` uses the checkpoint's calibrated curve (a·exp(b·s)); ``threshold``
+    is the calibration-free absolute skip threshold. They are mutually exclusive.
+    """
+
+    target_sparsity: float | None = None
+    threshold: float | None = None
+    disabled_until_timestep: float = 0.0
+
+    def __post_init__(self) -> None:
+        self.target_sparsity = _in_range(self.target_sparsity, "skip_softmax.target_sparsity", 0.0, 1.0)
+        self.threshold = _in_range(self.threshold, "skip_softmax.threshold", 0.0, None)
+        self.disabled_until_timestep = (
+            _in_range(self.disabled_until_timestep, "skip_softmax.disabled_until_timestep", 0.0, 1.0) or 0.0
+        )
+        if self.target_sparsity is not None and self.threshold is not None:
+            raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.")
+
+
 @dataclass
 class AttentionSpec:
-    """Specifies a backend and its backend-specific parameters for one attention role."""
+    """Specifies a backend and its typed backend-specific config for one attention role."""
 
-    backend: str  # registry name, e.g. "FLASH_ATTN"
-    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific kwargs
+    backend: str
+    skip_softmax: SkipSoftmaxSpec | None = None
+    skip_calibration: dict | None = field(default=None, repr=False)
 
     def __post_init__(self) -> None:
         if not isinstance(self.backend, str):
             raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
+        self.skip_softmax = self._coerce_skip_softmax(self.skip_softmax)
+        if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
+            raise ValueError(
+                f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
+                f"Remove skip_softmax or set backend to TRTLLM_ATTN."
+            )
 
-        if self.extra is None:
-            self.extra = {}
-        elif isinstance(self.extra, Mapping):
-            self.extra = dict(self.extra)
-        else:
-            raise TypeError(f"Expected dict for AttentionSpec.extra, got {type(self.extra)!r}")
+    @staticmethod
+    def _coerce_skip_softmax(value: Any) -> SkipSoftmaxSpec | None:
+        if value is None or isinstance(value, SkipSoftmaxSpec):
+            return value
+        if isinstance(value, Mapping):
+            return SkipSoftmaxSpec(**dict(value))
+        raise TypeError(f"Expected dict or SkipSoftmaxSpec for skip_softmax, got {type(value)!r}")
+
+    def backend_kwargs(self) -> dict[str, Any] | None:
+        """Serialize typed backend config into the kwargs dict the backend impl consumes."""
+        if self.skip_softmax is None:
+            return None
+        ss = self.skip_softmax
+        kw: dict[str, Any] = {}
+        if ss.threshold is not None:
+            kw["skip_softmax_threshold"] = ss.threshold
+        if ss.target_sparsity is not None:
+            kw["target_sparsity"] = ss.target_sparsity
+        if ss.disabled_until_timestep:
+            kw["disabled_until_timestep"] = ss.disabled_until_timestep
+        return kw or None
 
 
 @dataclass
@@ -1408,13 +1477,13 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "extra"}
+        spec_keys = {"backend", "skip_softmax"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys:
             if not node_keys <= spec_keys:
                 raise ValueError(
-                    f"Invalid per_role entry for role {role!r}: cannot mix backend/extra with nested role keys."
+                    f"Invalid per_role entry for role {role!r}: cannot mix backend/skip_softmax with nested role keys."
                 )
             normalized[role] = node_dict
             return

vllm_omni/diffusion/forward_context.py

@@ -28,6 +28,7 @@ class ForwardContext:
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
+    denoise_timestep: float | None = None
     # Per-request reference latent for img2img DiT models (e.g. Ming)
     ref_latent: torch.Tensor | None = None
     # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
@@ -214,6 +215,21 @@ def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
         _forward_context.denoise_step_idx = step_idx
 
 
+class DenoiseProgressMixin:
+    def record_denoise_step(
+        self,
+        step_idx: int | None,
+        timestep=None,
+        scheduler=None,
+    ) -> None:
+        set_forward_context_denoise_step_idx(step_idx)
+        if timestep is None or _forward_context is None:
+            return
+        scheduler = scheduler if scheduler is not None else getattr(self, "scheduler", None)
+        ntt = getattr(getattr(scheduler, "config", None), "num_train_timesteps", None)
+        _forward_context.denoise_timestep = float(timestep) / ntt if ntt else None
+
+
 def set_forward_context_ref_latent(ref_latent: torch.Tensor | None) -> None:
     """Set the per-request reference latent on the active ForwardContext.
 

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -378,6 +378,7 @@ class DiffusersPipelineLoader:
                 model.to("cpu")
                 logger.info("Quantization complete, offloaded model back to CPU")
 
+        self._apply_skip_softmax_calibration(model)
         return model.eval()
 
     @staticmethod
@@ -391,6 +392,14 @@ class DiffusersPipelineLoader:
                 return True
         return False
 
+    def _apply_skip_softmax_calibration(self, model: nn.Module) -> None:
+        from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
+            apply_skip_softmax_calibration,
+        )
+
+        cfg = getattr(self.od_config, "diffusion_attention_config", None)
+        apply_skip_softmax_calibration(cfg, model)
+
     def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
         """Process weights after loading for quantization methods.
 

vllm_omni/diffusion/model_metadata.py

@@ -10,6 +10,7 @@ class DiffusionModelMetadata:
     # config/model plumbing can read it without importing concrete pipelines.
     supports_multimodal_inputs: bool = False
     max_multimodal_image_inputs: int | None = None
+    attention_mask_free: bool = False
 
 
 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
@@ -35,6 +36,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES,
     ),
+    "WanPipeline": DiffusionModelMetadata(attention_mask_free=True),
+    "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True),
+    "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True),
+    "WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True),
 }
 
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -24,7 +24,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -264,6 +264,7 @@ class Wan22Pipeline(
     PipelineParallelMixin,
     CFGParallelMixin,
     ProgressBarMixin,
+    DenoiseProgressMixin,
     DiffusionPipelineProfilerMixin,
     SupportsComponentDiscovery,
 ):
@@ -460,7 +461,7 @@ class Wan22Pipeline(
         with self.progress_bar(total=len(timesteps)) as pbar:
             for step_idx, t in enumerate(timesteps):
                 self._current_timestep = t
-                set_forward_context_denoise_step_idx(step_idx)
+                self.record_denoise_step(step_idx, t)
 
                 # Select model based on timestep and boundary_ratio
                 # High noise stage (t >= boundary_timestep): use transformer

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -24,7 +24,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -144,6 +144,7 @@ class Wan22I2VPipeline(
     CFGParallelMixin,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
+    DenoiseProgressMixin,
     SupportsComponentDiscovery,
 ):
     """
@@ -354,7 +355,7 @@ class Wan22I2VPipeline(
                     current_model = self.transformer_2
                     current_guidance_scale = guidance_high
 
-                set_forward_context_denoise_step_idx(step_idx)
+                self.record_denoise_step(step_idx, t)
 
                 # Prepare latent input
                 if self.expand_timesteps:

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
@@ -462,7 +463,13 @@ def _load_wan_t5_as_umt5(
 
 
 class Wan22S2VPipeline(
-    nn.Module, SupportImageInput, SupportAudioInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    SupportImageInput,
+    SupportAudioInput,
+    CFGParallelMixin,
+    DenoiseProgressMixin,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
 ):
     """
     Wan2.2 Speech-to-Video Pipeline.
@@ -993,8 +1000,9 @@ class Wan22S2VPipeline(
         do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None
 
         with self.progress_bar(total=len(timesteps)) as pbar:
-            for t in timesteps:
+            for step_idx, t in enumerate(timesteps):
                 self._current_timestep = t
+                self.record_denoise_step(step_idx, t)
 
                 latent_model_input = [latents.to(device)]
                 timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -25,7 +25,6 @@ from vllm.model_executor.layers.quantization.base_config import QuantizationConf
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
     Wan22Pipeline,
@@ -204,7 +203,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         with self.progress_bar(total=len(timesteps)) as pbar:
             for step_idx, t in enumerate(timesteps):
                 self._current_timestep = t
-                set_forward_context_denoise_step_idx(step_idx)
+                self.record_denoise_step(step_idx, t)
 
                 if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None:
                     current_model = self.transformer_2

vllm_omni/entrypoints/cli/serve.py

@@ -561,7 +561,8 @@ class OmniServeCommand(CLISubcommand):
             help="Diffusion attention config. Accepts JSON or vLLM-style dotted flags. "
             "Examples: "
             "--diffusion-attention-config.default.backend FLASH_ATTN, "
-            "--diffusion-attention-config.per_role.self.backend SPARSE_BLOCK, "
+            "--diffusion-attention-config.default.backend TRTLLM_ATTN "
+            "--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5, "
             "--diffusion-attention-config.per_role.cross.backend SAGE_ATTN, "
             '--diffusion-attention-config \'{"default": {"backend": "FLASH_ATTN"}, '
             '"per_role": {"cross": {"backend": "SAGE_ATTN"}}}\'.',

vllm_omni/platforms/cuda/platform.py

@@ -57,6 +57,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         cls,
         selected_backend: str | None,
         head_size: int,
+        allow_trtllm_default: bool = True,
     ) -> str:
         from vllm_omni.diffusion.envs import PACKAGES_CHECKER
 
@@ -119,6 +120,18 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
             # not abort startup — just treat FlashInfer as unavailable.
             logger.debug("FlashInfer probe failed (%s); treating as unavailable", e)
 
+        # TRTLLM_ATTN needs the trtllm-gen kernel specifically, not just any FlashInfer
+        # wheel. Probe the actual symbol so a released wheel lacking it does not get
+        # auto-routed to TRTLLM_ATTN and then crash on the first forward.
+        trtllm_gen_available = False
+        if flashinfer_available:
+            try:
+                from flashinfer.prefill import trtllm_ragged_attention_deepseek  # noqa: F401
+
+                trtllm_gen_available = True
+            except Exception as e:
+                logger.debug("trtllm-gen kernel probe failed (%s); treating as unavailable", e)
+
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
@@ -172,10 +185,37 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
                         "SageAttention/sageattention3_blackwell. Falling back to TORCH_SDPA backend."
                     )
                     return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            if backend_upper == "TRTLLM_ATTN":
+                trtllm_attn_supported = compute_capability is not None and compute_capability.major == 10
+                if not trtllm_attn_supported:
+                    raise ValueError(
+                        "TRTLLM_ATTN diffusion attention backend requires a datacenter "
+                        "Blackwell GPU (SM100 / SM103, compute capability 10.x). Select a "
+                        "different --diffusion-attention-backend."
+                    )
+                if not flashinfer_available:
+                    raise ValueError(
+                        "TRTLLM_ATTN diffusion attention backend requires flashinfer, which is not installed."
+                    )
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
+        trtllm_attn_default_ok = (
+            allow_trtllm_default
+            and compute_capability is not None
+            and compute_capability.major == 10
+            and head_size == 128
+            and trtllm_gen_available
+        )
+        if trtllm_attn_default_ok:
+            logger.info(
+                "Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell %s, head_dim %d)",
+                sm_str,
+                head_size,
+            )
+            return DiffusionAttentionBackendEnum.TRTLLM_ATTN.get_path()
+
         if is_blackwell and cudnn_blackwell_ready:
             logger.info(
                 "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d)",