[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>
Jim Ban committed
Sep 8, 2026 at 16:49 UTC
c0d8d5413e233466e09d4f2d80867204d72acb60
8 files changed
+425
-126
tests/models/transformers/fusers/test_linear.py
+164
-2
@@ -10,14 +10,16 @@ import torch
10
import torch.nn as nn
11
import torch.nn.functional as F
12
13
-from vllm.model_executor.models.transformers.fuser import get_fuser, get_fusers
13
+from vllm.model_executor.models.transformers.fuser import Fusers, get_fuser, get_fusers
14
from vllm.model_executor.models.transformers.fusers import (
15
GLUFuser,
16
+ MergedColumnParallelFuser,
17
PackedQKVFuser,
18
QKVFuser,
19
packed_qkv,
20
qkv,
21
)
22
+from vllm.model_executor.models.transformers.fx_utils import trace
23
24
25
class SiluAndMulStub(nn.Module):
@@ -79,6 +81,26 @@ class NotAnActGLUMLP(GLUMLP):
81
self.act_fn = nn.Dropout()
82
83
84
+class BroadcastGLU(NoDownGLU):
85
+ """A broadcast multiply cannot use an equal-halves AndMul kernel."""
86
+
87
+ def __init__(self):
88
+ super().__init__()
89
+ self.gate_proj = nn.Linear(16, 1, bias=False)
90
+
91
+
92
+class ExtraProjGLU(GLUMLP):
93
+ """A third linear on the same input (e.g. a router) -> gate/up still fuse."""
94
+
95
+ def __init__(self):
96
+ super().__init__()
97
+ self.router = nn.Linear(16, 4, bias=False)
98
+
99
+ def forward(self, x):
100
+ gated = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
101
+ return gated, self.router(x)
102
+
103
+
104
class UntraceableMLP(GLUMLP):
105
"""Data-dependent control flow *before* the GLU -> no match."""
106
@@ -166,6 +188,42 @@ class ReversedFakeAttention(FakeAttention):
188
return self.o_proj(attn_output.reshape(*input_shape, -1)), None
189
190
191
+class FourParallelLinears(nn.Module):
192
+ """Four same-input projections of different widths -> one merged linear."""
193
+
194
+ def __init__(self):
195
+ super().__init__()
196
+ self.proj_a = nn.Linear(8, 4)
197
+ self.proj_b = nn.Linear(8, 6)
198
+ self.proj_c = nn.Linear(8, 8)
199
+ self.proj_d = nn.Linear(8, 10)
200
+
201
+ def forward(self, x):
202
+ return self.proj_a(x), self.proj_b(x), self.proj_c(x), self.proj_d(x)
203
+
204
+
205
+class MutatingParallelLinears(FourParallelLinears):
206
+ def forward(self, x):
207
+ a = self.proj_a(x)
208
+ x.add_(1)
209
+ return a, self.proj_b(x), self.proj_c(x), self.proj_d(x)
210
+
211
+
212
+class PlainQKV(nn.Module):
213
+ """q/k/v returned as a bare tuple: no reshape stands between the calls and
214
+ the return, so both QKVFuser and MergedColumnParallelFuser can rewrite it."""
215
+
216
+ def __init__(self, hidden: int = 32, heads: int = 4, head_dim: int = 8):
217
+ super().__init__()
218
+ self.q_proj = nn.Linear(hidden, heads * head_dim, bias=False)
219
+ self.k_proj = nn.Linear(hidden, heads * head_dim, bias=False)
220
+ self.v_proj = nn.Linear(hidden, heads * head_dim, bias=False)
221
+ self.head_dim = head_dim
222
+
223
+ def forward(self, x):
224
+ return self.q_proj(x), self.k_proj(x), self.v_proj(x)
225
+
226
+
227
class ExtraProjAttention(FakeAttention):
228
"""A second non-qkv linear of a different width -> `o_proj` still found."""
229
@@ -420,6 +478,19 @@ def test_detects_and_rewrites_glu(mlp_cls, bias):
478
torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
479
480
481
+def test_glu_fuses_alongside_unrelated_projections():
482
+ """Sibling linears the GLU does not consume must not block the fusion."""
483
+ with torch.device("meta"):
484
+ fuser = get_fuser(ExtraProjGLU(), GLUFuser)
485
+ assert fuser is not None
486
+ assert (fuser.gate_name, fuser.up_name) == ("gate_proj", "up_proj")
487
+ real = ExtraProjGLU()
488
+ x = torch.randn(4, 16)
489
+ expected = real(x)
490
+ fused = _apply_glu_fuser_with_stubs(real, fuser)
491
+ torch.testing.assert_close(fused(x), expected)
492
+
493
+
494
def test_glu_identifies_down_projection():
495
"""The row projection consuming `act(gate(x)) * up(x)` is identified.
496
@@ -490,6 +561,97 @@ def test_qkv_identifies_output_projection():
561
assert get_fuser(ResidDropoutAttention(), QKVFuser).o_name == "o_proj"
562
563
564
+def test_merged_column_fuser_rejects_input_mutation():
565
+ """The later projections must not be moved ahead of an input mutation."""
566
+ with torch.device("meta"):
567
+ module = MutatingParallelLinears()
568
+ fuser = MergedColumnParallelFuser.match(trace(module), module)
569
+ assert fuser is not None
570
+ with pytest.raises(ValueError, match="cross other operations"):
571
+ fuser.update_forward(module)
572
+
573
+
574
+def test_merged_column_fuser_supports_any_number_of_linears(
575
+ default_vllm_config, monkeypatch
576
+):
577
+ from vllm.model_executor import parameter
578
+ from vllm.model_executor.layers import linear
579
+ from vllm.model_executor.layers.linear import MergedColumnParallelLinear
580
+ from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
581
+ from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
582
+
583
+ monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0)
584
+ monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1)
585
+ monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
586
+ monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
587
+ with torch.device("meta"):
588
+ meta = FourParallelLinears()
589
+ fuser = MergedColumnParallelFuser.match(trace(meta), meta)
590
+ assert fuser is not None
591
+ assert fuser.linear_names == ("proj_a", "proj_b", "proj_c", "proj_d")
592
+ assert fuser.shards == [
593
+ ("proj_a", 0),
594
+ ("proj_b", 1),
595
+ ("proj_c", 2),
596
+ ("proj_d", 3),
597
+ ]
598
+ fuser.update_forward(meta)
599
+
600
+ real = FourParallelLinears()
601
+ x = torch.randn(2, 8)
602
+ expected = real(x)
603
+ weights = [
604
+ (name, weight.detach().clone()) for name, weight in real.named_parameters()
605
+ ]
606
+ fused = fuser.fuse(real, "parallel", default_vllm_config)
607
+ root = nn.Module()
608
+ root.parallel = fused
609
+ mapper = WeightsMapper(orig_to_new_stacked=fuser.orig_to_new_stacked("parallel"))
610
+ AutoWeightsLoader(root).load_weights(
611
+ ((f"parallel.{name}", weight) for name, weight in weights), mapper=mapper
612
+ )
613
+
614
+ merged = getattr(fused, fuser.merged_name)
615
+ assert isinstance(merged, MergedColumnParallelLinear)
616
+ dispatch_cpu_unquantized_gemm(merged, remove_weight=False)
617
+ for actual, reference in zip(fused(x), expected):
618
+ torch.testing.assert_close(actual, reference)
619
+
620
+
621
+def test_merged_name_is_unique_per_linear_names():
622
+ """The computed name must not collide across different projections, so
623
+ stacking two unrelated fusions can never clobber packed_modules_mapping."""
624
+ ab_c = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c"))
625
+ a_bc = MergedColumnParallelFuser(source_cls="M", linear_names=("a", "b_c"))
626
+ assert ab_c.merged_name != a_bc.merged_name
627
+ # Same names -> same (and therefore safely re-mergeable) name.
628
+ dup = MergedColumnParallelFuser(source_cls="M", linear_names=("a_b", "c"))
629
+ assert ab_c.merged_name == dup.merged_name
630
+
631
+
632
+def test_heterogeneous_instances_fall_back_to_merged_column_fuser(
633
+ default_vllm_config,
634
+):
635
+ """Same class and structure -> both fusers are cached as candidates. A
636
+ layer whose out_features aren't a multiple of its own head_dim (as a
637
+ misconfigured or heterogeneous checkpoint might have) fails QKVFuser's
638
+ validation and must still fuse via the generic fallback, while a
639
+ well-formed sibling of the same class keeps using QKVFuser."""
640
+ with torch.device("meta"):
641
+ container = nn.Module()
642
+ container.compatible = PlainQKV(head_dim=8)
643
+ container.incompatible = PlainQKV(head_dim=8)
644
+ container.incompatible.head_dim = 5 # 32 % 5 != 0, unlike q_proj's 32 % 8
645
+
646
+ candidates = get_fusers(container.compatible)
647
+ assert any(isinstance(f, QKVFuser) for f in candidates)
648
+ assert any(isinstance(f, MergedColumnParallelFuser) for f in candidates)
649
+
650
+ fusers = Fusers(container, default_vllm_config)
651
+ assert isinstance(fusers[container.compatible][0], QKVFuser)
652
+ assert isinstance(fusers[container.incompatible][0], MergedColumnParallelFuser)
653
+
654
+
655
@pytest.mark.parametrize("kv_heads", [1, 2])
656
def test_detects_and_rewrites_packed_qkv(kv_heads):
657
"""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():
778
}
779
780
619
-@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
781
+@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP, BroadcastGLU])
782
def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
783
with torch.device("meta"):
784
module = cls()
vllm/model_executor/models/transformers/fuser.py
+24
-20
@@ -1,13 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Fuser detection for the Transformers modeling backend.
4
-
5
-`get_fusers` traces a module class once (see `fx_utils`) and matches it against each
6
-concrete fuser in `fusers`; `Fusers` caches the result per class for a whole model.
7
-At most one match may give the module a different forward, and any number of fusers that
8
-leave the forward alone apply alongside. `base.recursive_replace` then applies them per
9
-instance. RMSNorm-shaped modules the tracer cannot match are warned about.
10
-"""
3
+"""Fuser detection for the Transformers modeling backend."""
4
5
from collections import UserDict
6
from typing import TYPE_CHECKING, TypeVar
@@ -20,6 +13,7 @@ from vllm.model_executor.models.transformers.fusers import (
13
AttentionFuser,
14
BaseFuser,
15
GLUFuser,
16
+ MergedColumnParallelFuser,
17
MLAFuser,
18
PackedQKVFuser,
19
QKVFuser,
@@ -42,23 +36,23 @@ def key(module: nn.Module) -> tuple:
36
37
38
FUSERS: tuple[type[BaseFuser], ...] = (
39
+ # Order these by priority, e.g. QKV would be preferred over MergedColumnParallel
40
MLAFuser,
41
GLUFuser,
42
QKVFuser,
43
PackedQKVFuser,
44
+ MergedColumnParallelFuser,
45
RMSNormFuser,
46
+ # Put fusers that don't redefine forward after this comment
47
AttentionFuser,
48
)
52
-"""Every fuser, in match order: those that redefine the forward first, then those
49
+"""Every fuser, in priority order: those that redefine the forward first, then those
50
that leave it alone. A new fuser is added here."""
51
52
53
@cached(cache={}, key=key)
54
def get_fusers(module: nn.Module) -> list[BaseFuser]:
58
- """Every fuser that applies to `module`'s class and shape (cached).
59
-
60
- The one that redefines the forward comes first, so applying them in order
61
- fuses before the rest read or adjust the result."""
55
+ """Every fuser that could apply to `module`'s class (cached), in `FUSERS` order."""
56
# Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a
57
# leaf module (raw tensor math, no submodules). Nothing else can match, and
58
# tracing is skipped for it.
@@ -66,13 +60,13 @@ def get_fusers(module: nn.Module) -> list[BaseFuser]:
60
is_leaf = next(module.children(), None) is None
61
graph = trace(module) if n_linear >= 2 or is_leaf else None
62
63
+ # Every redefining candidate that matches is cached, not just the first:
64
+ # instances of this class can be shaped heterogeneously, so which one is
65
+ # usable is decided per instance in Fusers.__getitem__, not here. A forward
66
+ # can only be rewritten from a trace of it, so each starts fresh regardless.
67
fusers: list[BaseFuser] = []
68
for fuser_cls in FUSERS:
71
- if fuser_cls.redefines_forward and (
72
- # A forward can only be rewritten from a trace of it, and only once,
73
- # because every rewrite starts from the original source.
74
- graph is None or any(fuser.redefines_forward for fuser in fusers)
75
- ):
69
+ if fuser_cls.redefines_forward and graph is None:
70
continue
71
if (fuser := fuser_cls.match(graph, module)) is None:
72
continue
@@ -118,5 +112,15 @@ class Fusers(UserDict):
112
113
def __getitem__(self, m: nn.Module) -> list[BaseFuser]:
114
"""The fusers this instance can take, in the order to apply them."""
121
- fusers = self.data.get(key(m), ())
122
- return [fuser for fuser in fusers if fuser.validate(m, self.vllm_config)]
115
+ chosen: list[BaseFuser] = []
116
+ redefined = False
117
+ for fuser in self.data.get(key(m), ()):
118
+ # Only one candidate may rewrite m's forward; take the first
119
+ # (highest-priority) one that validates for this instance.
120
+ if fuser.redefines_forward and redefined:
121
+ continue
122
+ if not fuser.validate(m, self.vllm_config):
123
+ continue
124
+ redefined = redefined or fuser.redefines_forward
125
+ chosen.append(fuser)
126
+ return chosen
vllm/model_executor/models/transformers/fusers/__init__.py
+4
@@ -9,6 +9,9 @@ from vllm.model_executor.models.transformers.fusers.base import (
9
StackedFuser,
10
)
11
from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
12
+from vllm.model_executor.models.transformers.fusers.merged_column import (
13
+ MergedColumnParallelFuser,
14
+)
15
from vllm.model_executor.models.transformers.fusers.mla import MLAFuser
16
from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
17
from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser
@@ -21,6 +24,7 @@ __all__ = [
24
"RewriteFuser",
25
"StackedFuser",
26
"GLUFuser",
27
+ "MergedColumnParallelFuser",
28
"MLAFuser",
29
"MoEBlockFuser",
30
"PackedQKVFuser",
vllm/model_executor/models/transformers/fusers/base.py
+37
-2
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Base classes for the Transformers backend fusers."""
4
5
+import ast
6
import types
7
from abc import ABC, abstractmethod
8
from collections.abc import Callable
@@ -10,6 +11,7 @@ from typing import TYPE_CHECKING, ClassVar
11
12
from torch import fx, nn
13
14
+from vllm.model_executor.models.transformers.fx_utils import replace_expr
15
from vllm.model_executor.models.utils import ShardId, maybe_prefix
16
17
if TYPE_CHECKING:
@@ -124,14 +126,14 @@ class StackedFuser(RewriteFuser):
126
127
merged_name: ClassVar[str]
128
"""Attribute name of the merged module created by `update_attrs`."""
127
- merged_cls: ClassVar[str]
129
+ merged_cls_name: ClassVar[str]
130
"""Name of the vLLM class the merged projection becomes (for logging)."""
131
132
def info(self, name: str) -> str:
133
sources = " + ".join(shard for shard, _ in self.shards)
134
return (
135
f"Fused: {sources} ({name}: {self.source_cls}) -> "
134
- f"{self.merged_name} ({self.merged_cls})"
136
+ f"{self.merged_name} ({self.merged_cls_name})"
137
)
138
139
@property
@@ -157,3 +159,36 @@ class StackedFuser(RewriteFuser):
159
"""`{merged_name: [projection names]}` so quantization can unpack the
160
fused layer into its per-shard configs."""
161
return {self.merged_name: [name for name, _ in self.shards]}
162
+
163
+ def _splice_merged_split(
164
+ self,
165
+ funcdef: ast.FunctionDef,
166
+ calls: list[ast.Call],
167
+ block: list[ast.stmt],
168
+ index: int,
169
+ ) -> None:
170
+ """Insert `temps = self.<merged_name>(arg).split(sizes, -1)` at
171
+ `block[index]` and replace each of `calls` with its temp name.
172
+
173
+ `calls` must share one input argument; `block[index]` must be where
174
+ they are (or would be) evaluated. Raises if a generated temporary
175
+ would shadow an existing name in `funcdef`.
176
+ """
177
+ temps = [f"_vllm_merged_{i}" for i in range(len(calls))]
178
+ names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
179
+ if names & set(temps):
180
+ raise ValueError("fused temporaries would shadow existing names")
181
+ targets = ", ".join(temps)
182
+ sections = local_output_sizes(self.merged_name)
183
+ source = f"{targets} = self.{self.merged_name}(__arg__).split({sections}, -1)"
184
+ assign = ast.parse(source).body[0]
185
+ arg = next(
186
+ node
187
+ for node in ast.walk(assign)
188
+ if isinstance(node, ast.Name) and node.id == "__arg__"
189
+ )
190
+ replace_expr(assign, arg, calls[0].args[0])
191
+ ast.copy_location(assign, block[index])
192
+ block.insert(index, assign)
193
+ for call, temp in zip(calls, temps):
194
+ replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
vllm/model_executor/models/transformers/fusers/glu.py
+30
-51
@@ -12,8 +12,9 @@ from transformers.activations import ACT2CLS
12
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.activation import _ACTIVATION_AND_MUL_REGISTRY
15
-from vllm.model_executor.layers.linear import MergedColumnParallelLinear
16
-from vllm.model_executor.models.transformers.fusers.base import StackedFuser
15
+from vllm.model_executor.models.transformers.fusers.merged_column import (
16
+ MergedColumnParallelFuser,
17
+)
18
from vllm.model_executor.models.transformers.fx_utils import (
19
compile_forward,
20
find_node,
@@ -28,7 +29,7 @@ from vllm.model_executor.models.transformers.utils import (
29
log_replacement,
30
replace_linear_class,
31
)
31
-from vllm.model_executor.models.utils import ShardId, maybe_prefix
32
+from vllm.model_executor.models.utils import maybe_prefix
33
34
if TYPE_CHECKING:
35
from vllm.config import VllmConfig
@@ -46,19 +47,20 @@ ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys())
47
48
49
@dataclass
49
-class GLUFuser(StackedFuser):
50
+class GLUFuser(MergedColumnParallelFuser):
51
"""Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
52
53
act_name: str
53
- gate_name: str
54
- up_name: str
54
down_name: str | None
55
merged_name: ClassVar[str] = "gate_up_proj"
57
- merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
56
57
@property
60
- def shards(self) -> list[tuple[str, ShardId]]:
61
- return [(self.gate_name, 0), (self.up_name, 1)]
58
+ def gate_name(self) -> str:
59
+ return self.linear_names[0]
60
+
61
+ @property
62
+ def up_name(self) -> str:
63
+ return self.linear_names[1]
64
65
@classmethod
66
def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
@@ -122,22 +124,14 @@ class GLUFuser(StackedFuser):
124
return None
125
act_node, gate_node, up_node, mul_node = glu_nodes
126
125
- gate = module.get_submodule(gate_node.target)
126
- up = module.get_submodule(up_node.target)
127
- # Shapes must be compatible for a single merged GEMM.
128
- if gate.in_features == up.in_features and (gate.bias is None) == (
129
- up.bias is None
130
- ):
131
- predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
132
- down_node = find_node(graph, predicate)
133
- return cls(
134
- source_cls=type(module).__name__,
135
- act_name=act_node.target,
136
- gate_name=gate_node.target,
137
- up_name=up_node.target,
138
- down_name=down_node.target if down_node is not None else None,
139
- )
140
- return None
127
+ predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
128
+ down_node = find_node(graph, predicate)
129
+ return cls(
130
+ source_cls=type(module).__name__,
131
+ act_name=act_node.target,
132
+ linear_names=(gate_node.target, up_node.target),
133
+ down_name=down_node.target if down_node is not None else None,
134
+ )
135
136
def update_forward(self, module: nn.Module) -> None:
137
"""Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
@@ -166,47 +160,32 @@ class GLUFuser(StackedFuser):
160
self.fused_forward = compile_forward(funcdef, fn)
161
162
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
163
+ if not super().validate(module, vllm_config):
164
+ return False
165
+ # An AndMul kernel splits its input in half, so gate and up must match.
166
+ gate = module.get_submodule(self.gate_name)
167
+ up = module.get_submodule(self.up_name)
168
+ if gate.out_features != up.out_features:
169
+ logger.debug("%s and %s differ in size; skipping fusion", gate, up)
170
+ return False
171
act = module.get_submodule(self.act_name)
172
if self._get_act_and_mul_name(act) is None:
171
- logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))
173
+ logger.debug("%s has no AndMul equivalent; skipping fusion", act)
174
return False
175
return True
176
177
def update_attrs(
178
self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
179
) -> None:
178
- quant_config = vllm_config.quant_config
180
+ super().update_attrs(module, prefix, vllm_config)
181
act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
180
- gate = module.get_submodule(self.gate_name)
181
- up = module.get_submodule(self.up_name)
182
- merged = MergedColumnParallelLinear(
183
- input_size=gate.in_features,
184
- output_sizes=[gate.out_features, up.out_features],
185
- bias=gate.bias is not None,
186
- quant_config=quant_config,
187
- prefix=maybe_prefix(prefix, self.merged_name),
188
- return_bias=False,
189
- )
190
- logger.debug(
191
- "%s: %s, %s: %s -> %s: %s",
192
- self.gate_name,
193
- gate,
194
- self.up_name,
195
- up,
196
- self.merged_name,
197
- merged,
198
- )
199
- setattr(module, self.merged_name, merged)
182
setattr(module, self.act_name, act_fn)
201
- # Drop the consumed submodules so their (meta) params are not expected.
202
- delattr(module, self.gate_name)
203
- delattr(module, self.up_name)
183
# If there is a down projection, we know it must be rowwise.
184
if self.down_name is not None:
185
down_prefix = maybe_prefix(prefix, self.down_name)
186
down = module.get_submodule(self.down_name)
187
new_down = replace_linear_class(
209
- down, "rowwise", quant_config, prefix=down_prefix
188
+ down, "rowwise", vllm_config.quant_config, prefix=down_prefix
189
)
190
setattr(module, self.down_name, new_down)
191
log_replacement(down_prefix, down, new_down)
vllm/model_executor/models/transformers/fusers/merged_column.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Fuser for parallel linear projections."""
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 MergedColumnParallelLinear
13
+from vllm.model_executor.models.transformers.fusers.base import StackedFuser
14
+from vllm.model_executor.models.transformers.fx_utils import (
15
+ compile_forward,
16
+ innermost_block,
17
+ is_linear,
18
+ recover_forward,
19
+ single_self_call,
20
+)
21
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
22
+
23
+if TYPE_CHECKING:
24
+ from vllm.config import VllmConfig
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+@dataclass
30
+class MergedColumnParallelFuser(StackedFuser):
31
+ """Fuser for merging column-parallel linear projections."""
32
+
33
+ linear_names: tuple[str, ...]
34
+ merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"
35
+
36
+ @property
37
+ def merged_name(self) -> str:
38
+ """Programmatic name for the merged projection, based on the original names."""
39
+ # len is used to disambiguate names like `a_b` + `c` vs `a` + `b_c`
40
+ parts = "_".join(f"{len(name)}_{name}" for name in self.linear_names)
41
+ return f"merged_proj_{parts}"
42
+
43
+ @property
44
+ def shards(self) -> list[tuple[str, ShardId]]:
45
+ return [(name, index) for index, name in enumerate(self.linear_names)]
46
+
47
+ @staticmethod
48
+ def _names(group: list[fx.Node]) -> tuple[str, ...] | None:
49
+ names = tuple(str(node.target) for node in group)
50
+ if len(set(names)) != len(names) or any("." in name for name in names):
51
+ return None
52
+ return names
53
+
54
+ @classmethod
55
+ def match(
56
+ cls, graph: fx.Graph, module: nn.Module
57
+ ) -> "MergedColumnParallelFuser | None":
58
+ """Fuse the module's sibling linears when there is only one such group."""
59
+ by_input: dict[fx.Node, list[fx.Node]] = {}
60
+ for node in graph.nodes:
61
+ if (
62
+ is_linear(node, module)
63
+ and len(node.args) == 1
64
+ and not node.kwargs
65
+ and isinstance(node.args[0], fx.Node)
66
+ # Like QKVFuser/PackedQKVFuser: this fuser has no head or
67
+ # TP-replication awareness, so it must not absorb a QKV-shaped
68
+ # pattern (e.g. behind a norm) a more specific fuser declined.
69
+ and node.args[0].op == "placeholder"
70
+ ):
71
+ by_input.setdefault(node.args[0], []).append(node)
72
+ # A group whose members are not distinct direct children is dropped:
73
+ # the source rewrite addresses each projection as `self.<name>` once.
74
+ groups = [
75
+ (group, names)
76
+ for group in by_input.values()
77
+ if len(group) >= 2 and (names := cls._names(group)) is not None
78
+ ]
79
+ if len(groups) > 1:
80
+ logger.debug(
81
+ "%s has %d fusable sibling-linear groups; skipping fusion "
82
+ "since which one to merge is ambiguous",
83
+ type(module).__name__,
84
+ len(groups),
85
+ )
86
+ if len(groups) != 1:
87
+ return None
88
+ _, names = groups[0]
89
+ fuser = cls(source_cls=type(module).__name__, linear_names=names)
90
+ return None if hasattr(module, fuser.merged_name) else fuser
91
+
92
+ def update_forward(self, module: nn.Module) -> None:
93
+ """Replace the parallel calls with one merged call and split."""
94
+ funcdef, fn = recover_forward(type(module))
95
+ calls = [single_self_call(funcdef, name) for name in self.linear_names]
96
+ if len(set(ast.dump(call.args[0]) for call in calls)) != 1:
97
+ raise ValueError("parallel linears read different inputs")
98
+ blocks = [innermost_block(funcdef.body, call) for call in calls]
99
+ if any(found is None for found in blocks):
100
+ raise ValueError("parallel linear calls not found in the function body")
101
+ if len(set(id(block) for block, _ in blocks)) != 1:
102
+ raise ValueError("parallel linear calls are in different blocks")
103
+
104
+ block = blocks[0][0]
105
+ index = min(index for _, index in blocks)
106
+ calls_have_name_arg = all(isinstance(call.args[0], ast.Name) for call in calls)
107
+ # Moving projections must not cross operations that can change their input.
108
+ for statement in block[index : max(index for _, index in blocks) + 1]:
109
+ if not isinstance(statement, (ast.Assign, ast.Return)):
110
+ raise ValueError("parallel projections cross other operations")
111
+ if isinstance(statement, ast.Assign) and any(
112
+ not isinstance(target, ast.Name) for target in statement.targets
113
+ ):
114
+ raise ValueError("parallel projection assignment has side effects")
115
+ value = statement.value
116
+ values = value.elts if isinstance(value, ast.Tuple) else [value]
117
+ if any(value not in calls for value in values) or not calls_have_name_arg:
118
+ raise ValueError("parallel projections cross other operations")
119
+
120
+ # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1)
121
+ self._splice_merged_split(funcdef, calls, block, index)
122
+ self.fused_forward = compile_forward(funcdef, fn)
123
+
124
+ def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
125
+ """Check that the parallel linears are compatible for merging."""
126
+ linear_layers = [module.get_submodule(name) for name in self.linear_names]
127
+ tp_size = vllm_config.parallel_config.tensor_parallel_size
128
+ return (
129
+ len(linear_layers) >= 2
130
+ and len(set(linear.in_features for linear in linear_layers)) == 1
131
+ and len(set(linear.bias is None for linear in linear_layers)) == 1
132
+ and all(linear.out_features % tp_size == 0 for linear in linear_layers)
133
+ )
134
+
135
+ def update_attrs(
136
+ self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
137
+ ) -> None:
138
+ """Replace the module's parallel linears with one merged projection."""
139
+ linear_modules = [module.get_submodule(name) for name in self.linear_names]
140
+ merged = MergedColumnParallelLinear(
141
+ input_size=linear_modules[0].in_features,
142
+ output_sizes=[linear.out_features for linear in linear_modules],
143
+ bias=linear_modules[0].bias is not None,
144
+ quant_config=vllm_config.quant_config,
145
+ prefix=maybe_prefix(prefix, self.merged_name),
146
+ return_bias=False,
147
+ )
148
+ setattr(module, self.merged_name, merged)
149
+ for name in self.linear_names:
150
+ delattr(module, name)
151
+ logger.debug(
152
+ "%s -> %s: %s",
153
+ ", ".join(f"{n}: {m}" for n, m in zip(self.linear_names, linear_modules)),
154
+ self.merged_name,
155
+ merged,
156
+ )
vllm/model_executor/models/transformers/fusers/mla.py
+6
-29
@@ -35,10 +35,6 @@ if TYPE_CHECKING:
35
36
logger = init_logger(__name__)
37
38
-# Temporaries the fused down-projection binds in the rewritten forward.
39
-_Q_A_TEMP = "__q_a_fused"
40
-_KV_A_TEMP = "__kv_a_fused"
41
-
38
39
def _consumes_placeholder(node: fx.Node) -> bool:
40
"""Whether `node` is a linear applied directly to a `forward` input."""
@@ -131,7 +127,7 @@ class MLAFuser(StackedFuser):
127
kv_b_proj_name: str
128
o_proj_name: str | None
129
merged_name: ClassVar[str] = "fused_qkv_a_proj"
134
- merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
130
+ merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear"
131
132
@property
133
def has_q_lora(self) -> bool:
@@ -242,31 +238,12 @@ class MLAFuser(StackedFuser):
238
if self.has_q_lora:
239
# q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`.
240
# The fused call is inserted at the top-level statement preceding both.
245
- q_call = single_self_call(funcdef, self.q_a_proj_name)
246
- kv_call = single_self_call(funcdef, self.kv_a_proj_name)
247
- if ast.dump(q_call.args[0]) != ast.dump(kv_call.args[0]):
241
+ names = [self.q_a_proj_name, self.kv_a_proj_name]
242
+ calls = [single_self_call(funcdef, name) for name in names]
243
+ if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]):
244
raise ValueError("down-projections read different inputs")
249
- names = {n.id for n in ast.walk(funcdef) if isinstance(n, ast.Name)}
250
- if names & {_Q_A_TEMP, _KV_A_TEMP}:
251
- raise ValueError("fused temporaries would shadow existing names")
252
-
253
- merged = f"self.{self.merged_name}"
254
- targets = f"{_Q_A_TEMP}, {_KV_A_TEMP}"
255
- sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
256
- source = f"{targets} = {merged}(__arg__).split({sections}, -1)"
257
- assign = ast.parse(source=source).body[0]
258
- placeholder = next(
259
- node
260
- for node in ast.walk(assign)
261
- if isinstance(node, ast.Name) and node.id == "__arg__"
262
- )
263
- replace_expr(assign, placeholder, q_call.args[0])
264
-
265
- index = min(_top_level_index(funcdef, call) for call in (q_call, kv_call))
266
- ast.copy_location(assign, funcdef.body[index])
267
- funcdef.body.insert(index, assign)
268
- replace_expr(funcdef, q_call, ast.Name(id=_Q_A_TEMP, ctx=ast.Load()))
269
- replace_expr(funcdef, kv_call, ast.Name(id=_KV_A_TEMP, ctx=ast.Load()))
245
+ index = min(_top_level_index(funcdef, call) for call in calls)
246
+ self._splice_merged_split(funcdef, calls, funcdef.body, index)
247
248
# Transformers expands the latent into full key/value in a dedicated method.
249
# `MLAAttention` consumes the latent directly (absorbing `kv_b_proj`),
vllm/model_executor/models/transformers/fusers/qkv.py
+4
-22
@@ -13,14 +13,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
- local_output_sizes,
16
)
17
from vllm.model_executor.models.transformers.fx_utils import (
18
compile_forward,
19
innermost_block,
20
is_linear,
21
recover_forward,
23
- replace_expr,
22
returned_linear,
23
single_self_call,
24
)
@@ -45,7 +43,7 @@ class QKVFuser(StackedFuser):
43
v_name: str
44
o_name: str | None
45
merged_name: ClassVar[str] = "qkv_proj"
48
- merged_cls: ClassVar[str] = "QKVParallelLinear"
46
+ merged_cls_name: ClassVar[str] = "QKVParallelLinear"
47
48
@property
49
def shards(self) -> list[tuple[str, ShardId]]:
@@ -133,25 +131,9 @@ class QKVFuser(StackedFuser):
131
raise ValueError("projection calls are in different blocks")
132
133
# q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
136
- names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
137
- temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
138
- if names & set(temps):
139
- raise ValueError("fused temporaries would shadow existing names")
140
- merged = f"self.{self.merged_name}"
141
- sections = local_output_sizes(self.merged_name)
142
- template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
143
- assign = ast.parse(template).body[0]
144
- arg = next(
145
- node
146
- for node in ast.walk(assign)
147
- if isinstance(node, ast.Name) and node.id == "__arg__"
148
- )
149
- replace_expr(assign, arg, calls[0].args[0])
150
- block, index = blocks[0]
151
- ast.copy_location(assign, block[index])
152
- block.insert(min(index for _, index in blocks), assign)
153
- for call, temp in zip(calls, temps):
154
- replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
134
+ block = blocks[0][0]
135
+ index = min(index for _, index in blocks)
136
+ self._splice_merged_split(funcdef, calls, block, index)
137
self.fused_forward = compile_forward(funcdef, fn)
138
139
def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: