Fix MQA with tensor parallelism on transformers modeling backend (#49987)
Signed-off-by: microslaw <milosz.grunwald@intel.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Miłosz Grunwald committed
Jul 28, 2026 at 00:51 UTC
ebcef33766f8bcce6344f2bafa245d47097d3c31
7 files changed
+396
-44
tests/models/transformers/fusers/test_linear.py
+162
-1
@@ -11,7 +11,11 @@ import torch.nn as nn
11
import torch.nn.functional as F
12
13
from vllm.model_executor.models.transformers.fuser import get_fuser
14
-from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
14
+from vllm.model_executor.models.transformers.fusers import (
15
+ GLUFuser,
16
+ PackedQKVFuser,
17
+ QKVFuser,
18
+)
19
20
21
class SiluAndMulStub(nn.Module):
@@ -203,6 +207,103 @@ class PerHeadQKNormAttention(FakeAttention):
207
return self.o_proj((q + k + v).flatten(-2)), None
208
209
210
+class ResidDropoutAttention(FakeAttention):
211
+ """GPT-style dropout after `o_proj` -> the output projection is still found."""
212
+
213
+ def __init__(self, **kwargs):
214
+ super().__init__(**kwargs)
215
+ self.resid_dropout = nn.Dropout(0.0)
216
+
217
+ def forward(
218
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
219
+ ):
220
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
221
+
222
+ input_shape = hidden_states.shape[:-1]
223
+ hidden_shape = (*input_shape, -1, self.head_dim)
224
+ q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
225
+ k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
226
+ v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
227
+ attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
228
+ self.config._attn_implementation, None
229
+ )
230
+ attn_output, _ = attention_interface(
231
+ self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
232
+ )
233
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
234
+ return self.resid_dropout(self.o_proj(attn_output)), None
235
+
236
+
237
+class PackedQKVAttention(nn.Module):
238
+ """GPTBigCode-style: one packed projection split into q/k/v in the forward."""
239
+
240
+ is_causal = True
241
+
242
+ def __init__(
243
+ self,
244
+ hidden: int = 32,
245
+ head_dim: int = 8,
246
+ heads: int = 4,
247
+ kv_heads: int = 1,
248
+ bias: bool = False,
249
+ layer_idx: int = 0,
250
+ ):
251
+ super().__init__()
252
+ self.config = SimpleNamespace(_attn_implementation="vllm")
253
+ self.layer_idx = layer_idx
254
+ self.head_dim = head_dim
255
+ self.scaling = head_dim**-0.5
256
+ self.embed_dim = heads * head_dim
257
+ self.kv_dim = kv_heads * head_dim
258
+ self.c_attn = nn.Linear(hidden, self.embed_dim + 2 * self.kv_dim, bias=bias)
259
+ self.c_proj = nn.Linear(self.embed_dim, hidden, bias=bias)
260
+ self.resid_dropout = nn.Dropout(0.0)
261
+
262
+ def forward(
263
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
264
+ ):
265
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
266
+
267
+ input_shape = hidden_states.shape[:-1]
268
+ q, k, v = (
269
+ self.c_attn(hidden_states)
270
+ .unsqueeze(1)
271
+ .split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3)
272
+ )
273
+ q = q.view(*input_shape, -1, self.head_dim).transpose(1, 2)
274
+ if past_key_values is not None:
275
+ k, v = past_key_values.update(k, v, self.layer_idx)
276
+ attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
277
+ self.config._attn_implementation, None
278
+ )
279
+ attn_output, attn_weights = attention_interface(
280
+ self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
281
+ )
282
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
283
+ return self.resid_dropout(self.c_proj(attn_output)), attn_weights
284
+
285
+
286
+class PerHeadSplitAttention(nn.Module):
287
+ """A packed projection reshaped and split *per head* -> not a q/k/v split."""
288
+
289
+ def __init__(self, hidden: int = 32, head_dim: int = 8, heads: int = 4):
290
+ super().__init__()
291
+ self.head_dim = head_dim
292
+ self.heads = heads
293
+ self.c_attn = nn.Linear(hidden, 3 * heads * head_dim)
294
+ self.c_proj = nn.Linear(heads * head_dim, hidden)
295
+
296
+ def forward(self, hidden_states):
297
+ shape = (*hidden_states.shape[:2], self.heads, 3 * self.head_dim)
298
+ q, k, v = (
299
+ self.c_attn(hidden_states)
300
+ .view(shape)
301
+ .transpose(1, 2)
302
+ .split((self.head_dim, self.head_dim, self.head_dim), dim=3)
303
+ )
304
+ return self.c_proj((q + k + v).transpose(1, 2).flatten(-2))
305
+
306
+
307
class FakeSelfAttn(nn.Module):
308
"""Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
309
@@ -215,6 +316,14 @@ class FakeSelfAttn(nn.Module):
316
return q + 2 * k + 3 * v
317
318
319
+class FakeMQASelfAttn(FakeSelfAttn):
320
+ """Stand-in for grouped/multi-query layouts, where `k`/`v` are narrower."""
321
+
322
+ def forward(self, q, k, v):
323
+ groups = q.shape[-1] // k.shape[-1]
324
+ return q + (2 * k + 3 * v).repeat(1, groups)
325
+
326
+
327
@pytest.fixture(autouse=True)
328
def _clear_fuser_cache():
329
get_fuser.cache_clear()
@@ -267,6 +376,15 @@ def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
376
return module
377
378
379
+def _apply_packed_qkv_fuser_with_stubs(module: nn.Module, fuser: PackedQKVFuser):
380
+ """Apply a fuser at `tp_size == 1`, where the rewritten split is unchanged."""
381
+ qkv = module.get_submodule(fuser.qkv_name)
382
+ qkv.output_sizes = [fuser.q_size, fuser.kv_size, fuser.kv_size]
383
+ qkv.tp_size = 1
384
+ module.forward = MethodType(fuser.fused_forward, module)
385
+ return module
386
+
387
+
388
@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
389
@pytest.mark.parametrize("bias", [False, True])
390
def test_detects_and_rewrites_glu(mlp_cls, bias):
@@ -366,6 +484,49 @@ def test_qkv_identifies_output_projection():
484
# Norm children (q_norm/k_norm) must not disturb o_proj identification.
485
assert get_fuser(QKNormAttention()).o_name == "o_proj"
486
assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
487
+ # A module between o_proj and the return is transparent.
488
+ assert get_fuser(ResidDropoutAttention()).o_name == "o_proj"
489
+
490
+
491
+@pytest.mark.parametrize("kv_heads", [1, 2])
492
+def test_detects_and_rewrites_packed_qkv(kv_heads):
493
+ """A single projection split into q/k/v must be re-sharded, not merged.
494
+
495
+ Only the split sizes change: `QKVParallelLinear` loads the packed
496
+ checkpoint weight as-is, and shards q by heads while replicating k/v."""
497
+ with torch.device("meta"):
498
+ meta = PackedQKVAttention(kv_heads=kv_heads)
499
+ fuser = get_fuser(meta)
500
+ assert isinstance(fuser, PackedQKVFuser)
501
+ assert (fuser.qkv_name, fuser.o_name) == ("c_attn", "c_proj")
502
+ assert (fuser.q_size, fuser.kv_size) == (32, 8 * kv_heads)
503
+
504
+ # The hard-coded widths become the per-rank widths of the sharded linear
505
+ names = fuser.fused_forward.__code__.co_names
506
+ assert "output_sizes" in names and "tp_size" in names
507
+ assert "kv_dim" not in names and "embed_dim" not in names
508
+
509
+ # Numerics: the rewritten forward must match the original on a real instance
510
+ real = PackedQKVAttention(kv_heads=kv_heads, layer_idx=3)
511
+ for p in real.parameters():
512
+ nn.init.normal_(p, std=0.05)
513
+ x = torch.randn(1, 5, 32)
514
+ attention_instances = {3: FakeMQASelfAttn()}
515
+ expected, _ = real(x, attention_instances=attention_instances)
516
+ fused = _apply_packed_qkv_fuser_with_stubs(real, fuser)
517
+
518
+ # Fusion is in place: the module keeps its class and other attributes
519
+ assert fused is real and type(fused) is PackedQKVAttention
520
+ assert fused.layer_idx == 3 and fused.is_causal
521
+ out, _ = fused(x, attention_instances=attention_instances)
522
+ torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
523
+
524
+
525
+def test_per_head_split_is_not_packed_qkv():
526
+ """The split must consume the whole projection, else its sizes are head
527
+ widths and re-sharding by them would be wrong."""
528
+ with torch.device("meta"):
529
+ assert get_fuser(PerHeadSplitAttention()) is None
530
531
532
def test_fuser_is_cached_per_class_and_structure():
vllm/model_executor/models/transformers/fuser.py
+4
-3
@@ -18,9 +18,10 @@ from vllm.logger import init_logger
18
from vllm.model_executor.models.transformers.fusers import (
19
BaseFuser,
20
GLUFuser,
21
+ PackedQKVFuser,
22
QKVFuser,
23
+ RewriteFuser,
24
RMSNormFuser,
23
- StackedFuser,
25
)
26
from vllm.model_executor.models.transformers.fx_utils import trace
27
@@ -47,9 +48,9 @@ def get_fuser(module: nn.Module) -> BaseFuser | None:
48
return None
49
if (graph := trace(module)) is None:
50
return None
50
- for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser):
51
+ for fuser_cls in (GLUFuser, QKVFuser, PackedQKVFuser, RMSNormFuser):
52
if (fuser := fuser_cls.match(graph, module)) is not None:
52
- if isinstance(fuser, StackedFuser):
53
+ if isinstance(fuser, RewriteFuser):
54
try:
55
fuser.update_forward(module)
56
except Exception as exc:
vllm/model_executor/models/transformers/fusers/__init__.py
+8
-1
@@ -2,17 +2,24 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Concrete fusers for the Transformers modeling backend."""
4
5
-from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser
5
+from vllm.model_executor.models.transformers.fusers.base import (
6
+ BaseFuser,
7
+ RewriteFuser,
8
+ StackedFuser,
9
+)
10
from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
11
from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
12
+from vllm.model_executor.models.transformers.fusers.packed_qkv import PackedQKVFuser
13
from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser
14
from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
15
16
__all__ = [
17
"BaseFuser",
18
+ "RewriteFuser",
19
"StackedFuser",
20
"GLUFuser",
21
"MoEBlockFuser",
22
+ "PackedQKVFuser",
23
"QKVFuser",
24
"RMSNormFuser",
25
]
vllm/model_executor/models/transformers/fusers/base.py
+45
-34
@@ -57,27 +57,61 @@ class BaseFuser(ABC):
57
return {}
58
59
60
+def local_output_sizes(merged_name: str) -> str:
61
+ """Source for the per-rank widths of the merged linear `self.<merged_name>`."""
62
+ merged = f"self.{merged_name}"
63
+ return f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
64
+
65
+
66
@dataclass
61
-class StackedFuser(BaseFuser):
62
- """A fuser that merges sibling projections into one stacked linear and
63
- rewrites the forward to call it.
67
+class RewriteFuser(BaseFuser):
68
+ """A fuser that rewrites the module's forward and rebinds it.
69
65
- `match` and `update_forward` analyse the class once; `fuse` builds the merged
66
- submodule and binds the compiled forward on an instance in place, so it keeps
67
- its class and any attribute the fusion does not consume.
70
+ `match` and `update_forward` analyse the class once; `fuse` swaps the
71
+ submodules and binds the compiled forward on an instance in place, so it
72
+ keeps its class and any attribute the fusion does not consume.
73
"""
74
70
- merged_name: ClassVar[str]
71
- """Attribute name of the merged module created by `update_attrs`."""
72
- merged_cls: ClassVar[str]
73
- """Name of the vLLM class the merged projection becomes (for logging)."""
74
-
75
source_cls: str
76
"""Class of the HF module the fused projections belonged to (for logging)."""
77
78
fused_forward: Callable = field(init=False, repr=False)
79
"""The compiled rewritten forward, set by `update_forward`."""
80
81
+ @abstractmethod
82
+ def update_forward(self, module: nn.Module) -> None:
83
+ """Rewrite and compile `type(module)`'s forward source.
84
+
85
+ Raises if the source does not admit the rewrite (fusion is then skipped).
86
+ """
87
+
88
+ @abstractmethod
89
+ def update_attrs(
90
+ self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
91
+ ) -> None:
92
+ """Replace `module`'s submodules with their vLLM equivalents."""
93
+
94
+ def fuse(
95
+ self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
96
+ ) -> nn.Module:
97
+ """Fuse an already-validated `module` in place (see `Fusers.__getitem__`).
98
+
99
+ Builds the merged submodule and binds the compiled forward."""
100
+ self.update_attrs(module, prefix, vllm_config)
101
+ module.forward = types.MethodType(self.fused_forward, module)
102
+ return module
103
+
104
+
105
+@dataclass
106
+class StackedFuser(RewriteFuser):
107
+ """A fuser that merges sibling projections into one stacked linear and
108
+ rewrites the forward to call it."""
109
+
110
+ merged_name: ClassVar[str]
111
+ """Attribute name of the merged module created by `update_attrs`."""
112
+ merged_cls: ClassVar[str]
113
+ """Name of the vLLM class the merged projection becomes (for logging)."""
114
+
115
def info(self, name: str) -> str:
116
sources = " + ".join(shard for shard, _ in self.shards)
117
return (
@@ -108,26 +142,3 @@ class StackedFuser(BaseFuser):
142
"""`{merged_name: [projection names]}` so quantization can unpack the
143
fused layer into its per-shard configs."""
144
return {self.merged_name: [name for name, _ in self.shards]}
111
-
112
- @abstractmethod
113
- def update_forward(self, module: nn.Module) -> None:
114
- """Rewrite and compile `type(module)`'s forward source.
115
-
116
- Raises if the source does not admit the rewrite (fusion is then skipped).
117
- """
118
-
119
- @abstractmethod
120
- def update_attrs(
121
- self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
122
- ) -> None:
123
- """Replace `module`'s submodules with the merged module."""
124
-
125
- def fuse(
126
- self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
127
- ) -> nn.Module:
128
- """Fuse an already-validated `module` in place (see `Fusers.__getitem__`).
129
-
130
- Builds the merged submodule and binds the compiled forward."""
131
- self.update_attrs(module, prefix, vllm_config)
132
- module.forward = types.MethodType(self.fused_forward, module)
133
- return module
vllm/model_executor/models/transformers/fusers/packed_qkv.py
new
+163
@@ -0,0 +1,163 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Packed-QKV fuser: `c_attn(x).split((q, kv, kv))` -> a `QKVParallelLinear`."""
4
+
5
+import ast
6
+from dataclasses import dataclass
7
+from typing import TYPE_CHECKING
8
+
9
+from torch import fx, nn
10
+
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.linear import QKVParallelLinear
13
+from vllm.model_executor.models.transformers.fusers.base import (
14
+ RewriteFuser,
15
+ local_output_sizes,
16
+)
17
+from vllm.model_executor.models.transformers.fx_utils import (
18
+ compile_forward,
19
+ is_method,
20
+ recover_forward,
21
+ returned_linear,
22
+ upstream_linear,
23
+)
24
+from vllm.model_executor.models.transformers.utils import (
25
+ log_replacement,
26
+ replace_linear_class,
27
+)
28
+from vllm.model_executor.models.utils import maybe_prefix
29
+
30
+if TYPE_CHECKING:
31
+ from vllm.config import VllmConfig
32
+
33
+logger = init_logger(__name__)
34
+
35
+
36
+@dataclass
37
+class PackedQKVFuser(RewriteFuser):
38
+ """Fuser for attention with q, k and v packed into one projection."""
39
+
40
+ qkv_name: str
41
+ o_name: str | None
42
+ q_size: int
43
+ kv_size: int
44
+
45
+ def info(self, name: str) -> str:
46
+ return (
47
+ f"Fused: {self.qkv_name} ({name}: {self.source_cls}) -> QKVParallelLinear"
48
+ )
49
+
50
+ @staticmethod
51
+ def _packed_sizes(node: fx.Node) -> tuple[int, int] | None:
52
+ """`(q, kv)` from a `split((q, kv, kv), ...)` call, if it is one."""
53
+ if not is_method(node, "split") or len(node.args) < 2:
54
+ return None
55
+ sizes = node.args[1]
56
+ if not isinstance(sizes, (tuple, list)) or len(sizes) != 3:
57
+ return None
58
+ if not all(isinstance(size, int) for size in sizes):
59
+ return None
60
+ q_size, k_size, v_size = sizes
61
+ if k_size != v_size or q_size < k_size:
62
+ return None
63
+ return q_size, k_size
64
+
65
+ @classmethod
66
+ def match(cls, graph: fx.Graph, module: nn.Module) -> "PackedQKVFuser | None":
67
+ for node in graph.nodes:
68
+ if (sizes := cls._packed_sizes(node)) is None:
69
+ continue
70
+ q_size, kv_size = sizes
71
+ qkv_node = upstream_linear(node.args[0], module)
72
+ if qkv_node is None:
73
+ continue
74
+ qkv_name = str(qkv_node.target)
75
+ # The split must consume the whole projection.
76
+ if module.get_submodule(qkv_name).out_features != q_size + 2 * kv_size:
77
+ continue
78
+ # o_proj produces the module's output and consumes the query width.
79
+ o_name = returned_linear(graph, module)
80
+ if o_name == qkv_name or (
81
+ o_name is not None
82
+ and module.get_submodule(o_name).in_features != q_size
83
+ ):
84
+ o_name = None
85
+ return cls(
86
+ source_cls=type(module).__name__,
87
+ qkv_name=qkv_name,
88
+ o_name=o_name,
89
+ q_size=q_size,
90
+ kv_size=kv_size,
91
+ )
92
+ return None
93
+
94
+ def _split_call(self, funcdef: ast.FunctionDef) -> ast.Call:
95
+ """The unique `self.<qkv_name>(...)....split((a, b, c), ...)` call."""
96
+ calls = [
97
+ node
98
+ for node in ast.walk(funcdef)
99
+ if isinstance(node, ast.Call)
100
+ and isinstance(node.func, ast.Attribute)
101
+ and node.func.attr == "split"
102
+ and node.args
103
+ and isinstance(node.args[0], (ast.Tuple, ast.List))
104
+ and len(node.args[0].elts) == 3
105
+ and any(
106
+ isinstance(inner, ast.Attribute) and inner.attr == self.qkv_name
107
+ for inner in ast.walk(node.func.value)
108
+ )
109
+ ]
110
+ if len(calls) != 1:
111
+ raise ValueError(f"{self.qkv_name} has {len(calls)} three-way splits")
112
+ return calls[0]
113
+
114
+ def update_forward(self, module: nn.Module) -> None:
115
+ """Rewrite the split sizes to the sharded projection's per-rank widths."""
116
+ funcdef, fn = recover_forward(type(module))
117
+ split = self._split_call(funcdef)
118
+ # (q, kv, kv) -> [s // qkv.tp_size for s in qkv.output_sizes]
119
+ sections = local_output_sizes(self.qkv_name)
120
+ split.args[0] = ast.parse(sections, mode="eval").body
121
+ self.fused_forward = compile_forward(funcdef, fn)
122
+
123
+ def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool:
124
+ """Shapes must be compatible with a head-sharded packed GEMM."""
125
+ head_size = vllm_config.model_config.get_head_size()
126
+ qkv = module.get_submodule(self.qkv_name)
127
+ compatible = (
128
+ self.q_size % head_size == 0
129
+ and self.kv_size % head_size == 0
130
+ and qkv.out_features == self.q_size + 2 * self.kv_size
131
+ )
132
+ if not compatible:
133
+ logger.debug("%s is not compatible with packed QKV fusion", type(module))
134
+ return compatible
135
+
136
+ def update_attrs(
137
+ self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
138
+ ) -> None:
139
+ quant_config = vllm_config.quant_config
140
+ head_size = vllm_config.model_config.get_head_size()
141
+ qkv_prefix = maybe_prefix(prefix, self.qkv_name)
142
+ qkv = module.get_submodule(self.qkv_name)
143
+ merged = QKVParallelLinear(
144
+ hidden_size=qkv.in_features,
145
+ head_size=head_size,
146
+ total_num_heads=self.q_size // head_size,
147
+ total_num_kv_heads=self.kv_size // head_size,
148
+ bias=qkv.bias is not None,
149
+ quant_config=quant_config,
150
+ prefix=qkv_prefix,
151
+ return_bias=False,
152
+ )
153
+ setattr(module, self.qkv_name, merged)
154
+ log_replacement(qkv_prefix, qkv, merged)
155
+ # If there is an output projection, we know it must be rowwise.
156
+ if self.o_name is not None:
157
+ o_prefix = maybe_prefix(prefix, self.o_name)
158
+ o_proj = module.get_submodule(self.o_name)
159
+ new_o = replace_linear_class(
160
+ o_proj, "rowwise", quant_config, prefix=o_prefix
161
+ )
162
+ setattr(module, self.o_name, new_o)
163
+ log_replacement(o_prefix, o_proj, new_o)
vllm/model_executor/models/transformers/fusers/qkv.py
+5
-2
@@ -10,7 +10,10 @@ from torch import fx, nn
10
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.linear import QKVParallelLinear
13
-from vllm.model_executor.models.transformers.fusers.base import StackedFuser
13
+from vllm.model_executor.models.transformers.fusers.base import (
14
+ StackedFuser,
15
+ local_output_sizes,
16
+)
17
from vllm.model_executor.models.transformers.fx_utils import (
18
compile_forward,
19
innermost_block,
@@ -134,7 +137,7 @@ class QKVFuser(StackedFuser):
137
if names & set(temps):
138
raise ValueError("fused temporaries would shadow existing names")
139
merged = f"self.{self.merged_name}"
137
- sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
140
+ sections = local_output_sizes(self.merged_name)
141
template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
142
assign = ast.parse(template).body[0]
143
arg = next(
vllm/model_executor/models/transformers/fx_utils.py
+9
-3
@@ -394,8 +394,10 @@ def output_value(graph: fx.Graph) -> object | None:
394
def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:
395
"""Nearest linear producing `node`, walking back through splits/reshapes.
396
397
- Never walks through a leaf call (e.g. an attention interface): its inputs
398
- are what attention consumes, not what produced the value."""
397
+ Non-linear submodules are transparent too (e.g. the dropout GPT-style
398
+ attentions apply after their output projection). Never walks through a leaf
399
+ call (e.g. an attention interface): its inputs are what attention consumes,
400
+ not what produced the value."""
401
stack = [node]
402
seen: set[fx.Node] = set()
403
while stack:
@@ -405,7 +407,11 @@ def upstream_linear(node: object, module: nn.Module) -> fx.Node | None:
407
seen.add(current)
408
if is_linear(current, module):
409
return current
408
- if current.op in ("call_function", "call_method") and not is_leaf_call(current):
410
+ if current.op in (
411
+ "call_function",
412
+ "call_method",
413
+ "call_module",
414
+ ) and not is_leaf_call(current):
415
stack.extend(current.args)
416
return None
417