From 7762d664d8550bd21cd71f092c56b53e679d5a74 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:48:11 +0800 Subject: [PATCH] [Attention] Add trtllm diffusion attention backend with Skip-Softmax (#5283) Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- docs/.nav.yml | 1 + docs/design/feature/skip_softmax.md | 88 +++++ docs/design/module/dit_module.md | 4 +- .../diffusion/attention_backends.md | 62 ++- recipes/README.md | 1 + recipes/Wan-AI/Wan2.2-T2V.md | 186 +++++++++ .../attention/test_attention_config.py | 45 ++- tests/diffusion/attention/test_trtllm_attn.py | 358 ++++++++++++++++++ .../attention/test_trtllm_calibration.py | 224 +++++++++++ .../diffusion/attention/backends/registry.py | 1 + .../attention/backends/trtllm_attn.py | 226 +++++++++++ .../attention/backends/trtllm_calibration.py | 152 ++++++++ vllm_omni/diffusion/attention/layer.py | 16 +- vllm_omni/diffusion/attention/selector.py | 11 +- vllm_omni/diffusion/data.py | 91 ++++- vllm_omni/diffusion/forward_context.py | 16 + .../model_loader/diffusers_loader.py | 9 + vllm_omni/diffusion/model_metadata.py | 5 + .../models/wan2_2/pipeline_wan2_2.py | 5 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 5 +- .../models/wan2_2/pipeline_wan2_2_s2v.py | 12 +- .../models/wan2_2/pipeline_wan2_2_vace.py | 3 +- vllm_omni/entrypoints/cli/serve.py | 3 +- vllm_omni/platforms/cuda/platform.py | 40 ++ 24 files changed, 1512 insertions(+), 52 deletions(-) create mode 100644 docs/design/feature/skip_softmax.md create mode 100644 recipes/Wan-AI/Wan2.2-T2V.md create mode 100644 tests/diffusion/attention/test_trtllm_attn.py create mode 100644 tests/diffusion/attention/test_trtllm_calibration.py create mode 100644 vllm_omni/diffusion/attention/backends/trtllm_attn.py create mode 100644 vllm_omni/diffusion/attention/backends/trtllm_calibration.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 4cb9d20ed..1da3deb94 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md new file mode 100644 index 000000000..8e6375ee8 --- /dev/null +++ b/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. diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index ace355427..aabdb78b0 100644 --- a/docs/design/module/dit_module.md +++ b/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). diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 50e02971f..e54955c04 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/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 \ --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`. diff --git a/recipes/README.md b/recipes/README.md index 94273caa3..e5e8473ee 100644 --- a/recipes/README.md +++ b/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) | diff --git a/recipes/Wan-AI/Wan2.2-T2V.md b/recipes/Wan-AI/Wan2.2-T2V.md new file mode 100644 index 000000000..640129a47 --- /dev/null +++ b/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: +- 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 ` — Unified (Ulysses) Sequence Parallelism degree. + - `--cfg-parallel-size ` — 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). diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index a4c67851a..26a93e901 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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()) diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py new file mode 100644 index 000000000..75d04c2b8 --- /dev/null +++ b/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) diff --git a/tests/diffusion/attention/test_trtllm_calibration.py b/tests/diffusion/attention/test_trtllm_calibration.py new file mode 100644 index 000000000..4db7eaa16 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index a822d78fb..92d3b6926 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/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). diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py new file mode 100644 index 000000000..9e3b305f1 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/trtllm_calibration.py b/vllm_omni/diffusion/attention/backends/trtllm_calibration.py new file mode 100644 index 000000000..71b26a53f --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index a11aecbb8..77d65f79e 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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( diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 6eede02c2..9d357b4f6 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/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, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9e049b211..d3880f900 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index ea230ed9b..a8e23756f 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/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. diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index bb048451a..9864d8dbf 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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. diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 7e02cf5d8..4952539b7 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/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), } diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index d3726e087..e3617f84d 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 8b042b595..d8c1d4a8c 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/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: diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 47b8e02d1..52a42cb79 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/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) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 4173b8601..d0cdb1250 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 98c086d86..ce0549597 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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"}}}\'.', diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 7be135352..4f203cfc2 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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)",