| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """MLA fuser: adapt a Transformers MLA attention module for vLLM's MLA layer.""" |
| 4 | |
| 5 | import ast |
| 6 | import inspect |
| 7 | import textwrap |
| 8 | from dataclasses import dataclass |
| 9 | from typing import TYPE_CHECKING, ClassVar |
| 10 | |
| 11 | from torch import fx, nn |
| 12 | |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 15 | from vllm.model_executor.models.transformers.fusers.base import StackedFuser |
| 16 | from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser |
| 17 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 18 | compile_forward, |
| 19 | downstream_linear, |
| 20 | is_linear, |
| 21 | recover_forward, |
| 22 | replace_expr, |
| 23 | returned_linear, |
| 24 | upstream_linear, |
| 25 | ) |
| 26 | from vllm.model_executor.models.transformers.utils import ( |
| 27 | Style, |
| 28 | log_replacement, |
| 29 | replace_linear_class, |
| 30 | ) |
| 31 | from vllm.model_executor.models.utils import ShardId, maybe_prefix |
| 32 | |
| 33 | if TYPE_CHECKING: |
| 34 | from vllm.config import VllmConfig |
| 35 | |
| 36 | logger = init_logger(__name__) |
| 37 | |
| 38 | |
| 39 | def _consumes_placeholder(node: fx.Node) -> bool: |
| 40 | """Whether `node` is a linear applied directly to a `forward` input.""" |
| 41 | return ( |
| 42 | len(node.args) == 1 |
| 43 | and isinstance(node.args[0], fx.Node) |
| 44 | and node.args[0].op == "placeholder" |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def _norm_size(norm: nn.Module) -> int: |
| 49 | weight = getattr(norm, "weight", None) |
| 50 | return weight.numel() if weight is not None else -1 |
| 51 | |
| 52 | |
| 53 | def _is_rms_norm(module: nn.Module) -> bool: |
| 54 | """Whether `module` computes an RMSNorm, verified by `RMSNormFuser`'s matcher.""" |
| 55 | # Go via `get_fuser` because it caches |
| 56 | from vllm.model_executor.models.transformers.fuser import get_fuser |
| 57 | |
| 58 | return get_fuser(module, RMSNormFuser) is not None |
| 59 | |
| 60 | |
| 61 | def _top_level_index(funcdef: ast.FunctionDef, node: ast.AST) -> int: |
| 62 | """Index in `funcdef.body` of the top-level statement containing `node`.""" |
| 63 | for index, stmt in enumerate(funcdef.body): |
| 64 | if any(child is node for child in ast.walk(stmt)): |
| 65 | return index |
| 66 | raise ValueError("node is not in the function body") |
| 67 | |
| 68 | |
| 69 | def _single_expand_call( |
| 70 | funcdef: ast.FunctionDef, module: nn.Module, kv_b_proj_name: str |
| 71 | ) -> ast.Call: |
| 72 | """The KV expansion call `self.<method>(kv_c_normed, k_pe)`: the unique |
| 73 | two-argument call to a method of `module` with the expansion's signature. |
| 74 | |
| 75 | The expansion is `kv_c_normed, k_pe -> key, value`: two inputs, every |
| 76 | `return` a pair, and it applies `kv_b_proj` -- the projection `MLAAttention` |
| 77 | owns after the bypass, absorbed into the attention computation. |
| 78 | """ |
| 79 | |
| 80 | def is_expansion_method(name: str) -> bool: |
| 81 | method = getattr(type(module), name, None) |
| 82 | if method is None: |
| 83 | return False |
| 84 | code = getattr(inspect.unwrap(method), "__code__", None) |
| 85 | if ( |
| 86 | code is None |
| 87 | or code.co_argcount != 3 # self, kv_c_normed, k_pe |
| 88 | or kv_b_proj_name not in code.co_names |
| 89 | ): |
| 90 | return False |
| 91 | try: |
| 92 | tree = ast.parse(textwrap.dedent(inspect.getsource(method))) |
| 93 | except (OSError, TypeError): |
| 94 | return False |
| 95 | returns = [node for node in ast.walk(tree) if isinstance(node, ast.Return)] |
| 96 | return bool(returns) and all( |
| 97 | isinstance(ret.value, ast.Tuple) and len(ret.value.elts) == 2 |
| 98 | for ret in returns |
| 99 | ) |
| 100 | |
| 101 | calls = [ |
| 102 | node |
| 103 | for node in ast.walk(funcdef) |
| 104 | if isinstance(node, ast.Call) |
| 105 | and isinstance(node.func, ast.Attribute) |
| 106 | and isinstance(node.func.value, ast.Name) |
| 107 | and node.func.value.id == "self" |
| 108 | and len(node.args) == 2 |
| 109 | and not node.keywords |
| 110 | and is_expansion_method(node.func.attr) |
| 111 | ] |
| 112 | if len(calls) != 1: |
| 113 | raise ValueError("expected exactly one KV expansion method call") |
| 114 | return calls[0] |
| 115 | |
| 116 | |
| 117 | @dataclass |
| 118 | class MLAFuser(StackedFuser): |
| 119 | """Fuser for the MLA attention pattern.""" |
| 120 | |
| 121 | q_proj_name: str | None |
| 122 | q_a_proj_name: str | None |
| 123 | q_a_layernorm_name: str | None |
| 124 | q_b_proj_name: str | None |
| 125 | kv_a_proj_name: str |
| 126 | kv_a_layernorm_name: str |
| 127 | kv_b_proj_name: str |
| 128 | o_proj_name: str | None |
| 129 | merged_name: ClassVar[str] = "fused_qkv_a_proj" |
| 130 | merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" |
| 131 | |
| 132 | def info(self, name: str) -> str: |
| 133 | info_str = ( |
| 134 | f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)" |
| 135 | ) |
| 136 | if self.q_a_proj_name is not None: |
| 137 | info_str += "; " + super().info(name).removeprefix("Fused: ") |
| 138 | return info_str |
| 139 | |
| 140 | @property |
| 141 | def shards(self) -> list[tuple[str, ShardId]]: |
| 142 | """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection.""" |
| 143 | if self.q_a_proj_name is None: |
| 144 | return [] |
| 145 | return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)] |
| 146 | |
| 147 | @property |
| 148 | def packed_modules_mapping(self) -> dict[str, list[str]]: |
| 149 | if self.q_a_proj_name is None: |
| 150 | return {} |
| 151 | return super().packed_modules_mapping |
| 152 | |
| 153 | @classmethod |
| 154 | def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None": |
| 155 | # Find all `rms_norm(linear(placeholder))` chains. |
| 156 | chains = [] |
| 157 | for node in graph.nodes: |
| 158 | if node.op != "call_module" or is_linear(node, module) or not node.args: |
| 159 | continue |
| 160 | source = upstream_linear(node.args[0], module) |
| 161 | if source is None or not _consumes_placeholder(source): |
| 162 | continue |
| 163 | if not _is_rms_norm(module.get_submodule(node.target)): |
| 164 | continue |
| 165 | chains.append((source, node)) |
| 166 | |
| 167 | # Tell the chains apart by width. |
| 168 | def is_kv_a(chain) -> bool: |
| 169 | source, rms_norm = chain |
| 170 | source_mod = module.get_submodule(source.target) |
| 171 | rms_norm_mod = module.get_submodule(rms_norm.target) |
| 172 | return source_mod.out_features != _norm_size(rms_norm_mod) |
| 173 | |
| 174 | kv_a_chains = [chain for chain in chains if is_kv_a(chain)] |
| 175 | q_a_chains = [chain for chain in chains if not is_kv_a(chain)] |
| 176 | # Exactly one KV chain is MLA's signature. The query chain is optional. |
| 177 | if len(kv_a_chains) != 1 or len(q_a_chains) > 1: |
| 178 | return None |
| 179 | kv_a_proj, kv_a_layernorm = kv_a_chains[0] |
| 180 | |
| 181 | # Linear children claimed for a role so far; the rest resolve by elimination. |
| 182 | claimed_linears = {kv_a_proj.target} |
| 183 | q_proj_name = q_a_proj = q_a_layernorm = q_b_proj = None |
| 184 | if q_a_chains: |
| 185 | # Find `q_b_proj(q_a_layernorm(q_a_proj(placeholder)))`. |
| 186 | q_a_proj, q_a_layernorm = q_a_chains[0] |
| 187 | q_b_proj = downstream_linear(q_a_layernorm, module) |
| 188 | if q_b_proj is None: |
| 189 | return None |
| 190 | claimed_linears |= {q_a_proj.target, q_b_proj.target} |
| 191 | else: |
| 192 | # Find `q_proj(placeholder)`. |
| 193 | placeholder_linears = { |
| 194 | node.target |
| 195 | for node in graph.nodes |
| 196 | if is_linear(node, module) and _consumes_placeholder(node) |
| 197 | } |
| 198 | if len(q_proj_candidates := placeholder_linears - claimed_linears) != 1: |
| 199 | return None |
| 200 | q_proj_name = next(iter(q_proj_candidates)) |
| 201 | claimed_linears.add(q_proj_name) |
| 202 | |
| 203 | # Find `kv_b_proj(kv_a_layernorm(...))`. |
| 204 | kv_b_proj = downstream_linear(kv_a_layernorm, module) |
| 205 | if kv_b_proj is None or kv_b_proj.target in claimed_linears: |
| 206 | return None |
| 207 | claimed_linears.add(kv_b_proj.target) |
| 208 | |
| 209 | # Find `o_proj` if it is returned by the forward graph. |
| 210 | o_proj_name = returned_linear(graph, module) |
| 211 | if o_proj_name in claimed_linears: |
| 212 | o_proj_name = None |
| 213 | |
| 214 | return cls( |
| 215 | source_cls=type(module).__name__, |
| 216 | q_proj_name=q_proj_name, |
| 217 | q_a_proj_name=q_a_proj.target if q_a_proj else None, |
| 218 | q_a_layernorm_name=q_a_layernorm.target if q_a_layernorm else None, |
| 219 | q_b_proj_name=q_b_proj.target if q_b_proj else None, |
| 220 | kv_a_proj_name=kv_a_proj.target, |
| 221 | kv_a_layernorm_name=kv_a_layernorm.target, |
| 222 | kv_b_proj_name=kv_b_proj.target, |
| 223 | o_proj_name=o_proj_name, |
| 224 | ) |
| 225 | |
| 226 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 227 | return vllm_config.model_config.use_mla |
| 228 | |
| 229 | def update_forward(self, module: nn.Module) -> None: |
| 230 | """Merge `q_a_proj` and `kv_a_proj` into one fused down proj then split. |
| 231 | Bypass the KV expansion method so the compressed latent reaches the `vllm_mla` |
| 232 | attention interface unexpanded.""" |
| 233 | funcdef, fn = recover_forward(type(module)) |
| 234 | if (q_a_proj_name := self.q_a_proj_name) is not None: |
| 235 | # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. |
| 236 | # The fused call is inserted at the top-level statement preceding both. |
| 237 | names = [q_a_proj_name, self.kv_a_proj_name] |
| 238 | calls = self._unguarded_calls(funcdef, names) |
| 239 | if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]): |
| 240 | raise ValueError("down-projections read different inputs") |
| 241 | indices = [_top_level_index(funcdef, call) for call in calls] |
| 242 | self._check_input_stable(funcdef, module, calls, funcdef.body, indices) |
| 243 | self._splice_merged_split(funcdef, calls, funcdef.body, min(indices)) |
| 244 | |
| 245 | # Transformers expands the latent into full key/value in a dedicated method. |
| 246 | # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`), |
| 247 | # so replace the expansion call with its own arguments so `kv_c_normed, k_pe` |
| 248 | # flow to the interface in place of `key, value`. |
| 249 | expand_call = _single_expand_call(funcdef, module, self.kv_b_proj_name) |
| 250 | replace_expr( |
| 251 | funcdef, expand_call, ast.Tuple(elts=list(expand_call.args), ctx=ast.Load()) |
| 252 | ) |
| 253 | self.fused_forward = compile_forward(funcdef, fn) |
| 254 | |
| 255 | def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"): |
| 256 | quant_config = vllm_config.quant_config |
| 257 | |
| 258 | def replace_linear_by_name(name: str, style: Style): |
| 259 | linear = module.get_submodule(name) |
| 260 | replacement = replace_linear_class( |
| 261 | linear, style, quant_config, prefix=maybe_prefix(prefix, name) |
| 262 | ) |
| 263 | setattr(module, name, replacement) |
| 264 | log_replacement(maybe_prefix(prefix, name), linear, replacement) |
| 265 | |
| 266 | if (q_a_proj_name := self.q_a_proj_name) is not None: |
| 267 | # `match` sets the q-LoRA names together, or none of them |
| 268 | assert self.q_b_proj_name is not None |
| 269 | q_a = module.get_submodule(q_a_proj_name) |
| 270 | kv_a = module.get_submodule(self.kv_a_proj_name) |
| 271 | merged = MergedColumnParallelLinear( |
| 272 | input_size=q_a.in_features, |
| 273 | output_sizes=[q_a.out_features, kv_a.out_features], |
| 274 | bias=q_a.bias is not None, |
| 275 | quant_config=quant_config, |
| 276 | prefix=maybe_prefix(prefix, self.merged_name), |
| 277 | return_bias=False, |
| 278 | disable_tp=True, |
| 279 | ) |
| 280 | logger.debug( |
| 281 | "%s: %s, %s: %s -> %s: %s", |
| 282 | q_a_proj_name, |
| 283 | q_a, |
| 284 | self.kv_a_proj_name, |
| 285 | kv_a, |
| 286 | self.merged_name, |
| 287 | merged, |
| 288 | ) |
| 289 | setattr(module, self.merged_name, merged) |
| 290 | # The rewritten forward calls the merged projection instead. |
| 291 | delattr(module, q_a_proj_name) |
| 292 | delattr(module, self.kv_a_proj_name) |
| 293 | replace_linear_by_name(self.q_b_proj_name, "colwise") |
| 294 | else: |
| 295 | assert self.q_proj_name is not None |
| 296 | replace_linear_by_name(self.kv_a_proj_name, "replicate") |
| 297 | replace_linear_by_name(self.q_proj_name, "colwise") |
| 298 | |
| 299 | replace_linear_by_name(self.kv_b_proj_name, "colwise") |
| 300 | # MLAAttention calls kv_b_proj and expects vLLM's default return_bias=True |
| 301 | module.get_submodule(self.kv_b_proj_name).return_bias = True |
| 302 | if self.o_proj_name is not None: |
| 303 | replace_linear_by_name(self.o_proj_name, "rowwise") |