From ef7f2f1cd2158bd55d00ee811eeb09608468841a Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 11 May 2026 09:13:07 +0800 Subject: [PATCH] [Docs] Refactor the attention backend docs/skill (#3475) Signed-off-by: gcanlin Co-authored-by: Hongsheng Liu --- .../references/parallelism-patterns.md | 1 + .../references/transformer-adaptation.md | 20 ++++ .../references/troubleshooting.md | 2 +- .../model/adding_diffusion_model.md | 37 ++++++- docs/design/module/dit_module.md | 102 +++++++++--------- .../diffusion/attention_backends.md | 90 ++++++++++++++-- 6 files changed, 186 insertions(+), 66 deletions(-) diff --git a/.claude/skills/add-diffusion-model/references/parallelism-patterns.md b/.claude/skills/add-diffusion-model/references/parallelism-patterns.md index 933e2d232..56a21d36b 100644 --- a/.claude/skills/add-diffusion-model/references/parallelism-patterns.md +++ b/.claude/skills/add-diffusion-model/references/parallelism-patterns.md @@ -92,6 +92,7 @@ class TPSelfAttention(nn.Module): softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False, num_kv_heads=self.to_qkv.num_kv_heads, # Local KV heads per GPU + role="self", ) def forward(self, x): diff --git a/.claude/skills/add-diffusion-model/references/transformer-adaptation.md b/.claude/skills/add-diffusion-model/references/transformer-adaptation.md index 6e344b6a6..64abc7712 100644 --- a/.claude/skills/add-diffusion-model/references/transformer-adaptation.md +++ b/.claude/skills/add-diffusion-model/references/transformer-adaptation.md @@ -61,6 +61,7 @@ class SelfAttentionBlock(nn.Module): softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False, num_kv_heads=num_heads, + role="self", ) def forward(self, x, attn_mask=None): @@ -75,6 +76,24 @@ class SelfAttentionBlock(nn.Module): return self.to_out(out) ``` +**Declaring the attention role.** Always pass a `role` string to `Attention(...)`. It is the lookup key users target with `--diffusion-attention-config.per_role..backend` to swap kernels per site without touching model code. + +| Convention | When to use | +|---|---| +| `role="self"` | Q/K/V come from the same hidden state | +| `role="cross"` | K/V come from a separate `encoder_hidden_states`; pair with `skip_sequence_parallel=True` if K/V is replicated across SP ranks | + +For multi-modal sites that don't fit `self` / `cross`, use a dot-namespaced role and provide `role_category` so user config can fall back to the generic category: + +```python +self.audio_to_video_attn = Attention( + num_heads=num_heads, head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False, + role="mymodel.audio_to_video", + role_category="cross", +) +``` + #### Fused QKV with TP (Advanced) For tensor parallelism, use vLLM's parallel linear layers: @@ -104,6 +123,7 @@ class TPSelfAttention(nn.Module): softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False, num_kv_heads=num_heads, + role="self", ) ``` diff --git a/.claude/skills/add-diffusion-model/references/troubleshooting.md b/.claude/skills/add-diffusion-model/references/troubleshooting.md index 27acdd8d1..a15e61c87 100644 --- a/.claude/skills/add-diffusion-model/references/troubleshooting.md +++ b/.claude/skills/add-diffusion-model/references/troubleshooting.md @@ -68,7 +68,7 @@ out = out.reshape(batch, seq_len, -1) ### Different Output vs Diffusers Reference **Common causes**: -1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers +1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Force SDPA via `--diffusion-attention-backend TORCH_SDPA` (or `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA`) to match diffusers; for finer-grained per-role overrides see `--diffusion-attention-config.per_role..backend` 2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations 3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved 4. **Scheduler rounding**: Some schedulers have numerical sensitivity diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 1c08b8bef..35ff6dae2 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -115,17 +115,20 @@ class YourAttentionBlock(nn.Module): from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata -class YourAttentionBlock(nn.Module): +class YourSelfAttentionBlock(nn.Module): def __init__(self, ...): super().__init__() - # Initialize vLLM-Omni's Attention layer + # Initialize vLLM-Omni's Attention layer. + # `role` lets users target this site with --diffusion-attention-config + # (e.g. --diffusion-attention-config.per_role.self.backend SAGE_ATTN). self.attn = Attention( num_heads=self.num_heads, head_size=self.head_dim, softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False, # Diffusion models typically use bidirectional attention num_kv_heads=self.num_kv_heads, + role="self", ) def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...): @@ -133,7 +136,6 @@ class YourAttentionBlock(nn.Module): # Create attention metadata attn_metadata = AttentionMetadata(attn_mask=attention_mask) hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata) - ``` **Key Points:** @@ -141,8 +143,35 @@ class YourAttentionBlock(nn.Module): - **Attention layer initialization:** Done in `__init__`, not per-forward - **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape - **AttentionMetadata:** Wraps attention mask and other metadata +- **Role:** Tag every `Attention` site with a `role` string so users can configure backends per role (see below) -**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models. +**Declaring attention roles** + +The `role` argument is a free-form string that identifies this attention site. Users can match it from `--diffusion-attention-config.per_role..*` to swap backends without touching model code. Two conventions cover the common cases: + +| Convention | When to use | Example | +|---|---|---| +| `"self"` | Q/K/V come from the same hidden state | DiT self-attention block | +| `"cross"` | K/V come from a separate `encoder_hidden_states` | Text-conditioned cross-attention | + +For multi-modal or unusual sites, use a dot-namespaced role and pair it with `role_category` so it can fall back to the generic config when nothing model-specific is set: + +```python +# A model-specific cross-attention site that user config can target +# either as 'mymodel.audio_to_video' (exact) or as 'cross' (category fallback). +self.audio_to_video_attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim ** 0.5), + causal=False, + role="mymodel.audio_to_video", + role_category="cross", +) +``` + +For cross-attention sites whose K/V are replicated across ranks (e.g. text encoder output), pass `skip_sequence_parallel=True` to opt this layer out of sequence-parallel sharding. + +**Attention backends:** When the user does not configure a backend, vLLM-Omni asks the current platform for its default (typically `FLASH_ATTN` on CUDA when available). Users override the default via `--diffusion-attention-backend`, the `DIFFUSION_ATTENTION_BACKEND` env var, or finer-grained `--diffusion-attention-config.per_role.*` flags. See [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md) for the full configuration surface. #### 1.3: Replace Imports and Utilities diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index e0fe00299..4390793e9 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -455,7 +455,7 @@ To learn more about the diffusion pipeline and how to add a new diffusion pipeli #### Architecture -The attention system uses a **backend selector pattern** that automatically chooses the optimal attention implementation based on hardware and model configuration. +The attention system uses a **role-aware backend selector pattern**. Each `Attention` site declares a semantic `role` (`"self"`, `"cross"`, or a model-specific string). The selector consults the user's `AttentionConfig` to pick a backend per role and falls back to the platform default when nothing is configured. #### Backend Selection @@ -463,92 +463,90 @@ The attention system uses a **backend selector pattern** that automatically choo ```python class Attention(nn.Module): - def __init__(self, num_heads, head_size, causal, softmax_scale, ...): - # Auto-select backend - self.attn_backend = get_attn_backend(-1) - self.attn_impl_cls = self.attn_backend.get_impl_cls() - self.attention = self.attn_impl_cls(...) + def __init__(self, num_heads, head_size, causal, softmax_scale, *, + role="self", role_category=None, ...): + # Resolve backend for this role from the active OmniDiffusionConfig + config = get_current_diffusion_config_or_none() + attention_config = config.attention_config if config is not None else None + attn_backend_cls, spec = get_attn_backend_for_role( + role=role, + head_size=head_size, + attention_config=attention_config, + 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) ``` **Available Backends**: -- **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling +- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU. -- **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform +- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback. -- **SageAttention**: Sparse attention implementation from SageAttention library - -- **AscendAttention**: NPU-optimized attention for Ascend hardware +- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library. These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention). #### Backend Selection Mechanism +Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order: + ```python -def get_attn_backend(head_size: int) -> type[AttentionBackend]: - # Check environment variable - backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND") - - if backend_name: - return load_backend(backend_name.upper()) - - # Default to SDPA - return SDPABackend +def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None): + # 1. attention_config.per_role[role] — exact match + # 2. attention_config.per_role[role_category] — category fallback + # 3. attention_config.default — global default + # 4. Platform default — hardware-specific + if attention_config is not None: + spec, source = attention_config.resolve_with_source( + role=role, role_category=role_category, + ) + if spec is not None: + return load_backend(spec.backend), spec + return current_omni_platform.get_diffusion_attn_backend_cls(...), None ``` **Selection Priority**: -1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override +1. **Per-role override** (`--diffusion-attention-config.per_role..backend`) — finest control; matched on the layer's exact `role` string. - - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND` +2. **Role category fallback** (`--diffusion-attention-config.per_role..backend`) — used when an exact match is missing and the layer declared `role_category` (e.g. a `"mymodel.audio_to_video"` site can fall back to `"cross"`). - - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN` +3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var). -2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable +4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware. -3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.) +`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. + +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). **Backend Availability**: - **SDPA**: Always available (PyTorch built-in) -- **FlashAttention**: Requires `flash-attn` package installed +- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU - **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub) -- **AscendAttention**: Only available on Ascend NPU hardware - #### Attention Backend Registry -**Location**: `vllm_omni/diffusion/attention/selector.py` +**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms//platform.py` -The attention system uses a **registry pattern** to manage and dynamically load attention backends. This allows for easy extension and runtime selection of backends. - - -**Registry Structure**: +Backend resolution is delegated to the active platform: `current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size)` returns a fully-qualified class path that `_load_backend_cls` imports lazily, with the result cached by `(backend_name, head_size)`. ```python -# Registry mapping backend names to their module paths and class names -_BACKEND_CONFIG = { - "FLASH_ATTN": { - "module": "vllm_omni.diffusion.attention.backends.flash_attn", - "class": "FlashAttentionBackend", - }, - "TORCH_SDPA": { - "module": "vllm_omni.diffusion.attention.backends.sdpa", - "class": "SDPABackend", - }, - "SAGE_ATTN": { - "module": "vllm_omni.diffusion.attention.backends.sage_attn", - "class": "SageAttentionBackend", - }, - "ASCEND": { - "module": "vllm_omni.diffusion.attention.backends.ascend_attn", - "class": "AscendAttentionBackend", - }, -} +@cache +def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]: + backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls( + selected_backend=backend_name, + head_size=head_size, + ) + return _load_backend_cls(backend_cls_path) ``` +Each platform (`cuda`, `rocm`, `xpu`, `musa`, `npu`) maps backend names like `"FLASH_ATTN"`, `"TORCH_SDPA"`, `"SAGE_ATTN"` to the right backend class path for that hardware, including head-size compatibility checks. Passing `selected_backend=None` lets the platform pick its own default. + #### Attention Backend Integration The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example: diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 692bcc0f9..4a19ad8c4 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -1,10 +1,10 @@ # Diffusion Attention Backends -This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention. +This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention. ## Overview -Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`. +Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured. This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. @@ -14,7 +14,7 @@ On CUDA, the practical choices today are: - `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`. - `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed. -If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`. +If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`. ## Backend Options @@ -24,19 +24,77 @@ If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform t | `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. | | `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. | -## Selection Priority +## Configuration -Diffusion attention backend selection follows this order: +Diffusion attention backends can be configured three ways, in priority order: -1. `DIFFUSION_ATTENTION_BACKEND` -2. Platform default +1. **`--diffusion-attention-config`** — structured per-role config (highest priority). +2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend. +3. **Platform default** — used when nothing is configured. -Example: +`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`. + +### Global default + +Set the default backend for every diffusion attention layer: ```bash +# CLI flag +vllm-omni serve --diffusion-attention-backend SAGE_ATTN + +# Environment variable (also recognized for backwards compatibility) export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN ``` +### Per-role configuration + +Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order: + +1. Exact `per_role[role]` match +2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`) +3. `default` +4. Platform default + +Use vLLM-style dotted flags or one JSON blob: + +```bash +# Dotted flags +vllm-omni serve \ + --diffusion-attention-config.default.backend FLASH_ATTN \ + --diffusion-attention-config.per_role.cross.backend TORCH_SDPA + +# JSON +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`: + +```bash +--diffusion-attention-config.per_role.self.backend SPARSE_BLOCK \ +--diffusion-attention-config.per_role.self.extra.block_size 128 +``` + +### Programmatic API + +When constructing `OmniDiffusionConfig` directly: + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig + +config = OmniDiffusionConfig( + attention_config=AttentionConfig( + default=AttentionSpec(backend="FLASH_ATTN"), + per_role={ + "cross": AttentionSpec(backend="TORCH_SDPA"), + }, + ), + ..., +) +``` + +A plain dict is also accepted and normalized to `AttentionConfig`. + ## SageAttention Installation vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni. @@ -85,7 +143,21 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_ --output outputs/wan22_sage.mp4 ``` -### Compare Against FlashAttention +### Mixed backends across roles + +Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + --prompt "A dog running across a field of golden wheat." \ + --diffusion-attention-config.per_role.self.backend FLASH_ATTN \ + --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \ + --tensor-parallel-size 2 \ + --output outputs/wan22_mixed.mp4 +``` + +### Compare against FlashAttention Unset the backend override, or explicitly use `FLASH_ATTN`: