| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention fuser: the module that dispatches to the attention interface.""" |
| 4 | |
| 5 | import ast |
| 6 | import inspect |
| 7 | import textwrap |
| 8 | from collections.abc import Callable |
| 9 | from dataclasses import dataclass |
| 10 | from functools import cache |
| 11 | from typing import TYPE_CHECKING, ClassVar |
| 12 | |
| 13 | import torch |
| 14 | from torch import fx, nn |
| 15 | |
| 16 | from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader |
| 17 | from vllm.model_executor.models.transformers.fusers.base import BaseFuser |
| 18 | from vllm.model_executor.utils import set_weight_attrs |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | from vllm.config import VllmConfig |
| 22 | |
| 23 | VLLM_ATTN_IMPL = "vllm" |
| 24 | VLLM_MLA_ATTN_IMPL = "vllm_mla" |
| 25 | |
| 26 | |
| 27 | def _is_interface_lookup(node: ast.expr | None) -> bool: |
| 28 | """Whether `node` reads an entry out of `ALL_ATTENTION_FUNCTIONS`.""" |
| 29 | # ALL_ATTENTION_FUNCTIONS.get_interface(...) or ALL_ATTENTION_FUNCTIONS[...] |
| 30 | if isinstance(node, ast.Call): |
| 31 | node = node.func |
| 32 | if not isinstance(node, (ast.Attribute, ast.Subscript)): |
| 33 | return False |
| 34 | return ( |
| 35 | isinstance(node.value, ast.Name) and node.value.id == "ALL_ATTENTION_FUNCTIONS" |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | @cache |
| 40 | def interface_call(forward: Callable) -> ast.Call | None: |
| 41 | """The attention interface call in `forward`, if it makes exactly one.""" |
| 42 | try: |
| 43 | source = inspect.getsource(inspect.unwrap(forward)) |
| 44 | tree = ast.parse(textwrap.dedent(source)) |
| 45 | except (AttributeError, OSError, SyntaxError, TypeError): |
| 46 | return None |
| 47 | |
| 48 | # Find the names of the local variables that read from the interface lookup |
| 49 | names: set[str] = set() |
| 50 | for node in ast.walk(tree): |
| 51 | if not isinstance(node, (ast.Assign, ast.AnnAssign)): |
| 52 | continue |
| 53 | if not _is_interface_lookup(node.value): |
| 54 | continue |
| 55 | targets = node.targets if isinstance(node, ast.Assign) else [node.target] |
| 56 | names.update(t.id for t in targets if isinstance(t, ast.Name)) |
| 57 | # Find the calls to those names, and return the one if there is exactly one |
| 58 | calls = [ |
| 59 | node |
| 60 | for node in ast.walk(tree) |
| 61 | if isinstance(node, ast.Call) |
| 62 | and isinstance(node.func, ast.Name) |
| 63 | and node.func.id in names |
| 64 | ] |
| 65 | return calls[0] if len(calls) == 1 else None |
| 66 | |
| 67 | |
| 68 | def _resolve(node: ast.expr, module: nn.Module) -> object: |
| 69 | """The value of `node` on `module`, for literals and `self.<attr>`.""" |
| 70 | if isinstance(node, ast.Constant): |
| 71 | return node.value |
| 72 | if ( |
| 73 | isinstance(node, ast.Attribute) |
| 74 | and isinstance(node.value, ast.Name) |
| 75 | and node.value.id == "self" |
| 76 | ): |
| 77 | return getattr(module, node.attr, None) |
| 78 | return None |
| 79 | |
| 80 | |
| 81 | @dataclass |
| 82 | class AttentionFuser(BaseFuser): |
| 83 | """A module that dispatches through the Transformers attention interface.""" |
| 84 | |
| 85 | redefines_forward: ClassVar[bool] = False |
| 86 | |
| 87 | source_cls: str |
| 88 | """Class of the HF module that dispatches (for logging).""" |
| 89 | scale_expr: ast.expr | None = None |
| 90 | """Source of the `scaling=` the module hands the interface, if it hands one.""" |
| 91 | s_aux_expr: ast.expr | None = None |
| 92 | """Source of the `s_aux=` the module hands the interface, if it hands one.""" |
| 93 | |
| 94 | def info(self, name: str) -> str: |
| 95 | return f"Found: {name} ({self.source_cls}) -> attention interface" |
| 96 | |
| 97 | @classmethod |
| 98 | def match( |
| 99 | cls, graph: fx.Graph | None, module: nn.Module |
| 100 | ) -> "AttentionFuser | None": |
| 101 | if (call := interface_call(type(module).forward)) is None: |
| 102 | return None |
| 103 | scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"] |
| 104 | scale_expr = scaling[0] if len(scaling) == 1 else None |
| 105 | s_aux = [kw.value for kw in call.keywords if kw.arg == "s_aux"] |
| 106 | s_aux_expr = s_aux[0] if len(s_aux) == 1 else None |
| 107 | return cls( |
| 108 | source_cls=type(module).__name__, |
| 109 | scale_expr=scale_expr, |
| 110 | s_aux_expr=s_aux_expr, |
| 111 | ) |
| 112 | |
| 113 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 114 | """Whether `module` will actually dispatch to vLLM.""" |
| 115 | config = getattr(module, "config", None) |
| 116 | # Only patched in the text config, this excludes attention based mm encoders |
| 117 | vllm_attn_impls = {VLLM_ATTN_IMPL, VLLM_MLA_ATTN_IMPL} |
| 118 | return getattr(config, "_attn_implementation", None) in vllm_attn_impls |
| 119 | |
| 120 | def fuse( |
| 121 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 122 | ) -> nn.Module: |
| 123 | if (sinks := self.sinks(module)) is not None: |
| 124 | assert isinstance(self.s_aux_expr, ast.Attribute) |
| 125 | size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size |
| 126 | device = vllm_config.device_config.device |
| 127 | data = torch.empty(size, dtype=sinks.dtype, device=device) |
| 128 | sinks_param = nn.Parameter(data, requires_grad=False) |
| 129 | set_weight_attrs(sinks_param, {"weight_loader": sharded_weight_loader(0)}) |
| 130 | setattr(module, self.s_aux_expr.attr, sinks_param) |
| 131 | return module |
| 132 | |
| 133 | def layer_index(self, module: nn.Module) -> int | None: |
| 134 | """The layer `module` computes attention for, if it declares one.""" |
| 135 | layer_idx = getattr(module, "layer_idx", None) |
| 136 | return layer_idx if isinstance(layer_idx, int) else None |
| 137 | |
| 138 | def scale(self, module: nn.Module) -> float | None: |
| 139 | """The softmax scale `module` passes to the interface, or `None`.""" |
| 140 | if self.scale_expr is None: |
| 141 | return None |
| 142 | scale = _resolve(self.scale_expr, module) |
| 143 | if not isinstance(scale, (int, float)) or isinstance(scale, bool): |
| 144 | expression = ast.unparse(self.scale_expr) |
| 145 | raise ValueError( |
| 146 | f"Cannot resolve attention scaling expression {expression!r} in " |
| 147 | f"{type(module).__name__}." |
| 148 | ) |
| 149 | return float(scale) |
| 150 | |
| 151 | def sinks(self, module: nn.Module) -> nn.Parameter | None: |
| 152 | """The per-head sink tensor `module` passes to the interface, or `None`.""" |
| 153 | if self.s_aux_expr is None: |
| 154 | return None |
| 155 | s_aux = _resolve(self.s_aux_expr, module) |
| 156 | if not isinstance(s_aux, (nn.Parameter, type(None))): |
| 157 | expression = ast.unparse(self.s_aux_expr) |
| 158 | raise ValueError( |
| 159 | f"Cannot resolve attention s_aux expression {expression!r} in " |
| 160 | f"{type(module).__name__}." |
| 161 | ) |
| 162 | return s_aux |