@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 77c4f45c6c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 11:15 UTC

[7/N][Attention][Docs] Add documentation for attention backends (#32477)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Jan 28, 2026 at 17:20 UTC 77c4f45c6c949c4e56f838c43b6907f6473ea65f
4 files changed +1481 -1
.pre-commit-config.yaml
+4
@@ -154,6 +154,10 @@ repos:
154 files: ^docker/(Dockerfile|versions\.json)$
155 pass_filenames: false
156 additional_dependencies: [dockerfile-parse]
157 + - id: attention-backend-docs
158 + name: Check attention backend documentation is up to date
159 + entry: python tools/pre_commit/generate_attention_backend_docs.py --check
160 + language: python
161 # Keep `suggestion` last
162 - id: suggestion
163 name: Suggestion
docs/configuration/optimization.md
+7 -1
@@ -288,5 +288,11 @@ Based on the configuration, the content of the multi-modal caches on `P0` and `P
288 | shm | Shared Memory Caching | K | N/A | V | `mm_processor_cache_gb * api_server_count` |
289 | N/A | Disabled | N/A | N/A | N/A | `0` |
290
291 -K: Stores the hashes of multi-modal items
291 +K: Stores the hashes of multi-modal items
292 V: Stores the processed tensor data of multi-modal items
293 +
294 +## Attention Backend Selection
295 +
296 +vLLM supports multiple attention backends optimized for different hardware and use cases. The backend is automatically selected based on your GPU architecture, model type, and configuration, but you can also manually specify one for optimal performance.
297 +
298 +For detailed information on available backends, their feature support, and how to configure them, see the [Attention Backend Feature Support](../design/attention_backends.md) documentation.
docs/design/attention_backends.md new
+212
@@ -0,0 +1,212 @@
1 +# Attention Backend Feature Support
2 +
3 +This document is auto-generated by `tools/pre_commit/generate_attention_backend_docs.py`.
4 +It shows the feature support for each registered attention backend
5 +based on the checks in `AttentionBackend.validate_configuration()`.
6 +
7 +**Do not edit this file manually.** Run the following command to
8 +regenerate it:
9 +
10 +```bash
11 +python tools/pre_commit/generate_attention_backend_docs.py
12 +```
13 +
14 +## Setting the Attention Backend
15 +
16 +### Command Line
17 +
18 +There are two ways to specify the backend from the command line:
19 +
20 +**Option 1: Using `--attention-backend` (simple)**
21 +
22 +```bash
23 +vllm serve <model> --attention-backend FLASH_ATTN
24 +```
25 +
26 +**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)**
27 +
28 +```bash
29 +# Dot notation
30 +vllm serve <model> --attention-config.backend FLASH_ATTN
31 +vllm serve <model> -ac.backend FLASH_ATTN
32 +
33 +# JSON format
34 +vllm serve <model> --attention-config '{"backend": "FLASH_ATTN"}'
35 +vllm serve <model> -ac '{"backend": "FLASH_ATTN"}'
36 +```
37 +
38 +> **Note:** `--attention-backend` and `--attention-config.backend` are mutually
39 +> exclusive. Use one or the other, not both.
40 +
41 +### Python API
42 +
43 +Use `AttentionConfig` with the `LLM` class:
44 +
45 +```python
46 +from vllm import LLM
47 +from vllm.config import AttentionConfig
48 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
49 +
50 +# Method 1: Using AttentionConfig with enum
51 +llm = LLM(
52 + model="Qwen/Qwen3-0.6B",
53 + attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN),
54 +)
55 +
56 +# Method 2: Using attention_backend parameter with string
57 +llm = LLM(
58 + model="Qwen/Qwen3-0.6B",
59 + attention_backend="FLASH_ATTN",
60 +)
61 +```
62 +
63 +## Backend Selection Behavior
64 +
65 +### Manual Selection
66 +
67 +When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
68 +
69 +1. The backend is **validated** against your configuration (model dtype, head
70 + size, compute capability, etc.)
71 +2. If the backend **doesn't support** your configuration, an error is raised
72 + with the specific reason
73 +3. If valid, the backend is used
74 +
75 +Example error when selecting an incompatible backend:
76 +
77 +```text
78 +ValueError: Selected backend FLASHMLA is not valid for this configuration.
79 +Reason: ['compute capability not supported']
80 +```
81 +
82 +### Automatic Selection
83 +
84 +When no backend is specified (the default):
85 +
86 +1. vLLM iterates through backends in **priority order** (see tables below)
87 +2. Each backend is validated against your configuration
88 +3. The **first compatible backend** is selected
89 +4. If no backend is compatible, an error is raised listing all backends and
90 + their incompatibility reasons
91 +
92 +## Backend Priority (CUDA)
93 +
94 +When no backend is explicitly selected, vLLM chooses the first
95 +compatible backend from these priority-ordered lists.
96 +
97 +Priority is **1 = highest** (tried first).
98 +
99 +### Standard Attention (MHA, MQA, GQA)
100 +
101 +**Blackwell (SM 10.x):**
102 +
103 +| Priority | Backend |
104 +|----------|---------|
105 +| 1 | `FLASHINFER` |
106 +| 2 | `FLASH_ATTN` |
107 +| 3 | `TRITON_ATTN` |
108 +| 4 | `FLEX_ATTENTION` |
109 +
110 +**Ampere/Hopper (SM 8.x-9.x):**
111 +
112 +| Priority | Backend |
113 +|----------|---------|
114 +| 1 | `FLASH_ATTN` |
115 +| 2 | `FLASHINFER` |
116 +| 3 | `TRITON_ATTN` |
117 +| 4 | `FLEX_ATTENTION` |
118 +
119 +### MLA Attention (DeepSeek-style)
120 +
121 +**Blackwell (SM 10.x):**
122 +
123 +| Priority | Backend |
124 +|----------|---------|
125 +| 1 | `FLASHINFER_MLA` |
126 +| 2 | `CUTLASS_MLA` |
127 +| 3 | `FLASH_ATTN_MLA` |
128 +| 4 | `FLASHMLA` |
129 +| 5 | `TRITON_MLA` |
130 +| 6 | `FLASHMLA_SPARSE` |
131 +
132 +**Ampere/Hopper (SM 8.x-9.x):**
133 +
134 +| Priority | Backend |
135 +|----------|---------|
136 +| 1 | `FLASH_ATTN_MLA` |
137 +| 2 | `FLASHMLA` |
138 +| 3 | `FLASHINFER_MLA` |
139 +| 4 | `TRITON_MLA` |
140 +| 5 | `FLASHMLA_SPARSE` |
141 +
142 +> **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
143 +
144 +## Legend
145 +
146 +| Column | Description |
147 +|--------|-------------|
148 +| **Dtypes** | Supported model data types (fp16, bf16, fp32) |
149 +| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) |
150 +| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
151 +| **Head Sizes** | Supported attention head sizes |
152 +| **Sink** | Attention sink support (for StreamingLLM) |
153 +| **Sparse** | Sparse attention support (MLA only) |
154 +| **MM Prefix** | Multimodal prefix full attention support |
155 +| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
156 +| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
157 +
158 +**Symbols:** ✅ = Supported, ❌ = Not supported
159 +
160 +## Standard Attention (MHA, MQA, GQA) Backends
161 +
162 +| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | Attention Types | Compute Cap. |
163 +|---------|---------|--------|-----------|-------------|------------|------|-----------|-----------------|--------------|
164 +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | All | N/A |
165 +| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | Decoder | 7.x-9.x |
166 +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | Decoder | 10.x |
167 +| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | All | ≥8.0 |
168 +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | All | 9.x |
169 +| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | Any |
170 +| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | Decoder, Encoder Only | Any |
171 +| `ROCM_AITER_FA` | | fp16, bf16 | `auto` | %16 | 64, 128, 256 | ❌ | ❌ | Decoder | N/A |
172 +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | Decoder | N/A |
173 +| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | N/A |
174 +| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | Decoder | Any |
175 +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | All | Any |
176 +
177 +> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
178 +>
179 +> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2` or `3`. Default is FA3 on SM90, FA2 otherwise.
180 +
181 +## MLA (Multi-head Latent Attention) Backends
182 +
183 +MLA uses separate backends for prefill and decode phases.
184 +
185 +### Prefill Backends
186 +
187 +The prefill backend is selected at runtime based on hardware and
188 +configuration.
189 +
190 +| Backend | Description | Compute Cap. | Enable | Disable | Notes |
191 +|---------|-------------|--------------|--------|---------|-------|
192 +| TRT-LLM Ragged‡ | TensorRT-LLM ragged attention | 10.x | Default on SM100 | `-ac.use_trtllm_ragged_deepseek_prefill=0` | DeepSeek R1 dims only |
193 +| FlashInfer | FlashInfer CUTLASS backend | 10.x | `-ac.disable_flashinfer_prefill=0` | `-ac.disable_flashinfer_prefill=1` | DeepSeek R1 dims only |
194 +| cuDNN | cuDNN-based attention | 10.x | `-ac.use_cudnn_prefill=1` | `-ac.use_cudnn_prefill=0` | |
195 +| FlashAttention | FlashAttention varlen (FA2/FA3) | Any | Default fallback | Use other backends | FA3 on SM90, FA2 otherwise |
196 +
197 +> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
198 +> On other GPUs, FlashAttention is used as the default.
199 +
200 +### Decode Backends
201 +
202 +| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |
203 +|---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------|
204 +| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
205 +| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
206 +| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x |
207 +| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x |
208 +| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x |
209 +| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | Decoder | N/A |
210 +| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | Decoder | N/A |
211 +| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
212 +| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | Decoder | Any |
tools/pre_commit/generate_attention_backend_docs.py new
+1258
@@ -0,0 +1,1258 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Generates documentation table for attention backends showing feature support.
5 +
6 +This script parses all registered attention backends using AST (no imports needed)
7 +and generates a markdown table showing what features each backend supports,
8 +based on the checks in AttentionBackend.validate_configuration().
9 +
10 +This approach avoids requiring CUDA/ROCm/GPU libraries to be installed.
11 +
12 +When used as a pre-commit hook, this script receives filenames as arguments
13 +and only runs the check if any of the relevant files were modified.
14 +"""
15 +
16 +import argparse
17 +import ast
18 +import fnmatch
19 +import sys
20 +from pathlib import Path
21 +from typing import Any
22 +
23 +REPO_ROOT = Path(__file__).parent.parent.parent
24 +
25 +RELEVANT_PATTERNS = [
26 + "vllm/v1/attention/backends/*.py",
27 + "vllm/v1/attention/backends/**/*.py",
28 + "vllm/v1/attention/backends/fa_utils.py",
29 + "vllm/model_executor/layers/attention/mla_attention.py",
30 + "vllm/platforms/cuda.py",
31 + "tools/pre_commit/generate_attention_backend_docs.py",
32 + "docs/design/attention_backends.md",
33 +]
34 +
35 +
36 +def is_relevant_file(filepath: str) -> bool:
37 + """Check if a file matches any of the relevant patterns."""
38 + path = Path(filepath)
39 + if path.is_absolute():
40 + try:
41 + path = path.relative_to(REPO_ROOT)
42 + except ValueError:
43 + return False
44 + path_str = str(path)
45 +
46 + return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)
47 +
48 +
49 +BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends"
50 +REGISTRY_FILE = BACKENDS_DIR / "registry.py"
51 +CUDA_PLATFORM_FILE = REPO_ROOT / "vllm" / "platforms" / "cuda.py"
52 +FA_UTILS_FILE = BACKENDS_DIR / "fa_utils.py"
53 +FLASHINFER_UTILS_FILE = REPO_ROOT / "vllm" / "utils" / "flashinfer.py"
54 +MLA_ATTENTION_FILE = (
55 + REPO_ROOT / "vllm" / "model_executor" / "layers" / "attention" / "mla_attention.py"
56 +)
57 +
58 +
59 +def parse_registry() -> dict[str, str]:
60 + """Parse the registry.py file to get backend names and their class paths."""
61 + tree = ast.parse(REGISTRY_FILE.read_text())
62 + for node in ast.walk(tree):
63 + if isinstance(node, ast.ClassDef) and node.name == "AttentionBackendEnum":
64 + return _extract_enum_values(node)
65 + return {}
66 +
67 +
68 +def _extract_enum_values(node: ast.ClassDef) -> dict[str, str]:
69 + """Extract enum name -> value mapping from a class definition."""
70 + result: dict[str, str] = {}
71 + for item in node.body:
72 + if not isinstance(item, ast.Assign):
73 + continue
74 + for target in item.targets:
75 + if not isinstance(target, ast.Name):
76 + continue
77 + if isinstance(item.value, ast.Constant) and item.value.value:
78 + result[target.id] = item.value.value
79 + return result
80 +
81 +
82 +def get_file_from_class_path(class_path: str) -> Path | None:
83 + """Convert a class path to a file path."""
84 + if not class_path:
85 + return None
86 + module_path = class_path.rsplit(".", 1)[0].replace(".", "/")
87 + py_file = REPO_ROOT / f"{module_path}.py"
88 + return py_file if py_file.exists() else None
89 +
90 +
91 +def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
92 + """Parse fa_utils.py to detect FA2 vs FA3 feature differences.
93 +
94 + Returns a dict with 'fa2' and 'fa3' keys containing their respective
95 + feature overrides for compute capability, KV cache dtypes, and sink support.
96 + """
97 + if not FA_UTILS_FILE.exists():
98 + return {}
99 +
100 + try:
101 + tree = ast.parse(FA_UTILS_FILE.read_text())
102 + except Exception:
103 + return {}
104 +
105 + # Analyze the functions to determine FA3-specific features
106 + fa3_supports_fp8 = False
107 + fa3_supports_sinks = False
108 + fa3_compute_cap: str | None = None
109 +
110 + for node in ast.walk(tree):
111 + if not isinstance(node, ast.FunctionDef):
112 + continue
113 +
114 + # Check flash_attn_supports_fp8 - looks for `get_flash_attn_version() == 3`
115 + if node.name == "flash_attn_supports_fp8":
116 + for n in ast.walk(node):
117 + if (
118 + isinstance(n, ast.Compare)
119 + and isinstance(n.left, ast.Call)
120 + and isinstance(n.left.func, ast.Name)
121 + and n.left.func.id == "get_flash_attn_version"
122 + ):
123 + fa3_supports_fp8 = True
124 + break
125 +
126 + # Check flash_attn_supports_sinks - looks for `get_flash_attn_version() == 3`
127 + if node.name == "flash_attn_supports_sinks":
128 + for n in ast.walk(node):
129 + if (
130 + isinstance(n, ast.Compare)
131 + and isinstance(n.left, ast.Call)
132 + and isinstance(n.left.func, ast.Name)
133 + and n.left.func.id == "get_flash_attn_version"
134 + ):
135 + fa3_supports_sinks = True
136 + break
137 +
138 + # Check get_flash_attn_version for FA3 compute capability
139 + # Look for the ternary: 3 if (device_capability.major == 9 ...) else 2
140 + if node.name == "get_flash_attn_version":
141 + for n in ast.walk(node):
142 + # Look for IfExp (ternary) with `device_capability.major == 9`
143 + if isinstance(n, ast.IfExp):
144 + test = n.test
145 + # Check if test is a BoolOp (and) containing the major check
146 + if isinstance(test, ast.BoolOp):
147 + for val in test.values:
148 + if (
149 + isinstance(val, ast.Compare)
150 + and isinstance(val.left, ast.Attribute)
151 + and val.left.attr == "major"
152 + and val.comparators
153 + and isinstance(val.comparators[0], ast.Constant)
154 + ):
155 + fa3_compute_cap = f"{val.comparators[0].value}.x"
156 + break
157 +
158 + return {
159 + "fa2": {
160 + "supports_fp8": False,
161 + "supports_sink": False,
162 + },
163 + "fa3": {
164 + "compute_capability": fa3_compute_cap,
165 + "supports_fp8": fa3_supports_fp8,
166 + "supports_sink": fa3_supports_sinks,
167 + },
168 + }
169 +
170 +
171 +def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
172 + """Parse flashinfer.py to detect TRTLLM-specific features.
173 +
174 + FLASHINFER uses TRTLLM attention on SM100 (Blackwell), which has different
175 + capabilities (e.g., sink support) than native FlashInfer on earlier GPUs.
176 + """
177 + if not FLASHINFER_UTILS_FILE.exists():
178 + return {}
179 +
180 + try:
181 + tree = ast.parse(FLASHINFER_UTILS_FILE.read_text())
182 + except Exception:
183 + return {}
184 +
185 + trtllm_compute_cap: str | None = None
186 +
187 + for node in ast.walk(tree):
188 + if not isinstance(node, ast.FunctionDef):
189 + continue
190 +
191 + # Parse supports_trtllm_attention for compute capability
192 + # Look for: current_platform.is_device_capability_family(100)
193 + if node.name == "supports_trtllm_attention":
194 + for n in ast.walk(node):
195 + if (
196 + isinstance(n, ast.Call)
197 + and isinstance(n.func, ast.Attribute)
198 + and n.func.attr == "is_device_capability_family"
199 + and n.args
200 + and isinstance(n.args[0], ast.Constant)
201 + and isinstance(n.args[0].value, int)
202 + ):
203 + cap = n.args[0].value
204 + # Convert 100 -> "10.x"
205 + trtllm_compute_cap = f"{cap // 10}.x"
206 + break
207 +
208 + if not trtllm_compute_cap:
209 + return {}
210 +
211 + return {
212 + "native": {
213 + # Native FlashInfer: everything except SM100
214 + "supports_sink": False,
215 + },
216 + "trtllm": {
217 + # TRTLLM pathway on Blackwell
218 + "compute_capability": trtllm_compute_cap,
219 + "supports_sink": True,
220 + },
221 + }
222 +
223 +
224 +def parse_mla_prefill_backends() -> list[dict[str, Any]]:
225 + """Parse MLA prefill backend options from mla_attention.py.
226 +
227 + MLA uses different backends for prefill vs decode. The decode backends are
228 + registered in the registry, but prefill backends are selected at runtime
229 + based on conditions in MLACommonImpl.__init__.
230 +
231 + Returns a list of prefill backend info dicts with their requirements.
232 + """
233 + if not MLA_ATTENTION_FILE.exists():
234 + return []
235 +
236 + try:
237 + tree = ast.parse(MLA_ATTENTION_FILE.read_text())
238 + except Exception:
239 + return []
240 +
241 + # Find compute capability requirements by parsing use_* functions
242 + flashinfer_cc: str | None = None
243 + cudnn_cc: str | None = None
244 + trtllm_cc: str | None = None
245 +
246 + for node in ast.walk(tree):
247 + if not isinstance(node, ast.FunctionDef):
248 + continue
249 +
250 + # Parse use_flashinfer_prefill for compute capability (SM100)
251 + if node.name == "use_flashinfer_prefill":
252 + for n in ast.walk(node):
253 + if (
254 + isinstance(n, ast.Call)
255 + and isinstance(n.func, ast.Attribute)
256 + and n.func.attr == "is_device_capability_family"
257 + and n.args
258 + and isinstance(n.args[0], ast.Constant)
259 + and isinstance(n.args[0].value, int)
260 + ):
261 + flashinfer_cc = f"{n.args[0].value // 10}.x"
262 +
263 + # Parse use_cudnn_prefill for compute capability (SM100)
264 + if node.name == "use_cudnn_prefill":
265 + for n in ast.walk(node):
266 + if (
267 + isinstance(n, ast.Call)
268 + and isinstance(n.func, ast.Attribute)
269 + and n.func.attr == "is_device_capability_family"
270 + and n.args
271 + and isinstance(n.args[0], ast.Constant)
272 + and isinstance(n.args[0].value, int)
273 + ):
274 + cudnn_cc = f"{n.args[0].value // 10}.x"
275 +
276 + # Parse use_trtllm_ragged_deepseek_prefill for compute capability
277 + if node.name == "use_trtllm_ragged_deepseek_prefill":
278 + for n in ast.walk(node):
279 + if (
280 + isinstance(n, ast.Call)
281 + and isinstance(n.func, ast.Attribute)
282 + and n.func.attr == "is_device_capability_family"
283 + and n.args
284 + and isinstance(n.args[0], ast.Constant)
285 + and isinstance(n.args[0].value, int)
286 + ):
287 + trtllm_cc = f"{n.args[0].value // 10}.x"
288 +
289 + # Build prefill backend list based on what we found
290 + # Order matches the priority in MLACommonImpl.__init__
291 + prefill_backends: list[dict[str, Any]] = []
292 +
293 + # TRT-LLM Ragged (highest priority if available)
294 + if trtllm_cc:
295 + prefill_backends.append(
296 + {
297 + "name": "TRT-LLM Ragged‡",
298 + "description": "TensorRT-LLM ragged attention",
299 + "compute_capability": trtllm_cc,
300 + "enable": "Default on SM100",
301 + "disable": "`-ac.use_trtllm_ragged_deepseek_prefill=0`",
302 + "notes": "DeepSeek R1 dims only",
303 + }
304 + )
305 +
306 + # FlashInfer prefill
307 + if flashinfer_cc:
308 + prefill_backends.append(
309 + {
310 + "name": "FlashInfer",
311 + "description": "FlashInfer CUTLASS backend",
312 + "compute_capability": flashinfer_cc,
313 + "enable": "`-ac.disable_flashinfer_prefill=0`",
314 + "disable": "`-ac.disable_flashinfer_prefill=1`",
315 + "notes": "DeepSeek R1 dims only",
316 + }
317 + )
318 +
319 + # cuDNN prefill
320 + if cudnn_cc:
321 + prefill_backends.append(
322 + {
323 + "name": "cuDNN",
324 + "description": "cuDNN-based attention",
325 + "compute_capability": cudnn_cc,
326 + "enable": "`-ac.use_cudnn_prefill=1`",
327 + "disable": "`-ac.use_cudnn_prefill=0`",
328 + "notes": "",
329 + }
330 + )
331 +
332 + # FlashAttention is always available as fallback
333 + prefill_backends.append(
334 + {
335 + "name": "FlashAttention",
336 + "description": "FlashAttention varlen (FA2/FA3)",
337 + "compute_capability": "Any",
338 + "enable": "Default fallback",
339 + "disable": "Use other backends",
340 + "notes": "FA3 on SM90, FA2 otherwise",
341 + }
342 + )
343 +
344 + return prefill_backends
345 +
346 +
347 +def find_class_in_ast(tree: ast.AST, class_name: str) -> ast.ClassDef | None:
348 + """Find a class definition in an AST."""
349 + for node in ast.walk(tree):
350 + if isinstance(node, ast.ClassDef) and node.name == class_name:
351 + return node
352 + return None
353 +
354 +
355 +def find_method(node: ast.ClassDef, method_name: str) -> ast.FunctionDef | None:
356 + """Find a method in a class definition."""
357 + for item in node.body:
358 + if isinstance(item, ast.FunctionDef) and item.name == method_name:
359 + return item
360 + return None
361 +
362 +
363 +def method_returns_true(method: ast.FunctionDef | None) -> bool:
364 + """Check if a method simply returns True."""
365 + if method is None:
366 + return False
367 + for node in ast.walk(method):
368 + if not isinstance(node, ast.Return):
369 + continue
370 + if isinstance(node.value, ast.Constant) and node.value.value is True:
371 + return True
372 + return False
373 +
374 +
375 +def _parse_list_class_var(node: ast.ClassDef, var_name: str) -> list[str] | None:
376 + """Parse a list-type class variable, returning None if not found."""
377 + for item in node.body:
378 + if not isinstance(item, ast.AnnAssign):
379 + continue
380 + if not isinstance(item.target, ast.Name):
381 + continue
382 + if item.target.id != var_name:
383 + continue
384 + if not (item.value and isinstance(item.value, ast.List)):
385 + continue
386 + result = []
387 + for elt in item.value.elts:
388 + if isinstance(elt, ast.Attribute):
389 + result.append(elt.attr)
390 + elif isinstance(elt, ast.Constant):
391 + result.append(str(elt.value))
392 + return result
393 + return None
394 +
395 +
396 +def parse_supported_dtypes(node: ast.ClassDef) -> str:
397 + """Parse supported_dtypes class variable."""
398 + dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"}
399 + dtypes = _parse_list_class_var(node, "supported_dtypes")
400 + if dtypes is None:
401 + return "fp16, bf16"
402 + return ", ".join(dtype_map.get(d, d) for d in dtypes)
403 +
404 +
405 +def parse_kv_cache_dtypes(node: ast.ClassDef) -> str:
406 + """Parse supported_kv_cache_dtypes class var or supports_kv_cache_dtype method."""
407 + # First try the class variable
408 + dtypes = _parse_list_class_var(node, "supported_kv_cache_dtypes")
409 + if dtypes:
410 + return ", ".join(dtypes)
411 +
412 + # Fall back to parsing the supports_kv_cache_dtype method
413 + # Look for `kv_cache_dtype in ["auto", "bfloat16"]` pattern
414 + method = find_method(node, "supports_kv_cache_dtype")
415 + if method:
416 + for n in ast.walk(method):
417 + if (
418 + isinstance(n, ast.Compare)
419 + and len(n.ops) == 1
420 + and isinstance(n.ops[0], ast.In)
421 + and len(n.comparators) == 1
422 + and isinstance(n.comparators[0], ast.List)
423 + ):
424 + dtypes = [
425 + e.value
426 + for e in n.comparators[0].elts
427 + if isinstance(e, ast.Constant) and isinstance(e.value, str)
428 + ]
429 + if dtypes:
430 + return ", ".join(dtypes)
431 +
432 + return "auto"
433 +
434 +
435 +def _parse_return_list(
436 + method: ast.FunctionDef | None, handle_multiple_of: bool = False
437 +) -> list[str]:
438 + """Extract list items from a method's return statement."""
439 + if method is None:
440 + return []
441 + for stmt in ast.walk(method):
442 + if not isinstance(stmt, ast.Return):
443 + continue
444 + if not isinstance(stmt.value, ast.List):
445 + continue
446 + sizes = []
447 + for elt in stmt.value.elts:
448 + if isinstance(elt, ast.Constant):
449 + sizes.append(str(elt.value))
450 + elif (
451 + handle_multiple_of
452 + and isinstance(elt, ast.Call)
453 + and isinstance(elt.func, ast.Name)
454 + and elt.func.id == "MultipleOf"
455 + and elt.args
456 + and isinstance(elt.args[0], ast.Constant)
457 + ):
458 + sizes.append(f"%{elt.args[0].value}")
459 + if sizes:
460 + return sizes
461 + return []
462 +
463 +
464 +def parse_block_sizes(node: ast.ClassDef) -> str:
465 + """Parse get_supported_kernel_block_sizes method."""
466 + method = find_method(node, "get_supported_kernel_block_sizes")
467 + sizes = _parse_return_list(method, handle_multiple_of=True)
468 + return ", ".join(sizes) if sizes else "Any"
469 +
470 +
471 +def parse_head_sizes(node: ast.ClassDef) -> str:
472 + """Parse get_supported_head_sizes method."""
473 + method = find_method(node, "get_supported_head_sizes")
474 + sizes = _parse_return_list(method)
475 + return ", ".join(sizes) if sizes else "Any"
476 +
477 +
478 +def parse_compute_capability(node: ast.ClassDef) -> str:
479 + """Parse supports_compute_capability method."""
480 + method = find_method(node, "supports_compute_capability")
481 + if method is None:
482 + return "Any"
483 +
484 + min_cap: tuple[int, int] | None = None
485 + max_cap: tuple[int, int] | None = None
486 + major_list: list[int] = []
487 +
488 + for n in ast.walk(method):
489 + if not isinstance(n, ast.Compare):
490 + continue
491 +
492 + # Handle `capability >= DeviceCapability(...)` or `capability <= ...`
493 + for op, comp in zip(n.ops, n.comparators):
494 + if not (
495 + isinstance(comp, ast.Call)
496 + and isinstance(comp.func, ast.Name)
497 + and comp.func.id == "DeviceCapability"
498 + and comp.args
499 + and isinstance(comp.args[0], ast.Constant)
500 + ):
501 + continue
502 + major = comp.args[0].value
503 + minor = 0
504 + if len(comp.args) > 1 and isinstance(comp.args[1], ast.Constant):
505 + minor = comp.args[1].value
506 + if isinstance(op, ast.GtE):
507 + min_cap = (major, minor)
508 + elif isinstance(op, ast.LtE):
509 + max_cap = (major, minor)
510 +
511 + # Handle `capability.major == N` or `capability.major in [N, M]`
512 + if (
513 + isinstance(n.left, ast.Attribute)
514 + and n.left.attr == "major"
515 + and len(n.ops) == 1
516 + and len(n.comparators) == 1
517 + ):
518 + comp = n.comparators[0]
519 + if isinstance(n.ops[0], ast.Eq) and isinstance(comp, ast.Constant):
520 + major_list.append(comp.value)
521 + elif isinstance(n.ops[0], ast.In) and isinstance(comp, ast.List):
522 + major_list.extend(
523 + e.value
524 + for e in comp.elts
525 + if isinstance(e, ast.Constant) and isinstance(e.value, int)
526 + )
527 +
528 + if major_list:
529 + major_list.sort()
530 + if len(major_list) == 1:
531 + return f"{major_list[0]}.x"
532 + return f"{major_list[0]}.x-{major_list[-1]}.x"
533 +
534 + if min_cap:
535 + if max_cap:
536 + return f"{min_cap[0]}.x-{max_cap[0]}.x"
537 + return f"≥{min_cap[0]}.{min_cap[1]}"
538 +
539 + return "Any"
540 +
541 +
542 +def parse_attention_types(node: ast.ClassDef) -> str:
543 + """Parse supports_attn_type method."""
544 + method = find_method(node, "supports_attn_type")
545 + if method is None:
546 + return "Decoder"
547 +
548 + type_map = {
549 + "DECODER": "Decoder",
550 + "ENCODER": "Encoder",
551 + "ENCODER_ONLY": "Encoder Only",
552 + "ENCODER_DECODER": "Enc-Dec",
553 + }
554 + types: set[str] = set()
555 +
556 + for n in ast.walk(method):
557 + # Handle `attn_type in (AttentionType.DECODER, ...)`
558 + if not (
559 + isinstance(n, ast.Compare)
560 + and len(n.ops) == 1
561 + and isinstance(n.ops[0], ast.In)
562 + and len(n.comparators) == 1
563 + and isinstance(n.comparators[0], ast.Tuple | ast.Set)
564 + ):
565 + continue
566 +
567 + for elt in n.comparators[0].elts:
568 + if isinstance(elt, ast.Attribute) and elt.attr in type_map:
569 + types.add(type_map[elt.attr])
570 +
571 + if not types:
572 + return "Decoder"
573 + return "All" if len(types) >= 3 else ", ".join(sorted(types))
574 +
575 +
576 +def check_method_overrides(node: ast.ClassDef, method_name: str) -> bool:
577 + """Check if a method is overridden and returns True."""
578 + method = find_method(node, method_name)
579 + return method_returns_true(method)
580 +
581 +
582 +def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None:
583 + """Analyze a backend class and extract feature information."""
584 + file_path = get_file_from_class_path(class_path)
585 + if file_path is None:
586 + return None
587 +
588 + try:
589 + tree = ast.parse(file_path.read_text())
590 + except Exception as e:
591 + print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr)
592 + return None
593 +
594 + class_name = class_path.rsplit(".", 1)[1]
595 + class_node = find_class_in_ast(tree, class_name)
596 + if class_node is None:
597 + return None
598 +
599 + # Check if this is an MLA backend by parent class or naming
600 + parent = None
601 + if class_node.bases:
602 + base = class_node.bases[0]
603 + parent = base.id if isinstance(base, ast.Name) else None
604 + mla_parents = {"MLACommonBackend", "FlashMLABackend", "FlashMLASparseBackend"}
605 + is_mla_backend = (
606 + parent in mla_parents
607 + or ".mla." in class_path.lower()
608 + or "_mla" in backend_name.lower()
609 + )
610 +
611 + # Determine compute capability - use N/A for non-CUDA backends
612 + is_non_cuda = backend_name.startswith(("CPU_", "ROCM_"))
613 + compute_cap = "N/A" if is_non_cuda else parse_compute_capability(class_node)
614 +
615 + return {
616 + "name": backend_name,
617 + "dtypes": parse_supported_dtypes(class_node),
618 + "kv_cache_dtypes": parse_kv_cache_dtypes(class_node),
619 + "block_sizes": parse_block_sizes(class_node),
620 + "head_sizes": parse_head_sizes(class_node),
621 + "attn_types": parse_attention_types(class_node),
622 + "compute_capability": compute_cap,
623 + "is_mla": is_mla_backend or check_method_overrides(class_node, "is_mla"),
624 + "supports_sink": check_method_overrides(class_node, "supports_sink"),
625 + "is_sparse": check_method_overrides(class_node, "is_sparse"),
626 + "supports_mm_prefix": check_method_overrides(class_node, "supports_mm_prefix"),
627 + }
628 +
629 +
630 +def add_literal_quotes(value: str) -> str:
631 + """Add literal backticks around all comma-separated items in a string."""
632 + items = [item.strip() for item in value.split(",")]
633 + quoted_items = [f"`{item}`" for item in items]
634 + return ", ".join(quoted_items)
635 +
636 +
637 +def bool_to_emoji(value: bool) -> str:
638 + """Convert a boolean to a checkmark or X emoji."""
639 + return "✅" if value else "❌"
640 +
641 +
642 +def generate_markdown_table(
643 + backends: list[dict[str, Any]], title: str, is_mla_table: bool = False
644 +) -> str:
645 + """Generate a markdown table from backend info.
646 +
647 + Args:
648 + backends: List of backend info dictionaries.
649 + title: Table title.
650 + is_mla_table: If True, include MLA and Sparse columns (for MLA table).
651 + If False, exclude them (for standard attention table).
652 + """
653 + if not backends:
654 + return f"## {title}\n\nNo backends found.\n"
655 +
656 + # Check if any backend has a version (for FA2/FA3 split)
657 + has_versions = any(b.get("version") for b in backends)
658 +
659 + if is_mla_table:
660 + header = (
661 + "| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes "
662 + "| Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |"
663 + )
664 + separator = (
665 + "|---------|--------|-----------|-------------|------------"
666 + "|------|--------|-----------|-----------------|--------------|"
667 + )
668 + elif has_versions:
669 + header = (
670 + "| Backend | Version | Dtypes | KV Dtypes | Block Sizes "
671 + "| Head Sizes | Sink | MM Prefix | Attention Types | Compute Cap. |"
672 + )
673 + separator = (
674 + "|---------|---------|--------|-----------|-------------"
675 + "|------------|------|-----------|-----------------|--------------|"
676 + )
677 + else:
678 + header = (
679 + "| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes "
680 + "| Sink | MM Prefix | Attention Types | Compute Cap. |"
681 + )
682 + separator = (
683 + "|---------|--------|-----------|-------------|------------"
684 + "|------|-----------|-----------------|--------------|"
685 + )
686 + lines = [f"## {title}", "", header, separator]
687 +
688 + def sort_key(x: dict[str, Any]) -> tuple[str, int]:
689 + """Sort key that keeps parent/child rows together in order."""
690 + return (x.get("_sort_key", x["name"]), x.get("_sort_order", 0))
691 +
692 + for info in sorted(backends, key=sort_key):
693 + if is_mla_table:
694 + row = "| `{}` | {} | {} | {} | {} | {} | {} | {} | {} | {} |".format(
695 + info["name"],
696 + info["dtypes"],
697 + add_literal_quotes(info["kv_cache_dtypes"]),
698 + info["block_sizes"],
699 + info["head_sizes"],
700 + bool_to_emoji(info["supports_sink"]),
701 + bool_to_emoji(info["is_sparse"]),
702 + bool_to_emoji(info["supports_mm_prefix"]),
703 + info["attn_types"],
704 + info["compute_capability"],
705 + )
706 + elif has_versions:
707 + row = "| `{}` | {} | {} | {} | {} | {} | {} | {} | {} | {} |".format(
708 + info["name"],
709 + info.get("version", ""),
710 + info["dtypes"],
711 + add_literal_quotes(info["kv_cache_dtypes"]),
712 + info["block_sizes"],
713 + info["head_sizes"],
714 + bool_to_emoji(info["supports_sink"]),
715 + bool_to_emoji(info["supports_mm_prefix"]),
716 + info["attn_types"],
717 + info["compute_capability"],
718 + )
719 + else:
720 + row = "| `{}` | {} | {} | {} | {} | {} | {} | {} | {} |".format(
721 + info["name"],
722 + info["dtypes"],
723 + add_literal_quotes(info["kv_cache_dtypes"]),
724 + info["block_sizes"],
725 + info["head_sizes"],
726 + bool_to_emoji(info["supports_sink"]),
727 + bool_to_emoji(info["supports_mm_prefix"]),
728 + info["attn_types"],
729 + info["compute_capability"],
730 + )
731 + lines.append(row)
732 +
733 + lines.append("")
734 + return "\n".join(lines)
735 +
736 +
737 +def parse_cuda_priority_lists() -> dict[str, list[str]]:
738 + """Parse priority lists from cuda.py using AST.
739 +
740 + The structure of _get_backend_priorities is:
741 + if use_mla:
742 + if device_capability.major == 10:
743 + return [MLA list for SM100]
744 + else:
745 + return [MLA list for default]
746 + else:
747 + if device_capability.major == 10:
748 + return [Standard list for SM100]
749 + else:
750 + return [Standard list for default]
751 + """
752 + if not CUDA_PLATFORM_FILE.exists():
753 + return {}
754 +
755 + try:
756 + source = CUDA_PLATFORM_FILE.read_text()
757 + tree = ast.parse(source)
758 + except Exception:
759 + return {}
760 +
761 + priorities: dict[str, list[str]] = {}
762 +
763 + # Find the _get_backend_priorities function
764 + for node in ast.walk(tree):
765 + if not isinstance(node, ast.FunctionDef):
766 + continue
767 + if node.name != "_get_backend_priorities":
768 + continue
769 +
770 + # Process the function body directly
771 + for stmt in node.body:
772 + if not isinstance(stmt, ast.If):
773 + continue
774 +
775 + # Check if this is the "if use_mla:" branch
776 + is_mla_branch = (
777 + isinstance(stmt.test, ast.Name) and stmt.test.id == "use_mla"
778 + )
779 +
780 + if is_mla_branch:
781 + _extract_priorities(stmt.body, priorities, "mla")
782 + if stmt.orelse:
783 + _extract_priorities(stmt.orelse, priorities, "standard")
784 + else:
785 + _extract_priorities([stmt], priorities, "standard")
786 +
787 + return priorities
788 +
789 +
790 +def _get_backends_from_return(stmts: list) -> list[str]:
791 + """Extract backend names from return statements in a list of statements."""
792 + for stmt in stmts:
793 + if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List):
794 + return [e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)]
795 + return []
796 +
797 +
798 +def _is_sm100_check(test: ast.expr) -> bool:
799 + """Check if test is `something.major == 10`."""
800 + return (
801 + isinstance(test, ast.Compare)
802 + and isinstance(test.left, ast.Attribute)
803 + and test.left.attr == "major"
804 + and len(test.ops) == 1
805 + and isinstance(test.ops[0], ast.Eq)
806 + and len(test.comparators) == 1
807 + and isinstance(test.comparators[0], ast.Constant)
808 + and test.comparators[0].value == 10
809 + )
810 +
811 +
812 +def _extract_priorities(body: list, priorities: dict[str, list[str]], prefix: str):
813 + """Extract priority lists from if/else statement body."""
814 + for stmt in body:
815 + if isinstance(stmt, ast.If):
816 + is_sm100 = _is_sm100_check(stmt.test)
817 + if_key = f"{prefix}_sm100" if is_sm100 else f"{prefix}_default"
818 + else_key = f"{prefix}_default" if is_sm100 else f"{prefix}_sm100"
819 +
820 + if backends := _get_backends_from_return(stmt.body):
821 + priorities[if_key] = backends
822 + if backends := _get_backends_from_return(stmt.orelse):
823 + priorities[else_key] = backends
824 +
825 + elif isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List):
826 + backends = [e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)]
827 + priorities[f"{prefix}_default"] = backends
828 +
829 +
830 +def generate_usage_section() -> str:
831 + """Generate the usage documentation section."""
832 + return """## Setting the Attention Backend
833 +
834 +### Command Line
835 +
836 +There are two ways to specify the backend from the command line:
837 +
838 +**Option 1: Using `--attention-backend` (simple)**
839 +
840 +```bash
841 +vllm serve <model> --attention-backend FLASH_ATTN
842 +```
843 +
844 +**Option 2: Using `--attention-config.backend` / `-ac.backend` (structured config)**
845 +
846 +```bash
847 +# Dot notation
848 +vllm serve <model> --attention-config.backend FLASH_ATTN
849 +vllm serve <model> -ac.backend FLASH_ATTN
850 +
851 +# JSON format
852 +vllm serve <model> --attention-config '{"backend": "FLASH_ATTN"}'
853 +vllm serve <model> -ac '{"backend": "FLASH_ATTN"}'
854 +```
855 +
856 +> **Note:** `--attention-backend` and `--attention-config.backend` are mutually
857 +> exclusive. Use one or the other, not both.
858 +
859 +### Python API
860 +
861 +Use `AttentionConfig` with the `LLM` class:
862 +
863 +```python
864 +from vllm import LLM
865 +from vllm.config import AttentionConfig
866 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
867 +
868 +# Method 1: Using AttentionConfig with enum
869 +llm = LLM(
870 + model="Qwen/Qwen3-0.6B",
871 + attention_config=AttentionConfig(backend=AttentionBackendEnum.FLASH_ATTN),
872 +)
873 +
874 +# Method 2: Using attention_backend parameter with string
875 +llm = LLM(
876 + model="Qwen/Qwen3-0.6B",
877 + attention_backend="FLASH_ATTN",
878 +)
879 +```
880 +
881 +## Backend Selection Behavior
882 +
883 +### Manual Selection
884 +
885 +When you explicitly set a backend via `--attention-backend` or `AttentionConfig`:
886 +
887 +1. The backend is **validated** against your configuration (model dtype, head
888 + size, compute capability, etc.)
889 +2. If the backend **doesn't support** your configuration, an error is raised
890 + with the specific reason
891 +3. If valid, the backend is used
892 +
893 +Example error when selecting an incompatible backend:
894 +
895 +```text
896 +ValueError: Selected backend FLASHMLA is not valid for this configuration.
897 +Reason: ['compute capability not supported']
898 +```
899 +
900 +### Automatic Selection
901 +
902 +When no backend is specified (the default):
903 +
904 +1. vLLM iterates through backends in **priority order** (see tables below)
905 +2. Each backend is validated against your configuration
906 +3. The **first compatible backend** is selected
907 +4. If no backend is compatible, an error is raised listing all backends and
908 + their incompatibility reasons
909 +"""
910 +
911 +
912 +def _priority_table(title: str, backends: list[str]) -> list[str]:
913 + """Generate a priority table for a list of backends."""
914 + return [
915 + f"**{title}:**",
916 + "",
917 + "| Priority | Backend |",
918 + "|----------|---------|",
919 + *[f"| {i} | `{b}` |" for i, b in enumerate(backends, 1)],
920 + "",
921 + ]
922 +
923 +
924 +def generate_priority_section(priorities: dict[str, list[str]]) -> str:
925 + """Generate the priority ranking section."""
926 + lines = [
927 + "## Backend Priority (CUDA)",
928 + "",
929 + "When no backend is explicitly selected, vLLM chooses the first",
930 + "compatible backend from these priority-ordered lists.",
931 + "",
932 + "Priority is **1 = highest** (tried first).",
933 + "",
934 + "### Standard Attention (MHA, MQA, GQA)",
935 + "",
936 + ]
937 +
938 + sm100 = "Blackwell (SM 10.x)"
939 + ampere = "Ampere/Hopper (SM 8.x-9.x)"
940 +
941 + if "standard_sm100" in priorities:
942 + lines.extend(_priority_table(sm100, priorities["standard_sm100"]))
943 + if "standard_default" in priorities:
944 + lines.extend(_priority_table(ampere, priorities["standard_default"]))
945 +
946 + lines.extend(["### MLA Attention (DeepSeek-style)", ""])
947 +
948 + if "mla_sm100" in priorities:
949 + lines.extend(_priority_table(sm100, priorities["mla_sm100"]))
950 + if "mla_default" in priorities:
951 + lines.extend(_priority_table(ampere, priorities["mla_default"]))
952 +
953 + lines.append(
954 + "> **Note:** ROCm and CPU platforms have their own selection logic. "
955 + "See the platform-specific documentation for details."
956 + )
957 + lines.append("")
958 +
959 + return "\n".join(lines)
960 +
961 +
962 +def generate_mla_section(
963 + prefill_backends: list[dict[str, Any]], decode_backends: list[dict[str, Any]]
964 +) -> str:
965 + """Generate the complete MLA section with prefill and decode tables."""
966 + lines = [
967 + "## MLA (Multi-head Latent Attention) Backends",
968 + "",
969 + "MLA uses separate backends for prefill and decode phases.",
970 + "",
971 + "### Prefill Backends",
972 + "",
973 + "The prefill backend is selected at runtime based on hardware and",
974 + "configuration.",
975 + "",
976 + "| Backend | Description | Compute Cap. | Enable | Disable | Notes |",
977 + "|---------|-------------|--------------|--------|---------|-------|",
978 + ]
979 +
980 + for backend in prefill_backends:
981 + row = "| {} | {} | {} | {} | {} | {} |".format(
982 + backend["name"],
983 + backend["description"],
984 + backend["compute_capability"],
985 + backend["enable"],
986 + backend["disable"],
987 + backend.get("notes", ""),
988 + )
989 + lines.append(row)
990 +
991 + lines.extend(
992 + [
993 + "",
994 + "> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).",
995 + "> On other GPUs, FlashAttention is used as the default.",
996 + "",
997 + "### Decode Backends",
998 + "",
999 + ]
1000 + )
1001 +
1002 + # Generate decode backends table
1003 + header = (
1004 + "| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes "
1005 + "| Sink | Sparse | MM Prefix | Attention Types | Compute Cap. |"
1006 + )
1007 + separator = (
1008 + "|---------|--------|-----------|-------------|------------"
1009 + "|------|--------|-----------|-----------------|--------------|"
1010 + )
1011 + lines.extend([header, separator])
1012 +
1013 + def sort_key(x: dict[str, Any]) -> tuple[str, int]:
1014 + return (x.get("_sort_key", x["name"]), x.get("_sort_order", 0))
1015 +
1016 + for info in sorted(decode_backends, key=sort_key):
1017 + row = "| `{}` | {} | {} | {} | {} | {} | {} | {} | {} | {} |".format(
1018 + info["name"],
1019 + info["dtypes"],
1020 + add_literal_quotes(info["kv_cache_dtypes"]),
1021 + info["block_sizes"],
1022 + info["head_sizes"],
1023 + bool_to_emoji(info["supports_sink"]),
1024 + bool_to_emoji(info["is_sparse"]),
1025 + bool_to_emoji(info["supports_mm_prefix"]),
1026 + info["attn_types"],
1027 + info["compute_capability"],
1028 + )
1029 + lines.append(row)
1030 +
1031 + lines.append("")
1032 + return "\n".join(lines)
1033 +
1034 +
1035 +def generate_legend() -> str:
1036 + """Generate a legend explaining the table columns."""
1037 + return """## Legend
1038 +
1039 +| Column | Description |
1040 +|--------|-------------|
1041 +| **Dtypes** | Supported model data types (fp16, bf16, fp32) |
1042 +| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) |
1043 +| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
1044 +| **Head Sizes** | Supported attention head sizes |
1045 +| **Sink** | Attention sink support (for StreamingLLM) |
1046 +| **Sparse** | Sparse attention support (MLA only) |
1047 +| **MM Prefix** | Multimodal prefix full attention support |
1048 +| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
1049 +| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
1050 +
1051 +**Symbols:** ✅ = Supported, ❌ = Not supported
1052 +"""
1053 +
1054 +
1055 +def generate_docs() -> str:
1056 + """Generate the complete documentation."""
1057 + attention_backends_map = parse_registry()
1058 +
1059 + # Parse priority lists from cuda.py
1060 + priorities = parse_cuda_priority_lists()
1061 +
1062 + # Parse FlashAttention FA2/FA3 feature differences
1063 + fa_features = parse_flash_attn_features()
1064 +
1065 + # Parse FlashInfer TRTLLM feature differences (native vs TRTLLM on Blackwell)
1066 + fi_features = parse_flashinfer_trtllm_features()
1067 +
1068 + # Parse MLA prefill backends
1069 + mla_prefill_backends = parse_mla_prefill_backends()
1070 +
1071 + # Collect backend info
1072 + all_backends = []
1073 + for backend_name, class_path in attention_backends_map.items():
1074 + if backend_name in ("CUSTOM", "TORCH_SDPA"):
1075 + continue
1076 + info = analyze_backend(backend_name, class_path)
1077 + if info:
1078 + all_backends.append(info)
1079 +
1080 + # Expand FLASH_ATTN into FA2 and FA3 variants with different capabilities
1081 + if fa_features:
1082 + expanded_backends = []
1083 + for backend in all_backends:
1084 + if backend["name"] == "FLASH_ATTN":
1085 + # Create FA2 entry (keeps base backend's compute_capability)
1086 + fa2 = backend.copy()
1087 + fa2["name"] = "FLASH_ATTN"
1088 + fa2["version"] = "FA2*"
1089 + fa2["_sort_key"] = "FLASH_ATTN"
1090 + fa2["_sort_order"] = 0
1091 + fa2["supports_sink"] = fa_features["fa2"]["supports_sink"]
1092 +
1093 + # Create FA3 entry (uses parsed compute_capability from fa_utils)
1094 + fa3 = backend.copy()
1095 + fa3["name"] = "FLASH_ATTN"
1096 + fa3["version"] = "FA3*"
1097 + fa3["_sort_key"] = "FLASH_ATTN"
1098 + fa3["_sort_order"] = 1
1099 + if fa_features["fa3"]["compute_capability"]:
1100 + fa3["compute_capability"] = fa_features["fa3"]["compute_capability"]
1101 + fa3["supports_sink"] = fa_features["fa3"]["supports_sink"]
1102 + if fa_features["fa3"]["supports_fp8"]:
1103 + # Add fp8 dtypes to the base backend's kv_cache_dtypes
1104 + base_dtypes = backend["kv_cache_dtypes"].split(", ")
1105 + fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"]
1106 + new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes]
1107 + fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes)
1108 +
1109 + # Add FA2 first, then FA3
1110 + expanded_backends.append(fa2)
1111 + expanded_backends.append(fa3)
1112 + else:
1113 + backend["_sort_key"] = backend["name"]
1114 + backend["_sort_order"] = 0
1115 + backend["version"] = "" # No version for other backends
1116 + expanded_backends.append(backend)
1117 + all_backends = expanded_backends
1118 +
1119 + # Expand FLASHINFER into native and TRTLLM variants
1120 + if fi_features:
1121 + expanded_backends = []
1122 + for backend in all_backends:
1123 + if backend["name"] == "FLASHINFER":
1124 + # Parse original compute capability to get min CC
1125 + orig_cap = backend["compute_capability"]
1126 + parts = orig_cap.replace(".x", "").split("-")
1127 + min_cc = parts[0] if parts else "7"
1128 + trtllm_cc = fi_features["trtllm"]["compute_capability"]
1129 +
1130 + # Create native entry (pre-Blackwell GPUs)
1131 + native = backend.copy()
1132 + native["name"] = "FLASHINFER"
1133 + native["version"] = "Native†"
1134 + native["_sort_key"] = "FLASHINFER"
1135 + native["_sort_order"] = 0
1136 + native["supports_sink"] = fi_features["native"]["supports_sink"]
1137 + # Native FlashInfer is used on GPUs before SM100 (Blackwell)
1138 + native["compute_capability"] = f"{min_cc}.x-9.x"
1139 +
1140 + # Create TRTLLM entry
1141 + trtllm = backend.copy()
1142 + trtllm["name"] = "FLASHINFER"
1143 + trtllm["version"] = "TRTLLM†"
1144 + trtllm["_sort_key"] = "FLASHINFER"
1145 + trtllm["_sort_order"] = 1
1146 + trtllm["compute_capability"] = trtllm_cc
1147 + trtllm["supports_sink"] = fi_features["trtllm"]["supports_sink"]
1148 +
1149 + expanded_backends.append(native)
1150 + expanded_backends.append(trtllm)
1151 + else:
1152 + expanded_backends.append(backend)
1153 + all_backends = expanded_backends
1154 +
1155 + # Split into MLA and non-MLA
1156 + mla_backends = [b for b in all_backends if b["is_mla"]]
1157 + non_mla_backends = [b for b in all_backends if not b["is_mla"]]
1158 +
1159 + # Generate documentation
1160 + script_path = "tools/pre_commit/generate_attention_backend_docs.py"
1161 + doc_lines = [
1162 + "# Attention Backend Feature Support",
1163 + "",
1164 + f"This document is auto-generated by `{script_path}`.",
1165 + "It shows the feature support for each registered attention backend",
1166 + "based on the checks in `AttentionBackend.validate_configuration()`.",
1167 + "",
1168 + "**Do not edit this file manually.** Run the following command to",
1169 + "regenerate it:",
1170 + "",
1171 + "```bash",
1172 + f"python {script_path}",
1173 + "```",
1174 + "",
1175 + ]
1176 +
1177 + # Add usage documentation
1178 + doc_lines.append(generate_usage_section())
1179 +
1180 + # Add priority section
1181 + doc_lines.append(generate_priority_section(priorities))
1182 +
1183 + # Add legend and feature tables
1184 + doc_lines.append(generate_legend())
1185 + standard_title = "Standard Attention (MHA, MQA, GQA) Backends"
1186 + doc_lines.append(
1187 + generate_markdown_table(non_mla_backends, standard_title, is_mla_table=False)
1188 + )
1189 + # Add footnotes for version/variant distinctions (in table order)
1190 + footnotes = []
1191 + if fi_features:
1192 + footnotes.append(
1193 + "> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which "
1194 + "supports sinks. Disable via `--attention-config.use_trtllm_attention=0`."
1195 + )
1196 + if fa_features:
1197 + footnotes.append(
1198 + "> **\\*** Specify the FlashAttention version via "
1199 + "`--attention-config.flash_attn_version=2` or `3`. Default is FA3 on SM90, "
1200 + "FA2 otherwise."
1201 + )
1202 + if footnotes:
1203 + doc_lines.append("\n>\n".join(footnotes) + "\n")
1204 +
1205 + # Add MLA section with prefill and decode backends
1206 + doc_lines.append(generate_mla_section(mla_prefill_backends, mla_backends))
1207 +
1208 + return "\n".join(doc_lines)
1209 +
1210 +
1211 +def main():
1212 + parser = argparse.ArgumentParser(
1213 + description="Generate attention backend documentation table"
1214 + )
1215 + parser.add_argument(
1216 + "--output",
1217 + "-o",
1218 + type=str,
1219 + default=str(REPO_ROOT / "docs" / "design" / "attention_backends.md"),
1220 + help="Output file path (default: docs/design/attention_backends.md)",
1221 + )
1222 + parser.add_argument(
1223 + "--check",
1224 + action="store_true",
1225 + help="Check if the documentation is up to date (for pre-commit)",
1226 + )
1227 + parser.add_argument(
1228 + "files",
1229 + nargs="*",
1230 + help="Files to check (passed by pre-commit). If none are relevant, skip.",
1231 + )
1232 + args = parser.parse_args()
1233 +
1234 + if args.files and not any(is_relevant_file(f) for f in args.files):
1235 + sys.exit(0)
1236 +
1237 + output_path = Path(args.output)
1238 + new_content = generate_docs()
1239 +
1240 + if args.check:
1241 + needs_update = (
1242 + not output_path.exists() or output_path.read_text() != new_content
1243 + )
1244 + if needs_update:
1245 + output_path.parent.mkdir(parents=True, exist_ok=True)
1246 + output_path.write_text(new_content)
1247 + print(f"🔄 Regenerated: {output_path}")
1248 + sys.exit(1)
1249 + print(f"✅ Up to date: {output_path}")
1250 + sys.exit(0)
1251 +
1252 + output_path.parent.mkdir(parents=True, exist_ok=True)
1253 + output_path.write_text(new_content)
1254 + print(f"Generated: {output_path}")
1255 +
1256 +
1257 +if __name__ == "__main__":
1258 + main()