| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split.""" |
| 4 | |
| 5 | import ast |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING, ClassVar |
| 8 | |
| 9 | from torch import fx, nn |
| 10 | |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.linear import QKVParallelLinear |
| 13 | from vllm.model_executor.models.transformers.fusers.base import ( |
| 14 | StackedFuser, |
| 15 | fused_head_size, |
| 16 | ) |
| 17 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 18 | block_chain, |
| 19 | compile_forward, |
| 20 | is_linear, |
| 21 | recover_forward, |
| 22 | returned_linear, |
| 23 | ) |
| 24 | from vllm.model_executor.models.transformers.utils import ( |
| 25 | log_replacement, |
| 26 | replace_linear_class, |
| 27 | ) |
| 28 | from vllm.model_executor.models.utils import ShardId, maybe_prefix |
| 29 | |
| 30 | if TYPE_CHECKING: |
| 31 | from vllm.config import VllmConfig |
| 32 | |
| 33 | logger = init_logger(__name__) |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class QKVFuser(StackedFuser): |
| 38 | """Fuser for the attention QKV pattern `q(x), k(x), v(x)`.""" |
| 39 | |
| 40 | q_name: str |
| 41 | k_name: str |
| 42 | v_name: str |
| 43 | o_name: str | None |
| 44 | merged_name: ClassVar[str] = "qkv_proj" |
| 45 | merged_cls_name: ClassVar[str] = "QKVParallelLinear" |
| 46 | |
| 47 | @property |
| 48 | def shards(self) -> list[tuple[str, ShardId]]: |
| 49 | return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")] |
| 50 | |
| 51 | @classmethod |
| 52 | def _get_qkv_nodes( |
| 53 | cls, graph: fx.Graph, module: nn.Module |
| 54 | ) -> tuple[fx.Node, fx.Node, fx.Node] | None: |
| 55 | """Search `graph` for the QKV pattern `q(x), k(x), v(x)`.""" |
| 56 | by_input: dict[fx.Node, list[fx.Node]] = {} |
| 57 | for node in graph.nodes: |
| 58 | if ( |
| 59 | is_linear(node, module) |
| 60 | and len(node.args) == 1 |
| 61 | and not node.kwargs |
| 62 | and isinstance(node.args[0], fx.Node) |
| 63 | and node.args[0].op == "placeholder" |
| 64 | ): |
| 65 | by_input.setdefault(node.args[0], []).append(node) |
| 66 | triples = [nodes for nodes in by_input.values() if len(nodes) == 3] |
| 67 | if len(triples) != 1: |
| 68 | return None |
| 69 | |
| 70 | q_node, k_node, v_node = nodes = triples[0] |
| 71 | outs = [module.get_submodule(node.target).out_features for node in nodes] |
| 72 | if len(set(outs)) == 2: |
| 73 | # q is identified as the larger projection (GQA) |
| 74 | (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1) |
| 75 | k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2) |
| 76 | if module.get_submodule(q_node.target).out_features != max(outs): |
| 77 | return None |
| 78 | elif len(set(outs)) != 1: |
| 79 | return None |
| 80 | return q_node, k_node, v_node |
| 81 | |
| 82 | @classmethod |
| 83 | def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None": |
| 84 | if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None: |
| 85 | return None |
| 86 | q, k, v = qkv_nodes |
| 87 | names = dict(q_name=q.target, k_name=k.target, v_name=v.target) |
| 88 | # o_proj produces the module's output. |
| 89 | o_name = returned_linear(graph, module) |
| 90 | # o_proj must be compatible with the q/k/v projections. |
| 91 | if o_name in names.values() or ( |
| 92 | o_name is not None |
| 93 | and module.get_submodule(o_name).in_features |
| 94 | != module.get_submodule(q.target).out_features |
| 95 | ): |
| 96 | o_name = None |
| 97 | names["o_name"] = o_name |
| 98 | return cls(source_cls=type(module).__name__, **names) |
| 99 | |
| 100 | def update_forward(self, module: nn.Module) -> None: |
| 101 | """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source. |
| 102 | |
| 103 | A projection may be guarded by an existence check -- a |
| 104 | `self.<proj> is (not) None` comparison (e.g. Gemma 4's `v_proj`) or a bare |
| 105 | truthiness test -- which is folded to its constant value (see |
| 106 | `bypass_existence_guard`). The calls may sit in different branches, so the |
| 107 | fused GEMM is inserted before the earliest of them, in the innermost block |
| 108 | that dominates all three. |
| 109 | """ |
| 110 | funcdef, fn = recover_forward(type(module)) |
| 111 | proj_names = (self.q_name, self.k_name, self.v_name) |
| 112 | calls = self._unguarded_calls(funcdef, proj_names) |
| 113 | arg_dumps = {ast.dump(call.args[0]) for call in calls} |
| 114 | if len(arg_dumps) != 1: |
| 115 | raise ValueError("projection inputs are written differently") |
| 116 | # The trace may be partial, so prove projection exclusivity in source: |
| 117 | # no other linear child may consume the same input (else the matched |
| 118 | # three may not be q, k and v) |
| 119 | other_linears = { |
| 120 | name |
| 121 | for name, child in module.named_children() |
| 122 | if isinstance(child, nn.Linear) |
| 123 | } - set(proj_names) |
| 124 | for node in ast.walk(funcdef): |
| 125 | if ( |
| 126 | isinstance(node, ast.Call) |
| 127 | and isinstance(node.func, ast.Attribute) |
| 128 | and node.func.attr in other_linears |
| 129 | and any(ast.dump(arg) in arg_dumps for arg in node.args) |
| 130 | ): |
| 131 | raise ValueError("another linear consumes the same input") |
| 132 | |
| 133 | # Insert the fused GEMM before the earliest call, in the innermost block |
| 134 | # common to all three (the calls may be split across branches). |
| 135 | chains = [block_chain(funcdef.body, call) for call in calls] |
| 136 | if any(not chain for chain in chains): |
| 137 | raise ValueError("projection calls not found in the function body") |
| 138 | depth = 0 |
| 139 | for level in zip(*chains): |
| 140 | if len({id(block) for block, _ in level}) != 1: |
| 141 | break |
| 142 | depth += 1 |
| 143 | block = chains[0][depth - 1][0] |
| 144 | indices = [chain[depth - 1][1] for chain in chains] |
| 145 | insert_index = min(indices) |
| 146 | self._check_input_stable(funcdef, module, calls, block, indices) |
| 147 | |
| 148 | # q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1) |
| 149 | self._splice_merged_split(funcdef, calls, block, insert_index) |
| 150 | self.fused_forward = compile_forward(funcdef, fn) |
| 151 | |
| 152 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 153 | """Shapes must be compatible for a single merged, head-sharded GEMM.""" |
| 154 | q = module.get_submodule(self.q_name) |
| 155 | k = module.get_submodule(self.k_name) |
| 156 | v = module.get_submodule(self.v_name) |
| 157 | head_size = fused_head_size(module, vllm_config) |
| 158 | compatible = ( |
| 159 | q.in_features == k.in_features == v.in_features |
| 160 | and len({proj.bias is None for proj in (q, k, v)}) == 1 |
| 161 | and k.out_features == v.out_features |
| 162 | and q.out_features % head_size == 0 |
| 163 | and k.out_features % head_size == 0 |
| 164 | ) |
| 165 | if not compatible: |
| 166 | logger.debug("%s is not compatible with QKV fusion", type(module)) |
| 167 | return compatible |
| 168 | |
| 169 | def update_attrs( |
| 170 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 171 | ) -> None: |
| 172 | quant_config = vllm_config.quant_config |
| 173 | head_size = fused_head_size(module, vllm_config) |
| 174 | q = module.get_submodule(self.q_name) |
| 175 | k = module.get_submodule(self.k_name) |
| 176 | merged = QKVParallelLinear( |
| 177 | hidden_size=q.in_features, |
| 178 | head_size=head_size, |
| 179 | total_num_heads=q.out_features // head_size, |
| 180 | total_num_kv_heads=k.out_features // head_size, |
| 181 | bias=q.bias is not None, |
| 182 | quant_config=quant_config, |
| 183 | prefix=maybe_prefix(prefix, self.merged_name), |
| 184 | return_bias=False, |
| 185 | ) |
| 186 | logger.debug( |
| 187 | "%s: %s, %s: %s, %s: %s -> %s: %s", |
| 188 | self.q_name, |
| 189 | q, |
| 190 | self.k_name, |
| 191 | k, |
| 192 | self.v_name, |
| 193 | module.get_submodule(self.v_name), |
| 194 | self.merged_name, |
| 195 | merged, |
| 196 | ) |
| 197 | setattr(module, self.merged_name, merged) |
| 198 | # Drop the consumed submodules so their (meta) params are not expected. |
| 199 | for name in (self.q_name, self.k_name, self.v_name): |
| 200 | delattr(module, name) |
| 201 | # If there is an output projection, we know it must be rowwise. |
| 202 | if self.o_name is not None: |
| 203 | o_proj_prefix = maybe_prefix(prefix, self.o_name) |
| 204 | o_proj = module.get_submodule(self.o_name) |
| 205 | new_o = replace_linear_class( |
| 206 | o_proj, "rowwise", quant_config, prefix=o_proj_prefix |
| 207 | ) |
| 208 | setattr(module, self.o_name, new_o) |
| 209 | log_replacement(o_proj_prefix, o_proj, new_o) |