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

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

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