[Docs] Refactor the attention backend docs/skill (#3475)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.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):

.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.<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",
         )
 ```
 

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

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

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.<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.<category>.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/<device>/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:

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 <model> --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 <model> \
+    --diffusion-attention-config.default.backend FLASH_ATTN \
+    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA
+
+# JSON
+vllm-omni serve <model> \
+    --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
+```
+
+Backends may also accept backend-specific parameters via `extra`:
+
+```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`: