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 8, 15:15 UTC
[Doc] Add DCP support to attention backend doc (#33936)
Michael Goin committed
Feb 9, 2026 at 18:33 UTC
5e75a14a667dccf7f48781568f19f1a6b9c8014a
2 files changed
+769
-644
docs/design/attention_backends.md
+26
-25
@@ -152,6 +152,7 @@ Priority is **1 = highest** (tried first).
152
| **Sink** | Attention sink support (for StreamingLLM) |
153
| **Sparse** | Sparse attention support (MLA only) |
154
| **MM Prefix** | Multimodal prefix full attention support |
155
+| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
156
| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
157
| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
158
@@ -159,20 +160,20 @@ Priority is **1 = highest** (tried first).
160
161
## Standard Attention (MHA, MQA, GQA) Backends
162
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, 32 | 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 |
163
+| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
164
+|---------|---------|--------|-----------|-------------|------------|------|-----------|-----|-----------------|--------------|
165
+| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
166
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
167
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
168
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
169
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
170
+| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
171
+| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
172
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
173
+| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
174
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
175
+| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
176
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
177
178
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
179
>
@@ -199,14 +200,14 @@ configuration.
200
201
### Decode Backends
202
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 |
203
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
204
+|---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|
205
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
206
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
207
+| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
208
+| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
209
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
210
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
211
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto` | Any | 576 | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
212
+| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
213
+| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
tools/pre_commit/generate_attention_backend_docs.py
+743
-619
@@ -17,9 +17,14 @@ import argparse
17
import ast
18
import fnmatch
19
import sys
20
+from collections.abc import Callable
21
from pathlib import Path
22
from typing import Any
23
24
+# ---------------------------------------------------------------------------
25
+# Constants and file paths
26
+# ---------------------------------------------------------------------------
27
+
28
REPO_ROOT = Path(__file__).parent.parent.parent
29
30
RELEVANT_PATTERNS = [
@@ -32,6 +37,18 @@ RELEVANT_PATTERNS = [
37
"docs/design/attention_backends.md",
38
]
39
40
+BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends"
41
+REGISTRY_FILE = BACKENDS_DIR / "registry.py"
42
+CUDA_PLATFORM_FILE = REPO_ROOT / "vllm" / "platforms" / "cuda.py"
43
+FA_UTILS_FILE = BACKENDS_DIR / "fa_utils.py"
44
+FLASHINFER_UTILS_FILE = REPO_ROOT / "vllm" / "utils" / "flashinfer.py"
45
+MLA_ATTENTION_FILE = (
46
+ REPO_ROOT / "vllm" / "model_executor" / "layers" / "attention" / "mla_attention.py"
47
+)
48
+
49
+# Backends to skip during doc generation
50
+SKIP_BACKENDS = {"CUSTOM", "TORCH_SDPA"}
51
+
52
53
def is_relevant_file(filepath: str) -> bool:
54
"""Check if a file matches any of the relevant patterns."""
@@ -46,351 +63,234 @@ def is_relevant_file(filepath: str) -> bool:
63
return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)
64
65
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
-)
66
+# ---------------------------------------------------------------------------
67
+# AST utility helpers
68
+# ---------------------------------------------------------------------------
69
70
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())
71
+def find_class_in_ast(tree: ast.AST, class_name: str) -> ast.ClassDef | None:
72
+ """Find a class definition in an AST."""
73
for node in ast.walk(tree):
63
- if isinstance(node, ast.ClassDef) and node.name == "AttentionBackendEnum":
64
- return _extract_enum_values(node)
65
- return {}
74
+ if isinstance(node, ast.ClassDef) and node.name == class_name:
75
+ return node
76
+ return None
77
78
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] = {}
79
+def find_method(node: ast.ClassDef, method_name: str) -> ast.FunctionDef | None:
80
+ """Find a method in a class definition."""
81
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.
82
+ if isinstance(item, ast.FunctionDef) and item.name == method_name:
83
+ return item
84
+ return None
85
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 {}
86
100
- try:
101
- tree = ast.parse(FA_UTILS_FILE.read_text())
102
- except Exception:
103
- return {}
87
+def method_returns_true(method: ast.FunctionDef | None) -> bool:
88
+ """Check if a method simply returns True."""
89
+ if method is None:
90
+ return False
91
+ for node in ast.walk(method):
92
+ if (
93
+ isinstance(node, ast.Return)
94
+ and isinstance(node.value, ast.Constant)
95
+ and node.value.value is True
96
+ ):
97
+ return True
98
+ return False
99
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
100
110
- for node in ast.walk(tree):
111
- if not isinstance(node, ast.FunctionDef):
112
- continue
101
+def check_method_overrides(node: ast.ClassDef, method_name: str) -> bool:
102
+ """Check if a method is overridden and returns True."""
103
+ return method_returns_true(find_method(node, method_name))
104
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
105
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):
106
+def _find_bool_class_var(class_node: ast.ClassDef, var_name: str) -> bool | None:
107
+ """Find a bool class variable in a class definition. Returns None if not found."""
108
+ for item in class_node.body:
109
+ # Check for annotated assignment: attr: bool = True/False
110
+ if (
111
+ isinstance(item, ast.AnnAssign)
112
+ and isinstance(item.target, ast.Name)
113
+ and item.target.id == var_name
114
+ and isinstance(item.value, ast.Constant)
115
+ and isinstance(item.value.value, bool)
116
+ ):
117
+ return item.value.value
118
+ # Check for plain assignment: attr = True/False
119
+ if isinstance(item, ast.Assign):
120
+ for target in item.targets:
121
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"
122
+ isinstance(target, ast.Name)
123
+ and target.id == var_name
124
+ and isinstance(item.value, ast.Constant)
125
+ and isinstance(item.value.value, bool)
126
):
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 {}
127
+ return item.value.value
128
+ return None
129
185
- trtllm_compute_cap: str | None = None
130
187
- for node in ast.walk(tree):
188
- if not isinstance(node, ast.FunctionDef):
131
+def _parse_list_class_var(node: ast.ClassDef, var_name: str) -> list[str] | None:
132
+ """Parse a list-type class variable, returning None if not found."""
133
+ for item in node.body:
134
+ if not isinstance(item, ast.AnnAssign):
135
continue
136
+ if not isinstance(item.target, ast.Name):
137
+ continue
138
+ if item.target.id != var_name:
139
+ continue
140
+ if not (item.value and isinstance(item.value, ast.List)):
141
+ continue
142
+ result = []
143
+ for elt in item.value.elts:
144
+ if isinstance(elt, ast.Attribute):
145
+ result.append(elt.attr)
146
+ elif isinstance(elt, ast.Constant):
147
+ result.append(str(elt.value))
148
+ return result
149
+ return None
150
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
- }
151
152
+def _parse_return_list(
153
+ method: ast.FunctionDef | None, handle_multiple_of: bool = False
154
+) -> list[str]:
155
+ """Extract list items from a method's return statement."""
156
+ if method is None:
157
+ return []
158
+ for stmt in ast.walk(method):
159
+ if not isinstance(stmt, ast.Return):
160
+ continue
161
+ if not isinstance(stmt.value, ast.List):
162
+ continue
163
+ sizes = []
164
+ for elt in stmt.value.elts:
165
+ if isinstance(elt, ast.Constant):
166
+ sizes.append(str(elt.value))
167
+ elif (
168
+ handle_multiple_of
169
+ and isinstance(elt, ast.Call)
170
+ and isinstance(elt.func, ast.Name)
171
+ and elt.func.id == "MultipleOf"
172
+ and elt.args
173
+ and isinstance(elt.args[0], ast.Constant)
174
+ ):
175
+ sizes.append(f"%{elt.args[0].value}")
176
+ if sizes:
177
+ return sizes
178
+ return []
179
224
-def parse_mla_prefill_backends() -> list[dict[str, Any]]:
225
- """Parse MLA prefill backend options from mla_attention.py.
180
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__.
181
+def _get_parent_class_name(class_node: ast.ClassDef) -> str | None:
182
+ """Get the first parent class name (simple name only).
183
231
- Returns a list of prefill backend info dicts with their requirements.
184
+ Handles both simple inheritance (class Foo(Bar)) and generic
185
+ inheritance (class Foo(Bar[T])).
186
"""
233
- if not MLA_ATTENTION_FILE.exists():
234
- return []
187
+ if not class_node.bases:
188
+ return None
189
+ base = class_node.bases[0]
190
+ if isinstance(base, ast.Name):
191
+ return base.id
192
+ if isinstance(base, ast.Subscript) and isinstance(base.value, ast.Name):
193
+ return base.value.id
194
+ return None
195
236
- try:
237
- tree = ast.parse(MLA_ATTENTION_FILE.read_text())
238
- except Exception:
239
- return []
196
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
197
+def _resolve_import_to_file(
198
+ tree: ast.AST, class_name: str, source_file: Path | None = None
199
+) -> Path | None:
200
+ """Try to resolve a class name to its source file via imports in the AST.
201
202
+ Handles both absolute imports (from vllm.foo import Bar) and relative
203
+ imports (from .foo import Bar) when source_file is provided.
204
+ """
205
for node in ast.walk(tree):
247
- if not isinstance(node, ast.FunctionDef):
206
+ if not isinstance(node, ast.ImportFrom):
207
continue
208
+ for alias in node.names:
209
+ actual_name = alias.asname or alias.name
210
+ if actual_name != class_name:
211
+ continue
212
+ if not node.module:
213
+ continue
214
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
- )
215
+ if node.level and node.level > 0 and source_file:
216
+ # Relative import: resolve from the source file's directory
217
+ base_dir = source_file.parent
218
+ for _ in range(node.level - 1):
219
+ base_dir = base_dir.parent
220
+ module_path = node.module.replace(".", "/")
221
+ py_file = base_dir / f"{module_path}.py"
222
+ else:
223
+ # Absolute import
224
+ module_path = node.module.replace(".", "/")
225
+ py_file = REPO_ROOT / f"{module_path}.py"
226
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
- )
227
+ if py_file.exists():
228
+ return py_file
229
+ return None
230
344
- return prefill_backends
231
232
+def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
233
+ """Find a compute capability from is_device_capability_family() calls in a function.
234
347
-def find_class_in_ast(tree: ast.AST, class_name: str) -> ast.ClassDef | None:
348
- """Find a class definition in an AST."""
235
+ Looks for the pattern: current_platform.is_device_capability_family(N)
236
+ and converts N (e.g. 100) to a CC string (e.g. "10.x").
237
+ """
238
for node in ast.walk(tree):
350
- if isinstance(node, ast.ClassDef) and node.name == class_name:
351
- return node
239
+ if not isinstance(node, ast.FunctionDef) or node.name != func_name:
240
+ continue
241
+ for n in ast.walk(node):
242
+ if (
243
+ isinstance(n, ast.Call)
244
+ and isinstance(n.func, ast.Attribute)
245
+ and n.func.attr == "is_device_capability_family"
246
+ and n.args
247
+ and isinstance(n.args[0], ast.Constant)
248
+ and isinstance(n.args[0].value, int)
249
+ ):
250
+ return f"{n.args[0].value // 10}.x"
251
return None
252
253
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
254
+# ---------------------------------------------------------------------------
255
+# Registry and file resolution
256
+# ---------------------------------------------------------------------------
257
258
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
259
+def parse_registry() -> dict[str, str]:
260
+ """Parse the registry.py file to get backend names and their class paths."""
261
+ tree = ast.parse(REGISTRY_FILE.read_text())
262
+ for node in ast.walk(tree):
263
+ if isinstance(node, ast.ClassDef) and node.name == "AttentionBackendEnum":
264
+ return _extract_enum_values(node)
265
+ return {}
266
267
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."""
268
+def _extract_enum_values(node: ast.ClassDef) -> dict[str, str]:
269
+ """Extract enum name -> value mapping from a class definition."""
270
+ result: dict[str, str] = {}
271
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)):
272
+ if not isinstance(item, ast.Assign):
273
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
274
+ for target in item.targets:
275
+ if not isinstance(target, ast.Name):
276
+ continue
277
+ if isinstance(item.value, ast.Constant) and item.value.value:
278
+ result[target.id] = item.value.value
279
+ return result
280
+
281
+
282
+def get_file_from_class_path(class_path: str) -> Path | None:
283
+ """Convert a class path to a file path."""
284
+ if not class_path:
285
+ return None
286
+ module_path = class_path.rsplit(".", 1)[0].replace(".", "/")
287
+ py_file = REPO_ROOT / f"{module_path}.py"
288
+ return py_file if py_file.exists() else None
289
+
290
+
291
+# ---------------------------------------------------------------------------
292
+# Backend feature extraction from AST
293
+# ---------------------------------------------------------------------------
294
295
296
def parse_supported_dtypes(node: ast.ClassDef) -> str:
@@ -432,35 +332,6 @@ def parse_kv_cache_dtypes(node: ast.ClassDef) -> str:
332
return "auto"
333
334
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
-
335
def parse_block_sizes(node: ast.ClassDef) -> str:
336
"""Parse get_supported_kernel_block_sizes method."""
337
method = find_method(node, "get_supported_kernel_block_sizes")
@@ -536,202 +407,444 @@ def parse_compute_capability(node: ast.ClassDef) -> str:
407
return f"{min_cap[0]}.x-{max_cap[0]}.x"
408
return f"≥{min_cap[0]}.{min_cap[1]}"
409
539
- return "Any"
410
+ return "Any"
411
+
412
+
413
+def parse_attention_types(node: ast.ClassDef) -> str:
414
+ """Parse supports_attn_type method."""
415
+ method = find_method(node, "supports_attn_type")
416
+ if method is None:
417
+ return "Decoder"
418
+
419
+ type_map = {
420
+ "DECODER": "Decoder",
421
+ "ENCODER": "Encoder",
422
+ "ENCODER_ONLY": "Encoder Only",
423
+ "ENCODER_DECODER": "Enc-Dec",
424
+ }
425
+ types: set[str] = set()
426
+
427
+ for n in ast.walk(method):
428
+ # Handle `attn_type in (AttentionType.DECODER, ...)`
429
+ if not (
430
+ isinstance(n, ast.Compare)
431
+ and len(n.ops) == 1
432
+ and isinstance(n.ops[0], ast.In)
433
+ and len(n.comparators) == 1
434
+ and isinstance(n.comparators[0], ast.Tuple | ast.Set)
435
+ ):
436
+ continue
437
+
438
+ for elt in n.comparators[0].elts:
439
+ if isinstance(elt, ast.Attribute) and elt.attr in type_map:
440
+ types.add(type_map[elt.attr])
441
+
442
+ if not types:
443
+ return "Decoder"
444
+ return "All" if len(types) >= 3 else ", ".join(sorted(types))
445
+
446
+
447
+def parse_impl_bool_attr(
448
+ tree: ast.AST,
449
+ class_name: str,
450
+ attr_name: str,
451
+ default: bool = False,
452
+ source_file: Path | None = None,
453
+ _visited: set[str] | None = None,
454
+) -> bool:
455
+ """Parse a boolean class attribute from an impl class, following inheritance.
456
+
457
+ Walks up the inheritance chain within the same file and across files
458
+ (by resolving imports) to find the attribute value.
459
+ """
460
+ if _visited is None:
461
+ _visited = set()
462
+ if class_name in _visited:
463
+ return default
464
+ _visited.add(class_name)
465
+
466
+ class_node = find_class_in_ast(tree, class_name)
467
+ if class_node is None:
468
+ return default
469
+
470
+ # Check directly on this class
471
+ value = _find_bool_class_var(class_node, attr_name)
472
+ if value is not None:
473
+ return value
474
+
475
+ # Check parent class
476
+ parent_name = _get_parent_class_name(class_node)
477
+ if parent_name:
478
+ # Try parent in same file first
479
+ parent_node = find_class_in_ast(tree, parent_name)
480
+ if parent_node is not None:
481
+ return parse_impl_bool_attr(
482
+ tree, parent_name, attr_name, default, source_file, _visited
483
+ )
484
+
485
+ # Try resolving cross-file import
486
+ parent_file = _resolve_import_to_file(tree, parent_name, source_file)
487
+ if parent_file:
488
+ try:
489
+ parent_tree = ast.parse(parent_file.read_text())
490
+ return parse_impl_bool_attr(
491
+ parent_tree,
492
+ parent_name,
493
+ attr_name,
494
+ default,
495
+ parent_file,
496
+ _visited,
497
+ )
498
+ except Exception:
499
+ pass
500
+
501
+ return default
502
+
503
+
504
+def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None:
505
+ """Analyze a backend class and extract feature information."""
506
+ file_path = get_file_from_class_path(class_path)
507
+ if file_path is None:
508
+ return None
509
+
510
+ try:
511
+ tree = ast.parse(file_path.read_text())
512
+ except Exception as e:
513
+ print(f" Warning: Could not parse {file_path}: {e}", file=sys.stderr)
514
+ return None
515
+
516
+ class_name = class_path.rsplit(".", 1)[1]
517
+ class_node = find_class_in_ast(tree, class_name)
518
+ if class_node is None:
519
+ return None
520
+
521
+ # Check if this is an MLA backend by parent class or naming
522
+ parent = _get_parent_class_name(class_node)
523
+ mla_parents = {"MLACommonBackend", "FlashMLABackend", "FlashMLASparseBackend"}
524
+ is_mla_backend = (
525
+ parent in mla_parents
526
+ or ".mla." in class_path.lower()
527
+ or "_mla" in backend_name.lower()
528
+ )
529
+
530
+ # Determine compute capability - use N/A for non-CUDA backends
531
+ is_non_cuda = backend_name.startswith(("CPU_", "ROCM_"))
532
+ compute_cap = "N/A" if is_non_cuda else parse_compute_capability(class_node)
533
+
534
+ # Parse impl class features (DCP support)
535
+ impl_method = find_method(class_node, "get_impl_cls")
536
+ impl_class_name = None
537
+ if impl_method:
538
+ for stmt in ast.walk(impl_method):
539
+ if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.Name):
540
+ impl_class_name = stmt.value.id
541
+ break
542
+
543
+ supports_dcp = False
544
+ if impl_class_name:
545
+ supports_dcp = parse_impl_bool_attr(
546
+ tree, impl_class_name, "can_return_lse_for_decode", False, file_path
547
+ )
548
+
549
+ return {
550
+ "name": backend_name,
551
+ "dtypes": parse_supported_dtypes(class_node),
552
+ "kv_cache_dtypes": parse_kv_cache_dtypes(class_node),
553
+ "block_sizes": parse_block_sizes(class_node),
554
+ "head_sizes": parse_head_sizes(class_node),
555
+ "attn_types": parse_attention_types(class_node),
556
+ "compute_capability": compute_cap,
557
+ "is_mla": is_mla_backend or check_method_overrides(class_node, "is_mla"),
558
+ "supports_sink": check_method_overrides(class_node, "supports_sink"),
559
+ "is_sparse": check_method_overrides(class_node, "is_sparse"),
560
+ "supports_mm_prefix": check_method_overrides(class_node, "supports_mm_prefix"),
561
+ "supports_dcp": supports_dcp,
562
+ }
563
+
564
+
565
+# ---------------------------------------------------------------------------
566
+# Special backend variant parsers (FA2/FA3, FlashInfer TRTLLM, MLA prefill)
567
+# ---------------------------------------------------------------------------
568
+
569
+
570
+def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
571
+ """Parse fa_utils.py to detect FA2 vs FA3 feature differences.
572
+
573
+ Returns a dict with 'fa2' and 'fa3' keys containing their respective
574
+ feature overrides for compute capability, KV cache dtypes, and sink support.
575
+ """
576
+ if not FA_UTILS_FILE.exists():
577
+ return {}
578
+
579
+ try:
580
+ tree = ast.parse(FA_UTILS_FILE.read_text())
581
+ except Exception:
582
+ return {}
583
+
584
+ # Analyze the functions to determine FA3-specific features
585
+ fa3_supports_fp8 = False
586
+ fa3_supports_sinks = False
587
+ fa3_compute_cap: str | None = None
588
+
589
+ for node in ast.walk(tree):
590
+ if not isinstance(node, ast.FunctionDef):
591
+ continue
592
+
593
+ # Check flash_attn_supports_fp8 - looks for `get_flash_attn_version() == 3`
594
+ if node.name == "flash_attn_supports_fp8":
595
+ for n in ast.walk(node):
596
+ if (
597
+ isinstance(n, ast.Compare)
598
+ and isinstance(n.left, ast.Call)
599
+ and isinstance(n.left.func, ast.Name)
600
+ and n.left.func.id == "get_flash_attn_version"
601
+ ):
602
+ fa3_supports_fp8 = True
603
+ break
604
+
605
+ # Check flash_attn_supports_sinks - looks for `get_flash_attn_version() == 3`
606
+ if node.name == "flash_attn_supports_sinks":
607
+ for n in ast.walk(node):
608
+ if (
609
+ isinstance(n, ast.Compare)
610
+ and isinstance(n.left, ast.Call)
611
+ and isinstance(n.left.func, ast.Name)
612
+ and n.left.func.id == "get_flash_attn_version"
613
+ ):
614
+ fa3_supports_sinks = True
615
+ break
616
+
617
+ # Check get_flash_attn_version for FA3 compute capability
618
+ # Look for the ternary: 3 if (device_capability.major == 9 ...) else 2
619
+ if node.name == "get_flash_attn_version":
620
+ for n in ast.walk(node):
621
+ # Look for IfExp (ternary) with `device_capability.major == 9`
622
+ if isinstance(n, ast.IfExp):
623
+ test = n.test
624
+ # Check if test is a BoolOp (and) containing the major check
625
+ if isinstance(test, ast.BoolOp):
626
+ for val in test.values:
627
+ if (
628
+ isinstance(val, ast.Compare)
629
+ and isinstance(val.left, ast.Attribute)
630
+ and val.left.attr == "major"
631
+ and val.comparators
632
+ and isinstance(val.comparators[0], ast.Constant)
633
+ ):
634
+ fa3_compute_cap = f"{val.comparators[0].value}.x"
635
+ break
636
+
637
+ return {
638
+ "fa2": {
639
+ "supports_fp8": False,
640
+ "supports_sink": False,
641
+ },
642
+ "fa3": {
643
+ "compute_capability": fa3_compute_cap,
644
+ "supports_fp8": fa3_supports_fp8,
645
+ "supports_sink": fa3_supports_sinks,
646
+ },
647
+ }
648
+
649
+
650
+def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
651
+ """Parse flashinfer.py to detect TRTLLM-specific features.
652
+
653
+ FLASHINFER uses TRTLLM attention on SM100 (Blackwell), which has different
654
+ capabilities (e.g., sink support) than native FlashInfer on earlier GPUs.
655
+ """
656
+ if not FLASHINFER_UTILS_FILE.exists():
657
+ return {}
658
+
659
+ try:
660
+ tree = ast.parse(FLASHINFER_UTILS_FILE.read_text())
661
+ except Exception:
662
+ return {}
663
664
+ trtllm_compute_cap = _find_cc_in_function(tree, "supports_trtllm_attention")
665
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"
666
+ if not trtllm_compute_cap:
667
+ return {}
668
548
- type_map = {
549
- "DECODER": "Decoder",
550
- "ENCODER": "Encoder",
551
- "ENCODER_ONLY": "Encoder Only",
552
- "ENCODER_DECODER": "Enc-Dec",
669
+ return {
670
+ "native": {
671
+ # Native FlashInfer: everything except SM100
672
+ "supports_sink": False,
673
+ },
674
+ "trtllm": {
675
+ # TRTLLM pathway on Blackwell
676
+ "compute_capability": trtllm_compute_cap,
677
+ "supports_sink": True,
678
+ },
679
}
554
- types: set[str] = set()
680
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
681
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])
682
+def parse_mla_prefill_backends() -> list[dict[str, Any]]:
683
+ """Parse MLA prefill backend options from mla_attention.py.
684
571
- if not types:
572
- return "Decoder"
573
- return "All" if len(types) >= 3 else ", ".join(sorted(types))
685
+ MLA uses different backends for prefill vs decode. The decode backends are
686
+ registered in the registry, but prefill backends are selected at runtime
687
+ based on conditions in MLACommonImpl.__init__.
688
689
+ Returns a list of prefill backend info dicts with their requirements.
690
+ """
691
+ if not MLA_ATTENTION_FILE.exists():
692
+ return []
693
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)
694
+ try:
695
+ tree = ast.parse(MLA_ATTENTION_FILE.read_text())
696
+ except Exception:
697
+ return []
698
699
+ # Find compute capability requirements by parsing use_* functions
700
+ trtllm_cc = _find_cc_in_function(tree, "use_trtllm_ragged_deepseek_prefill")
701
+ flashinfer_cc = _find_cc_in_function(tree, "use_flashinfer_prefill")
702
+ cudnn_cc = _find_cc_in_function(tree, "use_cudnn_prefill")
703
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
704
+ # Build prefill backend list based on what we found
705
+ # Order matches the priority in MLACommonImpl.__init__
706
+ prefill_backends: list[dict[str, Any]] = []
707
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
708
+ # TRT-LLM Ragged (highest priority if available)
709
+ if trtllm_cc:
710
+ prefill_backends.append(
711
+ {
712
+ "name": "TRT-LLM Ragged‡",
713
+ "description": "TensorRT-LLM ragged attention",
714
+ "compute_capability": trtllm_cc,
715
+ "enable": "Default on SM100",
716
+ "disable": "`-ac.use_trtllm_ragged_deepseek_prefill=0`",
717
+ "notes": "DeepSeek R1 dims only",
718
+ }
719
+ )
720
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
721
+ # FlashInfer prefill
722
+ if flashinfer_cc:
723
+ prefill_backends.append(
724
+ {
725
+ "name": "FlashInfer",
726
+ "description": "FlashInfer CUTLASS backend",
727
+ "compute_capability": flashinfer_cc,
728
+ "enable": "`-ac.disable_flashinfer_prefill=0`",
729
+ "disable": "`-ac.disable_flashinfer_prefill=1`",
730
+ "notes": "DeepSeek R1 dims only",
731
+ }
732
+ )
733
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()
734
+ # cuDNN prefill
735
+ if cudnn_cc:
736
+ prefill_backends.append(
737
+ {
738
+ "name": "cuDNN",
739
+ "description": "cuDNN-based attention",
740
+ "compute_capability": cudnn_cc,
741
+ "enable": "`-ac.use_cudnn_prefill=1`",
742
+ "disable": "`-ac.use_cudnn_prefill=0`",
743
+ "notes": "",
744
+ }
745
+ )
746
+
747
+ # FlashAttention is always available as fallback
748
+ prefill_backends.append(
749
+ {
750
+ "name": "FlashAttention",
751
+ "description": "FlashAttention varlen (FA2/FA3)",
752
+ "compute_capability": "Any",
753
+ "enable": "Default fallback",
754
+ "disable": "Use other backends",
755
+ "notes": "FA3 on SM90, FA2 otherwise",
756
+ }
757
)
758
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)
759
+ return prefill_backends
760
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
- }
761
762
+# ---------------------------------------------------------------------------
763
+# Backend variant expansion (FA2/FA3, FlashInfer native/TRTLLM)
764
+# ---------------------------------------------------------------------------
765
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)
766
767
+def _expand_flash_attn_variants(
768
+ all_backends: list[dict[str, Any]],
769
+ fa_features: dict[str, dict[str, Any]],
770
+) -> list[dict[str, Any]]:
771
+ """Expand FLASH_ATTN into FA2 and FA3 variants with different capabilities."""
772
+ expanded = []
773
+ for backend in all_backends:
774
+ if backend["name"] != "FLASH_ATTN":
775
+ backend.setdefault("_sort_key", backend["name"])
776
+ backend.setdefault("_sort_order", 0)
777
+ backend.setdefault("version", "")
778
+ expanded.append(backend)
779
+ continue
780
637
-def bool_to_emoji(value: bool) -> str:
638
- """Convert a boolean to a checkmark or X emoji."""
639
- return "✅" if value else "❌"
781
+ # Create FA2 entry (keeps base backend's compute_capability)
782
+ fa2 = backend.copy()
783
+ fa2["version"] = "FA2*"
784
+ fa2["_sort_key"] = "FLASH_ATTN"
785
+ fa2["_sort_order"] = 0
786
+ fa2["supports_sink"] = fa_features["fa2"]["supports_sink"]
787
+
788
+ # Create FA3 entry (uses parsed compute_capability from fa_utils)
789
+ fa3 = backend.copy()
790
+ fa3["version"] = "FA3*"
791
+ fa3["_sort_key"] = "FLASH_ATTN"
792
+ fa3["_sort_order"] = 1
793
+ if fa_features["fa3"]["compute_capability"]:
794
+ fa3["compute_capability"] = fa_features["fa3"]["compute_capability"]
795
+ fa3["supports_sink"] = fa_features["fa3"]["supports_sink"]
796
+ if fa_features["fa3"]["supports_fp8"]:
797
+ base_dtypes = backend["kv_cache_dtypes"].split(", ")
798
+ fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"]
799
+ new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes]
800
+ fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes)
801
+
802
+ expanded.append(fa2)
803
+ expanded.append(fa3)
804
+ return expanded
805
+
806
+
807
+def _expand_flashinfer_variants(
808
+ all_backends: list[dict[str, Any]],
809
+ fi_features: dict[str, dict[str, Any]],
810
+) -> list[dict[str, Any]]:
811
+ """Expand FLASHINFER into native and TRTLLM variants."""
812
+ expanded = []
813
+ for backend in all_backends:
814
+ if backend["name"] != "FLASHINFER":
815
+ expanded.append(backend)
816
+ continue
817
818
+ # Parse original compute capability to get min CC
819
+ orig_cap = backend["compute_capability"]
820
+ parts = orig_cap.replace(".x", "").split("-")
821
+ min_cc = parts[0] if parts else "7"
822
+ trtllm_cc = fi_features["trtllm"]["compute_capability"]
823
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.
824
+ # Create native entry (pre-Blackwell GPUs)
825
+ native = backend.copy()
826
+ native["version"] = "Native†"
827
+ native["_sort_key"] = "FLASHINFER"
828
+ native["_sort_order"] = 0
829
+ native["supports_sink"] = fi_features["native"]["supports_sink"]
830
+ native["compute_capability"] = f"{min_cc}.x-9.x"
831
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"
832
+ # Create TRTLLM entry
833
+ trtllm = backend.copy()
834
+ trtllm["version"] = "TRTLLM†"
835
+ trtllm["_sort_key"] = "FLASHINFER"
836
+ trtllm["_sort_order"] = 1
837
+ trtllm["compute_capability"] = trtllm_cc
838
+ trtllm["supports_sink"] = fi_features["trtllm"]["supports_sink"]
839
656
- # Check if any backend has a version (for FA2/FA3 split)
657
- has_versions = any(b.get("version") for b in backends)
840
+ expanded.append(native)
841
+ expanded.append(trtllm)
842
+ return expanded
843
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)
844
733
- lines.append("")
734
- return "\n".join(lines)
845
+# ---------------------------------------------------------------------------
846
+# CUDA priority list parsing
847
+# ---------------------------------------------------------------------------
848
849
850
def parse_cuda_priority_lists() -> dict[str, list[str]]:
@@ -827,6 +940,105 @@ def _extract_priorities(body: list, priorities: dict[str, list[str]], prefix: st
940
priorities[f"{prefix}_default"] = backends
941
942
943
+# ---------------------------------------------------------------------------
944
+# Data-driven table rendering
945
+#
946
+# Each column is a (header, formatter) pair. The formatter takes a backend
947
+# info dict and returns the cell string. Tables are assembled by selecting
948
+# which columns to include, then calling _render_table().
949
+# ---------------------------------------------------------------------------
950
+
951
+# Column type alias for readability
952
+TableColumn = tuple[str, Callable[[dict[str, Any]], str]]
953
+
954
+# Shared column definitions -- order here matches the output table order
955
+_COL_BACKEND: TableColumn = ("Backend", lambda b: f"`{b['name']}`")
956
+_COL_VERSION: TableColumn = ("Version", lambda b: b.get("version", ""))
957
+_COL_DTYPES: TableColumn = ("Dtypes", lambda b: b["dtypes"])
958
+_COL_KV_DTYPES: TableColumn = (
959
+ "KV Dtypes",
960
+ lambda b: add_literal_quotes(b["kv_cache_dtypes"]),
961
+)
962
+_COL_BLOCK_SIZES: TableColumn = ("Block Sizes", lambda b: b["block_sizes"])
963
+_COL_HEAD_SIZES: TableColumn = ("Head Sizes", lambda b: b["head_sizes"])
964
+_COL_SINK: TableColumn = ("Sink", lambda b: bool_to_emoji(b["supports_sink"]))
965
+_COL_SPARSE: TableColumn = ("Sparse", lambda b: bool_to_emoji(b["is_sparse"]))
966
+_COL_MM_PREFIX: TableColumn = (
967
+ "MM Prefix",
968
+ lambda b: bool_to_emoji(b["supports_mm_prefix"]),
969
+)
970
+_COL_DCP: TableColumn = ("DCP", lambda b: bool_to_emoji(b["supports_dcp"]))
971
+_COL_ATTN_TYPES: TableColumn = ("Attention Types", lambda b: b["attn_types"])
972
+_COL_COMPUTE_CAP: TableColumn = ("Compute Cap.", lambda b: b["compute_capability"])
973
+
974
+
975
+def add_literal_quotes(value: str) -> str:
976
+ """Add literal backticks around all comma-separated items in a string."""
977
+ items = [item.strip() for item in value.split(",")]
978
+ return ", ".join(f"`{item}`" for item in items)
979
+
980
+
981
+def bool_to_emoji(value: bool) -> str:
982
+ """Convert a boolean to a checkmark or X emoji."""
983
+ return "✅" if value else "❌"
984
+
985
+
986
+def _build_columns(is_mla: bool, has_versions: bool) -> list[TableColumn]:
987
+ """Build the column list for a backend feature table.
988
+
989
+ The column selection depends on whether it's an MLA table (includes
990
+ Sparse column) and whether any backend has version variants (includes
991
+ Version column).
992
+ """
993
+ cols: list[TableColumn] = [_COL_BACKEND]
994
+ if has_versions:
995
+ cols.append(_COL_VERSION)
996
+ cols.extend([_COL_DTYPES, _COL_KV_DTYPES, _COL_BLOCK_SIZES, _COL_HEAD_SIZES])
997
+ cols.append(_COL_SINK)
998
+ if is_mla:
999
+ cols.append(_COL_SPARSE)
1000
+ cols.extend([_COL_MM_PREFIX, _COL_DCP, _COL_ATTN_TYPES, _COL_COMPUTE_CAP])
1001
+ return cols
1002
+
1003
+
1004
+def _sort_key(x: dict[str, Any]) -> tuple[str, int]:
1005
+ """Sort key that keeps parent/child rows together in order."""
1006
+ return (x.get("_sort_key", x["name"]), x.get("_sort_order", 0))
1007
+
1008
+
1009
+def _render_table(
1010
+ columns: list[TableColumn],
1011
+ backends: list[dict[str, Any]],
1012
+) -> list[str]:
1013
+ """Render a markdown table from column specs and backend data."""
1014
+ header = "| " + " | ".join(name for name, _ in columns) + " |"
1015
+ sep = "|" + "|".join("-" * (len(name) + 2) for name, _ in columns) + "|"
1016
+ lines = [header, sep]
1017
+ for info in sorted(backends, key=_sort_key):
1018
+ row = "| " + " | ".join(fmt(info) for _, fmt in columns) + " |"
1019
+ lines.append(row)
1020
+ return lines
1021
+
1022
+
1023
+def generate_markdown_table(
1024
+ backends: list[dict[str, Any]], title: str, is_mla_table: bool = False
1025
+) -> str:
1026
+ """Generate a titled markdown table from backend info."""
1027
+ if not backends:
1028
+ return f"## {title}\n\nNo backends found.\n"
1029
+ has_versions = any(b.get("version") for b in backends)
1030
+ columns = _build_columns(is_mla_table, has_versions)
1031
+ lines = [f"## {title}", ""]
1032
+ lines.extend(_render_table(columns, backends))
1033
+ lines.append("")
1034
+ return "\n".join(lines)
1035
+
1036
+
1037
+# ---------------------------------------------------------------------------
1038
+# Markdown section generators (usage, priority, legend, MLA)
1039
+# ---------------------------------------------------------------------------
1040
+
1041
+
1042
def generate_usage_section() -> str:
1043
"""Generate the usage documentation section."""
1044
return """## Setting the Attention Backend
@@ -959,6 +1171,27 @@ def generate_priority_section(priorities: dict[str, list[str]]) -> str:
1171
return "\n".join(lines)
1172
1173
1174
+def generate_legend() -> str:
1175
+ """Generate a legend explaining the table columns."""
1176
+ return """## Legend
1177
+
1178
+| Column | Description |
1179
+|--------|-------------|
1180
+| **Dtypes** | Supported model data types (fp16, bf16, fp32) |
1181
+| **KV Dtypes** | Supported KV cache data types (`auto`, `fp8`, `fp8_e4m3`, etc.) |
1182
+| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
1183
+| **Head Sizes** | Supported attention head sizes |
1184
+| **Sink** | Attention sink support (for StreamingLLM) |
1185
+| **Sparse** | Sparse attention support (MLA only) |
1186
+| **MM Prefix** | Multimodal prefix full attention support |
1187
+| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
1188
+| **Attention Types** | Supported attention patterns (Decoder, Encoder, Enc-Dec) |
1189
+| **Compute Cap.** | Required CUDA compute capability (N/A for non-CUDA backends) |
1190
+
1191
+**Symbols:** ✅ = Supported, ❌ = Not supported
1192
+"""
1193
+
1194
+
1195
def generate_mla_section(
1196
prefill_backends: list[dict[str, Any]], decode_backends: list[dict[str, Any]]
1197
) -> str:
@@ -999,57 +1232,17 @@ def generate_mla_section(
1232
]
1233
)
1234
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)
1235
+ # Reuse data-driven table rendering for decode backends
1236
+ columns = _build_columns(is_mla=True, has_versions=False)
1237
+ lines.extend(_render_table(columns, decode_backends))
1238
1239
lines.append("")
1240
return "\n".join(lines)
1241
1242
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
-"""
1243
+# ---------------------------------------------------------------------------
1244
+# Top-level orchestration
1245
+# ---------------------------------------------------------------------------
1246
1247
1248
def generate_docs() -> str:
@@ -1071,86 +1264,17 @@ def generate_docs() -> str:
1264
# Collect backend info
1265
all_backends = []
1266
for backend_name, class_path in attention_backends_map.items():
1074
- if backend_name in ("CUSTOM", "TORCH_SDPA"):
1267
+ if backend_name in SKIP_BACKENDS:
1268
continue
1269
info = analyze_backend(backend_name, class_path)
1270
if info:
1271
all_backends.append(info)
1272
1080
- # Expand FLASH_ATTN into FA2 and FA3 variants with different capabilities
1273
+ # Expand backends into version variants
1274
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
1275
+ all_backends = _expand_flash_attn_variants(all_backends, fa_features)
1276
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
1277
+ all_backends = _expand_flashinfer_variants(all_backends, fi_features)
1278
1279
# Split into MLA and non-MLA
1280
mla_backends = [b for b in all_backends if b["is_mla"]]