| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Packed-QKV fuser: `c_attn(x).split((q, kv, kv))` -> a `QKVParallelLinear`.""" |
| 4 | |
| 5 | import ast |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING |
| 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 | RewriteFuser, |
| 15 | fused_head_size, |
| 16 | local_output_sizes, |
| 17 | ) |
| 18 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 19 | compile_forward, |
| 20 | is_method, |
| 21 | recover_forward, |
| 22 | returned_linear, |
| 23 | upstream_linear, |
| 24 | ) |
| 25 | from vllm.model_executor.models.transformers.utils import ( |
| 26 | log_replacement, |
| 27 | replace_linear_class, |
| 28 | ) |
| 29 | from vllm.model_executor.models.utils import maybe_prefix |
| 30 | |
| 31 | if TYPE_CHECKING: |
| 32 | from vllm.config import VllmConfig |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | @dataclass |
| 38 | class PackedQKVFuser(RewriteFuser): |
| 39 | """Fuser for attention with q, k and v packed into one projection.""" |
| 40 | |
| 41 | qkv_name: str |
| 42 | o_name: str | None |
| 43 | q_size: int |
| 44 | kv_size: int |
| 45 | |
| 46 | def info(self, name: str) -> str: |
| 47 | return ( |
| 48 | f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear" |
| 49 | ) |
| 50 | |
| 51 | @staticmethod |
| 52 | def _packed_sizes(node: fx.Node) -> tuple[int, int] | None: |
| 53 | """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one.""" |
| 54 | if not is_method(node, "split") or len(node.args) < 2: |
| 55 | return None |
| 56 | sizes = node.args[1] |
| 57 | if not isinstance(sizes, (tuple, list)) or len(sizes) != 3: |
| 58 | return None |
| 59 | if not all(isinstance(size, int) for size in sizes): |
| 60 | return None |
| 61 | q_size, k_size, v_size = sizes |
| 62 | if k_size != v_size or q_size < k_size: |
| 63 | return None |
| 64 | return q_size, k_size |
| 65 | |
| 66 | @classmethod |
| 67 | def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None": |
| 68 | for node in graph.nodes: |
| 69 | if (sizes := cls._packed_sizes(node)) is None: |
| 70 | continue |
| 71 | q_size, kv_size = sizes |
| 72 | qkv_node = upstream_linear(node.args[0], module) |
| 73 | if qkv_node is None: |
| 74 | continue |
| 75 | qkv_name = str(qkv_node.target) |
| 76 | # The split must consume the whole projection. |
| 77 | if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size: |
| 78 | continue |
| 79 | # o_proj produces the module's output and consumes the query width. |
| 80 | o_name = returned_linear(graph, module) |
| 81 | if o_name == qkv_name or ( |
| 82 | o_name is not None |
| 83 | and module.get_submodule(o_name).in_features != q_size |
| 84 | ): |
| 85 | o_name = None |
| 86 | return cls( |
| 87 | source_cls=type(module).__name__, |
| 88 | qkv_name=qkv_name, |
| 89 | o_name=o_name, |
| 90 | q_size=q_size, |
| 91 | kv_size=kv_size, |
| 92 | ) |
| 93 | return None |
| 94 | |
| 95 | def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call: |
| 96 | """The unique `self.<qkv_name>(...)....split((a, b, c), ...)` call.""" |
| 97 | calls = [ |
| 98 | node |
| 99 | for node in ast.walk(funcdef) |
| 100 | if isinstance(node, ast.Call) |
| 101 | and isinstance(node.func, ast.Attribute) |
| 102 | and node.func.attr == "split" |
| 103 | and node.args |
| 104 | and isinstance(node.args[0], (ast.Tuple, ast.List)) |
| 105 | and len(node.args[0].elts) == 3 |
| 106 | and any( |
| 107 | isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name |
| 108 | for inner in ast.walk(node.func.value) |
| 109 | ) |
| 110 | ] |
| 111 | if len(calls) != 1: |
| 112 | raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits") |
| 113 | return calls[0] |
| 114 | |
| 115 | def update_forward(self, module: nn.Module) -> None: |
| 116 | """Rewrite the split sizes to the sharded projection's per-rank widths.""" |
| 117 | funcdef, fn = recover_forward(type(module)) |
| 118 | split = self._split_call(funcdef) |
| 119 | # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes] |
| 120 | sections = local_output_sizes(self.qkv_name) |
| 121 | split.args[0] = ast.parse(sections, mode="eval").body |
| 122 | self.fused_forward = compile_forward(funcdef, fn) |
| 123 | |
| 124 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 125 | """Shapes must be compatible with a head-sharded packed GEMM.""" |
| 126 | head_size = fused_head_size(module, vllm_config) |
| 127 | qkv = module.get_submodule(self.qkv_name) |
| 128 | compatible = ( |
| 129 | self.q_size % head_size == 0 |
| 130 | and self.kv_size % head_size == 0 |
| 131 | and qkv.out_features == self.q_size + 2 * self.kv_size |
| 132 | ) |
| 133 | if not compatible: |
| 134 | logger.debug("%s is not compatible with packed QKV fusion", type(module)) |
| 135 | return compatible |
| 136 | |
| 137 | def update_attrs( |
| 138 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 139 | ) -> None: |
| 140 | quant_config = vllm_config.quant_config |
| 141 | head_size = fused_head_size(module, vllm_config) |
| 142 | qkv_prefix = maybe_prefix(prefix, self.qkv_name) |
| 143 | qkv = module.get_submodule(self.qkv_name) |
| 144 | merged = QKVParallelLinear( |
| 145 | hidden_size=qkv.in_features, |
| 146 | head_size=head_size, |
| 147 | total_num_heads=self.q_size // head_size, |
| 148 | total_num_kv_heads=self.kv_size // head_size, |
| 149 | bias=qkv.bias is not None, |
| 150 | quant_config=quant_config, |
| 151 | prefix=qkv_prefix, |
| 152 | return_bias=False, |
| 153 | ) |
| 154 | setattr(module, self.qkv_name, merged) |
| 155 | log_replacement(qkv_prefix, qkv, merged) |
| 156 | # If there is an output projection, we know it must be rowwise. |
| 157 | if self.o_name is not None: |
| 158 | o_prefix = maybe_prefix(prefix, self.o_name) |
| 159 | o_proj = module.get_submodule(self.o_name) |
| 160 | new_o = replace_linear_class( |
| 161 | o_proj, "rowwise", quant_config, prefix=o_prefix |
| 162 | ) |
| 163 | setattr(module, self.o_name, new_o) |
| 164 | log_replacement(o_prefix, o_proj, new_o) |