wehub-resource-sync
213 lines
8.5 KiB
Python
213 lines
8.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split."""
|
|
|
|
import ast
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, ClassVar
|
|
|
|
from torch import fx, nn
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.linear import QKVParallelLinear
|
|
from vllm.model_executor.models.transformers.fusers.base import StackedFuser
|
|
from vllm.model_executor.models.transformers.fx_utils import (
|
|
compile_forward,
|
|
innermost_block,
|
|
is_linear,
|
|
recover_forward,
|
|
replace_expr,
|
|
single_self_call,
|
|
)
|
|
from vllm.model_executor.models.transformers.utils import (
|
|
log_replacement,
|
|
replace_linear_class,
|
|
)
|
|
from vllm.model_executor.models.utils import ShardId, maybe_prefix
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config.model import ModelConfig
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class QKVFuser(StackedFuser):
|
|
"""Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
|
|
|
|
q_name: str
|
|
k_name: str
|
|
v_name: str
|
|
o_name: str | None
|
|
merged_name: ClassVar[str] = "qkv_proj"
|
|
merged_cls: ClassVar[str] = "QKVParallelLinear"
|
|
|
|
@property
|
|
def shards(self) -> list[tuple[str, ShardId]]:
|
|
return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]
|
|
|
|
@classmethod
|
|
def _get_qkv_nodes(
|
|
cls, graph: fx.Graph, module: nn.Module
|
|
) -> tuple[fx.Node, fx.Node, fx.Node] | None:
|
|
"""Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
|
|
by_input: dict[fx.Node, list[fx.Node]] = {}
|
|
for node in graph.nodes:
|
|
if (
|
|
is_linear(node, module)
|
|
and len(node.args) == 1
|
|
and not node.kwargs
|
|
and isinstance(node.args[0], fx.Node)
|
|
and node.args[0].op == "placeholder"
|
|
):
|
|
by_input.setdefault(node.args[0], []).append(node)
|
|
triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
|
|
if len(triples) != 1:
|
|
return None
|
|
|
|
q_node, k_node, v_node = nodes = triples[0]
|
|
outs = [module.get_submodule(node.target).out_features for node in nodes]
|
|
if len(set(outs)) == 2:
|
|
# q is identified as the larger projection (GQA)
|
|
(q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
|
|
k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
|
|
if module.get_submodule(q_node.target).out_features != max(outs):
|
|
return None
|
|
elif len(set(outs)) != 1:
|
|
return None
|
|
return q_node, k_node, v_node
|
|
|
|
@classmethod
|
|
def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
|
|
if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
|
|
return None
|
|
q, k, v = qkv_nodes
|
|
names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
|
|
attn_width = module.get_submodule(q.target).out_features
|
|
candidates = [
|
|
name
|
|
for name, child in module.named_children()
|
|
if isinstance(child, nn.Linear)
|
|
and name not in names.values()
|
|
and child.in_features == attn_width
|
|
]
|
|
names["o_name"] = candidates[0] if len(candidates) == 1 else None
|
|
return cls(source_cls=type(module).__name__, **names)
|
|
|
|
def update_forward(self, module: nn.Module) -> None:
|
|
"""Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
|
|
funcdef, fn = recover_forward(type(module))
|
|
calls = [
|
|
single_self_call(funcdef, name)
|
|
for name in (self.q_name, self.k_name, self.v_name)
|
|
]
|
|
arg_dumps = {ast.dump(call.args[0]) for call in calls}
|
|
if len(arg_dumps) != 1:
|
|
raise ValueError("projection inputs are written differently")
|
|
# The trace may be partial, so prove projection exclusivity in source:
|
|
# no other linear child may consume the same input (else the matched
|
|
# three may not be q, k and v)
|
|
other_linears = {
|
|
name
|
|
for name, child in module.named_children()
|
|
if isinstance(child, nn.Linear)
|
|
} - {self.q_name, self.k_name, self.v_name}
|
|
for node in ast.walk(funcdef):
|
|
if (
|
|
isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Attribute)
|
|
and node.func.attr in other_linears
|
|
and any(ast.dump(arg) in arg_dumps for arg in node.args)
|
|
):
|
|
raise ValueError("another linear consumes the same input")
|
|
blocks = [innermost_block(funcdef.body, call) for call in calls]
|
|
if any(found is None for found in blocks):
|
|
raise ValueError("projection calls not found in the function body")
|
|
if len({id(block) for block, _ in blocks}) != 1:
|
|
raise ValueError("projection calls are in different blocks")
|
|
|
|
# q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
|
|
names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
|
|
temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
|
|
if names & set(temps):
|
|
raise ValueError("fused temporaries would shadow existing names")
|
|
merged = f"self.{self.merged_name}"
|
|
sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
|
|
template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
|
|
assign = ast.parse(template).body[0]
|
|
arg = next(
|
|
node
|
|
for node in ast.walk(assign)
|
|
if isinstance(node, ast.Name) and node.id == "__arg__"
|
|
)
|
|
replace_expr(assign, arg, calls[0].args[0])
|
|
block, index = blocks[0]
|
|
ast.copy_location(assign, block[index])
|
|
block.insert(min(index for _, index in blocks), assign)
|
|
for call, temp in zip(calls, temps):
|
|
replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
|
|
self.fused_forward = compile_forward(funcdef, fn)
|
|
|
|
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
|
|
"""Shapes must be compatible for a single merged, head-sharded GEMM."""
|
|
q = module.get_submodule(self.q_name)
|
|
k = module.get_submodule(self.k_name)
|
|
v = module.get_submodule(self.v_name)
|
|
head_size = model_config.get_head_size()
|
|
compatible = (
|
|
q.in_features == k.in_features == v.in_features
|
|
and len({proj.bias is None for proj in (q, k, v)}) == 1
|
|
and k.out_features == v.out_features
|
|
and q.out_features % head_size == 0
|
|
and k.out_features % head_size == 0
|
|
)
|
|
if not compatible:
|
|
logger.debug("%s is not compatible with QKV fusion", type(module))
|
|
return compatible
|
|
|
|
def update_attrs(
|
|
self,
|
|
module: nn.Module,
|
|
prefix: str,
|
|
model_config: "ModelConfig",
|
|
quant_config: "QuantizationConfig",
|
|
) -> None:
|
|
head_size = model_config.get_head_size()
|
|
q = module.get_submodule(self.q_name)
|
|
k = module.get_submodule(self.k_name)
|
|
merged = QKVParallelLinear(
|
|
hidden_size=q.in_features,
|
|
head_size=head_size,
|
|
total_num_heads=q.out_features // head_size,
|
|
total_num_kv_heads=k.out_features // head_size,
|
|
bias=q.bias is not None,
|
|
quant_config=quant_config,
|
|
prefix=maybe_prefix(prefix, self.merged_name),
|
|
return_bias=False,
|
|
)
|
|
logger.debug(
|
|
"%s: %s, %s: %s, %s: %s -> %s: %s",
|
|
self.q_name,
|
|
q,
|
|
self.k_name,
|
|
k,
|
|
self.v_name,
|
|
module.get_submodule(self.v_name),
|
|
self.merged_name,
|
|
merged,
|
|
)
|
|
setattr(module, self.merged_name, merged)
|
|
# Drop the consumed submodules so their (meta) params are not expected.
|
|
for name in (self.q_name, self.k_name, self.v_name):
|
|
delattr(module, name)
|
|
# If there is an output projection, we know it must be rowwise.
|
|
if self.o_name is not None:
|
|
o_proj_prefix = maybe_prefix(prefix, self.o_name)
|
|
o_proj = module.get_submodule(self.o_name)
|
|
new_o = replace_linear_class(
|
|
o_proj, "rowwise", quant_config, prefix=o_proj_prefix
|
|
)
|
|
setattr(module, self.o_name, new_o)
|
|
log_replacement(o_proj_prefix, o_proj, new_o)
|