# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Generates documentation table for attention backends showing feature support. This script parses all registered attention backends using AST (no imports needed) and generates a markdown table showing what features each backend supports, based on the checks in AttentionBackend.validate_configuration(). This approach avoids requiring CUDA/ROCm/GPU libraries to be installed. It runs as an mkdocs-gen-files script, so the page is generated at docs build time rather than being committed to the repository. """ import ast import logging import sys from collections.abc import Callable from pathlib import Path from typing import Any sys.path.insert(0, str(Path(__file__).parent)) from generated_content import fill_markers # noqa: E402 logger = logging.getLogger("mkdocs") # --------------------------------------------------------------------------- # Constants and file paths # --------------------------------------------------------------------------- REPO_ROOT = Path(__file__).parent.parent.parent.parent BACKENDS_DIR = REPO_ROOT / "vllm" / "v1" / "attention" / "backends" REGISTRY_FILE = BACKENDS_DIR / "registry.py" CUDA_PLATFORM_FILE = REPO_ROOT / "vllm" / "platforms" / "cuda.py" FA_UTILS_FILE = BACKENDS_DIR / "fa_utils.py" FLASHINFER_UTILS_FILE = REPO_ROOT / "vllm" / "utils" / "flashinfer.py" MLA_ATTENTION_FILE = ( REPO_ROOT / "vllm" / "model_executor" / "layers" / "attention" / "mla_attention.py" ) # Backends to skip during doc generation SKIP_BACKENDS = {"CUSTOM", "TORCH_SDPA"} BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = { # fp8 is an alias for fp8_ds_mla for FlashMLA Sparse "FLASHMLA_SPARSE": {"fp8"}, } MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" # --------------------------------------------------------------------------- # AST utility helpers # --------------------------------------------------------------------------- def find_class_in_ast(tree: ast.AST, class_name: str) -> ast.ClassDef | None: """Find a class definition in an AST.""" for node in ast.walk(tree): if isinstance(node, ast.ClassDef) and node.name == class_name: return node return None def find_method(node: ast.ClassDef, method_name: str) -> ast.FunctionDef | None: """Find a method in a class definition.""" for item in node.body: if isinstance(item, ast.FunctionDef) and item.name == method_name: return item return None def method_returns_true(method: ast.FunctionDef | None) -> bool: """Check if a method simply returns True.""" if method is None: return False for node in ast.walk(method): if ( isinstance(node, ast.Return) and isinstance(node.value, ast.Constant) and node.value.value is True ): return True return False def check_method_overrides(node: ast.ClassDef, method_name: str) -> bool: """Check if a method is overridden and returns True.""" return method_returns_true(find_method(node, method_name)) def _find_bool_class_var(class_node: ast.ClassDef, var_name: str) -> bool | None: """Find a bool class variable in a class definition. Returns None if not found.""" for item in class_node.body: # Check for annotated assignment: attr: bool = True/False if ( isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name) and item.target.id == var_name and isinstance(item.value, ast.Constant) and isinstance(item.value.value, bool) ): return item.value.value # Check for plain assignment: attr = True/False if isinstance(item, ast.Assign): for target in item.targets: if ( isinstance(target, ast.Name) and target.id == var_name and isinstance(item.value, ast.Constant) and isinstance(item.value.value, bool) ): return item.value.value return None def _parse_list_class_var(node: ast.ClassDef, var_name: str) -> list[str] | None: """Parse a list-type class variable, returning None if not found.""" for item in node.body: if not isinstance(item, ast.AnnAssign): continue if not isinstance(item.target, ast.Name): continue if item.target.id != var_name: continue if not (item.value and isinstance(item.value, ast.List)): continue result = [] for elt in item.value.elts: if isinstance(elt, ast.Attribute): result.append(elt.attr) elif isinstance(elt, ast.Constant): result.append(str(elt.value)) return result return None def _parse_return_list( method: ast.FunctionDef | None, handle_multiple_of: bool = False ) -> list[str]: """Extract list items from a method's return statement.""" if method is None: return [] for stmt in ast.walk(method): if not isinstance(stmt, ast.Return): continue if not isinstance(stmt.value, ast.List): continue sizes = [] for elt in stmt.value.elts: if isinstance(elt, ast.Constant): sizes.append(str(elt.value)) elif ( handle_multiple_of and isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == "MultipleOf" and elt.args and isinstance(elt.args[0], ast.Constant) ): sizes.append(f"%{elt.args[0].value}") if sizes: return sizes return [] def _get_parent_class_name(class_node: ast.ClassDef) -> str | None: """Get the first parent class name (simple name only). Handles both simple inheritance (class Foo(Bar)) and generic inheritance (class Foo(Bar[T])). """ if not class_node.bases: return None base = class_node.bases[0] if isinstance(base, ast.Name): return base.id if isinstance(base, ast.Subscript) and isinstance(base.value, ast.Name): return base.value.id return None def _resolve_import_to_file( tree: ast.AST, class_name: str, source_file: Path | None = None ) -> Path | None: """Try to resolve a class name to its source file via imports in the AST. Handles both absolute imports (from vllm.foo import Bar) and relative imports (from .foo import Bar) when source_file is provided. """ for node in ast.walk(tree): if not isinstance(node, ast.ImportFrom): continue for alias in node.names: actual_name = alias.asname or alias.name if actual_name != class_name: continue if not node.module: continue if node.level and node.level > 0 and source_file: # Relative import: resolve from the source file's directory base_dir = source_file.parent for _ in range(node.level - 1): base_dir = base_dir.parent module_path = node.module.replace(".", "/") py_file = base_dir / f"{module_path}.py" else: # Absolute import module_path = node.module.replace(".", "/") py_file = REPO_ROOT / f"{module_path}.py" if py_file.exists(): return py_file return None def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None: """Find a compute capability from is_device_capability_family() calls in a function. Looks for the pattern: current_platform.is_device_capability_family(N) and converts N (e.g. 100) to a CC string (e.g. "10.x"). """ for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef) or node.name != func_name: continue for n in ast.walk(node): if ( isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "is_device_capability_family" and n.args and isinstance(n.args[0], ast.Constant) and isinstance(n.args[0].value, int) ): return f"{n.args[0].value // 10}.x" return None def _find_exact_cc_in_function(tree: ast.AST, func_name: str) -> str | None: """Find a compute capability from is_device_capability() calls in a function.""" for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef) or node.name != func_name: continue for n in ast.walk(node): if ( isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "is_device_capability" and n.args and isinstance(n.args[0], ast.Constant) and isinstance(n.args[0].value, int) ): capability = n.args[0].value return f"{capability // 10}.{capability % 10}" return None # --------------------------------------------------------------------------- # Registry and file resolution # --------------------------------------------------------------------------- def parse_registry() -> dict[str, str]: """Parse the registry.py file to get backend names and their class paths.""" tree = ast.parse(REGISTRY_FILE.read_text()) for node in ast.walk(tree): if isinstance(node, ast.ClassDef) and node.name == "AttentionBackendEnum": return _extract_enum_values(node) return {} def _extract_enum_values(node: ast.ClassDef) -> dict[str, str]: """Extract enum name -> value mapping from a class definition.""" result: dict[str, str] = {} for item in node.body: if not isinstance(item, ast.Assign): continue for target in item.targets: if not isinstance(target, ast.Name): continue if isinstance(item.value, ast.Constant) and item.value.value: result[target.id] = item.value.value return result def get_file_from_class_path(class_path: str) -> Path | None: """Convert a class path to a file path.""" if not class_path: return None module_path = class_path.rsplit(".", 1)[0].replace(".", "/") py_file = REPO_ROOT / f"{module_path}.py" return py_file if py_file.exists() else None def parse_mla_prefill_registry() -> dict[str, str]: """Parse MLAPrefillBackendEnum from the prefill registry. Returns: A dict mapping backend names to their class paths. """ if not MLA_PREFILL_REGISTRY_FILE.exists(): return {} try: tree = ast.parse(MLA_PREFILL_REGISTRY_FILE.read_text()) except Exception: return {} for node in ast.walk(tree): if isinstance(node, ast.ClassDef) and node.name == "MLAPrefillBackendEnum": return _extract_enum_values(node) return {} def parse_mla_prefill_priorities() -> dict[str, list[str]]: """Parse MLA prefill backend priorities from selector.py. Returns: A dict with keys like 'blackwell' and 'default' containing lists of backend enum names in priority order. """ if not MLA_PREFILL_SELECTOR_FILE.exists(): return {} try: tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text()) except Exception: return {} priorities: dict[str, list[str]] = {} for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef): continue if node.name != "_get_mla_prefill_backend_priorities": continue # Look for if statements checking device_capability.major for stmt in ast.walk(node): if not isinstance(stmt, ast.If): continue # Check if it's a capability.major == 10 check (Blackwell) is_blackwell = ( isinstance(stmt.test, ast.Compare) and isinstance(stmt.test.left, ast.Attribute) and stmt.test.left.attr == "major" and stmt.test.comparators and isinstance(stmt.test.comparators[0], ast.Constant) and stmt.test.comparators[0].value == 10 ) # Extract backends from return statements for body_stmt in stmt.body: if isinstance(body_stmt, ast.Return) and isinstance( body_stmt.value, ast.List ): backends = [] for elt in body_stmt.value.elts: if isinstance(elt, ast.Attribute): backends.append(elt.attr) if is_blackwell: priorities["blackwell"] = backends else: priorities["default"] = backends # Extract from else branch for else_stmt in stmt.orelse: if isinstance(else_stmt, ast.Return) and isinstance( else_stmt.value, ast.List ): backends = [] for elt in else_stmt.value.elts: if isinstance(elt, ast.Attribute): backends.append(elt.attr) priorities["default"] = backends return priorities def parse_mla_dimensions_call(node: ast.AST) -> str | None: """Parse an MLADimensions(...) call into a compact display string.""" if not isinstance(node, ast.Call): return None func = node.func if not isinstance(func, ast.Name) or func.id != "MLADimensions": return None dimensions: dict[str, int] = {} for keyword in node.keywords: if ( keyword.arg is not None and isinstance(keyword.value, ast.Constant) and isinstance(keyword.value.value, int) ): dimensions[keyword.arg] = keyword.value.value qk_nope_head_dim = dimensions.get("qk_nope_head_dim") qk_rope_head_dim = dimensions.get("qk_rope_head_dim") v_head_dim = dimensions.get("v_head_dim") if qk_nope_head_dim is None or qk_rope_head_dim is None or v_head_dim is None: return None return ( f"(qk_nope_head_dim={qk_nope_head_dim}, " f"qk_rope_head_dim={qk_rope_head_dim}, v_head_dim={v_head_dim})" ) def parse_supported_mla_dimensions(node: ast.AST | None) -> list[str]: """Parse a supported_mla_dimensions class variable.""" if not isinstance(node, ast.List): return [] supported_dimensions = [] for element in node.elts: dimensions = parse_mla_dimensions_call(element) if dimensions is not None: supported_dimensions.append(dimensions) return supported_dimensions def _parse_mla_dimension_reference( node: ast.AST, named_dimensions: dict[str, str], ) -> str | None: if isinstance(node, ast.Name): return named_dimensions.get(node.id) return parse_mla_dimensions_call(node) def _parse_mla_dimensions_return( statements: list[ast.stmt], named_dimensions: dict[str, str], ) -> list[str]: for statement in statements: if not isinstance(statement, ast.Return): continue value = statement.value if not ( isinstance(value, ast.Compare) and isinstance(value.left, ast.Name) and value.left.id == "mla_dimensions" and len(value.ops) == 1 and len(value.comparators) == 1 ): continue comparator = value.comparators[0] if isinstance(value.ops[0], ast.Eq): dimension = _parse_mla_dimension_reference( comparator, named_dimensions, ) return [dimension] if dimension is not None else [] if isinstance(value.ops[0], ast.In) and isinstance( comparator, ast.List | ast.Tuple | ast.Set ): dimensions = [ _parse_mla_dimension_reference(element, named_dimensions) for element in comparator.elts ] return [dimension for dimension in dimensions if dimension is not None] return [] def _parse_fa_version_condition(node: ast.AST) -> int | None: if not ( isinstance(node, ast.Compare) and len(node.ops) == 1 and isinstance(node.ops[0], ast.Eq) and len(node.comparators) == 1 and isinstance(node.comparators[0], ast.Constant) and isinstance(node.comparators[0].value, int) ): return None left = node.left is_fa_version = isinstance(left, ast.Name) and left.id == "fa_version" is_get_fa_version = ( isinstance(left, ast.Call) and isinstance(left.func, ast.Name) and left.func.id == "get_flash_attn_version" ) if is_fa_version or is_get_fa_version: return node.comparators[0].value return None def parse_supports_mla_dimensions( method: ast.FunctionDef | None, ) -> tuple[list[str], dict[int | None, list[str]]]: """Parse dimensions and FA-version branches from a support method.""" if method is None: return [], {} named_dimensions: dict[str, str] = {} for statement in method.body: if not ( isinstance(statement, ast.Assign) and len(statement.targets) == 1 and isinstance(statement.targets[0], ast.Name) ): continue dimension = parse_mla_dimensions_call(statement.value) if dimension is not None: named_dimensions[statement.targets[0].id] = dimension support_by_version: dict[int | None, list[str]] = {} for statement in method.body: if isinstance(statement, ast.Return): dimensions = _parse_mla_dimensions_return([statement], named_dimensions) if dimensions: support_by_version[None] = dimensions elif isinstance(statement, ast.If): fa_version = _parse_fa_version_condition(statement.test) if fa_version is None: continue dimensions = _parse_mla_dimensions_return(statement.body, named_dimensions) if dimensions: support_by_version[fa_version] = dimensions default_dimensions = _parse_mla_dimensions_return( statement.orelse, named_dimensions ) if default_dimensions: support_by_version[None] = default_dimensions supported_dimensions: list[str] = [] for dimensions in support_by_version.values(): for dimension in dimensions: if dimension not in supported_dimensions: supported_dimensions.append(dimension) return supported_dimensions, support_by_version def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None: """Parse a single MLA prefill backend file to extract its properties. Args: class_path: The fully qualified class path. Returns: A dict with backend properties, or None if parsing fails. """ file_path = get_file_from_class_path(class_path) if file_path is None: return None try: tree = ast.parse(file_path.read_text()) except Exception: return None class_name = class_path.rsplit(".", 1)[1] class_node = find_class_in_ast(tree, class_name) if class_node is None: return None info: dict[str, Any] = { "compute_capability": "Any", "supported_mla_dimensions": [], "mla_dimension_support_by_fa_version": {}, "dtypes": "fp16, bf16", # Default from base class } # Parse class variables for item in class_node.body: if ( isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name) and item.target.id == "supported_mla_dimensions" ): info["supported_mla_dimensions"] = parse_supported_mla_dimensions( item.value ) # Parse supported_dtypes class variable if ( isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name) and item.target.id == "supported_dtypes" and isinstance(item.value, ast.List) ): dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"} dtypes = [] for elt in item.value.elts: if isinstance(elt, ast.Attribute): dtypes.append(dtype_map.get(elt.attr, elt.attr)) if dtypes: info["dtypes"] = ", ".join(dtypes) dimensions, support_by_version = parse_supports_mla_dimensions( find_method(class_node, "supports_mla_dimensions") ) if dimensions: info["supported_mla_dimensions"] = dimensions info["mla_dimension_support_by_fa_version"] = support_by_version # Parse get_name static method get_name_method = find_method(class_node, "get_name") if get_name_method: for n in ast.walk(get_name_method): if isinstance(n, ast.Return) and isinstance(n.value, ast.Constant): info["name"] = n.value.value # Parse supports_compute_capability classmethod cc_method = find_method(class_node, "supports_compute_capability") if cc_method: for n in ast.walk(cc_method): # Look for capability.major == 10 style checks if ( isinstance(n, ast.Compare) and isinstance(n.left, ast.Attribute) and n.left.attr == "major" and n.comparators and isinstance(n.comparators[0], ast.Constant) ): major = n.comparators[0].value info["compute_capability"] = f"{major}.x" return info def parse_mla_prefill_backends() -> list[dict[str, Any]]: """Parse MLA prefill backend options from the prefill registry. MLA uses different backends for prefill vs decode. The decode backends are registered in the main registry, but prefill backends have their own registry at vllm/v1/attention/backends/mla/prefill/registry.py. Returns a list of prefill backend info dicts with their requirements. """ registry = parse_mla_prefill_registry() priorities = parse_mla_prefill_priorities() if not registry: return [] # Get the priority order (Blackwell order shows all backends) priority_order = priorities.get("blackwell", list(registry.keys())) prefill_backends: list[dict[str, Any]] = [] # Backend-specific metadata that can't be easily parsed from code backend_metadata = { "TRTLLM_RAGGED": { "description": "TensorRT-LLM ragged attention", }, "FLASHINFER": { "description": "FlashInfer CUTLASS backend", }, "FLASH_ATTN": { "description": "FlashAttention varlen (FA2/FA3/FA4)", }, } for backend_name in priority_order: if backend_name not in registry: continue class_path = registry[backend_name] backend_info = parse_mla_prefill_backend_file(class_path) if backend_info is None: continue metadata = backend_metadata.get(backend_name, {}) display_name = backend_info.get("name", backend_name) # Add marker for the highest-priority automatic backend. marker = "" if backend_name == priority_order[0] and priorities.get("blackwell"): marker = "‡" notes = "" supported_mla_dimensions = backend_info.get("supported_mla_dimensions", []) if supported_mla_dimensions: notes = " or ".join(supported_mla_dimensions) + " only" support_by_version = backend_info.get( "mla_dimension_support_by_fa_version", {} ) if any(version is not None for version in support_by_version): default_dimensions = support_by_version.get(None, []) dimension_notes = [] for dimension in supported_mla_dimensions: versions = [ f"FA{version}" for version in (2, 3, 4) if dimension in support_by_version.get(version, default_dimensions) ] version_limit = "" if len(versions) == 3 else " only" dimension_notes.append( f"{dimension} ({'/'.join(versions)}{version_limit})" ) notes = " or ".join(dimension_notes) elif backend_name == "FLASH_ATTN": notes = "FA4 on SM100+, FA3 on SM90, FA2 otherwise" prefill_backends.append( { "name": display_name, "marker": marker, "description": metadata.get("description", ""), "dtypes": backend_info.get("dtypes", "fp16, bf16"), "compute_capability": backend_info.get("compute_capability", "Any"), "notes": notes, } ) return prefill_backends # --------------------------------------------------------------------------- # Backend feature extraction from AST # --------------------------------------------------------------------------- def parse_supported_dtypes(node: ast.ClassDef) -> str: """Parse supported_dtypes class variable.""" dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"} dtypes = _parse_list_class_var(node, "supported_dtypes") if dtypes is None: return "fp16, bf16" return ", ".join(dtype_map.get(d, d) for d in dtypes) def parse_kv_cache_dtypes(node: ast.ClassDef) -> str: """Parse supported_kv_cache_dtypes class var or supports_kv_cache_dtype method.""" # First try the class variable dtypes = _parse_list_class_var(node, "supported_kv_cache_dtypes") if dtypes: return ", ".join(dtypes) # Fall back to parsing the supports_kv_cache_dtype method # Look for `kv_cache_dtype in ["auto", "bfloat16"]` pattern method = find_method(node, "supports_kv_cache_dtype") if method: for n in ast.walk(method): if ( isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], ast.In) and len(n.comparators) == 1 and isinstance(n.comparators[0], ast.List) ): dtypes = [ e.value for e in n.comparators[0].elts if isinstance(e, ast.Constant) and isinstance(e.value, str) ] if dtypes: return ", ".join(dtypes) return "auto" def parse_block_sizes(node: ast.ClassDef) -> str: """Parse get_supported_kernel_block_sizes method.""" method = find_method(node, "get_supported_kernel_block_sizes") sizes = _parse_return_list(method, handle_multiple_of=True) return ", ".join(sizes) if sizes else "Any" def parse_head_sizes(node: ast.ClassDef) -> str: """Parse get_supported_head_sizes method.""" method = find_method(node, "get_supported_head_sizes") sizes = _parse_return_list(method) return ", ".join(sizes) if sizes else "Any" def parse_compute_capability(node: ast.ClassDef) -> str: """Parse supports_compute_capability method.""" method = find_method(node, "supports_compute_capability") if method is None: return "Any" min_cap: tuple[int, int] | None = None max_cap: tuple[int, int] | None = None major_list: list[int] = [] for n in ast.walk(method): if not isinstance(n, ast.Compare): continue # Handle `capability >= DeviceCapability(...)` or `capability <= ...` for op, comp in zip(n.ops, n.comparators): if not ( isinstance(comp, ast.Call) and isinstance(comp.func, ast.Name) and comp.func.id == "DeviceCapability" and comp.args and isinstance(comp.args[0], ast.Constant) ): continue major = comp.args[0].value minor = 0 if len(comp.args) > 1 and isinstance(comp.args[1], ast.Constant): minor = comp.args[1].value if isinstance(op, ast.GtE): min_cap = (major, minor) elif isinstance(op, ast.LtE): max_cap = (major, minor) # Handle `capability.major == N` or `capability.major in [N, M]` if ( isinstance(n.left, ast.Attribute) and n.left.attr == "major" and len(n.ops) == 1 and len(n.comparators) == 1 ): comp = n.comparators[0] if isinstance(n.ops[0], ast.Eq) and isinstance(comp, ast.Constant): major_list.append(comp.value) elif isinstance(n.ops[0], ast.In) and isinstance(comp, ast.List): major_list.extend( e.value for e in comp.elts if isinstance(e, ast.Constant) and isinstance(e.value, int) ) if major_list: major_list.sort() if len(major_list) == 1: return f"{major_list[0]}.x" if major_list == list(range(major_list[0], major_list[-1] + 1)): return f"{major_list[0]}.x-{major_list[-1]}.x" return ", ".join(f"{major}.x" for major in major_list) if min_cap: if max_cap: return f"{min_cap[0]}.x-{max_cap[0]}.x" return f"≥{min_cap[0]}.{min_cap[1]}" return "Any" def parse_attention_types(node: ast.ClassDef) -> str: """Parse supports_attn_type method.""" method = find_method(node, "supports_attn_type") if method is None: return "Decoder" type_map = { "DECODER": "Decoder", "ENCODER": "Encoder", "ENCODER_ONLY": "Encoder Only", "ENCODER_DECODER": "Enc-Dec", } types: set[str] = set() for n in ast.walk(method): # Handle `attn_type in (AttentionType.DECODER, ...)` if not ( isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], ast.In) and len(n.comparators) == 1 and isinstance(n.comparators[0], ast.Tuple | ast.Set) ): continue for elt in n.comparators[0].elts: if isinstance(elt, ast.Attribute) and elt.attr in type_map: types.add(type_map[elt.attr]) if not types: return "Decoder" return "All" if types >= set(type_map.values()) else ", ".join(sorted(types)) def parse_impl_bool_attr( tree: ast.AST, class_name: str, attr_name: str, default: bool = False, source_file: Path | None = None, _visited: set[str] | None = None, ) -> bool: """Parse a boolean class attribute from an impl class, following inheritance. Walks up the inheritance chain within the same file and across files (by resolving imports) to find the attribute value. """ if _visited is None: _visited = set() if class_name in _visited: return default _visited.add(class_name) class_node = find_class_in_ast(tree, class_name) if class_node is None: return default # Check directly on this class value = _find_bool_class_var(class_node, attr_name) if value is not None: return value # Check parent class parent_name = _get_parent_class_name(class_node) if parent_name: # Try parent in same file first parent_node = find_class_in_ast(tree, parent_name) if parent_node is not None: return parse_impl_bool_attr( tree, parent_name, attr_name, default, source_file, _visited ) # Try resolving cross-file import parent_file = _resolve_import_to_file(tree, parent_name, source_file) if parent_file: try: parent_tree = ast.parse(parent_file.read_text()) return parse_impl_bool_attr( parent_tree, parent_name, attr_name, default, parent_file, _visited, ) except Exception: pass return default def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None: """Analyze a backend class and extract feature information.""" file_path = get_file_from_class_path(class_path) if file_path is None: return None try: tree = ast.parse(file_path.read_text()) except Exception as e: logger.warning("Could not parse %s: %s", file_path, e) return None class_name = class_path.rsplit(".", 1)[1] class_node = find_class_in_ast(tree, class_name) if class_node is None: return None # Check if this is an MLA backend by parent class or naming parent = _get_parent_class_name(class_node) mla_parents = {"MLACommonBackend", "FlashMLABackend", "FlashMLASparseBackend"} is_mla_backend = ( parent in mla_parents or ".mla." in class_path.lower() or "_mla" in backend_name.lower() ) # Determine compute capability - use N/A for non-CUDA backends is_non_cuda = backend_name.startswith(("CPU_", "ROCM_")) compute_cap = "N/A" if is_non_cuda else parse_compute_capability(class_node) # Parse impl class features (DCP support) impl_method = find_method(class_node, "get_impl_cls") impl_class_name = None if impl_method: for stmt in ast.walk(impl_method): if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.Name): impl_class_name = stmt.value.id break supports_dcp = False if impl_class_name: supports_dcp = parse_impl_bool_attr( tree, impl_class_name, "can_return_lse_for_decode", False, file_path ) kv_cache_dtypes = parse_kv_cache_dtypes(class_node) if backend_name in BACKEND_KV_DTYPE_EXCLUDES: excluded = BACKEND_KV_DTYPE_EXCLUDES[backend_name] kv_cache_dtypes = ", ".join( d for d in (d.strip() for d in kv_cache_dtypes.split(",")) if d not in excluded ) return { "name": backend_name, "dtypes": parse_supported_dtypes(class_node), "kv_cache_dtypes": kv_cache_dtypes, "block_sizes": parse_block_sizes(class_node), "head_sizes": parse_head_sizes(class_node), "attn_types": parse_attention_types(class_node), "compute_capability": compute_cap, "is_mla": is_mla_backend or check_method_overrides(class_node, "is_mla"), "supports_sink": check_method_overrides(class_node, "supports_sink"), "supports_non_causal": check_method_overrides( class_node, "supports_non_causal" ), "is_sparse": check_method_overrides(class_node, "is_sparse"), "supports_mm_prefix": check_method_overrides(class_node, "supports_mm_prefix"), "supports_dcp": supports_dcp, } # --------------------------------------------------------------------------- # Special backend variant parsers (FA2/FA3/FA4, FlashInfer TRTLLM, MLA prefill) # --------------------------------------------------------------------------- def _parse_fa4_supported_caps() -> str | None: """Parse flash_attn_interface.py for FA4 supported compute capabilities. Looks for `cc not in [9, 10, 11]` pattern in _is_fa4_supported(). """ fa_interface_file = ( REPO_ROOT / "vllm" / "vllm_flash_attn" / "flash_attn_interface.py" ) if not fa_interface_file.exists(): return None try: tree = ast.parse(fa_interface_file.read_text()) except Exception: return None for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef) or node.name != "_is_fa4_supported": continue for n in ast.walk(node): if not ( isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], ast.NotIn) and isinstance(n.comparators[0], ast.List) ): continue caps: list[int] = [ e.value for e in n.comparators[0].elts if isinstance(e, ast.Constant) and isinstance(e.value, int) ] if caps: caps.sort() return f"{caps[0]}.x-{caps[-1]}.x" return None def parse_flash_attn_features() -> dict[str, dict[str, Any]]: """Parse fa_utils.py to detect FA2 vs FA3 vs FA4 feature differences. Returns a dict with 'fa2', 'fa3', and 'fa4' keys containing their respective feature overrides for compute capability, KV cache dtypes, and sink support. """ if not FA_UTILS_FILE.exists(): return {} try: tree = ast.parse(FA_UTILS_FILE.read_text()) except Exception: return {} # Analyze the functions to determine FA3/FA4-specific features fa3_supports_fp8 = False fa4_supports_fp8 = False fa3_supports_sinks = False fa4_supports_sinks = False fa3_compute_cap: str | None = None fa4_compute_cap: str | None = None for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef): continue # Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version. # Accept both equality checks and membership checks such as `in (3, 4)`. if node.name == "flash_attn_supports_kv_cache_dtype": for n in ast.walk(node): if not ( isinstance(n, ast.Compare) and len(n.ops) == 1 and len(n.comparators) == 1 ): continue is_version_compare = ( isinstance(n.left, ast.Name) and n.left.id == "fa_version" ) or ( isinstance(n.left, ast.Call) and isinstance(n.left.func, ast.Name) and n.left.func.id == "get_flash_attn_version" ) if not is_version_compare: continue versions: list[Any] = [] comparator = n.comparators[0] if isinstance(n.ops[0], ast.Eq) and isinstance( comparator, ast.Constant ): versions = [comparator.value] elif isinstance(n.ops[0], ast.In) and isinstance( comparator, (ast.Tuple, ast.List, ast.Set) ): versions = [ elt.value for elt in comparator.elts if isinstance(elt, ast.Constant) ] fa3_supports_fp8 |= 3 in versions fa4_supports_fp8 |= 4 in versions # Check flash_attn_supports_sinks - looks for `fa_version == 3/4` # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`) if node.name == "flash_attn_supports_sinks": for n in ast.walk(node): if ( isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], ast.Eq) and isinstance(n.comparators[0], ast.Constant) ): is_version_compare = ( isinstance(n.left, ast.Name) and n.left.id == "fa_version" ) or ( isinstance(n.left, ast.Call) and isinstance(n.left.func, ast.Name) and n.left.func.id == "get_flash_attn_version" ) if is_version_compare: val = n.comparators[0].value if val == 3: fa3_supports_sinks = True elif val == 4: fa4_supports_sinks = True elif ( isinstance(n, ast.Compare) and len(n.ops) == 1 and isinstance(n.ops[0], ast.In) and isinstance(n.comparators[0], (ast.Tuple, ast.List, ast.Set)) ): is_version_compare = ( isinstance(n.left, ast.Name) and n.left.id == "fa_version" ) or ( isinstance(n.left, ast.Call) and isinstance(n.left.func, ast.Name) and n.left.func.id == "get_flash_attn_version" ) if is_version_compare: for elt in n.comparators[0].elts: if isinstance(elt, ast.Constant): if elt.value == 3: fa3_supports_sinks = True elif elt.value == 4: fa4_supports_sinks = True # Check get_flash_attn_version for FA3/FA4 compute capability if node.name == "get_flash_attn_version": for n in ast.walk(node): # Handle IfExp (ternary) with `device_capability.major == 9` if isinstance(n, ast.IfExp): test = n.test if isinstance(test, ast.BoolOp): for val in test.values: if ( isinstance(val, ast.Compare) and isinstance(val.left, ast.Attribute) and val.left.attr == "major" and val.comparators and isinstance(val.comparators[0], ast.Constant) ): fa3_compute_cap = f"{val.comparators[0].value}.x" break # Handle If statements for FA3/FA4 detection # e.g. `if device_capability.major == 9` -> FA3 # `elif device_capability.major >= 10` -> FA4 if isinstance(n, ast.If): test = n.test comparisons = ( [v for v in test.values if isinstance(v, ast.Compare)] if isinstance(test, ast.BoolOp) else [test] if isinstance(test, ast.Compare) else [] ) for comp in comparisons: if not ( isinstance(comp.left, ast.Attribute) and comp.left.attr == "major" and comp.comparators and isinstance(comp.comparators[0], ast.Constant) and isinstance(comp.comparators[0].value, int) ): continue op = comp.ops[0] val = comp.comparators[0].value if isinstance(op, ast.Eq) and fa3_compute_cap is None: fa3_compute_cap = f"{val}.x" elif isinstance(op, ast.GtE) and fa4_compute_cap is None: fa4_compute_cap = f"≥{val}.0" # Fallback: try to parse FA4 compute caps from flash_attn_interface.py if fa4_compute_cap is None: fa4_compute_cap = _parse_fa4_supported_caps() return { "fa2": { "supports_fp8": False, "supports_sink": False, }, "fa3": { "compute_capability": fa3_compute_cap, "supports_fp8": fa3_supports_fp8, "supports_sink": fa3_supports_sinks, }, "fa4": { "compute_capability": fa4_compute_cap, "supports_fp8": fa4_supports_fp8, "supports_sink": fa4_supports_sinks, }, } def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]: """Parse flashinfer.py to detect FlashInfer TRTLLM API variants. FLASHINFER uses XQA on SM90 and trtllm-gen on SM100 through FlashInfer's TRTLLM decode API. These variants have different capabilities than native FlashInfer. """ if not FLASHINFER_UTILS_FILE.exists(): return {} try: tree = ast.parse(FLASHINFER_UTILS_FILE.read_text()) except Exception: return {} xqa_compute_cap = _find_exact_cc_in_function(tree, "supports_trtllm_attention") trtllm_gen_compute_cap = _find_cc_in_function(tree, "supports_trtllm_attention") if not xqa_compute_cap and not trtllm_gen_compute_cap: return {} # KV cache dtypes that only work with a dedicated kernel (e.g. nvfp4 # requires the SM100 NVFP4 MHA kernel) and should not appear in the # generic attention-backend feature matrix. kernel_only_kv_dtypes = ["nvfp4"] return { "native": { # Native FlashInfer path. "supports_sink": False, }, "xqa": { # XQA decode path on Hopper. "compute_capability": xqa_compute_cap, "supports_sink": False, }, "trtllm_gen": { # trtllm-gen pathway on Blackwell. "compute_capability": trtllm_gen_compute_cap, "supports_sink": True, }, "exclude_kv_dtypes": kernel_only_kv_dtypes, } # --------------------------------------------------------------------------- # Backend variant expansion (FA2/FA3/FA4, FlashInfer native/XQA/trtllm-gen) # --------------------------------------------------------------------------- def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str: """Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string. The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``, but actual support varies by FA version (e.g. FA2 has no fp8 support), so each variant's row is normalized to match its own capability. """ fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"} dtypes = kv_cache_dtypes.split(", ") base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes] if supports_fp8: advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes] return ", ".join(base_dtypes + advertised_fp8_dtypes) return ", ".join(base_dtypes) def _expand_flash_attn_variants( all_backends: list[dict[str, Any]], fa_features: dict[str, dict[str, Any]], ) -> list[dict[str, Any]]: """Expand FLASH_ATTN into FA2, FA3, and FA4 variants.""" expanded = [] for backend in all_backends: if backend["name"] != "FLASH_ATTN": backend.setdefault("_sort_key", backend["name"]) backend.setdefault("_sort_order", 0) backend.setdefault("version", "") expanded.append(backend) continue # Create FA2 entry (keeps base backend's compute_capability) fa2 = backend.copy() fa2["version"] = "FA2*" fa2["_sort_key"] = "FLASH_ATTN" fa2["_sort_order"] = 0 fa2["supports_sink"] = fa_features["fa2"]["supports_sink"] fa2["kv_cache_dtypes"] = _apply_fp8_support( backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"] ) # Create FA3 entry (uses parsed compute_capability from fa_utils) fa3 = backend.copy() fa3["version"] = "FA3*" fa3["_sort_key"] = "FLASH_ATTN" fa3["_sort_order"] = 1 if fa_features["fa3"]["compute_capability"]: fa3["compute_capability"] = fa_features["fa3"]["compute_capability"] fa3["supports_sink"] = fa_features["fa3"]["supports_sink"] fa3["kv_cache_dtypes"] = _apply_fp8_support( backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"] ) expanded.append(fa2) expanded.append(fa3) # Create FA4 entry if FA4 features are available if "fa4" in fa_features: fa4 = backend.copy() fa4["version"] = "FA4*" fa4["_sort_key"] = "FLASH_ATTN" fa4["_sort_order"] = 2 if fa_features["fa4"].get("compute_capability"): fa4["compute_capability"] = fa_features["fa4"]["compute_capability"] fa4["supports_sink"] = fa_features["fa4"]["supports_sink"] fa4["kv_cache_dtypes"] = _apply_fp8_support( backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"] ) expanded.append(fa4) return expanded def _expand_flashinfer_variants( all_backends: list[dict[str, Any]], fi_features: dict[str, dict[str, Any]], ) -> list[dict[str, Any]]: """Expand FLASHINFER into native, XQA, and trtllm-gen variants.""" expanded = [] for backend in all_backends: if backend["name"] != "FLASHINFER": expanded.append(backend) continue # Parse original compute capability to get min CC orig_cap = backend["compute_capability"] parts = orig_cap.replace(".x", "").split("-") min_cc = parts[0] if parts else "7" # Create native entry. native = backend.copy() native["version"] = "Native†" native["_sort_key"] = "FLASHINFER" native["_sort_order"] = 0 native["supports_sink"] = fi_features["native"]["supports_sink"] native["compute_capability"] = f"{min_cc}.x-9.x" # Remove KV dtypes only supported by SM100 kernels (e.g. nvfp4) exclude = fi_features.get("exclude_kv_dtypes", []) if exclude: native["kv_cache_dtypes"] = ", ".join( d for d in (d.strip() for d in native["kv_cache_dtypes"].split(",")) if d not in exclude ) # Create XQA entry. xqa = backend.copy() xqa["version"] = "XQA†" xqa["_sort_key"] = "FLASHINFER" xqa["_sort_order"] = 1 xqa["compute_capability"] = fi_features["xqa"]["compute_capability"] xqa["supports_sink"] = fi_features["xqa"]["supports_sink"] xqa["supports_non_causal"] = False if exclude: xqa["kv_cache_dtypes"] = ", ".join( d for d in (d.strip() for d in xqa["kv_cache_dtypes"].split(",")) if d not in exclude ) # Create trtllm-gen entry. trtllm_gen = backend.copy() trtllm_gen["version"] = "trtllm-gen†" trtllm_gen["_sort_key"] = "FLASHINFER" trtllm_gen["_sort_order"] = 2 trtllm_gen["compute_capability"] = fi_features["trtllm_gen"][ "compute_capability" ] trtllm_gen["supports_sink"] = fi_features["trtllm_gen"]["supports_sink"] expanded.append(native) if fi_features["xqa"]["compute_capability"]: expanded.append(xqa) if fi_features["trtllm_gen"]["compute_capability"]: expanded.append(trtllm_gen) return expanded # --------------------------------------------------------------------------- # CUDA priority list parsing # --------------------------------------------------------------------------- def parse_cuda_priority_lists() -> dict[str, list[str]]: """Parse priority lists from cuda.py using AST. The structure of _get_backend_priorities is: if use_mla: if device_capability.major == 10: return [MLA list for SM100] else: return [MLA list for default] else: if device_capability.major == 10: return [Standard list for SM100] else: return [Standard list for default] """ if not CUDA_PLATFORM_FILE.exists(): return {} try: source = CUDA_PLATFORM_FILE.read_text() tree = ast.parse(source) except Exception: return {} priorities: dict[str, list[str]] = {} # Find the _get_backend_priorities function for node in ast.walk(tree): if not isinstance(node, ast.FunctionDef): continue if node.name != "_get_backend_priorities": continue # Process the function body directly for stmt in node.body: if not isinstance(stmt, ast.If): continue # Check if this is the "if use_mla:" branch is_mla_branch = ( isinstance(stmt.test, ast.Name) and stmt.test.id == "use_mla" ) if is_mla_branch: _extract_priorities(stmt.body, priorities, "mla") if stmt.orelse: _extract_priorities(stmt.orelse, priorities, "standard") else: _extract_priorities([stmt], priorities, "standard") return priorities def _get_backends_from_return(stmts: list) -> list[str]: """Extract backend names from return statements in a list of statements. Handles starred unpacking (e.g. ``*sparse_backends``) by resolving the variable from assignments found in the same statement list. When the variable is conditionally assigned (inside an ``if/else``), the ``else`` branch value is used as the representative default. """ # Collect variable assignments so we can resolve starred expressions. # For conditional assignments, last-written (else branch) wins. var_assigns: dict[str, list[str]] = {} for stmt in stmts: if isinstance(stmt, ast.Assign) and isinstance(stmt.value, ast.List): for target in stmt.targets: if isinstance(target, ast.Name): var_assigns[target.id] = [ e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute) ] elif isinstance(stmt, ast.If): for branch in (stmt.body, stmt.orelse): for branch_stmt in branch: if isinstance(branch_stmt, ast.Assign) and isinstance( branch_stmt.value, ast.List ): for target in branch_stmt.targets: if isinstance(target, ast.Name): var_assigns[target.id] = [ e.attr for e in branch_stmt.value.elts if isinstance(e, ast.Attribute) ] for stmt in stmts: if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List): backends: list[str] = [] for e in stmt.value.elts: if isinstance(e, ast.Attribute): backends.append(e.attr) elif ( isinstance(e, ast.Starred) and isinstance(e.value, ast.Name) and e.value.id in var_assigns ): backends.extend(var_assigns[e.value.id]) return backends return [] def _is_sm100_check(test: ast.expr) -> bool: """Check if test is `something.major == 10`, possibly inside an `and`.""" if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And): return any(_is_sm100_check(value) for value in test.values) return ( isinstance(test, ast.Compare) and isinstance(test.left, ast.Attribute) and test.left.attr == "major" and len(test.ops) == 1 and isinstance(test.ops[0], ast.Eq) and len(test.comparators) == 1 and isinstance(test.comparators[0], ast.Constant) and test.comparators[0].value == 10 ) def _extract_priorities(body: list, priorities: dict[str, list[str]], prefix: str): """Extract priority lists from if/else statement body.""" for stmt in body: if isinstance(stmt, ast.If): is_sm100 = _is_sm100_check(stmt.test) if_key = f"{prefix}_sm100" if is_sm100 else f"{prefix}_default" else_key = f"{prefix}_default" if is_sm100 else f"{prefix}_sm100" if backends := _get_backends_from_return(stmt.body): priorities[if_key] = backends if backends := _get_backends_from_return(stmt.orelse): priorities[else_key] = backends elif isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List): backends = [e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)] priorities[f"{prefix}_default"] = backends # --------------------------------------------------------------------------- # Data-driven table rendering # # Each column is a (header, formatter) pair. The formatter takes a backend # info dict and returns the cell string. Tables are assembled by selecting # which columns to include, then calling _render_table(). # --------------------------------------------------------------------------- # Column type alias for readability TableColumn = tuple[str, Callable[[dict[str, Any]], str]] # Shared column definitions -- order here matches the output table order _COL_BACKEND: TableColumn = ("Backend", lambda b: f"`{b['name']}`") _COL_VERSION: TableColumn = ("Version", lambda b: b.get("version", "")) _COL_DTYPES: TableColumn = ("Dtypes", lambda b: b["dtypes"]) _COL_KV_DTYPES: TableColumn = ( "KV Dtypes", lambda b: add_literal_quotes(b["kv_cache_dtypes"]), ) _COL_BLOCK_SIZES: TableColumn = ("Block Sizes", lambda b: b["block_sizes"]) _COL_HEAD_SIZES: TableColumn = ("Head Sizes", lambda b: b["head_sizes"]) _COL_SINK: TableColumn = ("Sink", lambda b: bool_to_emoji(b["supports_sink"])) _COL_NON_CAUSAL: TableColumn = ( "Non-Causal", lambda b: bool_to_emoji(b["supports_non_causal"]), ) _COL_SPARSE: TableColumn = ("Sparse", lambda b: bool_to_emoji(b["is_sparse"])) _COL_MM_PREFIX: TableColumn = ( "MM Prefix", lambda b: bool_to_emoji(b["supports_mm_prefix"]), ) _COL_DCP: TableColumn = ("DCP", lambda b: bool_to_emoji(b["supports_dcp"])) _COL_ATTN_TYPES: TableColumn = ("Attention Types", lambda b: b["attn_types"]) _COL_COMPUTE_CAP: TableColumn = ("Compute Cap.", lambda b: b["compute_capability"]) def add_literal_quotes(value: str) -> str: """Add literal backticks around all comma-separated items in a string.""" items = [item.strip() for item in value.split(",")] return ", ".join(f"`{item}`" for item in items) def bool_to_emoji(value: bool) -> str: """Convert a boolean to a checkmark or X emoji.""" return "✅" if value else "❌" def _build_columns(is_mla: bool, has_versions: bool) -> list[TableColumn]: """Build the column list for a backend feature table. The column selection depends on whether it's an MLA table (includes Sparse column) and whether any backend has version variants (includes Version column). """ cols: list[TableColumn] = [_COL_BACKEND] if has_versions: cols.append(_COL_VERSION) cols.extend([_COL_DTYPES, _COL_KV_DTYPES, _COL_BLOCK_SIZES, _COL_HEAD_SIZES]) cols.append(_COL_SINK) cols.append(_COL_NON_CAUSAL) if is_mla: cols.append(_COL_SPARSE) cols.extend([_COL_MM_PREFIX, _COL_DCP, _COL_ATTN_TYPES, _COL_COMPUTE_CAP]) return cols def _sort_key(x: dict[str, Any]) -> tuple[str, int]: """Sort key that keeps parent/child rows together in order.""" return (x.get("_sort_key", x["name"]), x.get("_sort_order", 0)) def _render_table( columns: list[TableColumn], backends: list[dict[str, Any]], ) -> list[str]: """Render a markdown table from column specs and backend data.""" header = "| " + " | ".join(name for name, _ in columns) + " |" sep = "| " + " | ".join("-" * len(name) for name, _ in columns) + " |" lines = [header, sep] for info in sorted(backends, key=_sort_key): row = "| " + " | ".join(fmt(info) for _, fmt in columns) + " |" lines.append(row.replace(" ", " ")) return lines def _priority_table( title: str, backends: list[str], annotations: dict[str, str] | None = None, ) -> list[str]: """Render a priority table for a list of backends.""" def _fmt(b: str) -> str: suffix = annotations.get(b, "") if annotations else "" return f"`{b}`{suffix}" return [ f"**{title}:**", "", "| Priority | Backend |", "| -------- | ------- |", *[f"| {i} | {_fmt(b)} |" for i, b in enumerate(backends, 1)], "", ] _SM100 = "Blackwell (SM 10.x)" _AMPERE = "Ampere/Hopper (SM 8.x-9.x)" def _priority_block( priorities: dict[str, list[str]], sm100_key: str, default_key: str, sm100_annotations: dict[str, str] | None = None, ) -> str: """Render whichever priority tables exist for one attention category.""" lines: list[str] = [] if sm100_key in priorities: lines += _priority_table(_SM100, priorities[sm100_key], sm100_annotations) if default_key in priorities: lines += _priority_table(_AMPERE, priorities[default_key]) return "\n".join(lines).strip() def _feature_table(backends: list[dict[str, Any]], is_mla: bool) -> str: """Render a backend feature table (header, separator, one row per backend).""" has_versions = any(b.get("version") for b in backends) columns = _build_columns(is_mla, has_versions) return "\n".join(_render_table(columns, backends)) def _mla_prefill_table(prefill_backends: list[dict[str, Any]]) -> str: """Render the MLA prefill backend table.""" lines = [ "| Backend | Description | Dtypes | Compute Cap. | Notes |", "| ------- | ----------- | ------ | ------------ | ----- |", ] for backend in prefill_backends: row = "| `{}`{} | {} | {} | {} | {} |".format( backend["name"], backend.get("marker", ""), backend["description"], backend.get("dtypes", "fp16, bf16"), backend["compute_capability"], backend.get("notes", ""), ) lines.append(row.replace(" ", " ")) return "\n".join(lines) def build_blocks() -> dict[str, str]: """Build the generated table blocks keyed by their `gen:` marker name. Only the tables are generated here; the surrounding prose lives in the handwritten ``docs/design/attention_backends.md`` page. """ attention_backends_map = parse_registry() priorities = parse_cuda_priority_lists() fa_features = parse_flash_attn_features() fi_features = parse_flashinfer_trtllm_features() mla_prefill_backends = parse_mla_prefill_backends() all_backends = [] for backend_name, class_path in attention_backends_map.items(): if backend_name in SKIP_BACKENDS: continue info = analyze_backend(backend_name, class_path) if info: all_backends.append(info) if fa_features: all_backends = _expand_flash_attn_variants(all_backends, fa_features) if fi_features: all_backends = _expand_flashinfer_variants(all_backends, fi_features) # DeepSeek V4 (*_DSV4) and MiniMax M3 sparse backends get their own tables # rather than mixing into the main MLA / standard tables (the ROCm V4 backend # isn't flagged is_mla by the AST heuristic, so filter purely on the name). def _is_v4(b: dict[str, Any]) -> bool: return b["name"].endswith("_DSV4") def _is_minimax(b: dict[str, Any]) -> bool: return not b["is_mla"] and not _is_v4(b) and b["name"].startswith("MINIMAX") v4_decode_backends = [b for b in all_backends if _is_v4(b)] minimax_backends = [b for b in all_backends if _is_minimax(b)] mla_backends = [b for b in all_backends if b["is_mla"] and not _is_v4(b)] non_mla_backends = [ b for b in all_backends if not b["is_mla"] and not _is_v4(b) and not _is_minimax(b) ] mla_sm100_annotations = {"FLASHINFER_MLA_SPARSE": "**\\***"} return { "priority-standard": _priority_block( priorities, "standard_sm100", "standard_default" ), "priority-mla": _priority_block( priorities, "mla_sm100", "mla_default", mla_sm100_annotations ), "table-standard": _feature_table(non_mla_backends, is_mla=False), "table-minimax": _feature_table(minimax_backends, is_mla=False), "table-mla-prefill": _mla_prefill_table(mla_prefill_backends), "table-mla-decode": _feature_table(mla_backends, is_mla=True), "table-mla-v4-decode": _feature_table(v4_decode_backends, is_mla=True), } logger.info("Generating attention backend documentation") fill_markers("design/attention_backends.md", build_blocks())