| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Base classes for the Transformers backend fusers.""" |
| 4 | |
| 5 | import ast |
| 6 | import types |
| 7 | from abc import ABC, abstractmethod |
| 8 | from collections.abc import Callable, Iterable |
| 9 | from dataclasses import dataclass, field |
| 10 | from typing import TYPE_CHECKING, ClassVar |
| 11 | |
| 12 | from torch import fx, nn |
| 13 | |
| 14 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 15 | aliasing_names, |
| 16 | bypass_existence_guard, |
| 17 | replace_expr, |
| 18 | self_call_and_refs, |
| 19 | written_names, |
| 20 | ) |
| 21 | from vllm.model_executor.models.utils import ShardId, maybe_prefix |
| 22 | |
| 23 | if TYPE_CHECKING: |
| 24 | from vllm.config import VllmConfig |
| 25 | |
| 26 | |
| 27 | @dataclass |
| 28 | class BaseFuser(ABC): |
| 29 | """A detected fusion and how to apply it. |
| 30 | |
| 31 | `match` analyses the module *class* once (cached, see `get_fusers`); `fuse` |
| 32 | then applies the fusion to an instance in `recursive_replace`, returning the |
| 33 | module to install in its place. |
| 34 | """ |
| 35 | |
| 36 | redefines_forward: ClassVar[bool] = True |
| 37 | """Whether `fuse` gives the module a different forward, |
| 38 | by rewriting its source or by returning a different module in its place.""" |
| 39 | |
| 40 | @abstractmethod |
| 41 | def info(self, name: str) -> str: |
| 42 | """A human-readable description of the fusion at `name`, for logging.""" |
| 43 | |
| 44 | @classmethod |
| 45 | @abstractmethod |
| 46 | def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None": |
| 47 | """Match the pattern in `graph`, returning a fuser if found.""" |
| 48 | |
| 49 | @abstractmethod |
| 50 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 51 | """Whether this fuser can be applied to this `module` instance.""" |
| 52 | |
| 53 | @abstractmethod |
| 54 | def fuse( |
| 55 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 56 | ) -> nn.Module: |
| 57 | """Apply the fusion to an already-validated `module`, returning the |
| 58 | module to install in its place (mutated in place, or freshly built).""" |
| 59 | |
| 60 | def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: |
| 61 | """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes |
| 62 | (none unless it stacks weights).""" |
| 63 | return {} |
| 64 | |
| 65 | @property |
| 66 | def packed_modules_mapping(self) -> dict[str, list[str]]: |
| 67 | """`packed_modules_mapping` entries this fuser contributes (none unless |
| 68 | it stacks weights).""" |
| 69 | return {} |
| 70 | |
| 71 | |
| 72 | def fused_head_size(module: nn.Module, vllm_config: "VllmConfig") -> int: |
| 73 | """The head size of `module`, which head counts are derived from. |
| 74 | |
| 75 | Prefer the module's own `head_dim`, which Transformers sets per instance: |
| 76 | the model-wide value is the largest head size across layers, so on a |
| 77 | heterogeneous checkpoint it would divide out the wrong number of heads, and |
| 78 | it is the text head size, so it does not describe a vision tower at all. |
| 79 | """ |
| 80 | return getattr(module, "head_dim", None) or vllm_config.model_config.get_head_size() |
| 81 | |
| 82 | |
| 83 | def local_output_sizes(merged_name: str) -> str: |
| 84 | """Source for the per-rank widths of the merged linear `self.<merged_name>`.""" |
| 85 | merged = f"self.{merged_name}" |
| 86 | return f"[s // {merged}.tp_size for s in {merged}.output_sizes]" |
| 87 | |
| 88 | |
| 89 | @dataclass |
| 90 | class RewriteFuser(BaseFuser): |
| 91 | """A fuser that rewrites the module's forward and rebinds it. |
| 92 | |
| 93 | `match` and `update_forward` analyse the class once; `fuse` swaps the |
| 94 | submodules and binds the compiled forward on an instance in place, so it |
| 95 | keeps its class and any attribute the fusion does not consume. |
| 96 | """ |
| 97 | |
| 98 | source_cls: str |
| 99 | """Class of the HF module the fused projections belonged to (for logging).""" |
| 100 | |
| 101 | fused_forward: Callable = field(init=False, repr=False) |
| 102 | """The compiled rewritten forward, set by `update_forward`.""" |
| 103 | |
| 104 | @abstractmethod |
| 105 | def update_forward(self, module: nn.Module) -> None: |
| 106 | """Rewrite and compile `type(module)`'s forward source. |
| 107 | |
| 108 | Raises if the source does not admit the rewrite (fusion is then skipped). |
| 109 | """ |
| 110 | |
| 111 | @abstractmethod |
| 112 | def update_attrs( |
| 113 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 114 | ) -> None: |
| 115 | """Replace `module`'s submodules with their vLLM equivalents.""" |
| 116 | |
| 117 | def fuse( |
| 118 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 119 | ) -> nn.Module: |
| 120 | """Fuse an already-validated `module` in place (see `Fusers.__getitem__`). |
| 121 | |
| 122 | Builds the merged submodule and binds the compiled forward.""" |
| 123 | self.update_attrs(module, prefix, vllm_config) |
| 124 | module.forward = types.MethodType(self.fused_forward, module) |
| 125 | return module |
| 126 | |
| 127 | |
| 128 | @dataclass |
| 129 | class StackedFuser(RewriteFuser): |
| 130 | """A fuser that merges sibling projections into one stacked linear and |
| 131 | rewrites the forward to call it.""" |
| 132 | |
| 133 | merged_cls_name: ClassVar[str] |
| 134 | """Name of the vLLM class the merged projection becomes (for logging).""" |
| 135 | |
| 136 | @property |
| 137 | @abstractmethod |
| 138 | def merged_name(self) -> str: |
| 139 | """Attribute name of the merged module created by `update_attrs`.""" |
| 140 | |
| 141 | def info(self, name: str) -> str: |
| 142 | sources = " + ".join(shard for shard, _ in self.shards) |
| 143 | return ( |
| 144 | f"Fused: {sources} ({name}: {self.source_cls}) -> " |
| 145 | f"{self.merged_name} ({self.merged_cls_name})" |
| 146 | ) |
| 147 | |
| 148 | @property |
| 149 | @abstractmethod |
| 150 | def shards(self) -> list[tuple[str, ShardId]]: |
| 151 | """Each projection's original name and its shard id in the merged module. |
| 152 | |
| 153 | Source for both `orig_to_new_stacked` and `packed_modules_mapping`.""" |
| 154 | |
| 155 | def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]: |
| 156 | """`WeightsMapper.orig_to_new_stacked` entries for one fused instance. |
| 157 | |
| 158 | Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname |
| 159 | so only this exact layer is remapped, never a same-named projection |
| 160 | elsewhere (e.g. an unfused MoE expert's `gate_proj`).""" |
| 161 | merged = maybe_prefix(prefix, self.merged_name) |
| 162 | return { |
| 163 | maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards |
| 164 | } |
| 165 | |
| 166 | @property |
| 167 | def packed_modules_mapping(self) -> dict[str, list[str]]: |
| 168 | """`{merged_name: [projection names]}` so quantization can unpack the |
| 169 | fused layer into its per-shard configs.""" |
| 170 | return {self.merged_name: [name for name, _ in self.shards]} |
| 171 | |
| 172 | def _unguarded_calls( |
| 173 | self, funcdef: ast.FunctionDef, names: Iterable[str] |
| 174 | ) -> list[ast.Call]: |
| 175 | """One `self.<name>(arg)` call per projection, existence guards folded. |
| 176 | |
| 177 | `update_attrs` deletes the projections it merges, so any reference to one beyond |
| 178 | its call site must be a guard on its existence; folding those to their constant |
| 179 | value keeps the rewritten forward off a name that no longer exists. A reference |
| 180 | that is not such a guard raises, so fusion is skipped.""" |
| 181 | calls = [] |
| 182 | for name in names: |
| 183 | call, refs = self_call_and_refs(funcdef, name) |
| 184 | for ref in refs: |
| 185 | bypass_existence_guard(funcdef, ref, name) |
| 186 | calls.append(call) |
| 187 | return calls |
| 188 | |
| 189 | def _check_input_stable( |
| 190 | self, |
| 191 | funcdef: ast.FunctionDef, |
| 192 | module: nn.Module, |
| 193 | calls: list[ast.Call], |
| 194 | block: list[ast.stmt], |
| 195 | indices: list[int], |
| 196 | ) -> None: |
| 197 | """Raise unless hoisting the merged GEMM preserves what it reads. |
| 198 | |
| 199 | Fusing moves every projection to one call at `min(indices)`, so the |
| 200 | merged GEMM reads the input once, up front, where the last of `calls` |
| 201 | would have read it later. That holds only if nothing in between changes |
| 202 | the input, and a change need not name it: writing any view that shares |
| 203 | its storage changes it too. Both halves of the check over-approximate, |
| 204 | since a false hit costs a fusion while a miss returns wrong numbers. |
| 205 | """ |
| 206 | arg_names = { |
| 207 | node.id |
| 208 | for node in ast.walk(calls[0].args[0]) |
| 209 | if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load) |
| 210 | } |
| 211 | # Protect the names the input reads and anything that may alias them. |
| 212 | tracked = aliasing_names(funcdef, arg_names, module) |
| 213 | # Any of them rebound, written through, or passed to a call that writes |
| 214 | # its argument would leave the hoisted GEMM reading a different value. |
| 215 | region = block[min(indices) : max(indices) + 1] |
| 216 | if tracked & written_names(region): |
| 217 | raise ValueError("projection input is rebound or mutated before all calls") |
| 218 | |
| 219 | def _splice_merged_split( |
| 220 | self, |
| 221 | funcdef: ast.FunctionDef, |
| 222 | calls: list[ast.Call], |
| 223 | block: list[ast.stmt], |
| 224 | index: int, |
| 225 | ) -> None: |
| 226 | """Insert `temps = self.<merged_name>(arg).split(sizes, -1)` at |
| 227 | `block[index]` and replace each of `calls` with its temp name. |
| 228 | |
| 229 | `calls` must share one input argument; `block[index]` must be where |
| 230 | they are (or would be) evaluated. Raises if a generated temporary |
| 231 | would shadow an existing name in `funcdef`. |
| 232 | """ |
| 233 | temps = [f"_vllm_merged_{i}" for i in range(len(calls))] |
| 234 | names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)} |
| 235 | if names & set(temps): |
| 236 | raise ValueError("fused temporaries would shadow existing names") |
| 237 | targets = ", ".join(temps) |
| 238 | sections = local_output_sizes(self.merged_name) |
| 239 | source = f"{targets} = self.{self.merged_name}(__arg__).split({sections}, -1)" |
| 240 | assign = ast.parse(source).body[0] |
| 241 | arg = next( |
| 242 | node |
| 243 | for node in ast.walk(assign) |
| 244 | if isinstance(node, ast.Name) and node.id == "__arg__" |
| 245 | ) |
| 246 | replace_expr(assign, arg, calls[0].args[0]) |
| 247 | ast.copy_location(assign, block[index]) |
| 248 | block.insert(index, assign) |
| 249 | for call, temp in zip(calls, temps): |
| 250 | replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load())) |