[Transformers] Generalize merged-column linear fusion (#55301)

Signed-off-by: BANANASJIM <bananasjim1@gmail.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

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()

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

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",

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.<merged_name>(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()))

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)

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.<name>` 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,
+        )

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`),

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: