From c0d8d5413e233466e09d4f2d80867204d72acb60 Mon Sep 17 00:00:00 2001 From: Jim Ban <77719403+BANANASJIM@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:49:37 -0700 Subject: [PATCH] [Transformers] Generalize merged-column linear fusion (#55301) Signed-off-by: BANANASJIM Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../models/transformers/fusers/test_linear.py | 166 +++++++++++++++++- .../models/transformers/fuser.py | 44 ++--- .../models/transformers/fusers/__init__.py | 4 + .../models/transformers/fusers/base.py | 39 +++- .../models/transformers/fusers/glu.py | 81 ++++----- .../transformers/fusers/merged_column.py | 156 ++++++++++++++++ .../models/transformers/fusers/mla.py | 35 +--- .../models/transformers/fusers/qkv.py | 26 +-- 8 files changed, 425 insertions(+), 126 deletions(-) create mode 100644 vllm/model_executor/models/transformers/fusers/merged_column.py diff --git a/tests/models/transformers/fusers/test_linear.py b/tests/models/transformers/fusers/test_linear.py index 020459d7d9..cc1af349fa 100644 --- a/tests/models/transformers/fusers/test_linear.py +++ b/tests/models/transformers/fusers/test_linear.py @@ -10,14 +10,16 @@ import torch import torch.nn as nn import torch.nn.functional as F -from vllm.model_executor.models.transformers.fuser import get_fuser, get_fusers +from vllm.model_executor.models.transformers.fuser import Fusers, get_fuser, get_fusers from vllm.model_executor.models.transformers.fusers import ( GLUFuser, + MergedColumnParallelFuser, PackedQKVFuser, QKVFuser, packed_qkv, qkv, ) +from vllm.model_executor.models.transformers.fx_utils import trace class SiluAndMulStub(nn.Module): @@ -79,6 +81,26 @@ class NotAnActGLUMLP(GLUMLP): self.act_fn = nn.Dropout() +class BroadcastGLU(NoDownGLU): + """A broadcast multiply cannot use an equal-halves AndMul kernel.""" + + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(16, 1, bias=False) + + +class ExtraProjGLU(GLUMLP): + """A third linear on the same input (e.g. a router) -> gate/up still fuse.""" + + def __init__(self): + super().__init__() + self.router = nn.Linear(16, 4, bias=False) + + def forward(self, x): + gated = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return gated, self.router(x) + + class UntraceableMLP(GLUMLP): """Data-dependent control flow *before* the GLU -> no match.""" @@ -166,6 +188,42 @@ class ReversedFakeAttention(FakeAttention): return self.o_proj(attn_output.reshape(*input_shape, -1)), None +class FourParallelLinears(nn.Module): + """Four same-input projections of different widths -> one merged linear.""" + + def __init__(self): + super().__init__() + self.proj_a = nn.Linear(8, 4) + self.proj_b = nn.Linear(8, 6) + self.proj_c = nn.Linear(8, 8) + self.proj_d = nn.Linear(8, 10) + + def forward(self, x): + return self.proj_a(x), self.proj_b(x), self.proj_c(x), self.proj_d(x) + + +class MutatingParallelLinears(FourParallelLinears): + def forward(self, x): + a = self.proj_a(x) + x.add_(1) + return a, self.proj_b(x), self.proj_c(x), self.proj_d(x) + + +class PlainQKV(nn.Module): + """q/k/v returned as a bare tuple: no reshape stands between the calls and + the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it.""" + + def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8): + super().__init__() + self.q_proj = nn.Linear(hidden, heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden, heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden, heads * head_dim, bias=False) + self.head_dim = head_dim + + def forward(self, x): + return self.q_proj(x), self.k_proj(x), self.v_proj(x) + + class ExtraProjAttention(FakeAttention): """A second non-qkv linear of a different width -> `o_proj` still found.""" @@ -420,6 +478,19 @@ def test_detects_and_rewrites_glu(mlp_cls, bias): torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5) +def test_glu_fuses_alongside_unrelated_projections(): + """Sibling linears the GLU does not consume must not block the fusion.""" + with torch.device("meta"): + fuser = get_fuser(ExtraProjGLU(), GLUFuser) + assert fuser is not None + assert (fuser.gate_name, fuser.up_name) == ("gate_proj", "up_proj") + real = ExtraProjGLU() + x = torch.randn(4, 16) + expected = real(x) + fused = _apply_glu_fuser_with_stubs(real, fuser) + torch.testing.assert_close(fused(x), expected) + + def test_glu_identifies_down_projection(): """The row projection consuming `act(gate(x)) * up(x)` is identified. @@ -490,6 +561,97 @@ def test_qkv_identifies_output_projection(): assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj" +def test_merged_column_fuser_rejects_input_mutation(): + """The later projections must not be moved ahead of an input mutation.""" + with torch.device("meta"): + module = MutatingParallelLinears() + fuser = MergedColumnParallelFuser.match(trace(module), module) + assert fuser is not None + with pytest.raises(ValueError, match="cross other operations"): + fuser.update_forward(module) + + +def test_merged_column_fuser_supports_any_number_of_linears( + default_vllm_config, monkeypatch +): + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + from vllm.model_executor.layers.linear import MergedColumnParallelLinear + from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm + from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper + + monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + with torch.device("meta"): + meta = FourParallelLinears() + fuser = MergedColumnParallelFuser.match(trace(meta), meta) + assert fuser is not None + assert fuser.linear_names == ("proj_a", "proj_b", "proj_c", "proj_d") + assert fuser.shards == [ + ("proj_a", 0), + ("proj_b", 1), + ("proj_c", 2), + ("proj_d", 3), + ] + fuser.update_forward(meta) + + real = FourParallelLinears() + x = torch.randn(2, 8) + expected = real(x) + weights = [ + (name, weight.detach().clone()) for name, weight in real.named_parameters() + ] + fused = fuser.fuse(real, "parallel", default_vllm_config) + root = nn.Module() + root.parallel = fused + mapper = WeightsMapper(orig_to_new_stacked=fuser.orig_to_new_stacked("parallel")) + AutoWeightsLoader(root).load_weights( + ((f"parallel.{name}", weight) for name, weight in weights), mapper=mapper + ) + + merged = getattr(fused, fuser.merged_name) + assert isinstance(merged, MergedColumnParallelLinear) + dispatch_cpu_unquantized_gemm(merged, remove_weight=False) + for actual, reference in zip(fused(x), expected): + torch.testing.assert_close(actual, reference) + + +def test_merged_name_is_unique_per_linear_names(): + """The computed name must not collide across different projections, so + stacking two unrelated fusions can never clobber packed_modules_mapping.""" + ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) + a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c")) + assert ab_c.merged_name != a_bc.merged_name + # Same names -> same (and therefore safely re-mergeable) name. + dup = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c")) + assert ab_c.merged_name == dup.merged_name + + +def test_heterogeneous_instances_fall_back_to_merged_column_fuser( + default_vllm_config, +): + """Same class and structure -> both fusers are cached as candidates. A + layer whose out_features aren't a multiple of its own head_dim (as a + misconfigured or heterogeneous checkpoint might have) fails QKVFuser's + validation and must still fuse via the generic fallback, while a + well-formed sibling of the same class keeps using QKVFuser.""" + with torch.device("meta"): + container = nn.Module() + container.compatible = PlainQKV(head_dim=8) + container.incompatible = PlainQKV(head_dim=8) + container.incompatible.head_dim = 5 # 32 % 5 != 0, unlike q_proj's 32 % 8 + + candidates = get_fusers(container.compatible) + assert any(isinstance(f, QKVFuser) for f in candidates) + assert any(isinstance(f, MergedColumnParallelFuser) for f in candidates) + + fusers = Fusers(container, default_vllm_config) + assert isinstance(fusers[container.compatible][0], QKVFuser) + assert isinstance(fusers[container.incompatible][0], MergedColumnParallelFuser) + + @pytest.mark.parametrize("kv_heads", [1, 2]) def test_detects_and_rewrites_packed_qkv(kv_heads): """A single projection split into q/k/v must be re-sharded, not merged. @@ -616,7 +778,7 @@ def test_weight_mappings_are_scoped_to_fused_prefixes(): } -@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP]) +@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP, BroadcastGLU]) def test_unfusable_modules_are_not_fused(cls, default_vllm_config): with torch.device("meta"): module = cls() diff --git a/vllm/model_executor/models/transformers/fuser.py b/vllm/model_executor/models/transformers/fuser.py index 6f211af0d3..91c232f6d5 100644 --- a/vllm/model_executor/models/transformers/fuser.py +++ b/vllm/model_executor/models/transformers/fuser.py @@ -1,13 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fuser detection for the Transformers modeling backend. - -`get_fusers` traces a module class once (see `fx_utils`) and matches it against each -concrete fuser in `fusers`; `Fusers` caches the result per class for a whole model. -At most one match may give the module a different forward, and any number of fusers that -leave the forward alone apply alongside. `base.recursive_replace` then applies them per -instance. RMSNorm-shaped modules the tracer cannot match are warned about. -""" +"""Fuser detection for the Transformers modeling backend.""" from collections import UserDict from typing import TYPE_CHECKING, TypeVar @@ -20,6 +13,7 @@ from vllm.model_executor.models.transformers.fusers import ( AttentionFuser, BaseFuser, GLUFuser, + MergedColumnParallelFuser, MLAFuser, PackedQKVFuser, QKVFuser, @@ -42,23 +36,23 @@ def key(module: nn.Module) -> tuple: FUSERS: tuple[type[BaseFuser], ...] = ( + # Order these by priority, e.g. QKV would be preferred over MergedColumnParallel MLAFuser, GLUFuser, QKVFuser, PackedQKVFuser, + MergedColumnParallelFuser, RMSNormFuser, + # Put fusers that don't redefine forward after this comment AttentionFuser, ) -"""Every fuser, in match order: those that redefine the forward first, then those +"""Every fuser, in priority order: those that redefine the forward first, then those that leave it alone. A new fuser is added here.""" @cached(cache={}, key=key) def get_fusers(module: nn.Module) -> list[BaseFuser]: - """Every fuser that applies to `module`'s class and shape (cached). - - The one that redefines the forward comes first, so applying them in order - fuses before the rest read or adjust the result.""" + """Every fuser that could apply to `module`'s class (cached), in `FUSERS` order.""" # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a # leaf module (raw tensor math, no submodules). Nothing else can match, and # tracing is skipped for it. @@ -66,13 +60,13 @@ def get_fusers(module: nn.Module) -> list[BaseFuser]: is_leaf = next(module.children(), None) is None graph = trace(module) if n_linear >= 2 or is_leaf else None + # Every redefining candidate that matches is cached, not just the first: + # instances of this class can be shaped heterogeneously, so which one is + # usable is decided per instance in Fusers.__getitem__, not here. A forward + # can only be rewritten from a trace of it, so each starts fresh regardless. fusers: list[BaseFuser] = [] for fuser_cls in FUSERS: - if fuser_cls.redefines_forward and ( - # A forward can only be rewritten from a trace of it, and only once, - # because every rewrite starts from the original source. - graph is None or any(fuser.redefines_forward for fuser in fusers) - ): + if fuser_cls.redefines_forward and graph is None: continue if (fuser := fuser_cls.match(graph, module)) is None: continue @@ -118,5 +112,15 @@ class Fusers(UserDict): def __getitem__(self, m: nn.Module) -> list[BaseFuser]: """The fusers this instance can take, in the order to apply them.""" - fusers = self.data.get(key(m), ()) - return [fuser for fuser in fusers if fuser.validate(m, self.vllm_config)] + chosen: list[BaseFuser] = [] + redefined = False + for fuser in self.data.get(key(m), ()): + # Only one candidate may rewrite m's forward; take the first + # (highest-priority) one that validates for this instance. + if fuser.redefines_forward and redefined: + continue + if not fuser.validate(m, self.vllm_config): + continue + redefined = redefined or fuser.redefines_forward + chosen.append(fuser) + return chosen diff --git a/vllm/model_executor/models/transformers/fusers/__init__.py b/vllm/model_executor/models/transformers/fusers/__init__.py index 735392a3ab..727dc0635f 100644 --- a/vllm/model_executor/models/transformers/fusers/__init__.py +++ b/vllm/model_executor/models/transformers/fusers/__init__.py @@ -9,6 +9,9 @@ from vllm.model_executor.models.transformers.fusers.base import ( StackedFuser, ) from vllm.model_executor.models.transformers.fusers.glu import GLUFuser +from vllm.model_executor.models.transformers.fusers.merged_column import ( + MergedColumnParallelFuser, +) from vllm.model_executor.models.transformers.fusers.mla import MLAFuser from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser @@ -21,6 +24,7 @@ __all__ = [ "RewriteFuser", "StackedFuser", "GLUFuser", + "MergedColumnParallelFuser", "MLAFuser", "MoEBlockFuser", "PackedQKVFuser", diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index e36b355739..7c7bfe273e 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Base classes for the Transformers backend fusers.""" +import ast import types from abc import ABC, abstractmethod from collections.abc import Callable @@ -10,6 +11,7 @@ from typing import TYPE_CHECKING, ClassVar from torch import fx, nn +from vllm.model_executor.models.transformers.fx_utils import replace_expr from vllm.model_executor.models.utils import ShardId, maybe_prefix if TYPE_CHECKING: @@ -124,14 +126,14 @@ class StackedFuser(RewriteFuser): merged_name: ClassVar[str] """Attribute name of the merged module created by `update_attrs`.""" - merged_cls: ClassVar[str] + merged_cls_name: ClassVar[str] """Name of the vLLM class the merged projection becomes (for logging).""" def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( f"Fused: {sources} ({name}: {self.source_cls}) -> " - f"{self.merged_name} ({self.merged_cls})" + f"{self.merged_name} ({self.merged_cls_name})" ) @property @@ -157,3 +159,36 @@ class StackedFuser(RewriteFuser): """`{merged_name: [projection names]}` so quantization can unpack the fused layer into its per-shard configs.""" return {self.merged_name: [name for name, _ in self.shards]} + + def _splice_merged_split( + self, + funcdef: ast.FunctionDef, + calls: list[ast.Call], + block: list[ast.stmt], + index: int, + ) -> None: + """Insert `temps = self.(arg).split(sizes, -1)` at + `block[index]` and replace each of `calls` with its temp name. + + `calls` must share one input argument; `block[index]` must be where + they are (or would be) evaluated. Raises if a generated temporary + would shadow an existing name in `funcdef`. + """ + temps = [f"_vllm_merged_{i}" for i in range(len(calls))] + names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)} + if names & set(temps): + raise ValueError("fused temporaries would shadow existing names") + targets = ", ".join(temps) + sections = local_output_sizes(self.merged_name) + source = f"{targets} = self.{self.merged_name}(__arg__).split({sections}, -1)" + assign = ast.parse(source).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]) + ast.copy_location(assign, block[index]) + block.insert(index, assign) + for call, temp in zip(calls, temps): + replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load())) diff --git a/vllm/model_executor/models/transformers/fusers/glu.py b/vllm/model_executor/models/transformers/fusers/glu.py index e19f5b71df..6b81777dc7 100644 --- a/vllm/model_executor/models/transformers/fusers/glu.py +++ b/vllm/model_executor/models/transformers/fusers/glu.py @@ -12,8 +12,9 @@ from transformers.activations import ACT2CLS from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_AND_MUL_REGISTRY -from vllm.model_executor.layers.linear import MergedColumnParallelLinear -from vllm.model_executor.models.transformers.fusers.base import StackedFuser +from vllm.model_executor.models.transformers.fusers.merged_column import ( + MergedColumnParallelFuser, +) from vllm.model_executor.models.transformers.fx_utils import ( compile_forward, find_node, @@ -28,7 +29,7 @@ from vllm.model_executor.models.transformers.utils import ( log_replacement, replace_linear_class, ) -from vllm.model_executor.models.utils import ShardId, maybe_prefix +from vllm.model_executor.models.utils import maybe_prefix if TYPE_CHECKING: from vllm.config import VllmConfig @@ -46,19 +47,20 @@ ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys()) @dataclass -class GLUFuser(StackedFuser): +class GLUFuser(MergedColumnParallelFuser): """Fuser for the GLU pattern `act(gate(x)) * up(x)`.""" act_name: str - gate_name: str - up_name: str down_name: str | None merged_name: ClassVar[str] = "gate_up_proj" - merged_cls: ClassVar[str] = "MergedColumnParallelLinear" @property - def shards(self) -> list[tuple[str, ShardId]]: - return [(self.gate_name, 0), (self.up_name, 1)] + def gate_name(self) -> str: + return self.linear_names[0] + + @property + def up_name(self) -> str: + return self.linear_names[1] @classmethod def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool: @@ -122,22 +124,14 @@ class GLUFuser(StackedFuser): return None act_node, gate_node, up_node, mul_node = glu_nodes - gate = module.get_submodule(gate_node.target) - up = module.get_submodule(up_node.target) - # Shapes must be compatible for a single merged GEMM. - if gate.in_features == up.in_features and (gate.bias is None) == ( - up.bias is None - ): - predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node - down_node = find_node(graph, predicate) - return cls( - source_cls=type(module).__name__, - act_name=act_node.target, - gate_name=gate_node.target, - up_name=up_node.target, - down_name=down_node.target if down_node is not None else None, - ) - return None + predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node + down_node = find_node(graph, predicate) + return cls( + source_cls=type(module).__name__, + act_name=act_node.target, + linear_names=(gate_node.target, up_node.target), + down_name=down_node.target if down_node is not None else None, + ) def update_forward(self, module: nn.Module) -> None: """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source.""" @@ -166,47 +160,32 @@ class GLUFuser(StackedFuser): self.fused_forward = compile_forward(funcdef, fn) def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: + if not super().validate(module, vllm_config): + return False + # An AndMul kernel splits its input in half, so gate and up must match. + gate = module.get_submodule(self.gate_name) + up = module.get_submodule(self.up_name) + if gate.out_features != up.out_features: + logger.debug("%s and %s differ in size; skipping fusion", gate, up) + return False act = module.get_submodule(self.act_name) if self._get_act_and_mul_name(act) is None: - logger.debug("No AndMul equivalent for %s; skipping fusion", type(act)) + logger.debug("%s has no AndMul equivalent; skipping fusion", act) return False return True def update_attrs( self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> None: - quant_config = vllm_config.quant_config + super().update_attrs(module, prefix, vllm_config) act_fn = self._get_act_and_mul(module.get_submodule(self.act_name)) - gate = module.get_submodule(self.gate_name) - up = module.get_submodule(self.up_name) - merged = MergedColumnParallelLinear( - input_size=gate.in_features, - output_sizes=[gate.out_features, up.out_features], - bias=gate.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", - self.gate_name, - gate, - self.up_name, - up, - self.merged_name, - merged, - ) - setattr(module, self.merged_name, merged) setattr(module, self.act_name, act_fn) - # Drop the consumed submodules so their (meta) params are not expected. - delattr(module, self.gate_name) - delattr(module, self.up_name) # If there is a down projection, we know it must be rowwise. if self.down_name is not None: down_prefix = maybe_prefix(prefix, self.down_name) down = module.get_submodule(self.down_name) new_down = replace_linear_class( - down, "rowwise", quant_config, prefix=down_prefix + down, "rowwise", vllm_config.quant_config, prefix=down_prefix ) setattr(module, self.down_name, new_down) log_replacement(down_prefix, down, new_down) diff --git a/vllm/model_executor/models/transformers/fusers/merged_column.py b/vllm/model_executor/models/transformers/fusers/merged_column.py new file mode 100644 index 0000000000..cd335438f0 --- /dev/null +++ b/vllm/model_executor/models/transformers/fusers/merged_column.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fuser for parallel linear projections.""" + +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 MergedColumnParallelLinear +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, + single_self_call, +) +from vllm.model_executor.models.utils import ShardId, maybe_prefix + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +logger = init_logger(__name__) + + +@dataclass +class MergedColumnParallelFuser(StackedFuser): + """Fuser for merging column-parallel linear projections.""" + + linear_names: tuple[str, ...] + merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" + + @property + def merged_name(self) -> str: + """Programmatic name for the merged projection, based on the original names.""" + # len is used to disambiguate names like `a_b` + `c` vs `a` + `b_c` + parts = "_".join(f"{len(name)}_{name}" for name in self.linear_names) + return f"merged_proj_{parts}" + + @property + def shards(self) -> list[tuple[str, ShardId]]: + return [(name, index) for index, name in enumerate(self.linear_names)] + + @staticmethod + def _names(group: list[fx.Node]) -> tuple[str, ...] | None: + names = tuple(str(node.target) for node in group) + if len(set(names)) != len(names) or any("." in name for name in names): + return None + return names + + @classmethod + def match( + cls, graph: fx.Graph, module: nn.Module + ) -> "MergedColumnParallelFuser | None": + """Fuse the module's sibling linears when there is only one such group.""" + 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) + # Like QKVFuser/PackedQKVFuser: this fuser has no head or + # TP-replication awareness, so it must not absorb a QKV-shaped + # pattern (e.g. behind a norm) a more specific fuser declined. + and node.args[0].op == "placeholder" + ): + by_input.setdefault(node.args[0], []).append(node) + # A group whose members are not distinct direct children is dropped: + # the source rewrite addresses each projection as `self.` once. + groups = [ + (group, names) + for group in by_input.values() + if len(group) >= 2 and (names := cls._names(group)) is not None + ] + if len(groups) > 1: + logger.debug( + "%s has %d fusable sibling-linear groups; skipping fusion " + "since which one to merge is ambiguous", + type(module).__name__, + len(groups), + ) + if len(groups) != 1: + return None + _, names = groups[0] + fuser = cls(source_cls=type(module).__name__, linear_names=names) + return None if hasattr(module, fuser.merged_name) else fuser + + def update_forward(self, module: nn.Module) -> None: + """Replace the parallel calls with one merged call and split.""" + funcdef, fn = recover_forward(type(module)) + calls = [single_self_call(funcdef, name) for name in self.linear_names] + if len(set(ast.dump(call.args[0]) for call in calls)) != 1: + raise ValueError("parallel linears read different inputs") + blocks = [innermost_block(funcdef.body, call) for call in calls] + if any(found is None for found in blocks): + raise ValueError("parallel linear calls not found in the function body") + if len(set(id(block) for block, _ in blocks)) != 1: + raise ValueError("parallel linear calls are in different blocks") + + block = blocks[0][0] + index = min(index for _, index in blocks) + calls_have_name_arg = all(isinstance(call.args[0], ast.Name) for call in calls) + # Moving projections must not cross operations that can change their input. + for statement in block[index : max(index for _, index in blocks) + 1]: + if not isinstance(statement, (ast.Assign, ast.Return)): + raise ValueError("parallel projections cross other operations") + if isinstance(statement, ast.Assign) and any( + not isinstance(target, ast.Name) for target in statement.targets + ): + raise ValueError("parallel projection assignment has side effects") + value = statement.value + values = value.elts if isinstance(value, ast.Tuple) else [value] + if any(value not in calls for value in values) or not calls_have_name_arg: + raise ValueError("parallel projections cross other operations") + + # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1) + self._splice_merged_split(funcdef, calls, block, index) + self.fused_forward = compile_forward(funcdef, fn) + + def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: + """Check that the parallel linears are compatible for merging.""" + linear_layers = [module.get_submodule(name) for name in self.linear_names] + tp_size = vllm_config.parallel_config.tensor_parallel_size + return ( + len(linear_layers) >= 2 + and len(set(linear.in_features for linear in linear_layers)) == 1 + and len(set(linear.bias is None for linear in linear_layers)) == 1 + and all(linear.out_features % tp_size == 0 for linear in linear_layers) + ) + + def update_attrs( + self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" + ) -> None: + """Replace the module's parallel linears with one merged projection.""" + linear_modules = [module.get_submodule(name) for name in self.linear_names] + merged = MergedColumnParallelLinear( + input_size=linear_modules[0].in_features, + output_sizes=[linear.out_features for linear in linear_modules], + bias=linear_modules[0].bias is not None, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, self.merged_name), + return_bias=False, + ) + setattr(module, self.merged_name, merged) + for name in self.linear_names: + delattr(module, name) + logger.debug( + "%s -> %s: %s", + ", ".join(f"{n}: {m}" for n, m in zip(self.linear_names, linear_modules)), + self.merged_name, + merged, + ) diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index 31bee9dc63..c0d4127830 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -35,10 +35,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -# Temporaries the fused down-projection binds in the rewritten forward. -_Q_A_TEMP = "__q_a_fused" -_KV_A_TEMP = "__kv_a_fused" - def _consumes_placeholder(node: fx.Node) -> bool: """Whether `node` is a linear applied directly to a `forward` input.""" @@ -131,7 +127,7 @@ class MLAFuser(StackedFuser): kv_b_proj_name: str o_proj_name: str | None merged_name: ClassVar[str] = "fused_qkv_a_proj" - merged_cls: ClassVar[str] = "MergedColumnParallelLinear" + merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" @property def has_q_lora(self) -> bool: @@ -242,31 +238,12 @@ class MLAFuser(StackedFuser): if self.has_q_lora: # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. # The fused call is inserted at the top-level statement preceding both. - q_call = single_self_call(funcdef, self.q_a_proj_name) - kv_call = single_self_call(funcdef, self.kv_a_proj_name) - if ast.dump(q_call.args[0]) != ast.dump(kv_call.args[0]): + names = [self.q_a_proj_name, self.kv_a_proj_name] + calls = [single_self_call(funcdef, name) for name in names] + if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]): raise ValueError("down-projections read different inputs") - names = {n.id for n in ast.walk(funcdef) if isinstance(n, ast.Name)} - if names & {_Q_A_TEMP, _KV_A_TEMP}: - raise ValueError("fused temporaries would shadow existing names") - - merged = f"self.{self.merged_name}" - targets = f"{_Q_A_TEMP}, {_KV_A_TEMP}" - sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]" - source = f"{targets} = {merged}(__arg__).split({sections}, -1)" - assign = ast.parse(source=source).body[0] - placeholder = next( - node - for node in ast.walk(assign) - if isinstance(node, ast.Name) and node.id == "__arg__" - ) - replace_expr(assign, placeholder, q_call.args[0]) - - index = min(_top_level_index(funcdef, call) for call in (q_call, kv_call)) - ast.copy_location(assign, funcdef.body[index]) - funcdef.body.insert(index, assign) - replace_expr(funcdef, q_call, ast.Name(id=_Q_A_TEMP, ctx=ast.Load())) - replace_expr(funcdef, kv_call, ast.Name(id=_KV_A_TEMP, ctx=ast.Load())) + index = min(_top_level_index(funcdef, call) for call in calls) + self._splice_merged_split(funcdef, calls, funcdef.body, index) # Transformers expands the latent into full key/value in a dedicated method. # `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`), diff --git a/vllm/model_executor/models/transformers/fusers/qkv.py b/vllm/model_executor/models/transformers/fusers/qkv.py index dcc8534c0c..e222d655e2 100644 --- a/vllm/model_executor/models/transformers/fusers/qkv.py +++ b/vllm/model_executor/models/transformers/fusers/qkv.py @@ -13,14 +13,12 @@ from vllm.model_executor.layers.linear import QKVParallelLinear from vllm.model_executor.models.transformers.fusers.base import ( StackedFuser, fused_head_size, - local_output_sizes, ) from vllm.model_executor.models.transformers.fx_utils import ( compile_forward, innermost_block, is_linear, recover_forward, - replace_expr, returned_linear, single_self_call, ) @@ -45,7 +43,7 @@ class QKVFuser(StackedFuser): v_name: str o_name: str | None merged_name: ClassVar[str] = "qkv_proj" - merged_cls: ClassVar[str] = "QKVParallelLinear" + merged_cls_name: ClassVar[str] = "QKVParallelLinear" @property def shards(self) -> list[tuple[str, ShardId]]: @@ -133,25 +131,9 @@ class QKVFuser(StackedFuser): 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 = local_output_sizes(self.merged_name) - 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())) + block = blocks[0][0] + index = min(index for _, index in blocks) + self._splice_merged_split(funcdef, calls, block, index) self.fused_forward = compile_forward(funcdef, fn) def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: