Make the Transformers modeling backend as fast as native vLLM (#47187)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Harry Mellor committed
Jul 6, 2026 at 16:59 UTC
5bce653e09ca62c870ea18d01a4180dc48d3bacb
23 files changed
+2679
-160
.buildkite/test-amd.yaml
+4
-4
@@ -472,7 +472,7 @@ steps:
472
commands:
473
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
474
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
475
- - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
475
+ - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
476
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
477
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
478
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
@@ -1637,10 +1637,10 @@ steps:
1637
source_file_dependencies:
1638
- vllm/
1639
- tests/models/test_terratorch.py
1640
- - tests/models/test_transformers.py
1640
+ - tests/models/transformers/test_backend.py
1641
- tests/models/test_registry.py
1642
commands:
1643
- - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
1643
+ - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
1644
1645
#----------------------------------------------------- mi300 · models / language -----------------------------------------------------#
1646
@@ -1860,7 +1860,7 @@ steps:
1860
- examples/
1861
commands:
1862
- pip install --upgrade git+https://github.com/huggingface/transformers
1863
- - pytest -v -s tests/models/test_transformers.py
1863
+ - pytest -v -s tests/models/transformers/test_backend.py
1864
- pytest -v -s tests/models/multimodal/test_mapping.py
1865
- python3 examples/basic/offline_inference/chat.py
1866
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
.buildkite/test_areas/models_basic.yaml
+4
-3
@@ -36,10 +36,10 @@ steps:
36
source_file_dependencies:
37
- vllm/
38
- tests/models/test_terratorch.py
39
- - tests/models/test_transformers.py
39
+ - tests/models/transformers/test_backend.py
40
- tests/models/test_registry.py
41
commands:
42
- - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
42
+ - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
43
mirror:
44
amd:
45
device: mi325_1
@@ -55,6 +55,7 @@ steps:
55
- vllm/
56
- tests/models/test_utils.py
57
- tests/models/test_vision.py
58
+ - tests/models/transformers/fusers/
59
device: cpu-small
60
commands:
60
- - pytest -v -s models/test_utils.py models/test_vision.py
61
+ - pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
.buildkite/test_areas/models_distributed.yaml
+1
-1
@@ -17,7 +17,7 @@ steps:
17
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
18
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
19
# Avoid importing model tests that cause CUDA reinitialization error
20
- - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
20
+ - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
21
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
22
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
23
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
.github/CODEOWNERS
+1
-1
@@ -119,7 +119,7 @@
119
120
# Transformers modeling backend
121
/vllm/model_executor/models/transformers @hmellor
122
-/tests/models/test_transformers.py @hmellor
122
+/tests/models/transformers @hmellor
123
124
# Docs
125
/docs/mkdocs @hmellor
docs/models/supported_models.md
+4
-2
@@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang
15
16
### Transformers
17
18
-vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend".
18
+vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation.
19
20
Currently, the Transformers modeling backend works for the following:
21
@@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded:
140
141
That's it!
142
143
-For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
143
+For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
144
145
<details class="code">
146
<summary>configuration_my_model.py</summary>
@@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
168
</details>
169
170
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported).
171
+ - vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated.
172
- `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s:
173
- You only need to do this for layers which are not present on all pipeline stages
174
- vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages
175
+ - When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage
176
- The `list` in the first element of the `tuple` contains the names of the input arguments
177
- The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
178
tests/models/transformers/__init__.py
tests/models/transformers/fusers/__init__.py
tests/models/transformers/fusers/test_linear.py
new
+480
@@ -0,0 +1,480 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the Transformers modeling backend's linear fusers."""
4
+
5
+import inspect
6
+from types import MethodType, SimpleNamespace
7
+
8
+import pytest
9
+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
14
+from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser
15
+
16
+
17
+class SiluAndMulStub(nn.Module):
18
+ """Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""
19
+
20
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
21
+ d = x.shape[-1] // 2
22
+ return F.silu(x[..., :d]) * x[..., d:]
23
+
24
+
25
+class NoDownGLU(nn.Module):
26
+ """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None."""
27
+
28
+ def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
29
+ super().__init__()
30
+ self.gate_proj = nn.Linear(hidden, inter, bias=bias)
31
+ self.up_proj = nn.Linear(hidden, inter, bias=bias)
32
+ self.act_fn = nn.SiLU()
33
+
34
+ def forward(self, x):
35
+ return self.act_fn(self.gate_proj(x)) * self.up_proj(x)
36
+
37
+
38
+class GLUMLP(NoDownGLU):
39
+ """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP."""
40
+
41
+ def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
42
+ super().__init__(hidden, inter, bias)
43
+ self.down_proj = nn.Linear(inter, hidden, bias=bias)
44
+
45
+ def forward(self, x):
46
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
47
+
48
+
49
+class ReversedGLUMLP(GLUMLP):
50
+ """`up(x) * act(gate(x))` — operands swapped (multiply is commutative)."""
51
+
52
+ def forward(self, x):
53
+ return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x)))
54
+
55
+
56
+class NotAnMLP(nn.Module):
57
+ """Two linears but no activation*linear multiply -> must not match."""
58
+
59
+ def __init__(self):
60
+ super().__init__()
61
+ self.fc1 = nn.Linear(8, 8)
62
+ self.fc2 = nn.Linear(8, 8)
63
+
64
+ def forward(self, x):
65
+ return self.fc2(self.fc1(x))
66
+
67
+
68
+class NotAnActGLUMLP(GLUMLP):
69
+ """GLU-shaped, but the "activation" is not a known activation module."""
70
+
71
+ def __init__(self):
72
+ super().__init__()
73
+ self.act_fn = nn.Dropout()
74
+
75
+
76
+class UntraceableMLP(GLUMLP):
77
+ """Data-dependent control flow *before* the GLU -> no match."""
78
+
79
+ def forward(self, x):
80
+ if x.sum() > 0: # noqa: SIM108 - intentionally untraceable
81
+ return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
82
+ return x
83
+
84
+
85
+class UntraceableTailGLUMLP(GLUMLP):
86
+ """Data-dependent control flow *after* the GLU -> still fusable."""
87
+
88
+ def forward(self, x):
89
+ y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
90
+ if y.sum() > torch.inf: # intentionally untraceable
91
+ y = y * 0
92
+ return y
93
+
94
+
95
+class FakeAttention(nn.Module):
96
+ """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""
97
+
98
+ is_causal = True
99
+
100
+ def __init__(
101
+ self,
102
+ hidden: int = 32,
103
+ head_dim: int = 8,
104
+ heads: int = 4,
105
+ kv_heads: int = 4,
106
+ bias: bool = False,
107
+ layer_idx: int = 0,
108
+ ):
109
+ super().__init__()
110
+ self.config = SimpleNamespace(_attn_implementation="vllm")
111
+ self.layer_idx = layer_idx
112
+ self.head_dim = head_dim
113
+ self.scaling = head_dim**-0.5
114
+ self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias)
115
+ self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
116
+ self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
117
+ self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias)
118
+
119
+ def forward(
120
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
121
+ ):
122
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
123
+
124
+ input_shape = hidden_states.shape[:-1]
125
+ hidden_shape = (*input_shape, -1, self.head_dim)
126
+ q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
127
+ k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
128
+ v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
129
+ if past_key_values is not None:
130
+ k, v = past_key_values.update(k, v, self.layer_idx)
131
+ attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
132
+ self.config._attn_implementation, None
133
+ )
134
+ attn_output, attn_weights = attention_interface(
135
+ self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
136
+ )
137
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
138
+ return self.o_proj(attn_output), attn_weights
139
+
140
+
141
+class ReversedFakeAttention(FakeAttention):
142
+ """Projections computed in (v, k, q) order — q must still be identified."""
143
+
144
+ def forward(
145
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
146
+ ):
147
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
148
+
149
+ input_shape = hidden_states.shape[:-1]
150
+ hidden_shape = (*input_shape, -1, self.head_dim)
151
+ v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
152
+ k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
153
+ q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
154
+ attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
155
+ self.config._attn_implementation, None
156
+ )
157
+ attn_output, _ = attention_interface(
158
+ self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
159
+ )
160
+ return self.o_proj(attn_output.reshape(*input_shape, -1)), None
161
+
162
+
163
+class ExtraProjAttention(FakeAttention):
164
+ """A second non-qkv linear of a different width -> `o_proj` still found."""
165
+
166
+ def __init__(self, **kwargs):
167
+ super().__init__(**kwargs)
168
+ self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
169
+
170
+
171
+class QKNormAttention(FakeAttention):
172
+ """OLMoE-style: a full-dim norm applied to the whole q/k projection output."""
173
+
174
+ def __init__(self, **kwargs):
175
+ super().__init__(**kwargs)
176
+ self.q_norm = nn.RMSNorm(self.q_proj.out_features)
177
+ self.k_norm = nn.RMSNorm(self.k_proj.out_features)
178
+
179
+ def forward(
180
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
181
+ ):
182
+ q = self.q_norm(self.q_proj(hidden_states))
183
+ k = self.k_norm(self.k_proj(hidden_states))
184
+ v = self.v_proj(hidden_states)
185
+ return self.o_proj(q + k + v), None
186
+
187
+
188
+class PerHeadQKNormAttention(FakeAttention):
189
+ """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape."""
190
+
191
+ def __init__(self, **kwargs):
192
+ super().__init__(**kwargs)
193
+ self.q_norm = nn.RMSNorm(self.head_dim)
194
+ self.k_norm = nn.RMSNorm(self.head_dim)
195
+
196
+ def forward(
197
+ self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
198
+ ):
199
+ shape = (*hidden_states.shape[:-1], -1, self.head_dim)
200
+ q = self.q_norm(self.q_proj(hidden_states).view(shape))
201
+ k = self.k_norm(self.k_proj(hidden_states).view(shape))
202
+ v = self.v_proj(hidden_states).view(shape)
203
+ return self.o_proj((q + k + v).flatten(-2)), None
204
+
205
+
206
+class FakeSelfAttn(nn.Module):
207
+ """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""
208
+
209
+ def __init__(self):
210
+ super().__init__()
211
+ self.impl = SimpleNamespace(scale=None)
212
+
213
+ def forward(self, q, k, v):
214
+ # MHA-shaped stub: any deterministic combination of q/k/v will do
215
+ return q + 2 * k + 3 * v
216
+
217
+
218
+@pytest.fixture(autouse=True)
219
+def _clear_fuser_cache():
220
+ get_fuser.cache_clear()
221
+ yield
222
+ get_fuser.cache_clear()
223
+
224
+
225
+def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
226
+ """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul)."""
227
+ gate = module.get_submodule(fuser.gate_name)
228
+ up = module.get_submodule(fuser.up_name)
229
+ merged = nn.Linear(
230
+ gate.in_features,
231
+ gate.out_features + up.out_features,
232
+ bias=gate.bias is not None,
233
+ )
234
+ with torch.no_grad():
235
+ merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0))
236
+ if gate.bias is not None:
237
+ merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0))
238
+ setattr(module, fuser.merged_name, merged)
239
+ setattr(module, fuser.act_name, SiluAndMulStub())
240
+ delattr(module, fuser.gate_name)
241
+ delattr(module, fuser.up_name)
242
+ module.forward = MethodType(fuser.fused_forward, module)
243
+ return module
244
+
245
+
246
+def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
247
+ """Apply a fuser using a plain merged `nn.Linear` (no TP sharding)."""
248
+ q, k, v = (
249
+ module.get_submodule(name)
250
+ for name in (fuser.q_name, fuser.k_name, fuser.v_name)
251
+ )
252
+ merged = nn.Linear(
253
+ q.in_features,
254
+ q.out_features + k.out_features + v.out_features,
255
+ bias=q.bias is not None,
256
+ )
257
+ with torch.no_grad():
258
+ merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
259
+ if q.bias is not None:
260
+ merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
261
+ merged.split_sizes = [q.out_features, k.out_features, v.out_features]
262
+ setattr(module, fuser.merged_name, merged)
263
+ for name in (fuser.q_name, fuser.k_name, fuser.v_name):
264
+ delattr(module, name)
265
+ module.forward = MethodType(fuser.fused_forward, module)
266
+ return module
267
+
268
+
269
+@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
270
+@pytest.mark.parametrize("bias", [False, True])
271
+def test_detects_and_rewrites_glu(mlp_cls, bias):
272
+ with torch.device("meta"):
273
+ meta = mlp_cls(bias=bias)
274
+ fuser = get_fuser(meta)
275
+ assert isinstance(fuser, GLUFuser)
276
+ assert (
277
+ fuser.gate_name,
278
+ fuser.up_name,
279
+ fuser.act_name,
280
+ fuser.down_name,
281
+ ) == ("gate_proj", "up_proj", "act_fn", "down_proj")
282
+
283
+ # The rewritten forward references the merged projection instead of the
284
+ # sources; the rest of the forward is untouched.
285
+ names = fuser.fused_forward.__code__.co_names
286
+ assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names
287
+ assert not {"gate_proj", "up_proj"} & set(names)
288
+
289
+ # Numerics: the fused forward must match the original on a real instance.
290
+ real = mlp_cls(bias=bias)
291
+ for p in real.parameters():
292
+ nn.init.normal_(p, std=0.05)
293
+ x = torch.randn(4, 16)
294
+ expected = real(x)
295
+ fused = _apply_glu_fuser_with_stubs(real, fuser)
296
+
297
+ # Fusion is in place: the module keeps its class and other attributes
298
+ assert fused is real and type(fused) is mlp_cls
299
+ torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
300
+
301
+
302
+def test_glu_identifies_down_projection():
303
+ """The row projection consuming `act(gate(x)) * up(x)` is identified.
304
+
305
+ It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
306
+ matches the column-parallel merged gate/up; `None` when there is no such
307
+ projection to force (fusion of gate/up still applies)."""
308
+ with torch.device("meta"):
309
+ assert get_fuser(GLUMLP()).down_name == "down_proj"
310
+ assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
311
+ assert get_fuser(NoDownGLU()).down_name is None
312
+
313
+
314
+@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
315
+@pytest.mark.parametrize("kv_heads", [4, 2])
316
+def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
317
+ if attn_cls is ReversedFakeAttention and kv_heads == 4:
318
+ pytest.skip("MHA q/k/v assignment is order-based by design")
319
+ with torch.device("meta"):
320
+ meta = attn_cls(kv_heads=kv_heads)
321
+ fuser = get_fuser(meta)
322
+ assert isinstance(fuser, QKVFuser)
323
+ # q (sharded differently under TP) must be identified exactly; k/v may be
324
+ # swapped for non-canonical compute order, which is numerically consistent
325
+ # because the weight mapping and the split indices follow the same
326
+ # assignment.
327
+ assert fuser.q_name == "q_proj"
328
+ assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"}
329
+ assert fuser.o_name == "o_proj"
330
+
331
+ # The projections are merged; everything else stays live Python with its
332
+ # original semantics (branches, kwargs, attribute reads)
333
+ code = fuser.fused_forward.__code__
334
+ names = code.co_names
335
+ assert "qkv_proj" in names and "split_sizes" in names and "o_proj" in names
336
+ assert not {"q_proj", "k_proj", "v_proj"} & set(names)
337
+ if attn_cls is FakeAttention:
338
+ assert "update" in names # the cache branch survives
339
+ assert code.co_flags & inspect.CO_VARKEYWORDS # **kwargs survives
340
+
341
+ # Numerics: the fused forward must match the original on a real instance,
342
+ # with a different layer_idx than the traced instance (kv_heads == heads so
343
+ # the q/k/v stub combination is shape-compatible).
344
+ real = attn_cls(kv_heads=4, layer_idx=3)
345
+ for p in real.parameters():
346
+ nn.init.normal_(p, std=0.05)
347
+ x = torch.randn(1, 5, 32)
348
+ attention_instances = {3: FakeSelfAttn()}
349
+ expected, _ = real(x, attention_instances=attention_instances)
350
+ fused = _apply_qkv_fuser_with_stubs(real, fuser)
351
+
352
+ # Fusion is in place: the module keeps its class and other attributes
353
+ assert fused is real and type(fused) is attn_cls
354
+ assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
355
+ out, _ = fused(x, attention_instances=attention_instances)
356
+ torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
357
+
358
+
359
+def test_qkv_identifies_output_projection():
360
+ with torch.device("meta"):
361
+ assert get_fuser(FakeAttention()).o_name == "o_proj"
362
+ assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
363
+ assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
364
+ # Norm children (q_norm/k_norm) must not disturb o_proj identification.
365
+ assert get_fuser(QKNormAttention()).o_name == "o_proj"
366
+ assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"
367
+
368
+
369
+def test_fuser_is_cached_per_class():
370
+ with torch.device("meta"):
371
+ fuser_a = get_fuser(GLUMLP())
372
+ fuser_b = get_fuser(GLUMLP())
373
+ assert fuser_a is fuser_b
374
+ assert GLUMLP in get_fuser.cache
375
+
376
+
377
+@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
378
+def test_non_matching_modules_return_none(cls):
379
+ with torch.device("meta"):
380
+ module = cls()
381
+ assert get_fuser(module) is None
382
+
383
+
384
+def test_untraceable_tail_still_fuses():
385
+ with torch.device("meta"):
386
+ meta = UntraceableTailGLUMLP()
387
+ fuser = get_fuser(meta)
388
+ assert isinstance(fuser, GLUFuser)
389
+
390
+ # Numerics: the live tail must survive the rewrite
391
+ real = UntraceableTailGLUMLP()
392
+ for p in real.parameters():
393
+ nn.init.normal_(p, std=0.05)
394
+ x = torch.randn(4, 16)
395
+ expected = real(x)
396
+ fused = _apply_glu_fuser_with_stubs(real, fuser)
397
+ torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)
398
+
399
+
400
+def test_weight_mappings_are_scoped_to_fused_prefixes():
401
+ from vllm.model_executor.models.utils import WeightsMapper
402
+
403
+ with torch.device("meta"):
404
+ glu_fuser = get_fuser(GLUMLP())
405
+ qkv_fuser = get_fuser(FakeAttention())
406
+
407
+ mapper = WeightsMapper()
408
+ for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
409
+ mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
410
+ mapper.orig_to_new_stacked.update(
411
+ qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn")
412
+ )
413
+
414
+ names = [
415
+ "model.layers.0.mlp.gate_proj.weight",
416
+ "model.layers.0.mlp.up_proj.weight",
417
+ "model.layers.1.mlp.gate_proj.weight",
418
+ "model.layers.0.self_attn.q_proj.weight",
419
+ "model.layers.0.self_attn.k_proj.weight",
420
+ "model.layers.0.self_attn.v_proj.weight",
421
+ # Unfused modules at other prefixes must be left untouched.
422
+ "model.layers.2.mlp.experts.0.gate_proj.weight",
423
+ "model.layers.1.self_attn.q_proj.weight",
424
+ ]
425
+ # `apply` rewrites the name and stamps the shard id onto each tensor.
426
+ weights = [(name, torch.empty(0)) for name in names]
427
+ mapped = list(mapper.apply(weights))
428
+ mapped_names = [name for name, _ in mapped]
429
+ shard_ids = [getattr(data, "shard_id", None) for _, data in mapped]
430
+
431
+ assert mapped_names == [
432
+ "model.layers.0.mlp.gate_up_proj.weight",
433
+ "model.layers.0.mlp.gate_up_proj.weight",
434
+ "model.layers.1.mlp.gate_up_proj.weight",
435
+ "model.layers.0.self_attn.qkv_proj.weight",
436
+ "model.layers.0.self_attn.qkv_proj.weight",
437
+ "model.layers.0.self_attn.qkv_proj.weight",
438
+ # Only the exact fused layers are remapped; everything else is untouched.
439
+ "model.layers.2.mlp.experts.0.gate_proj.weight",
440
+ "model.layers.1.self_attn.q_proj.weight",
441
+ ]
442
+ assert shard_ids == [0, 1, 0, "q", "k", "v", None, None]
443
+
444
+ # The fused layers are exposed to the quantization machinery via their
445
+ # original constituent projection names (what the checkpoint stores).
446
+ assert glu_fuser.packed_modules_mapping == {
447
+ "gate_up_proj": ["gate_proj", "up_proj"],
448
+ }
449
+ assert qkv_fuser.packed_modules_mapping == {
450
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
451
+ }
452
+
453
+
454
+@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
455
+def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
456
+ with torch.device("meta"):
457
+ module = cls()
458
+ fuser = get_fuser(module)
459
+ # Either no pattern matches the class, or this instance fails validation
460
+ # (`recursive_replace` gates fusion and its weight mappings on `validate`)
461
+ model_config = default_vllm_config.model_config
462
+ assert fuser is None or not fuser.validate(module, model_config)
463
+
464
+
465
+def test_act_and_mul_derived_from_module(default_vllm_config):
466
+ from transformers.activations import GELUTanh, SiLUActivation
467
+
468
+ from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul
469
+
470
+ assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
471
+ assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
472
+ gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh())
473
+ assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
474
+ gelu = GLUFuser._get_act_and_mul(nn.GELU())
475
+ assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
476
+ # Not activations at all -> no fusion
477
+ assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
478
+ assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
479
+ with pytest.raises(ValueError, match="No AndMul equivalent"):
480
+ GLUFuser._get_act_and_mul(nn.Dropout())
tests/models/transformers/fusers/test_moe.py
new
+299
@@ -0,0 +1,299 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the Transformers modeling backend's MoE fuser."""
4
+
5
+import pytest
6
+import torch
7
+import torch.nn as nn
8
+import torch.nn.functional as F
9
+
10
+from vllm.model_executor.models.transformers.fusers import MoEBlockFuser
11
+
12
+from .test_linear import GLUMLP
13
+
14
+
15
+class TopKRouter(nn.Module):
16
+ """HF v5 top-k router: `linear -> softmax -> topk (-> renorm)`."""
17
+
18
+ def __init__(self, num_experts=8, hidden=16, top_k=2, sigmoid=False):
19
+ super().__init__()
20
+ self.top_k = top_k
21
+ self.sigmoid = sigmoid
22
+ self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
23
+
24
+ def forward(self, hidden_states):
25
+ logits = F.linear(hidden_states, self.weight)
26
+ scores = torch.sigmoid(logits) if self.sigmoid else F.softmax(logits, dim=-1)
27
+ value, index = torch.topk(scores, self.top_k, dim=-1)
28
+ value = value / value.sum(dim=-1, keepdim=True)
29
+ return logits, value, index
30
+
31
+
32
+class CorrectionRouter(nn.Module):
33
+ """Grouped router with a score-correction bias buffer (DeepSeek-V3) -> declined."""
34
+
35
+ def __init__(self, num_experts=8, hidden=16):
36
+ super().__init__()
37
+ self.weight = nn.Parameter(torch.zeros(num_experts, hidden))
38
+ self.register_buffer("e_score_correction_bias", torch.zeros(num_experts))
39
+
40
+ def forward(self, hidden_states):
41
+ logits = F.linear(hidden_states, self.weight)
42
+ scores = torch.sigmoid(logits) + self.e_score_correction_bias
43
+ _, index = torch.topk(scores, 2, dim=-1)
44
+ return logits, scores, index
45
+
46
+
47
+class BiasedRouter(TopKRouter):
48
+ """A valid top-k router but not `weight`-only (extra `bias` param) -> declined."""
49
+
50
+ def __init__(self):
51
+ super().__init__()
52
+ self.bias = nn.Parameter(torch.zeros(8))
53
+
54
+ def forward(self, hidden_states):
55
+ logits = F.linear(hidden_states, self.weight) + self.bias
56
+ scores = F.softmax(logits, dim=-1)
57
+ value, index = torch.topk(scores, self.top_k, dim=-1)
58
+ return logits, value, index
59
+
60
+
61
+class DisconnectedRouter(TopKRouter):
62
+ """linear+softmax+top-k present but top-k ignores the logits -> not a router."""
63
+
64
+ def forward(self, hidden_states):
65
+ logits = F.linear(hidden_states, self.weight)
66
+ _ = F.softmax(logits, dim=-1) # scored, but not consumed by top-k
67
+ value, index = torch.topk(hidden_states, self.top_k, dim=-1)
68
+ return logits, value, index
69
+
70
+
71
+class MoEExperts(nn.Module):
72
+ """Packed experts (3D weights); only its name (`experts`) matters here."""
73
+
74
+ def __init__(self, num_experts=8, hidden=16, inter=32):
75
+ super().__init__()
76
+ self.gate_up_proj = nn.Parameter(torch.zeros(num_experts, 2 * inter, hidden))
77
+ self.down_proj = nn.Parameter(torch.zeros(num_experts, hidden, inter))
78
+
79
+ def forward(self, hidden_states, index, weights):
80
+ return hidden_states
81
+
82
+
83
+class MoEBlock(nn.Module):
84
+ """Single-tensor MoE block (Qwen3-style); subclasses override `_shared`."""
85
+
86
+ def __init__(self, router_cls=TopKRouter):
87
+ super().__init__()
88
+ self.experts = MoEExperts()
89
+ self.gate = router_cls()
90
+
91
+ def _shared(self, x, logits):
92
+ """The term added to the experts' output (none for a plain block)."""
93
+ return 0
94
+
95
+ def forward(self, hidden_states):
96
+ x = hidden_states.reshape(-1, hidden_states.shape[-1])
97
+ logits, weights, index = self.gate(x)
98
+ out = self.experts(x, index, weights) + self._shared(x, logits)
99
+ return out.reshape(hidden_states.shape)
100
+
101
+
102
+class MoEBlockNoShared(MoEBlock):
103
+ """No shared-expert child but a gate-derived add -> trace skipped, still fuses."""
104
+
105
+ def _shared(self, x, logits):
106
+ return logits.sum()
107
+
108
+
109
+class MoEBlockShared(MoEBlock):
110
+ """A block with a shared expert and its sigmoid gate (Qwen2-style)."""
111
+
112
+ def __init__(self):
113
+ super().__init__()
114
+ self.shared_expert = GLUMLP()
115
+ self.shared_expert_gate = nn.Linear(16, 1, bias=False)
116
+
117
+ def _shared(self, x, logits):
118
+ return torch.sigmoid(self.shared_expert_gate(x)) * self.shared_expert(x)
119
+
120
+
121
+class MoEBlockSharedNoGate(MoEBlock):
122
+ """A block with an ungated shared expert -> native, shared passed through."""
123
+
124
+ def __init__(self):
125
+ super().__init__()
126
+ self.shared_expert = GLUMLP()
127
+
128
+ def _shared(self, x, logits):
129
+ return self.shared_expert(x)
130
+
131
+
132
+class MoEBlockTuple(MoEBlock):
133
+ """A tuple-returning block (gpt-oss-style) -> must decline."""
134
+
135
+ def forward(self, hidden_states):
136
+ x = hidden_states.reshape(-1, hidden_states.shape[-1])
137
+ _, weights, index = self.gate(x)
138
+ return self.experts(x, index, weights), index
139
+
140
+
141
+class MoEBlockTupleVar(MoEBlock):
142
+ """Returns a name bound to a tuple, not a literal tuple -> must still decline."""
143
+
144
+ def forward(self, hidden_states):
145
+ x = hidden_states.reshape(-1, hidden_states.shape[-1])
146
+ _, weights, index = self.gate(x)
147
+ result = self.experts(x, index, weights), index
148
+ return result
149
+
150
+
151
+class MoEBlockNestedTupleReturn(MoEBlock):
152
+ """Tuple `return` in a nested helper; block returns one tensor -> still fuses."""
153
+
154
+ def forward(self, hidden_states):
155
+ def keep(a, b):
156
+ return a, b
157
+
158
+ x = hidden_states.reshape(-1, hidden_states.shape[-1])
159
+ _, weights, index = self.gate(x)
160
+ out, _ = keep(self.experts(x, index, weights), index)
161
+ return out.reshape(hidden_states.shape)
162
+
163
+
164
+class PlainMLP(nn.Module):
165
+ """A non-GLU FFN: `down(act(up(x)))`, no gating multiply."""
166
+
167
+ def __init__(self, hidden: int = 16, inter: int = 32):
168
+ super().__init__()
169
+ self.up_proj = nn.Linear(hidden, inter, bias=False)
170
+ self.down_proj = nn.Linear(inter, hidden, bias=False)
171
+ self.act_fn = nn.SiLU()
172
+
173
+ def forward(self, x):
174
+ return self.down_proj(self.act_fn(self.up_proj(x)))
175
+
176
+
177
+class MoEBlockSharedNonGLU(MoEBlock):
178
+ """A non-GLU shared expert -> detected by dataflow (no gate/up merge)."""
179
+
180
+ def __init__(self):
181
+ super().__init__()
182
+ self.shared_expert = PlainMLP()
183
+
184
+ def _shared(self, x, logits):
185
+ return self.shared_expert(x)
186
+
187
+
188
+class MoEBlockUnaccounted(MoEBlock):
189
+ """A weight-bearing child outside the fused dataflow (pre-router) -> declined."""
190
+
191
+ def __init__(self):
192
+ super().__init__()
193
+ self.extra = nn.Linear(16, 16, bias=False)
194
+
195
+ def forward(self, hidden_states):
196
+ x = self.extra(hidden_states.reshape(-1, hidden_states.shape[-1]))
197
+ _, weights, index = self.gate(x)
198
+ return self.experts(x, index, weights).reshape(hidden_states.shape)
199
+
200
+
201
+class BufferScale(nn.Module):
202
+ """A stateful child carrying only a buffer (no parameters)."""
203
+
204
+ def __init__(self, hidden: int = 16):
205
+ super().__init__()
206
+ self.register_buffer("scale", torch.ones(hidden))
207
+
208
+ def forward(self, x):
209
+ return x * self.scale
210
+
211
+
212
+class MoEBlockUnaccountedBuffer(MoEBlockUnaccounted):
213
+ """Like `MoEBlockUnaccounted`, but the extra child holds only a buffer."""
214
+
215
+ def __init__(self):
216
+ super().__init__()
217
+ self.extra = BufferScale()
218
+
219
+
220
+@pytest.mark.parametrize("sigmoid", [False, True])
221
+def test_moe_fuser_detects_router(sigmoid):
222
+ with torch.device("meta"):
223
+ block = MoEBlock(lambda: TopKRouter(sigmoid=sigmoid))
224
+ fuser = MoEBlockFuser.match(block, "experts")
225
+ assert isinstance(fuser, MoEBlockFuser)
226
+ assert fuser.gate_name == "gate"
227
+ assert fuser.scoring_func == ("sigmoid" if sigmoid else "softmax")
228
+ assert fuser.shared_name is None and fuser.shared_gate_name is None
229
+
230
+
231
+def test_moe_fuser_detects_shared_experts():
232
+ with torch.device("meta"):
233
+ block = MoEBlockShared()
234
+ fuser = MoEBlockFuser.match(block, "experts")
235
+ assert isinstance(fuser, MoEBlockFuser)
236
+ assert fuser.shared_name == "shared_expert"
237
+ assert fuser.shared_gate_name == "shared_expert_gate"
238
+
239
+
240
+def test_moe_fuser_skips_shared_detection_without_extra_children():
241
+ """With only experts and gate, shared-expert detection (and its block trace)
242
+ is skipped, so a gate-derived add is not misread as a shared expert."""
243
+ with torch.device("meta"):
244
+ block = MoEBlockNoShared()
245
+ fuser = MoEBlockFuser.match(block, "experts")
246
+ assert isinstance(fuser, MoEBlockFuser)
247
+ assert fuser.shared_name is None and fuser.shared_gate_name is None
248
+
249
+
250
+def test_moe_fuser_shared_without_gate():
251
+ with torch.device("meta"):
252
+ block = MoEBlockSharedNoGate()
253
+ fuser = MoEBlockFuser.match(block, "experts")
254
+ assert isinstance(fuser, MoEBlockFuser)
255
+ assert fuser.shared_name == "shared_expert"
256
+ assert fuser.shared_gate_name is None
257
+
258
+
259
+def test_moe_fuser_detects_non_glu_shared_expert():
260
+ with torch.device("meta"):
261
+ block = MoEBlockSharedNonGLU()
262
+ fuser = MoEBlockFuser.match(block, "experts")
263
+ assert isinstance(fuser, MoEBlockFuser)
264
+ # Recognised by dataflow (added to the experts' output), though not a GLU.
265
+ assert fuser.shared_name == "shared_expert"
266
+ assert fuser.shared_gate_name is None
267
+
268
+
269
+@pytest.mark.parametrize(
270
+ "block_cls",
271
+ [
272
+ lambda: MoEBlock(CorrectionRouter), # score-correction buffer (grouped)
273
+ lambda: MoEBlock(BiasedRouter), # router not weight-only (extra param)
274
+ MoEBlockTuple, # tuple-returning block (e.g. gpt-oss)
275
+ MoEBlockTupleVar, # tuple returned via a name binding, not a literal
276
+ MoEBlockUnaccounted, # weight-bearing child outside the fused dataflow
277
+ MoEBlockUnaccountedBuffer, # buffer-only child outside the fused dataflow
278
+ ],
279
+)
280
+def test_moe_fuser_declines_unsupported(block_cls):
281
+ with torch.device("meta"):
282
+ block = block_cls()
283
+ assert MoEBlockFuser.match(block, "experts") is None
284
+
285
+
286
+def test_moe_fuser_ignores_nested_returns():
287
+ """A tuple `return` inside a nested helper must not decline a block whose own
288
+ forward returns a single tensor."""
289
+ with torch.device("meta"):
290
+ block = MoEBlockNestedTupleReturn()
291
+ assert isinstance(MoEBlockFuser.match(block, "experts"), MoEBlockFuser)
292
+
293
+
294
+def test_moe_fuser_router_requires_connected_dataflow():
295
+ """A gate with linear + softmax + top-k present but not wired as a router
296
+ (top-k selects over the input, not the scored logits) is not detected."""
297
+ with torch.device("meta"):
298
+ block = MoEBlock(DisconnectedRouter)
299
+ assert MoEBlockFuser.match(block, "experts") is None
tests/models/transformers/fusers/test_rms_norm.py
new
+227
@@ -0,0 +1,227 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the Transformers modeling backend's RMSNorm fuser."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+import torch.nn as nn
10
+import torch.nn.functional as F
11
+
12
+from vllm.model_executor.models.transformers.fuser import get_fuser
13
+from vllm.model_executor.models.transformers.fusers import RMSNormFuser
14
+
15
+
16
+class RMSNorm(nn.Module):
17
+ """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`."""
18
+
19
+ def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True):
20
+ super().__init__()
21
+ if weight:
22
+ self.weight = nn.Parameter(torch.ones(hidden))
23
+ self.variance_epsilon = eps
24
+
25
+ def _rms(self, x):
26
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
27
+
28
+ def forward(self, x):
29
+ return self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
30
+
31
+
32
+class GemmaRMSNorm(RMSNorm):
33
+ """Zero-centered weight: `(1 + weight) * normalized`."""
34
+
35
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
36
+ super().__init__(hidden, eps)
37
+ self.weight = nn.Parameter(torch.zeros(hidden))
38
+
39
+ def forward(self, x):
40
+ return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype)
41
+
42
+
43
+class WeightlessRMSNorm(RMSNorm):
44
+ """No scale parameter (e.g. Gemma3n `with_scale=False`)."""
45
+
46
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
47
+ super().__init__(hidden, eps, weight=False)
48
+
49
+ def forward(self, x):
50
+ return self._rms(x.to(torch.float32)).to(x.dtype)
51
+
52
+
53
+class LayerNorm(RMSNorm):
54
+ """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast)."""
55
+
56
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
57
+ super().__init__(hidden, eps)
58
+
59
+ def forward(self, x):
60
+ return self.weight * self._rms(x)
61
+
62
+
63
+class NotAnRMSNorm(RMSNorm):
64
+ """Mean-subtracting LayerNorm-like math -> not an RMSNorm."""
65
+
66
+ def __init__(self, hidden: int = 16, eps: float = 1e-6):
67
+ super().__init__(hidden, eps)
68
+
69
+ def forward(self, x):
70
+ x = x - x.mean(-1, keepdim=True)
71
+ variance = x.var(-1, keepdim=True)
72
+ return self.weight * x / torch.sqrt(variance + self.variance_epsilon)
73
+
74
+
75
+class GatedRMSNorm(RMSNorm):
76
+ """Second input and tail compute -> not an RMSNorm."""
77
+
78
+ def forward(self, x, gate=None):
79
+ normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
80
+ return normed * F.silu(gate)
81
+
82
+
83
+class GatedFusedRMSNorm(nn.Module):
84
+ """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm."""
85
+
86
+ def __init__(self, hidden: int = 16, eps: float = 1e-5):
87
+ super().__init__()
88
+ self.weight = nn.Parameter(torch.ones(hidden))
89
+ self.eps = eps
90
+
91
+ def forward(self, x, gate=None):
92
+ return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate)
93
+
94
+
95
+class UntraceableGatedRMSNorm(RMSNorm):
96
+ """Tracer can't see tail compute in forward, but still has a second input (gate)."""
97
+
98
+ def forward(self, x, gate=None):
99
+ normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype)
100
+ if gate.sum() > 0: # untraceable -> partial graph, no visible tail
101
+ normed = normed * F.silu(gate)
102
+ return normed
103
+
104
+
105
+@pytest.mark.parametrize(
106
+ "cls,eps,zero_centered",
107
+ [
108
+ (RMSNorm, 1e-5, False),
109
+ (GemmaRMSNorm, 1e-6, True),
110
+ (WeightlessRMSNorm, 1e-6, False),
111
+ (LayerNorm, 1e-6, False),
112
+ (torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op
113
+ ],
114
+)
115
+def test_detects_rms_norm_variants(cls, eps, zero_centered):
116
+ with torch.device("meta"):
117
+ fuser = get_fuser(cls(16, eps=eps))
118
+ assert isinstance(fuser, RMSNormFuser)
119
+ assert fuser.zero_centered == zero_centered
120
+
121
+
122
+@pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU])
123
+def test_non_rms_norms_are_not_matched(cls):
124
+ with torch.device("meta"):
125
+ module = cls(16) if cls is nn.LayerNorm else cls()
126
+ assert not isinstance(get_fuser(module), RMSNormFuser)
127
+
128
+
129
+@pytest.mark.parametrize(
130
+ "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm]
131
+)
132
+def test_gated_rms_norm_is_not_fused(cls):
133
+ with torch.device("meta"):
134
+ assert not isinstance(get_fuser(cls()), RMSNormFuser)
135
+
136
+
137
+@pytest.mark.parametrize(
138
+ "cls,expected,zero_centered",
139
+ [
140
+ (RMSNorm, "RMSNorm", False),
141
+ (GemmaRMSNorm, "GemmaRMSNorm", True),
142
+ (WeightlessRMSNorm, "RMSNorm", False),
143
+ ],
144
+)
145
+def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config):
146
+ from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm
147
+ from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
148
+
149
+ # `default_vllm_config` supplies the config context the CustomOp needs; the
150
+ # weightless path reads hidden size from the model config, so stub it.
151
+ model_config = SimpleNamespace(get_hidden_size=lambda: 16)
152
+ with torch.device("meta"):
153
+ module = cls()
154
+ fuser = get_fuser(module)
155
+ built = fuser.fuse(module, "norm", model_config, None)
156
+ from vllm.model_executor.models.transformers.fusers.rms_norm import (
157
+ TPAwareNormMixin,
158
+ )
159
+
160
+ types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
161
+ assert isinstance(built, types_by_name[expected])
162
+ assert isinstance(built, TPAwareNormMixin) # fused norms self-correct under TP
163
+ assert built.variance_epsilon == module.variance_epsilon
164
+ assert isinstance(built.weight, nn.Parameter) == (
165
+ getattr(module, "weight", None) is not None
166
+ )
167
+
168
+
169
+def test_fused_rms_norm_op_default_eps(default_vllm_config):
170
+ """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path;
171
+ its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`."""
172
+ from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
173
+
174
+ with torch.device("meta"):
175
+ module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call
176
+ fuser = get_fuser(module)
177
+ assert isinstance(fuser, RMSNormFuser)
178
+ assert not fuser.zero_centered
179
+ model_config = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
180
+ built = fuser.fuse(module, "norm", model_config, None)
181
+ assert isinstance(built, VLLMRMSNorm)
182
+ assert built.variance_epsilon == torch.finfo(torch.float32).eps
183
+
184
+
185
+def test_eps_is_derived_per_instance(default_vllm_config):
186
+ """Two instances of the same norm class with different eps must fuse to their
187
+ own eps: the type-cached fuser holds only structure, not this value."""
188
+ model_config = SimpleNamespace(get_hidden_size=lambda: 16)
189
+ with torch.device("meta"):
190
+ for eps in (1e-5, 1e-6):
191
+ module = RMSNorm(16, eps=eps)
192
+ built = get_fuser(module).fuse(module, "norm", model_config, None)
193
+ assert built.variance_epsilon == eps
194
+
195
+
196
+def test_fused_norm_is_gather_capable(default_vllm_config):
197
+ """Every fused norm is emitted gather-capable, so a norm on a head-sharded
198
+ projection (OLMoE-style) self-corrects at runtime with no QKV-specific
199
+ plumbing. A full-width input skips the gather and equals a plain norm."""
200
+ from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
201
+ from vllm.model_executor.models.transformers.fusers import rms_norm
202
+
203
+ torch.manual_seed(0)
204
+ x = torch.randn(4, 16)
205
+ for gathered_cls, plain_cls in [
206
+ (rms_norm.TPAwareRMSNorm, RMSNorm),
207
+ (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm),
208
+ ]:
209
+ gathered = gathered_cls(hidden_size=16, eps=1e-6)
210
+ assert isinstance(gathered, rms_norm.TPAwareNormMixin)
211
+ plain = plain_cls(hidden_size=16, eps=1e-6)
212
+ with torch.no_grad():
213
+ weight = torch.randn(16)
214
+ gathered.weight.copy_(weight)
215
+ plain.weight.copy_(weight)
216
+ torch.testing.assert_close(gathered(x), plain(x))
217
+
218
+
219
+def test_gathered_norm_rejects_uneven_sharding(default_vllm_config):
220
+ """A sharded input (narrower than the full-width weight) that does not tile
221
+ the weight evenly across ranks is rejected before any collective."""
222
+ from vllm.model_executor.models.transformers.fusers import rms_norm
223
+
224
+ norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6)
225
+ norm.tp_size = 2 # emulate TP=2 without a real process group
226
+ with pytest.raises(ValueError, match="does not tile it evenly"):
227
+ norm(torch.randn(2, 3)) # 3 * 2 != 8
tests/models/transformers/test_backend.py
renamed
+39
-9
@@ -6,10 +6,16 @@ from typing import Any
6
7
import pytest
8
9
-from ..conftest import HfRunner, VllmRunner
10
-from ..utils import multi_gpu_test, prep_prompts
11
-from .registry import HF_EXAMPLE_MODELS
12
-from .utils import check_embeddings_close, check_logprobs_close
9
+from ...conftest import HfRunner, VllmRunner
10
+from ...utils import multi_gpu_test, prep_prompts
11
+from ..registry import HF_EXAMPLE_MODELS
12
+from ..utils import check_embeddings_close, check_logprobs_close
13
+
14
+
15
+@pytest.fixture(scope="function", autouse=True)
16
+def enable_pickle(monkeypatch):
17
+ """`LLM.apply_model` requires pickling a function."""
18
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
19
20
21
def get_model(arch: str) -> str:
@@ -18,6 +24,17 @@ def get_model(arch: str) -> str:
24
return model_info.default
25
26
27
+def get_num_fused(model) -> tuple[int, int]:
28
+ from vllm.model_executor.layers.linear import (
29
+ MergedColumnParallelLinear,
30
+ QKVParallelLinear,
31
+ )
32
+
33
+ glu = sum(isinstance(m, MergedColumnParallelLinear) for m in model.modules())
34
+ qkv = sum(isinstance(m, QKVParallelLinear) for m in model.modules())
35
+ return glu, qkv
36
+
37
+
38
def check_implementation(
39
runner_ref: type[HfRunner | VllmRunner],
40
runner_test: type[VllmRunner],
@@ -25,6 +42,7 @@ def check_implementation(
42
model: str,
43
kwargs_ref: dict[str, Any] | None = None,
44
kwargs_test: dict[str, Any] | None = None,
45
+ num_fused: tuple[int, int] = (1, 1),
46
**kwargs,
47
):
48
if kwargs_ref is None:
@@ -41,6 +59,12 @@ def check_implementation(
59
model_config = model_test.llm.llm_engine.model_config
60
assert model_config.using_transformers_backend()
61
62
+ num_layers = model_config.hf_config.get_text_config().num_hidden_layers
63
+ expected_glu, expected_qkv = num_fused
64
+ for num_glu, num_qkv in model_test.apply_model(get_num_fused):
65
+ assert num_glu == expected_glu * num_layers
66
+ assert num_qkv == expected_qkv * num_layers
67
+
68
outputs_test = model_test.generate_greedy_logprobs(*args)
69
70
with runner_ref(model, **kwargs_ref) as model_ref:
@@ -58,11 +82,11 @@ def check_implementation(
82
83
84
@pytest.mark.parametrize(
61
- "model,model_impl",
85
+ "model,model_impl,num_fused",
86
[
63
- ("meta-llama/Llama-3.2-1B-Instruct", "transformers"),
64
- ("hmellor/Ilama-3.2-1B", "auto"), # CUSTOM CODE
65
- ("allenai/OLMoE-1B-7B-0924", "transformers"), # MoE
87
+ ("meta-llama/Llama-3.2-1B-Instruct", "transformers", (1, 1)),
88
+ ("hmellor/Ilama-3.2-1B", "auto", (1, 1)), # CUSTOM CODE
89
+ ("allenai/OLMoE-1B-7B-0924", "transformers", (0, 1)), # MoE
90
],
91
) # trust_remote_code=True by default
92
def test_models(
@@ -71,6 +95,7 @@ def test_models(
95
example_prompts: list[str],
96
model: str,
97
model_impl: str,
98
+ num_fused: tuple[int, int],
99
) -> None:
100
import transformers
101
from packaging.version import Version
@@ -84,7 +109,12 @@ def test_models(
109
)
110
111
check_implementation(
87
- hf_runner, vllm_runner, example_prompts, model, model_impl=model_impl
112
+ hf_runner,
113
+ vllm_runner,
114
+ example_prompts,
115
+ model,
116
+ num_fused=num_fused,
117
+ model_impl=model_impl,
118
)
119
120
vllm/model_executor/models/transformers/__init__.py
+34
@@ -16,6 +16,10 @@
16
# limitations under the License.
17
"""Wrapper around `transformers` models"""
18
19
+from typing import TYPE_CHECKING
20
+
21
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
22
+
23
from vllm.model_executor.models.transformers.base import Base
24
from vllm.model_executor.models.transformers.causal import CausalMixin
25
from vllm.model_executor.models.transformers.legacy import LegacyMixin
@@ -32,6 +36,36 @@ from vllm.model_executor.models.transformers.pooling import (
36
)
37
from vllm.multimodal import MULTIMODAL_REGISTRY
38
39
+if TYPE_CHECKING:
40
+ import torch
41
+
42
+ from vllm.model_executor.layers.attention import Attention
43
+
44
+
45
+def vllm_attention_forward(
46
+ # Transformers args
47
+ module: "torch.nn.Module",
48
+ query: "torch.Tensor",
49
+ key: "torch.Tensor",
50
+ value: "torch.Tensor",
51
+ attention_mask: "torch.Tensor",
52
+ # Transformers kwargs
53
+ scaling: float | None = None,
54
+ # vLLM kwargs
55
+ attention_instances: dict[int, "Attention"] | None = None,
56
+ **kwargs,
57
+):
58
+ self_attn = attention_instances[module.layer_idx]
59
+ if scaling is not None:
60
+ self_attn.impl.scale = float(scaling)
61
+ hidden = query.shape[-2]
62
+ query, key, value = (x.transpose(1, 2) for x in (query, key, value))
63
+ query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))
64
+ return self_attn.forward(query, key, value), None
65
+
66
+
67
+ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward
68
+
69
70
# Text only models
71
class TransformersForCausalLM(CausalMixin, Base): ...
vllm/model_executor/models/transformers/base.py
+65
-53
@@ -31,7 +31,6 @@ from transformers.conversion_mapping import (
31
WeightRenaming,
32
get_model_conversion_mapping,
33
)
34
-from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
34
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config.utils import getattr_iter
@@ -42,6 +41,7 @@ from vllm.model_executor.layers.attention import (
41
Attention,
42
EncoderOnlyAttention,
43
)
44
+from vllm.model_executor.layers.fused_moe import MoERunner
45
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
46
from vllm.model_executor.models.interfaces import (
47
SupportsEagle,
@@ -51,6 +51,7 @@ from vllm.model_executor.models.interfaces import (
51
SupportsQuant,
52
)
53
from vllm.model_executor.models.interfaces_base import VllmModel
54
+from vllm.model_executor.models.transformers.fuser import BaseFuser, Fusers
55
from vllm.model_executor.models.transformers.utils import (
56
can_enable_torch_compile,
57
get_feature_request_tip,
@@ -58,7 +59,6 @@ from vllm.model_executor.models.transformers.utils import (
59
log_replacement,
60
replace_conv_class,
61
replace_linear_class,
61
- replace_rms_norm_class,
62
)
63
from vllm.model_executor.models.utils import (
64
AutoWeightsLoader,
@@ -74,37 +74,10 @@ if TYPE_CHECKING:
74
from transformers import PreTrainedModel
75
76
from vllm.config import VllmConfig
77
-else:
78
- PreTrainedModel = object
77
78
logger = init_logger(__name__)
79
80
83
-def vllm_flash_attention_forward(
84
- # Transformers args
85
- module: torch.nn.Module,
86
- query: torch.Tensor,
87
- key: torch.Tensor,
88
- value: torch.Tensor,
89
- attention_mask: torch.Tensor,
90
- # Transformers kwargs
91
- scaling: float | None = None,
92
- # vLLM kwargs
93
- attention_instances: dict[int, Attention] | None = None,
94
- **kwargs,
95
-):
96
- self_attn = attention_instances[module.layer_idx]
97
- if scaling is not None:
98
- self_attn.impl.scale = float(scaling)
99
- hidden = query.shape[-2]
100
- query, key, value = (x.transpose(1, 2) for x in (query, key, value))
101
- query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))
102
- return self_attn.forward(query, key, value), None
103
-
104
-
105
-ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_flash_attention_forward
106
-
107
-
81
class Base(
82
nn.Module,
83
VllmModel,
@@ -141,6 +114,9 @@ class Base(
114
"""Ignore unexpected weights whose qualname starts with these prefixes."""
115
self.ignore_unexpected_suffixes: list[str] = []
116
"""Ignore unexpected weights whose qualname ends with these suffixes."""
117
+ self.packed_modules_mapping: dict[str, list[str]] = {}
118
+ """Fused module -> constituent projections, populated by `recursive_replace`
119
+ for the quantization machinery and loaders (e.g. bitsandbytes)."""
120
121
# Attrs for Eagle3 (see self.set_aux_hidden_state_layers)
122
self._target_class: type[nn.Module] = nn.Module
@@ -217,7 +193,7 @@ class Base(
193
self.text_config._attn_implementation = "vllm"
194
self.config.dtype = torch.get_default_dtype()
195
220
- def _get_decoder_cls(self, **kwargs: dict) -> type[PreTrainedModel]:
196
+ def _get_decoder_cls(self, **kwargs: dict) -> type["PreTrainedModel"]:
197
"""
198
Get the decoder class from the model.
199
@@ -236,7 +212,7 @@ class Base(
212
213
def _decorate_cls_for_torch_compile(
214
self,
239
- cls: type[PreTrainedModel],
215
+ cls: type["PreTrainedModel"],
216
dynamic_arg_dims: dict[str, int] | None,
217
enable_if: Callable[["VllmConfig"], bool],
218
is_encoder: bool,
@@ -356,12 +332,24 @@ class Base(
332
if self.pp_group.world_size <= 1:
333
return
334
359
- if not self.model.supports_pp_plan:
335
+ if self.model.supports_pp_plan:
336
+ module = self.model
337
+ names = list(module._pp_plan.keys())
338
+ else:
339
+ module = self.model.get_decoder()
340
+ has_parameters = lambda m: next(m.parameters(), None) is not None
341
+ names = [n for n, c in module.named_children() if has_parameters(c)]
342
tip = get_feature_request_tip(
343
self.model_config.model, self.model_config.trust_remote_code
344
)
363
- raise ValueError(
364
- f"{type(self.model)} does not support pipeline parallel. {tip}"
345
+ logger.warning(
346
+ "%s does not define a pipeline parallel plan. The Transformers "
347
+ "modeling backend will infer the split from the layers of %s in order "
348
+ "of declaration and keep parameter-free modules on every rank. This "
349
+ "may fail if the model's structure is non-standard. %s",
350
+ type(self.model),
351
+ type(module),
352
+ tip,
353
)
354
355
def attrsetter(attr: str) -> Callable[[object, object], None]:
@@ -376,10 +364,9 @@ class Base(
364
365
module_lists = []
366
module_list_idx = None
379
- pp_plan = list(self.model._pp_plan.keys())
380
- for i, name in enumerate(pp_plan):
367
+ for i, name in enumerate(names):
368
# attrgetter in case the module is nested (e.g. "text_model.layers")
382
- if isinstance(attrgetter(name)(self.model), nn.ModuleList):
369
+ if isinstance(attrgetter(name)(module), nn.ModuleList):
370
module_lists.append(name)
371
module_list_idx = i
372
@@ -389,16 +376,16 @@ class Base(
376
"in the base model are not supported yet!"
377
)
378
if module_list_idx is None:
392
- raise ValueError(f"Could not find `ModuleList` in {type(self.model)}")
379
+ raise ValueError(f"Could not find `ModuleList` in {type(module)}")
380
381
# Layers before module list
395
- for name in pp_plan[:module_list_idx]:
382
+ for name in names[:module_list_idx]:
383
if self.pp_group.is_first_rank or (
384
self._get_tie_word_embeddings() and self.pp_group.is_last_rank
385
):
386
continue
387
# attrsetter in case the module is nested (e.g. "text_model.embed_tokens")
401
- attrsetter(name)(self.model, PPMissingLayer())
388
+ attrsetter(name)(module, PPMissingLayer())
389
390
# Module list
391
start_layer, end_layer = get_pp_indices(
@@ -406,41 +393,61 @@ class Base(
393
self.pp_group.rank_in_group,
394
self.pp_group.world_size,
395
)
409
- layers_name = pp_plan[module_list_idx]
396
+ layers_name = names[module_list_idx]
397
# attrgetter in case the module is nested (e.g. "text_model.layers")
411
- layers = attrgetter(layers_name)(self.model)
398
+ layers = attrgetter(layers_name)(module)
399
for i in range(len(layers)):
400
if start_layer <= i and i < end_layer:
401
continue
402
layers[i] = PPMissingLayer()
403
404
# Layers after module list
418
- for name in pp_plan[module_list_idx + 1 :]:
405
+ for name in names[module_list_idx + 1 :]:
406
# Modules that should be on last rank
407
if not self.pp_group.is_last_rank:
408
# attrsetter in case the module is nested (e.g. "text_model.norm")
422
- attrsetter(name)(self.model, PPMissingLayer())
409
+ attrsetter(name)(module, PPMissingLayer())
410
411
def recursive_replace(self):
412
"""Recursively replace modules in the model as needed.
413
414
Currently, this replaces:
415
416
+ - GLUs with a fused `MergedColumnParallelLinear` + `...AndMul`
417
+ - Attention QKV projections with a fused `QKVParallelLinear` + split
418
- `nn.Linear` with vLLM's tensor parallel linear classes
430
- - `*RMSNorm` with vLLM's `RMSNorm`
419
+ - `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
420
+ - RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
421
"""
432
- tp_plan = self.model.tp_plan
422
+ tp_plan = self.model.tp_plan or {}
423
424
if not tp_plan and self.tp_group.world_size > 1:
425
tip = get_feature_request_tip(
426
self.model_config.model, self.model_config.trust_remote_code
427
)
438
- raise ValueError(
439
- f"{type(self.model)} does not support tensor parallel. {tip}"
428
+ logger.warning_once(
429
+ "%s does not define a tensor parallel plan. The Transformers modeling "
430
+ "backend will shard the model the best it can during graph fusion and "
431
+ "replicate the rest. This may be suboptimal or fail if the model does "
432
+ "not fuse cleanly. %s",
433
+ type(self.model),
434
+ tip,
435
)
436
437
# Prefix the patterns because we always start from `self.model`
438
tp_plan = {maybe_prefix("model", k): v for k, v in tp_plan.items()}
439
+ # Detect fusable patterns once per module class (cached, so this is cheap)
440
+ fusers = Fusers(self.model, self.model_config)
441
+
442
+ def register_fusion(fuser: BaseFuser, prefix: str):
443
+ """Register a fused layer's mappings just before it is built."""
444
+ orig_to_new_stacked = fuser.orig_to_new_stacked(prefix)
445
+ self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked)
446
+
447
+ packed_modules_mapping = fuser.packed_modules_mapping
448
+ self.packed_modules_mapping.update(packed_modules_mapping)
449
+ if self.quant_config is not None:
450
+ self.quant_config.packed_modules_mapping.update(packed_modules_mapping)
451
452
def _recursive_replace(module: nn.Module, prefix: str):
453
for child_name, child_module in module.named_children():
@@ -479,11 +486,16 @@ class Base(
486
)
487
elif isinstance(child_module, (nn.Conv2d, nn.Conv3d)):
488
new_module = replace_conv_class(child_module)
482
- elif child_module.__class__.__name__.endswith("RMSNorm"):
483
- new_module = replace_rms_norm_class(
484
- child_module, self.text_config.hidden_size
489
+ elif (fuser := fusers[child_module]) is not None:
490
+ register_fusion(fuser, qual_name)
491
+ new_module = fuser.fuse(
492
+ child_module, qual_name, self.model_config, self.quant_config
493
)
486
- else:
494
+ logger.info_once(fuser.info(child_name))
495
+ _recursive_replace(new_module, prefix=qual_name)
496
+ elif not isinstance(child_module, MoERunner):
497
+ # MoERunner can contain aliases of shared experts and gates,
498
+ # so we don't want to recurse into it and break weight loading.
499
_recursive_replace(child_module, prefix=qual_name)
500
501
if new_module is not child_module:
@@ -538,7 +550,7 @@ class Base(
550
num_heads=num_heads,
551
head_size=head_size,
552
# NOTE: We use Llama scale as default, if it's set by
541
- # Transformers, it's updated in vllm_flash_attention_forward
553
+ # Transformers, it's updated in vllm_attention_forward
554
scale=head_size**-0.5,
555
num_kv_heads=num_kv_heads,
556
cache_config=self.cache_config,
vllm/model_executor/models/transformers/fuser.py
new
+78
@@ -0,0 +1,78 @@
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_fuser` traces a module class once (see `fx_utils`) and matches it against
6
+each concrete fuser in `fusers`; `Fusers` caches the result per class for a
7
+whole model. `base.recursive_replace` then applies the matched fuser per
8
+instance. RMSNorm-shaped modules the tracer cannot match are warned about.
9
+"""
10
+
11
+from collections import UserDict
12
+from typing import TYPE_CHECKING
13
+
14
+from cachetools import cached
15
+from torch import nn
16
+
17
+from vllm.logger import init_logger
18
+from vllm.model_executor.models.transformers.fusers import (
19
+ BaseFuser,
20
+ GLUFuser,
21
+ QKVFuser,
22
+ RMSNormFuser,
23
+ StackedFuser,
24
+)
25
+from vllm.model_executor.models.transformers.fx_utils import trace
26
+
27
+if TYPE_CHECKING:
28
+ from vllm.config.model import ModelConfig
29
+
30
+logger = init_logger(__name__)
31
+
32
+
33
+@cached(cache={}, key=type)
34
+def get_fuser(module: nn.Module) -> BaseFuser | None:
35
+ """The fuser for `type(module)` (cached per class), or `None` if no match."""
36
+ # Projection fusions need >=2 sibling linears; the RMSNorm fusion needs a
37
+ # leaf module (raw tensor math, no submodules). Nothing else can match, and
38
+ # tracing is skipped for it.
39
+ n_linear = sum(isinstance(c, nn.Linear) for c in module.children())
40
+ is_leaf = next(module.children(), None) is None
41
+ if n_linear < 2 and not is_leaf:
42
+ return None
43
+ if (graph := trace(module)) is None:
44
+ return None
45
+ for fuser_cls in (GLUFuser, QKVFuser, RMSNormFuser):
46
+ if (fuser := fuser_cls.match(graph, module)) is not None:
47
+ if isinstance(fuser, StackedFuser):
48
+ try:
49
+ fuser.update_forward(module)
50
+ except Exception as exc:
51
+ # An unrecognised source just means we cannot fuse here.
52
+ logger.debug(
53
+ "Could not rewrite %s for fusion: %s", type(module), exc
54
+ )
55
+ return None
56
+ return fuser
57
+ # A norm we could not match structurally is left unfused; flag likely misses.
58
+ if module.__class__.__name__.endswith("RMSNorm"):
59
+ logger.warning_once(
60
+ "%s looks like an RMSNorm but its computation did not match the "
61
+ "expected pattern, so it was left unfused.",
62
+ module.__class__.__name__,
63
+ )
64
+ return None
65
+
66
+
67
+class Fusers(UserDict):
68
+ """Mapping from module class to fuser, for all fusable classes in a model."""
69
+
70
+ def __init__(self, model: nn.Module, model_config: "ModelConfig"):
71
+ self.model_config = model_config
72
+ super().__init__({type(m): get_fuser(m) for m in model.modules()})
73
+
74
+ def __getitem__(self, m: nn.Module) -> BaseFuser | None:
75
+ fuser = self.data.get(type(m))
76
+ if fuser is not None and fuser.validate(m, self.model_config):
77
+ return fuser
78
+ return None
vllm/model_executor/models/transformers/fusers/__init__.py
new
+18
@@ -0,0 +1,18 @@
1
+# SPDX-License-Identifier: Apache-2.0
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
6
+from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
7
+from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
8
+from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser
9
+from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
10
+
11
+__all__ = [
12
+ "BaseFuser",
13
+ "StackedFuser",
14
+ "GLUFuser",
15
+ "MoEBlockFuser",
16
+ "QKVFuser",
17
+ "RMSNormFuser",
18
+]
vllm/model_executor/models/transformers/fusers/base.py
new
+146
@@ -0,0 +1,146 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Base classes for the Transformers backend fusers."""
4
+
5
+import types
6
+from abc import ABC, abstractmethod
7
+from collections.abc import Callable
8
+from dataclasses import dataclass, field
9
+from typing import TYPE_CHECKING, ClassVar
10
+
11
+from torch import fx, nn
12
+
13
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
14
+
15
+if TYPE_CHECKING:
16
+ from vllm.config.model import ModelConfig
17
+ from vllm.model_executor.layers.quantization import QuantizationConfig
18
+
19
+
20
+@dataclass
21
+class BaseFuser(ABC):
22
+ """A detected fusion and how to apply it.
23
+
24
+ `match` analyses the module *class* once (cached, see `get_fuser`); `fuse`
25
+ then applies the fusion to an instance in `recursive_replace`, returning the
26
+ module to install in its place.
27
+ """
28
+
29
+ @abstractmethod
30
+ def info(self, name: str) -> str:
31
+ """A human-readable description of the fusion at `name`, for logging."""
32
+
33
+ @classmethod
34
+ @abstractmethod
35
+ def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
36
+ """Match the pattern in `graph`, returning a fuser if found."""
37
+
38
+ @abstractmethod
39
+ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
40
+ """Whether this fuser can be applied to this `module` instance."""
41
+
42
+ @abstractmethod
43
+ def fuse(
44
+ self,
45
+ module: nn.Module,
46
+ prefix: str,
47
+ model_config: "ModelConfig",
48
+ quant_config: "QuantizationConfig",
49
+ ) -> nn.Module:
50
+ """Apply the fusion to an already-validated `module`, returning the
51
+ module to install in its place (mutated in place, or freshly built)."""
52
+
53
+ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
54
+ """`WeightsMapper.orig_to_new_stacked` entries this fuser contributes
55
+ (none unless it stacks weights)."""
56
+ return {}
57
+
58
+ @property
59
+ def packed_modules_mapping(self) -> dict[str, list[str]]:
60
+ """`packed_modules_mapping` entries this fuser contributes (none unless
61
+ it stacks weights)."""
62
+ return {}
63
+
64
+
65
+@dataclass
66
+class StackedFuser(BaseFuser):
67
+ """A fuser that merges sibling projections into one stacked linear and
68
+ rewrites the forward to call it.
69
+
70
+ `match` and `update_forward` analyse the class once; `fuse` builds the merged
71
+ submodule and binds the compiled forward on an instance in place, so it keeps
72
+ its class and any attribute the fusion does not consume.
73
+ """
74
+
75
+ merged_name: ClassVar[str]
76
+ """Attribute name of the merged module created by `update_attrs`."""
77
+ merged_cls: ClassVar[str]
78
+ """Name of the vLLM class the merged projection becomes (for logging)."""
79
+
80
+ source_cls: str
81
+ """Class of the HF module the fused projections belonged to (for logging)."""
82
+
83
+ fused_forward: Callable = field(init=False, repr=False)
84
+ """The compiled rewritten forward, set by `update_forward`."""
85
+
86
+ def info(self, name: str) -> str:
87
+ sources = " + ".join(shard for shard, _ in self.shards)
88
+ return (
89
+ f"Fused: {sources} ({name}: {self.source_cls}) -> "
90
+ f"{self.merged_name} ({self.merged_cls})"
91
+ )
92
+
93
+ @property
94
+ @abstractmethod
95
+ def shards(self) -> list[tuple[str, ShardId]]:
96
+ """Each projection's original name and its shard id in the merged module.
97
+
98
+ Source for both `orig_to_new_stacked` and `packed_modules_mapping`."""
99
+
100
+ def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
101
+ """`WeightsMapper.orig_to_new_stacked` entries for one fused instance.
102
+
103
+ Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname
104
+ so only this exact layer is remapped, never a same-named projection
105
+ elsewhere (e.g. an unfused MoE expert's `gate_proj`)."""
106
+ merged = maybe_prefix(prefix, self.merged_name)
107
+ return {
108
+ maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards
109
+ }
110
+
111
+ @property
112
+ def packed_modules_mapping(self) -> dict[str, list[str]]:
113
+ """`{merged_name: [projection names]}` so quantization can unpack the
114
+ fused layer into its per-shard configs."""
115
+ return {self.merged_name: [name for name, _ in self.shards]}
116
+
117
+ @abstractmethod
118
+ def update_forward(self, module: nn.Module) -> None:
119
+ """Rewrite and compile `type(module)`'s forward source.
120
+
121
+ Raises if the source does not admit the rewrite (fusion is then skipped).
122
+ """
123
+
124
+ @abstractmethod
125
+ def update_attrs(
126
+ self,
127
+ module: nn.Module,
128
+ prefix: str,
129
+ model_config: "ModelConfig",
130
+ quant_config: "QuantizationConfig",
131
+ ) -> None:
132
+ """Replace `module`'s submodules with the merged module."""
133
+
134
+ def fuse(
135
+ self,
136
+ module: nn.Module,
137
+ prefix: str,
138
+ model_config: "ModelConfig",
139
+ quant_config: "QuantizationConfig",
140
+ ) -> nn.Module:
141
+ """Fuse an already-validated `module` in place (see `Fusers.__getitem__`).
142
+
143
+ Builds the merged submodule and binds the compiled forward."""
144
+ self.update_attrs(module, prefix, model_config, quant_config)
145
+ module.forward = types.MethodType(self.fused_forward, module)
146
+ return module
vllm/model_executor/models/transformers/fusers/glu.py
new
+218
@@ -0,0 +1,218 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear."""
4
+
5
+import ast
6
+import operator
7
+from dataclasses import dataclass
8
+from typing import TYPE_CHECKING, ClassVar
9
+
10
+from torch import fx, nn
11
+from transformers.activations import ACT2CLS
12
+
13
+from vllm.logger import init_logger
14
+from vllm.model_executor.layers.activation import (
15
+ _ACTIVATION_AND_MUL_REGISTRY,
16
+ get_act_and_mul_fn,
17
+)
18
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
19
+from vllm.model_executor.models.transformers.fusers.base import StackedFuser
20
+from vllm.model_executor.models.transformers.fx_utils import (
21
+ compile_forward,
22
+ find_node,
23
+ is_linear,
24
+ peel,
25
+ recover_forward,
26
+ replace_expr,
27
+ single_self_call,
28
+)
29
+from vllm.model_executor.models.transformers.utils import (
30
+ log_replacement,
31
+ replace_linear_class,
32
+)
33
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
34
+
35
+if TYPE_CHECKING:
36
+ from vllm.config.model import ModelConfig
37
+ from vllm.model_executor.layers.quantization import QuantizationConfig
38
+
39
+logger = init_logger(__name__)
40
+
41
+
42
+CLS2ACT: dict[type, list[str]] = {}
43
+for _act_name, _act_cls in ACT2CLS.items():
44
+ if isinstance(_act_cls, tuple):
45
+ _act_cls = _act_cls[0]
46
+ CLS2ACT.setdefault(_act_cls, []).append(_act_name)
47
+
48
+ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys())
49
+
50
+
51
+@dataclass
52
+class GLUFuser(StackedFuser):
53
+ """Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
54
+
55
+ act_name: str
56
+ gate_name: str
57
+ up_name: str
58
+ down_name: str | None
59
+ merged_name: ClassVar[str] = "gate_up_proj"
60
+ merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
61
+
62
+ @property
63
+ def shards(self) -> list[tuple[str, ShardId]]:
64
+ return [(self.gate_name, 0), (self.up_name, 1)]
65
+
66
+ @classmethod
67
+ def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
68
+ """Is node `act(gate(x))` where `gate` is linear and `act` is not linear."""
69
+ return (
70
+ node.op == "call_module"
71
+ and not is_linear(node, module)
72
+ and len(node.args) == 1
73
+ and isinstance(node.args[0], fx.Node)
74
+ and is_linear(node.args[0], module)
75
+ )
76
+
77
+ @classmethod
78
+ def _get_glu_nodes(
79
+ cls, graph: fx.Graph, module: nn.Module
80
+ ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
81
+ """Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
82
+ for mul in graph.nodes:
83
+ if (
84
+ mul.op == "call_function"
85
+ and mul.target == operator.mul
86
+ and len(mul.args) == 2
87
+ and all(isinstance(arg, fx.Node) for arg in mul.args)
88
+ ):
89
+ a, b = mul.args
90
+ if cls._is_act_of_gate(a, module) and is_linear(b, module):
91
+ act, gate, up = a, a.args[0], b
92
+ elif cls._is_act_of_gate(b, module) and is_linear(a, module):
93
+ act, gate, up = b, b.args[0], a
94
+ else:
95
+ continue
96
+ if (
97
+ all(len(args) == 1 for args in (gate.args, up.args))
98
+ and isinstance(x := gate.args[0], fx.Node)
99
+ and x is up.args[0]
100
+ ):
101
+ return act, gate, up, mul
102
+ return None
103
+
104
+ @staticmethod
105
+ def _get_act_and_mul_name(act: nn.Module) -> str | None:
106
+ """Get the name of `act` if it has an `...AndMul` equivalent."""
107
+ for name in CLS2ACT.get(type(act), []):
108
+ if name in ACT_AND_MUL_NAMES:
109
+ return name
110
+ # nn.GELU is not in ACT2CLS, but could be in model code
111
+ if type(act) is nn.GELU:
112
+ return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu"
113
+ return None
114
+
115
+ @classmethod
116
+ def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
117
+ """Get the `...AndMul` equivalent of a Transformers activation module."""
118
+ if name := cls._get_act_and_mul_name(act):
119
+ return get_act_and_mul_fn(name)
120
+ raise ValueError(f"No AndMul equivalent for {type(act)}")
121
+
122
+ @classmethod
123
+ def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
124
+ if (glu_nodes := cls._get_glu_nodes(graph, module)) is None:
125
+ return None
126
+ act_node, gate_node, up_node, mul_node = glu_nodes
127
+
128
+ gate = module.get_submodule(gate_node.target)
129
+ up = module.get_submodule(up_node.target)
130
+ # Shapes must be compatible for a single merged GEMM.
131
+ if gate.in_features == up.in_features and (gate.bias is None) == (
132
+ up.bias is None
133
+ ):
134
+ predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
135
+ down_node = find_node(graph, predicate)
136
+ return cls(
137
+ source_cls=type(module).__name__,
138
+ act_name=act_node.target,
139
+ gate_name=gate_node.target,
140
+ up_name=up_node.target,
141
+ down_name=down_node.target if down_node is not None else None,
142
+ )
143
+ return None
144
+
145
+ def update_forward(self, module: nn.Module) -> None:
146
+ """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
147
+ funcdef, fn = recover_forward(type(module))
148
+ act_call = single_self_call(funcdef, self.act_name)
149
+ gate_call = single_self_call(funcdef, self.gate_name)
150
+ up_call = single_self_call(funcdef, self.up_name)
151
+ if act_call.args[0] is not gate_call:
152
+ raise ValueError("activation does not directly wrap the gate")
153
+ if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):
154
+ raise ValueError("gate and up inputs are written differently")
155
+ muls = [
156
+ node
157
+ for node in ast.walk(funcdef)
158
+ if isinstance(node, ast.BinOp)
159
+ and isinstance(node.op, ast.Mult)
160
+ and {id(node.left), id(node.right)} == {id(act_call), id(up_call)}
161
+ ]
162
+ if len(muls) != 1:
163
+ raise ValueError("no multiply of the activation and up projection")
164
+
165
+ # act(gate(x)) * up(x) -> act(gate_up(x))
166
+ assert isinstance(gate_call.func, ast.Attribute)
167
+ gate_call.func.attr = self.merged_name
168
+ replace_expr(funcdef, muls[0], act_call)
169
+ self.fused_forward = compile_forward(funcdef, fn)
170
+
171
+ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
172
+ act = module.get_submodule(self.act_name)
173
+ if self._get_act_and_mul_name(act) is None:
174
+ logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))
175
+ return False
176
+ return True
177
+
178
+ def update_attrs(
179
+ self,
180
+ module: nn.Module,
181
+ prefix: str,
182
+ model_config: "ModelConfig",
183
+ quant_config: "QuantizationConfig",
184
+ ) -> None:
185
+ act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
186
+ gate = module.get_submodule(self.gate_name)
187
+ up = module.get_submodule(self.up_name)
188
+ merged = MergedColumnParallelLinear(
189
+ input_size=gate.in_features,
190
+ output_sizes=[gate.out_features, up.out_features],
191
+ bias=gate.bias is not None,
192
+ quant_config=quant_config,
193
+ prefix=maybe_prefix(prefix, self.merged_name),
194
+ return_bias=False,
195
+ )
196
+ logger.debug(
197
+ "%s: %s, %s: %s -> %s: %s",
198
+ self.gate_name,
199
+ gate,
200
+ self.up_name,
201
+ up,
202
+ self.merged_name,
203
+ merged,
204
+ )
205
+ setattr(module, self.merged_name, merged)
206
+ setattr(module, self.act_name, act_fn)
207
+ # Drop the consumed submodules so their (meta) params are not expected.
208
+ delattr(module, self.gate_name)
209
+ delattr(module, self.up_name)
210
+ # If there is a down projection, we know it must be rowwise.
211
+ if self.down_name is not None:
212
+ down_prefix = maybe_prefix(prefix, self.down_name)
213
+ down = module.get_submodule(self.down_name)
214
+ new_down = replace_linear_class(
215
+ down, "rowwise", quant_config, prefix=down_prefix
216
+ )
217
+ setattr(module, self.down_name, new_down)
218
+ log_replacement(down_prefix, down, new_down)
vllm/model_executor/models/transformers/fusers/moe.py
new
+268
@@ -0,0 +1,268 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing."""
4
+
5
+import ast
6
+import inspect
7
+import textwrap
8
+import types
9
+from collections.abc import Iterator
10
+from dataclasses import dataclass
11
+from itertools import chain
12
+
13
+import torch
14
+from torch import fx, nn
15
+
16
+from vllm.distributed import tensor_model_parallel_all_gather
17
+from vllm.model_executor.layers.linear import ReplicatedLinear
18
+from vllm.model_executor.models.transformers.fx_utils import (
19
+ find_node,
20
+ is_op,
21
+ peel,
22
+ trace,
23
+)
24
+from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk
25
+
26
+
27
+def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
28
+ """`module`'s own state (i.e. named parameters and buffers)."""
29
+ return chain(module.named_parameters(), module.named_buffers())
30
+
31
+
32
+def _own_returns(node: ast.AST) -> Iterator[ast.Return]:
33
+ """`return` statements in `node`'s own scope, not in nested functions."""
34
+ stack = list(ast.iter_child_nodes(node))
35
+ while stack:
36
+ child = stack.pop()
37
+ if isinstance(child, ast.Return):
38
+ yield child
39
+ elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
40
+ stack.extend(ast.iter_child_nodes(child))
41
+
42
+
43
+def _returns_tuple(cls: type[nn.Module]) -> bool:
44
+ """Does `cls.forward()` return a tuple?"""
45
+ try:
46
+ source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward)))
47
+ forward = ast.parse(source).body[0]
48
+ except (OSError, SyntaxError, TypeError, IndexError):
49
+ return True
50
+ # Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`.
51
+ tuple_names = {
52
+ target.id
53
+ for node in ast.walk(forward)
54
+ if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple)
55
+ for target in node.targets
56
+ if isinstance(target, ast.Name)
57
+ }
58
+
59
+ def yields_tuple(value: ast.expr | None) -> bool:
60
+ if isinstance(value, ast.Tuple):
61
+ return True
62
+ if isinstance(value, ast.Name):
63
+ return value.id in tuple_names
64
+ if isinstance(value, ast.IfExp):
65
+ return yields_tuple(value.body) or yields_tuple(value.orelse)
66
+ return False
67
+
68
+ return any(yields_tuple(ret.value) for ret in _own_returns(forward))
69
+
70
+
71
+def _is_scalar_gate(module: nn.Module) -> bool:
72
+ """A linear projecting to a single logit (the shared-expert sigmoid gate)."""
73
+ weight = getattr(module, "weight", None)
74
+ return (
75
+ isinstance(module, nn.Linear)
76
+ and weight is not None
77
+ and weight.ndim == 2
78
+ and weight.shape[0] == 1
79
+ )
80
+
81
+
82
+def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
83
+ """Returns the set of nodes reachable from `node` by following `key` edges."""
84
+ seen: set[fx.Node] = set()
85
+ stack = [node]
86
+ while stack:
87
+ n = stack.pop()
88
+ if n in seen:
89
+ continue
90
+ seen.add(n)
91
+ stack.extend(getattr(n, key))
92
+ return seen
93
+
94
+
95
+class SharedExpertMLP(nn.Module):
96
+ """Wraps an HF shared expert, applying the output gating it is paired with."""
97
+
98
+ def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None):
99
+ super().__init__()
100
+ self.shared_experts = shared_experts
101
+ self.gate = gate
102
+
103
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
104
+ out = self.shared_experts(hidden_states)
105
+ if self.gate is not None:
106
+ out = torch.sigmoid(self.gate(hidden_states)[0]) * out
107
+ return out
108
+
109
+
110
+def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
111
+ """Standard MoE block forward.
112
+
113
+ Routing and any shared experts are handled inside `self.experts: MoERunner`."""
114
+ orig_shape = hidden_states.shape
115
+ hidden_states = hidden_states.reshape(-1, orig_shape[-1])
116
+ num_tokens = hidden_states.shape[0]
117
+ is_sequence_parallel = self.experts.moe_config.is_sequence_parallel
118
+ if is_sequence_parallel:
119
+ hidden_states = sequence_parallel_chunk(hidden_states)
120
+ out = self.experts(hidden_states, router_logits=hidden_states)
121
+ if is_sequence_parallel:
122
+ out = tensor_model_parallel_all_gather(out, 0)[:num_tokens]
123
+ return out.reshape(orig_shape)
124
+
125
+
126
+@dataclass
127
+class MoEBlockFuser:
128
+ """Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""
129
+
130
+ gate_name: str
131
+ scoring_func: str
132
+ shared_name: str | None
133
+ shared_gate_name: str | None
134
+
135
+ @staticmethod
136
+ def _match_router(gate: nn.Module) -> str | None:
137
+ """Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`."""
138
+ if [name for name, _ in named_state(gate)] != ["weight"]:
139
+ return None
140
+ graph = trace(gate)
141
+ if graph is None:
142
+ return None
143
+ topk = find_node(graph, lambda n: is_op(n, "topk"))
144
+ if topk is None:
145
+ return None
146
+ # Exactly one scoring op upstream of the top-k, fed (transitively) by a linear.
147
+ scorers = [
148
+ n
149
+ for n in _reaches(topk, "all_input_nodes")
150
+ if is_op(n, "softmax") or is_op(n, "sigmoid")
151
+ ]
152
+ if len(scorers) != 1:
153
+ return None
154
+ scorer = scorers[0]
155
+ if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")):
156
+ return None
157
+ return "softmax" if is_op(scorer, "softmax") else "sigmoid"
158
+
159
+ @staticmethod
160
+ def _match_shared_experts(
161
+ graph: fx.Graph, experts: str
162
+ ) -> tuple[str | None, str | None]:
163
+ """Detects the shared expert and its optional gate by dataflow."""
164
+ experts_predicate = lambda n: n.op == "call_module" and n.target == experts
165
+ if (experts_node := find_node(graph, experts_predicate)) is None:
166
+ return None, None
167
+ from_experts = _reaches(experts_node, "users")
168
+ for add in graph.nodes:
169
+ if not is_op(add, "add"):
170
+ continue
171
+ operands = [a for a in add.args if isinstance(a, fx.Node)]
172
+ # Exactly one side is the experts' output; the other is the shared path.
173
+ sides = [a in from_experts for a in operands]
174
+ if len(operands) != 2 or sides.count(True) != 1:
175
+ continue
176
+ cone = _reaches(operands[sides.index(False)], "all_input_nodes")
177
+ modules = [n for n in cone if n.op == "call_module" and n.target != experts]
178
+ # A sigmoid wrapping one of those modules marks the shared-expert gate.
179
+ gate = next(
180
+ (
181
+ src
182
+ for n in cone
183
+ if is_op(n, "sigmoid")
184
+ and isinstance(src := peel(n.args[0]), fx.Node)
185
+ and src in modules
186
+ ),
187
+ None,
188
+ )
189
+ shared = [n for n in modules if n is not gate]
190
+ if len(shared) != 1:
191
+ return None, None
192
+ return shared[0].target, (gate.target if gate is not None else None)
193
+ return None, None
194
+
195
+ @classmethod
196
+ def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None":
197
+ # Standard MoE block returns a single tensor.
198
+ if _returns_tuple(type(moe_block)):
199
+ return None
200
+ # Router: the child that scores + top-k selects.
201
+ gate_name = scoring_func = None
202
+ for name, child in moe_block.named_children():
203
+ if name != experts_name and (func := cls._match_router(child)) is not None:
204
+ gate_name, scoring_func = name, func
205
+ break
206
+ if gate_name is None or scoring_func is None:
207
+ return None
208
+ # Shared expert: a child the block adds to the experts' output.
209
+ shared_name = shared_gate_name = None
210
+ others = [
211
+ n
212
+ for n, _ in moe_block.named_children()
213
+ if n not in {experts_name, gate_name}
214
+ ]
215
+ if others:
216
+ graph = trace(moe_block)
217
+ if graph is None:
218
+ return None
219
+ shared_name, shared_gate_name = cls._match_shared_experts(
220
+ graph, experts_name
221
+ )
222
+ if shared_gate_name is not None and not _is_scalar_gate(
223
+ getattr(moe_block, shared_gate_name)
224
+ ):
225
+ return None
226
+ # Fail closed: `rewrite_forward` runs only the experts and the detected
227
+ # shared expert, so any other stateful child would be dropped.
228
+ accounted = {experts_name, gate_name, shared_name, shared_gate_name}
229
+ for name, child in moe_block.named_children():
230
+ if name not in accounted and next(named_state(child), None) is not None:
231
+ return None
232
+ return cls(gate_name, scoring_func, shared_name, shared_gate_name)
233
+
234
+ def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
235
+ """Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE."""
236
+ num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape
237
+ gate = ReplicatedLinear(
238
+ hidden_size,
239
+ num_experts,
240
+ bias=False,
241
+ prefix=maybe_prefix(prefix, self.gate_name),
242
+ )
243
+ setattr(moe_block, self.gate_name, gate)
244
+ return gate
245
+
246
+ def shared_experts(
247
+ self, moe_block: nn.Module, prefix: str
248
+ ) -> SharedExpertMLP | None:
249
+ """Build the HF shared expert (and its optional gate)
250
+ as a `SharedExpertMLP` for vLLM's fused MoE."""
251
+ if self.shared_name is None:
252
+ return None
253
+ shared_experts = getattr(moe_block, self.shared_name)
254
+ gate = None
255
+ if self.shared_gate_name is not None:
256
+ hf_gate = getattr(moe_block, self.shared_gate_name)
257
+ gate = ReplicatedLinear(
258
+ hf_gate.in_features,
259
+ hf_gate.out_features,
260
+ bias=hf_gate.bias is not None,
261
+ prefix=maybe_prefix(prefix, self.shared_gate_name),
262
+ )
263
+ setattr(moe_block, self.shared_gate_name, gate)
264
+ return SharedExpertMLP(shared_experts, gate)
265
+
266
+ def rewrite_forward(self, moe_block: nn.Module) -> None:
267
+ """Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
268
+ moe_block.forward = types.MethodType(_moe_block_forward, moe_block)
vllm/model_executor/models/transformers/fusers/qkv.py
new
+220
@@ -0,0 +1,220 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split."""
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 QKVParallelLinear
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
+ replace_expr,
20
+ single_self_call,
21
+)
22
+from vllm.model_executor.models.transformers.utils import (
23
+ log_replacement,
24
+ replace_linear_class,
25
+)
26
+from vllm.model_executor.models.utils import ShardId, maybe_prefix
27
+
28
+if TYPE_CHECKING:
29
+ from vllm.config.model import ModelConfig
30
+ from vllm.model_executor.layers.quantization import QuantizationConfig
31
+
32
+logger = init_logger(__name__)
33
+
34
+
35
+@dataclass
36
+class QKVFuser(StackedFuser):
37
+ """Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
38
+
39
+ q_name: str
40
+ k_name: str
41
+ v_name: str
42
+ o_name: str | None
43
+ merged_name: ClassVar[str] = "qkv_proj"
44
+ merged_cls: ClassVar[str] = "QKVParallelLinear"
45
+
46
+ @property
47
+ def shards(self) -> list[tuple[str, ShardId]]:
48
+ return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]
49
+
50
+ @classmethod
51
+ def _get_qkv_nodes(
52
+ cls, graph: fx.Graph, module: nn.Module
53
+ ) -> tuple[fx.Node, fx.Node, fx.Node] | None:
54
+ """Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
55
+ by_input: dict[fx.Node, list[fx.Node]] = {}
56
+ for node in graph.nodes:
57
+ if (
58
+ is_linear(node, module)
59
+ and len(node.args) == 1
60
+ and not node.kwargs
61
+ and isinstance(node.args[0], fx.Node)
62
+ and node.args[0].op == "placeholder"
63
+ ):
64
+ by_input.setdefault(node.args[0], []).append(node)
65
+ triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
66
+ if len(triples) != 1:
67
+ return None
68
+
69
+ q_node, k_node, v_node = nodes = triples[0]
70
+ outs = [module.get_submodule(node.target).out_features for node in nodes]
71
+ if len(set(outs)) == 2:
72
+ # q is identified as the larger projection (GQA)
73
+ (q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
74
+ k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
75
+ if module.get_submodule(q_node.target).out_features != max(outs):
76
+ return None
77
+ elif len(set(outs)) != 1:
78
+ return None
79
+ return q_node, k_node, v_node
80
+
81
+ @classmethod
82
+ def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
83
+ if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
84
+ return None
85
+ q, k, v = qkv_nodes
86
+ names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
87
+ attn_width = module.get_submodule(q.target).out_features
88
+ candidates = [
89
+ name
90
+ for name, child in module.named_children()
91
+ if isinstance(child, nn.Linear)
92
+ and name not in names.values()
93
+ and child.in_features == attn_width
94
+ ]
95
+ names["o_name"] = candidates[0] if len(candidates) == 1 else None
96
+ return cls(source_cls=type(module).__name__, **names)
97
+
98
+ def update_forward(self, module: nn.Module) -> None:
99
+ """Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
100
+ funcdef, fn = recover_forward(type(module))
101
+ calls = [
102
+ single_self_call(funcdef, name)
103
+ for name in (self.q_name, self.k_name, self.v_name)
104
+ ]
105
+ arg_dumps = {ast.dump(call.args[0]) for call in calls}
106
+ if len(arg_dumps) != 1:
107
+ raise ValueError("projection inputs are written differently")
108
+ # The trace may be partial, so prove projection exclusivity in source:
109
+ # no other linear child may consume the same input (else the matched
110
+ # three may not be q, k and v)
111
+ other_linears = {
112
+ name
113
+ for name, child in module.named_children()
114
+ if isinstance(child, nn.Linear)
115
+ } - {self.q_name, self.k_name, self.v_name}
116
+ for node in ast.walk(funcdef):
117
+ if (
118
+ isinstance(node, ast.Call)
119
+ and isinstance(node.func, ast.Attribute)
120
+ and node.func.attr in other_linears
121
+ and any(ast.dump(arg) in arg_dumps for arg in node.args)
122
+ ):
123
+ raise ValueError("another linear consumes the same input")
124
+ blocks = [innermost_block(funcdef.body, call) for call in calls]
125
+ if any(found is None for found in blocks):
126
+ raise ValueError("projection calls not found in the function body")
127
+ if len({id(block) for block, _ in blocks}) != 1:
128
+ raise ValueError("projection calls are in different blocks")
129
+
130
+ # q(x), k(x), v(x) -> q, k, v = qkv(x).split(self.qkv.split_sizes, -1)
131
+ names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
132
+ temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
133
+ if names & set(temps):
134
+ raise ValueError("fused temporaries would shadow existing names")
135
+ merged = f"self.{self.merged_name}"
136
+ template = (
137
+ f"{', '.join(temps)} = {merged}(__arg__).split({merged}.split_sizes, -1)"
138
+ )
139
+ assign = ast.parse(template).body[0]
140
+ arg = next(
141
+ node
142
+ for node in ast.walk(assign)
143
+ if isinstance(node, ast.Name) and node.id == "__arg__"
144
+ )
145
+ replace_expr(assign, arg, calls[0].args[0])
146
+ block, index = blocks[0]
147
+ ast.copy_location(assign, block[index])
148
+ block.insert(min(index for _, index in blocks), assign)
149
+ for call, temp in zip(calls, temps):
150
+ replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
151
+ self.fused_forward = compile_forward(funcdef, fn)
152
+
153
+ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
154
+ """Shapes must be compatible for a single merged, head-sharded GEMM."""
155
+ q = module.get_submodule(self.q_name)
156
+ k = module.get_submodule(self.k_name)
157
+ v = module.get_submodule(self.v_name)
158
+ head_size = model_config.get_head_size()
159
+ compatible = (
160
+ q.in_features == k.in_features == v.in_features
161
+ and len({proj.bias is None for proj in (q, k, v)}) == 1
162
+ and k.out_features == v.out_features
163
+ and q.out_features % head_size == 0
164
+ and k.out_features % head_size == 0
165
+ )
166
+ if not compatible:
167
+ logger.debug("%s is not compatible with QKV fusion", type(module))
168
+ return compatible
169
+
170
+ def update_attrs(
171
+ self,
172
+ module: nn.Module,
173
+ prefix: str,
174
+ model_config: "ModelConfig",
175
+ quant_config: "QuantizationConfig",
176
+ ) -> None:
177
+ head_size = model_config.get_head_size()
178
+ q = module.get_submodule(self.q_name)
179
+ k = module.get_submodule(self.k_name)
180
+ merged = QKVParallelLinear(
181
+ hidden_size=q.in_features,
182
+ head_size=head_size,
183
+ total_num_heads=q.out_features // head_size,
184
+ total_num_kv_heads=k.out_features // head_size,
185
+ bias=q.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 -> %s: %s",
192
+ self.q_name,
193
+ q,
194
+ self.k_name,
195
+ k,
196
+ self.v_name,
197
+ module.get_submodule(self.v_name),
198
+ self.merged_name,
199
+ merged,
200
+ )
201
+ # The rewritten forward splits the merged projection into this rank's
202
+ # shard sizes (see `update_forward`)
203
+ merged.split_sizes = [
204
+ merged.num_heads * merged.head_size,
205
+ merged.num_kv_heads * merged.head_size,
206
+ merged.num_kv_heads * merged.v_head_size,
207
+ ]
208
+ setattr(module, self.merged_name, merged)
209
+ # Drop the consumed submodules so their (meta) params are not expected.
210
+ for name in (self.q_name, self.k_name, self.v_name):
211
+ delattr(module, name)
212
+ # If there is an output projection, we know it must be rowwise.
213
+ if self.o_name is not None:
214
+ o_proj_prefix = maybe_prefix(prefix, self.o_name)
215
+ o_proj = module.get_submodule(self.o_name)
216
+ new_o = replace_linear_class(
217
+ o_proj, "rowwise", quant_config, prefix=o_proj_prefix
218
+ )
219
+ setattr(module, self.o_name, new_o)
220
+ log_replacement(o_proj_prefix, o_proj, new_o)
vllm/model_executor/models/transformers/fusers/rms_norm.py
new
+217
@@ -0,0 +1,217 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm."""
4
+
5
+from dataclasses import dataclass
6
+from typing import TYPE_CHECKING
7
+
8
+import torch
9
+from torch import fx, nn
10
+
11
+from vllm.distributed import (
12
+ get_tensor_model_parallel_rank,
13
+ get_tensor_model_parallel_world_size,
14
+ tensor_model_parallel_all_gather,
15
+)
16
+from vllm.distributed.parallel_state import model_parallel_is_initialized
17
+from vllm.distributed.utils import split_tensor_along_last_dim
18
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
19
+from vllm.model_executor.models.transformers.fusers.base import BaseFuser
20
+from vllm.model_executor.models.transformers.fx_utils import (
21
+ find_node,
22
+ forward_input_count,
23
+ is_op,
24
+ peel,
25
+ trace,
26
+)
27
+
28
+if TYPE_CHECKING:
29
+ from vllm.config.model import ModelConfig
30
+ from vllm.model_executor.layers.quantization import QuantizationConfig
31
+
32
+
33
+def _is_squared(node: object, x: fx.Node) -> bool:
34
+ """`x**2`, `x.square()` or `x * x`, through any dtype casts."""
35
+ node = peel(node)
36
+ if is_op(node, "pow"):
37
+ base, exp = node.args
38
+ return peel(base) is x and exp == 2
39
+ if is_op(node, "square"):
40
+ return peel(node.args[0]) is x
41
+ if is_op(node, "mul"):
42
+ a, b = node.args
43
+ return peel(a) is x and peel(b) is x
44
+ return False
45
+
46
+
47
+def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
48
+ """eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
49
+ add = peel(rsqrt.args[0])
50
+ if not is_op(add, "add"):
51
+ return None
52
+ consts = [a for a in add.args if isinstance(a, (int, float))]
53
+ nodes = [a for a in add.args if isinstance(a, fx.Node)]
54
+ if len(consts) != 1 or len(nodes) != 1:
55
+ return None
56
+ mean = peel(nodes[0])
57
+ if not is_op(mean, "mean"):
58
+ return None
59
+ if not _is_squared(mean.args[0], x):
60
+ return None
61
+ return float(consts[0])
62
+
63
+
64
+def _is_one_plus(node: object) -> bool:
65
+ """`1 + weight` in either operand order (marks a zero-centered weight)."""
66
+ node = peel(node)
67
+ if not is_op(node, "add"):
68
+ return False
69
+ return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
70
+
71
+
72
+def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:
73
+ """Does the forward compute anything after `node` before returning?"""
74
+ output = find_node(graph, lambda n: n.op == "output")
75
+ if output is None or not output.args:
76
+ return False
77
+ return peel(output.args[0]) is not node
78
+
79
+
80
+class TPAwareNormMixin(nn.Module):
81
+ """Mixin for RMSNorms that reconstructs a TP-sharded input before normalizing."""
82
+
83
+ def __init__(self, *args, **kwargs):
84
+ super().__init__(*args, **kwargs)
85
+ if model_parallel_is_initialized():
86
+ self.tp_size = get_tensor_model_parallel_world_size()
87
+ self.tp_rank = get_tensor_model_parallel_rank()
88
+ else:
89
+ self.tp_size, self.tp_rank = 1, 0
90
+
91
+ def forward(
92
+ self, x: torch.Tensor, residual: torch.Tensor | None = None
93
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
94
+ if self.tp_size > 1 and x.shape[-1] < (full := self.weight.shape[0]):
95
+ if x.shape[-1] * self.tp_size != full:
96
+ raise ValueError(
97
+ f"Cannot gather norm of width {full}: a TP-sharded input of "
98
+ f"width {x.shape[-1]} does not tile it evenly across "
99
+ f"{self.tp_size} ranks (replicated or uneven sharding)."
100
+ )
101
+ x = tensor_model_parallel_all_gather(x.contiguous())
102
+ x = super().forward(x)
103
+ splits = split_tensor_along_last_dim(x, num_partitions=self.tp_size)
104
+ return splits[self.tp_rank]
105
+ return super().forward(x, residual)
106
+
107
+
108
+class TPAwareRMSNorm(TPAwareNormMixin, RMSNorm):
109
+ """`RMSNorm` that reconstructs a TP-sharded input before normalizing."""
110
+
111
+
112
+class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm):
113
+ """`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing."""
114
+
115
+
116
+@dataclass
117
+class RMSNormFuser(BaseFuser):
118
+ """Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""
119
+
120
+ zero_centered: bool
121
+ """Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
122
+ source_cls: str
123
+ """Class name of the norm this was matched from (for logging)."""
124
+
125
+ def info(self, name: str) -> str:
126
+ norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
127
+ return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)"
128
+
129
+ @classmethod
130
+ def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
131
+ """Match a graph to the RMSNorm pattern, returning a fuser if found."""
132
+ if forward_input_count(type(module)) != 1:
133
+ return None
134
+ x = find_node(graph, lambda n: n.op == "placeholder")
135
+ if x is None:
136
+ return None
137
+ # Handle native torch `rms_norm` op.
138
+ rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
139
+ if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
140
+ if _has_trailing_compute(graph, rms_norm):
141
+ return None
142
+ return cls(zero_centered=False, source_cls=type(module).__name__)
143
+ # Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
144
+ # The rsqrt over the mean-square variance is the spine of the norm.
145
+ rsqrt = None
146
+ for node in graph.nodes:
147
+ if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
148
+ rsqrt = node
149
+ break
150
+ if rsqrt is None:
151
+ return None
152
+ # The `x * rsqrt(...)` normalize multiply.
153
+ normalize = find_node(
154
+ graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
155
+ )
156
+ if normalize is None:
157
+ return None
158
+ # An optional later `weight * normalized` (or `(1 + weight) * normalized`).
159
+ tail, zero_centered = normalize, False
160
+ for node in graph.nodes:
161
+ if not is_op(node, "mul") or node is normalize:
162
+ continue
163
+ operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
164
+ if len(operands) == 2 and normalize in operands:
165
+ weight = next(o for o in operands if o is not normalize)
166
+ tail, zero_centered = node, _is_one_plus(weight)
167
+ break
168
+ # The norm must be the last compute in forward, or it is not a pure norm.
169
+ if _has_trailing_compute(graph, tail):
170
+ return None
171
+ return cls(zero_centered=zero_centered, source_cls=type(module).__name__)
172
+
173
+ @staticmethod
174
+ def _eps_from_graph(graph: fx.Graph) -> float | None:
175
+ """Extract the `eps` constant from the graph, if present."""
176
+ if (x := find_node(graph, lambda n: n.op == "placeholder")) is None:
177
+ return None
178
+ fused = find_node(graph, lambda n: is_op(n, "rms_norm"))
179
+ if fused is not None and fused.args and peel(fused.args[0]) is x:
180
+ args, kwargs = fused.args, fused.kwargs
181
+ eps = args[3] if len(args) > 3 else kwargs.get("eps")
182
+ return eps if isinstance(eps, (int, float)) else None
183
+ for node in graph.nodes:
184
+ if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
185
+ return eps
186
+ return None
187
+
188
+ def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
189
+ return True
190
+
191
+ def fuse(
192
+ self,
193
+ module: nn.Module,
194
+ prefix: str,
195
+ model_config: "ModelConfig",
196
+ quant_config: "QuantizationConfig",
197
+ ) -> nn.Module:
198
+ """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
199
+ weight = getattr(module, "weight", None)
200
+ hidden_size = (
201
+ weight.size(0) if weight is not None else model_config.get_hidden_size()
202
+ )
203
+ graph = trace(module)
204
+ eps = self._eps_from_graph(graph) if graph is not None else None
205
+ if eps is None:
206
+ # If eps not in graph, match torch behaviour.
207
+ dtype = weight.dtype if weight is not None else model_config.dtype
208
+ eps = torch.finfo(dtype).eps
209
+ if self.zero_centered:
210
+ return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
211
+ has_weight = weight is not None
212
+ return TPAwareRMSNorm(
213
+ hidden_size=hidden_size,
214
+ eps=eps,
215
+ has_weight=has_weight,
216
+ dtype=weight.dtype if has_weight else None,
217
+ )
vllm/model_executor/models/transformers/fx_utils.py
new
+273
@@ -0,0 +1,273 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""fx tracing and forward-source rewriting for the Transformers backend fusers.
4
+
5
+A small engine, independent of any particular pattern: trace a module's forward
6
+with `torch.fx` (tolerating a partial graph), inspect the resulting nodes, and
7
+rewrite the forward's *source* (AST) so only matched calls change while the rest
8
+stays live Python. `fusion.py` builds the concrete fusion patterns on top.
9
+"""
10
+
11
+import ast
12
+import inspect
13
+import operator
14
+import textwrap
15
+from collections.abc import Callable
16
+
17
+import torch
18
+from torch import fx, nn
19
+from torch.nn import functional as F
20
+
21
+from vllm.logger import init_logger
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+def _infer_len(node: fx.Node) -> int | None:
27
+ """Concrete length of a proxy's value, inferred from its node chain.
28
+
29
+ Lets tracing pass through the shape unpacks and `*`-splats (e.g.
30
+ `(*input_shape, -1, head_dim)`) that precede the patterns in HF attention.
31
+ """
32
+ # `x.shape` has the rank of `x`, when known
33
+ if (
34
+ node.op == "call_function"
35
+ and node.target is getattr
36
+ and node.args[1] == "shape"
37
+ and (rank := _rank(node.args[0])) is not None
38
+ ):
39
+ return rank
40
+ # Slices of known-length values
41
+ if node.op == "call_function" and node.target is operator.getitem:
42
+ src_len = _infer_len(node.args[0])
43
+ index = node.args[1]
44
+ if src_len is not None and isinstance(index, slice):
45
+ return len(range(*index.indices(src_len)))
46
+ return None
47
+
48
+
49
+def _rank(node: fx.Node) -> int | None:
50
+ """The tensor rank of `node`'s value, if known."""
51
+ # vLLM always feeds the model [1, seq_len, hidden_size] hidden states
52
+ if node.op == "placeholder" and node.target == "hidden_states":
53
+ return 3
54
+ return None
55
+
56
+
57
+class _SizedProxy(fx.Proxy):
58
+ """Proxy whose `len` is inferred from the graph (see `_infer_len`)."""
59
+
60
+ def __len__(self) -> int:
61
+ length = _infer_len(self.node)
62
+ if length is None:
63
+ return super().__len__()
64
+ return length
65
+
66
+
67
+class _AllLeafTracer(fx.Tracer):
68
+ """Tracer that treats every submodule as a leaf.
69
+
70
+ Each child stays one `call_module` node, so matching sees the module's own
71
+ forward structure (activations aren't decomposed into e.g. `sigmoid * x`).
72
+ `iter` traces through the leading shape unpacks (see `_infer_len`); anything
73
+ else untraceable ends the trace early and the partial graph is matched.
74
+ """
75
+
76
+ def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool:
77
+ return True
78
+
79
+ def proxy(self, node: fx.Node) -> fx.Proxy:
80
+ return _SizedProxy(node, self)
81
+
82
+ def iter(self, obj: fx.Proxy):
83
+ length = _infer_len(obj.node)
84
+ if length is None:
85
+ return super().iter(obj)
86
+ return iter([obj[i] for i in range(length)])
87
+
88
+
89
+def trace(module: nn.Module) -> fx.Graph | None:
90
+ """Trace `module.forward`, returning the partial graph on failure.
91
+
92
+ The graph is only evidence for matching, and the patterns sit at the top of
93
+ their forwards, so a trace that fails partway can still be matched."""
94
+ tracer = _AllLeafTracer()
95
+ try:
96
+ return tracer.trace(module)
97
+ except Exception as exc:
98
+ logger.debug("Could not fully trace %s: %s", type(module), exc)
99
+ return getattr(tracer, "graph", None)
100
+
101
+
102
+def recover_forward(cls: type[nn.Module]) -> tuple[ast.FunctionDef, Callable]:
103
+ """Parse the source of `cls.forward`, ready for rewriting."""
104
+ fn = inspect.unwrap(cls.forward)
105
+ if fn.__code__.co_freevars:
106
+ raise ValueError("forward is a closure")
107
+ tree = ast.parse(textwrap.dedent(inspect.getsource(fn)))
108
+ funcdef = tree.body[0]
109
+ if not isinstance(funcdef, ast.FunctionDef):
110
+ raise ValueError("source is not a plain function definition")
111
+ # `fn` is already unwrapped; don't re-apply its decorators
112
+ funcdef.decorator_list.clear()
113
+ # Annotations may not evaluate outside the defining module (e.g. with
114
+ # postponed evaluation); they're not needed at runtime
115
+ funcdef.returns = None
116
+ args = funcdef.args
117
+ for arg in (
118
+ *args.posonlyargs,
119
+ *args.args,
120
+ *args.kwonlyargs,
121
+ *filter(None, (args.vararg, args.kwarg)),
122
+ ):
123
+ arg.annotation = None
124
+ # Recompiling outside the class body would break name mangling
125
+ for node in ast.walk(funcdef):
126
+ name = getattr(node, "attr", None) or getattr(node, "id", None)
127
+ if name and name.startswith("__") and not name.endswith("__"):
128
+ raise ValueError(f"{name} would be name mangled")
129
+ return funcdef, fn
130
+
131
+
132
+def forward_input_count(cls: type[nn.Module]) -> int:
133
+ """The number of tensor inputs `cls.forward` declares, excluding `self` and
134
+ any `*args`/`**kwargs`. Read from the signature, so it is independent of
135
+ whether the trace completes (unlike counting placeholders)."""
136
+ try:
137
+ params = list(inspect.signature(cls.forward).parameters.values())[1:]
138
+ except (ValueError, TypeError):
139
+ return 1 # uninspectable: assume a single input and let matching decide
140
+ fixed = (
141
+ inspect.Parameter.POSITIONAL_ONLY,
142
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
143
+ inspect.Parameter.KEYWORD_ONLY,
144
+ )
145
+ return sum(1 for p in params if p.kind in fixed)
146
+
147
+
148
+def compile_forward(funcdef: ast.FunctionDef, fn: Callable) -> Callable:
149
+ """Compile `funcdef` in `fn`'s module so tracebacks point at the source."""
150
+ module = ast.Module(body=[funcdef], type_ignores=[])
151
+ ast.fix_missing_locations(module)
152
+ ast.increment_lineno(module, fn.__code__.co_firstlineno - 1)
153
+ code = compile(module, fn.__code__.co_filename, "exec")
154
+ namespace: dict = {}
155
+ exec(code, fn.__globals__, namespace)
156
+ return namespace[funcdef.name]
157
+
158
+
159
+def single_self_call(funcdef: ast.FunctionDef, name: str) -> ast.Call:
160
+ """The unique `self.<name>(arg)` call in `funcdef`.
161
+
162
+ Raises unless `name` appears exactly once, as such a call, so the source
163
+ rewrite agrees with the fx match.
164
+ """
165
+ uses = [
166
+ node
167
+ for node in ast.walk(funcdef)
168
+ if isinstance(node, ast.Attribute) and node.attr == name
169
+ ]
170
+ if len(uses) != 1:
171
+ raise ValueError(f"{name} is referenced {len(uses)} times")
172
+ calls = [
173
+ node
174
+ for node in ast.walk(funcdef)
175
+ if isinstance(node, ast.Call)
176
+ and node.func is uses[0]
177
+ and len(node.args) == 1
178
+ and not isinstance(node.args[0], ast.Starred)
179
+ and not node.keywords
180
+ ]
181
+ if (
182
+ len(calls) != 1
183
+ or not isinstance(uses[0].value, ast.Name)
184
+ or uses[0].value.id != "self"
185
+ ):
186
+ raise ValueError(f"{name} is not a single-argument call on self")
187
+ return calls[0]
188
+
189
+
190
+def innermost_block(
191
+ block: list[ast.stmt], node: ast.AST
192
+) -> tuple[list[ast.stmt], int] | None:
193
+ """The innermost statement list containing `node`, and the index within."""
194
+ for index, stmt in enumerate(block):
195
+ if not any(child is node for child in ast.walk(stmt)):
196
+ continue
197
+ child_blocks = [
198
+ getattr(stmt, fld, None) for fld in ("body", "orelse", "finalbody")
199
+ ]
200
+ child_blocks += [h.body for h in getattr(stmt, "handlers", [])]
201
+ child_blocks += [c.body for c in getattr(stmt, "cases", [])]
202
+ for child_block in child_blocks:
203
+ if (
204
+ isinstance(child_block, list)
205
+ and child_block
206
+ and (found := innermost_block(child_block, node)) is not None
207
+ ):
208
+ return found
209
+ return block, index
210
+ return None
211
+
212
+
213
+def replace_expr(module: ast.AST, old: ast.expr, new: ast.expr) -> None:
214
+ """Replace the expression `old` (by identity) with `new` within `module`."""
215
+
216
+ class _Replacer(ast.NodeTransformer):
217
+ def visit(self, node: ast.AST) -> ast.AST:
218
+ if node is old:
219
+ return new
220
+ return super().generic_visit(node)
221
+
222
+ _Replacer().visit(module)
223
+
224
+
225
+def find_node(graph: fx.Graph, predicate: Callable[[fx.Node], bool]) -> fx.Node | None:
226
+ """The first node in `graph` matching `predicate`, or `None`."""
227
+ return next((n for n in graph.nodes if predicate(n)), None)
228
+
229
+
230
+def is_linear(node: fx.Node, module: nn.Module) -> bool:
231
+ """Is node `nn.Linear.__call__()`."""
232
+ return node.op == "call_module" and isinstance(
233
+ module.get_submodule(node.target), nn.Linear
234
+ )
235
+
236
+
237
+_DTYPE_CASTS = frozenset({"to", "float", "double", "half", "bfloat16", "type_as"})
238
+
239
+
240
+def peel(node: object) -> object:
241
+ """Strip dtype-cast wrappers (`.to(...)`, `.float()`, `.type_as(...)`)."""
242
+ while (
243
+ isinstance(node, fx.Node)
244
+ and node.op == "call_method"
245
+ and node.target in _DTYPE_CASTS
246
+ ):
247
+ node = node.args[0]
248
+ return node
249
+
250
+
251
+def is_fn(node: object, target: Callable) -> bool:
252
+ """Is node `<target>()`."""
253
+ return (
254
+ isinstance(node, fx.Node)
255
+ and node.op == "call_function"
256
+ and node.target is target
257
+ )
258
+
259
+
260
+def is_method(node: object, name: str) -> bool:
261
+ """Is node `.<name>()`."""
262
+ return (
263
+ isinstance(node, fx.Node) and node.op == "call_method" and node.target == name
264
+ )
265
+
266
+
267
+def is_op(node: object, name: str) -> bool:
268
+ """
269
+ Is node `torch.<name>()`, `F.<name>()`, `operator.<name>()`, or `Tensor.<name>()`.
270
+ """
271
+ return any(
272
+ is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)
273
+ ) or (hasattr(torch.Tensor, name) and is_method(node, name))
vllm/model_executor/models/transformers/moe.py
+83
-46
@@ -26,9 +26,11 @@ import torch.nn as nn
26
from vllm.config.utils import getattr_iter
27
from vllm.distributed import get_dp_group, get_ep_group
28
from vllm.forward_context import ForwardContext, get_forward_context
29
+from vllm.logger import init_logger
30
from vllm.model_executor.custom_op import PluggableLayer
31
from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner, RoutedExperts
32
from vllm.model_executor.models.interfaces import MixtureOfExperts
33
+from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
34
from vllm.model_executor.models.utils import maybe_prefix
35
from vllm.utils.torch_utils import direct_register_custom_op
36
@@ -37,6 +39,8 @@ from .utils import log_replacement
39
if TYPE_CHECKING:
40
from vllm.config import VllmConfig
41
42
+logger = init_logger(__name__)
43
+
44
45
@dataclass
46
class TransformersMoEState:
@@ -142,14 +146,15 @@ class MoEMixin(MixtureOfExperts):
146
self.num_physical_experts = num_physical_experts
147
self.num_local_physical_experts = num_local_physical_experts
148
self.num_redundant_experts = num_physical_experts - self.num_logical_experts
145
- for mlp in self.mlp_layers:
146
- mlp.n_local_physical_experts = num_local_physical_experts
147
- mlp.n_physical_experts = num_physical_experts
148
- mlp.n_redundant_experts = self.num_redundant_experts
149
- mlp.experts.update_expert_map()
149
+ for moe_block in self.mlp_layers:
150
+ moe_block.n_local_physical_experts = num_local_physical_experts
151
+ moe_block.n_physical_experts = num_physical_experts
152
+ moe_block.n_redundant_experts = self.num_redundant_experts
153
+ moe_block.experts.update_expert_map()
154
155
def recursive_replace(self):
156
"""Initialize the MoE layers."""
157
+ experts_name = "experts"
158
text_config = self.text_config
159
160
# Positional arguments
@@ -217,10 +222,13 @@ class MoEMixin(MixtureOfExperts):
222
# down_proj = (num_experts, intermediate_size, hidden_size)
223
params = list(child_module.parameters())
224
is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
220
- if child_name == "experts" and (is_modulelist or is_3d):
225
+ if child_name == experts_name and (is_modulelist or is_3d):
226
# Alias for readability
222
- mlp = module
227
+ moe_block = module
228
experts = child_module
229
+ # Class of the fused block (parent of gate/experts/shared)
230
+ moe_block_cls = type(moe_block).__name__
231
+ experts_cls = type(experts).__name__
232
# Do the experts have biases
233
has_bias = False
234
for experts_param_name, _ in experts.named_parameters():
@@ -230,64 +238,93 @@ class MoEMixin(MixtureOfExperts):
238
# If the config does not specify num_shared_experts, but
239
# the model has shared experts, we assume there is one.
240
if self.num_shared_experts == 0:
233
- for mlp_param_name, _ in mlp.named_parameters():
234
- if "shared_expert" in mlp_param_name:
241
+ for moe_block_param_name, _ in moe_block.named_parameters():
242
+ if "shared_expert" in moe_block_param_name:
243
self.num_shared_experts = 1
244
break
245
238
- # Replace experts module with FusedMoE
239
- moe_state = TransformersMoEState()
240
-
241
- def custom_routing_function(
242
- hidden_states: torch.Tensor,
243
- gating_output: torch.Tensor,
244
- topk: int,
245
- renormalize: bool,
246
- moe_state: TransformersMoEState,
247
- ):
248
- """Return `topk_weights` from `gating_output` and the
249
- `topk_ids` we stored in the layer earlier."""
250
- topk_weights = gating_output
251
- topk_ids = moe_state.topk_ids
252
- assert topk_ids is not None
253
- # Handle all gather in expert parallel
254
- if topk_ids.size(0) != hidden_states.size(0):
255
- dp_metadata = get_forward_context().dp_metadata
256
- sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
257
- is_sp = moe_state.is_sequence_parallel
258
- dist_group = get_ep_group() if is_sp else get_dp_group()
259
- assert sizes[dist_group.rank_in_group] == topk_ids.shape[0]
260
- (topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes)
261
- return topk_weights, topk_ids
262
-
263
- fused_experts = FusedMoE(
246
+ kwargs: dict[str, Any] = dict(
247
num_experts=num_experts,
248
top_k=top_k,
249
hidden_size=hidden_size,
250
intermediate_size=intermediate_size,
251
renormalize=renormalize,
269
- # Hard coded because topk happens in Transformers
252
use_grouped_topk=False,
271
- num_expert_group=num_expert_group,
272
- topk_group=topk_group,
253
quant_config=self.quant_config,
254
prefix=qual_name,
255
activation=activation,
256
enable_eplb=enable_eplb,
257
num_redundant_experts=num_redundant_experts,
258
has_bias=has_bias,
279
- custom_routing_function=partial(
280
- custom_routing_function,
281
- moe_state=moe_state,
282
- ),
283
- runner_cls=TransformersMoERunner,
259
routed_experts_cls=TransformersRoutedExperts,
285
- runner_args={"moe_state": moe_state},
260
)
287
- mlp.experts = fused_experts
261
+ fuser = MoEBlockFuser.match(moe_block, experts_name)
262
+ if self.num_expert_groups <= 1 and fuser is not None:
263
+ # MoE block forward is fully replaced.
264
+ # gate/router and shared expert (if any) runs in FusedMoE.
265
+ kwargs |= dict(
266
+ scoring_func=fuser.scoring_func,
267
+ is_sequence_parallel=(
268
+ self.parallel_config.use_sequence_parallel_moe
269
+ ),
270
+ gate=fuser.gate(moe_block, prefix),
271
+ shared_experts=fuser.shared_experts(moe_block, prefix),
272
+ )
273
+ fuser.rewrite_forward(moe_block)
274
+ routed = "gate + experts"
275
+ if fuser.shared_name:
276
+ routed += " + shared experts"
277
+ logger.info_once(
278
+ "Fused: %s (%s) -> FusedMoE (internal routing)",
279
+ routed,
280
+ moe_block_cls,
281
+ )
282
+ else:
283
+ # MoE block forward is unmodified.
284
+ # gate/router and shared expert (if any) runs in Transformers.
285
+ # We then smuggle the topk_ids in using a custom op.
286
+ moe_state = TransformersMoEState()
287
+
288
+ def custom_routing_function(
289
+ hidden_states: torch.Tensor,
290
+ gating_output: torch.Tensor,
291
+ topk: int,
292
+ renormalize: bool,
293
+ moe_state: TransformersMoEState,
294
+ ):
295
+ """Return `topk_weights` from `gating_output` and the
296
+ `topk_ids` we stored in the layer earlier."""
297
+ topk_weights = gating_output
298
+ topk_ids = moe_state.topk_ids
299
+ assert topk_ids is not None
300
+ # Handle all gather in expert parallel
301
+ if topk_ids.size(0) != hidden_states.size(0):
302
+ dp_metadata = get_forward_context().dp_metadata
303
+ sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
304
+ is_sp = moe_state.is_sequence_parallel
305
+ group = get_ep_group() if is_sp else get_dp_group()
306
+ assert sizes[group.rank_in_group] == topk_ids.shape[0]
307
+ (topk_ids,) = group.all_gatherv([topk_ids], 0, sizes)
308
+ return topk_weights, topk_ids
309
+
310
+ kwargs |= dict(
311
+ num_expert_group=num_expert_group,
312
+ topk_group=topk_group,
313
+ custom_routing_function=partial(
314
+ custom_routing_function, moe_state=moe_state
315
+ ),
316
+ runner_cls=TransformersMoERunner,
317
+ runner_args={"moe_state": moe_state},
318
+ )
319
+ logger.info_once(
320
+ "Fused: experts (%s) -> FusedMoE (external routing)",
321
+ experts_cls,
322
+ )
323
+ fused_experts = FusedMoE(**kwargs)
324
+ moe_block.experts = fused_experts
325
log_replacement(qual_name, experts, fused_experts)
326
# Update MixtureOfExperts mixin state
290
- self.mlp_layers.append(mlp)
327
+ self.mlp_layers.append(moe_block)
328
self.moe_layers.append(fused_experts)
329
else:
330
_recursive_replace(child_module, prefix=qual_name)
vllm/model_executor/models/transformers/utils.py
-41
@@ -23,10 +23,8 @@ from typing import TYPE_CHECKING, Literal
23
import torch
24
from torch import nn
25
26
-from vllm.config.utils import getattr_iter
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
29
-from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
28
from vllm.model_executor.layers.linear import (
29
ColumnParallelLinear,
30
ReplicatedLinear,
@@ -183,45 +181,6 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
181
)
182
183
186
-def replace_rms_norm_class(rms_norm: nn.Module, hidden_size: int) -> RMSNorm:
187
- """Replace a Transformers RMSNorm with vLLM's RMSNorm.
188
-
189
- This method assumes:
190
- - Weight is stored as `weight`.
191
- - Epsilon is stored as `eps` or `variance_epsilon`.
192
- - `with_scale` indicates whether the layer has a weight (Gemma3n only).
193
- - `var_hidden_size` is only ever used for Intern vision encoder in vLLM
194
- and Transformers doesn't appear to have the same concept.
195
- """
196
- eps = getattr_iter(rms_norm, ("eps", "variance_epsilon"), 1e-6)
197
- kwargs = {"hidden_size": hidden_size, "eps": eps}
198
- # Update hidden size if weight is available
199
- weight_meta = getattr(rms_norm, "weight", None)
200
- if weight_meta is not None:
201
- kwargs["hidden_size"] = weight_meta.size(0)
202
- # Check if weight is all zeros, which indicates GemmaRMSNorm
203
- # We must create a new instance because rms_norm is on meta
204
- try:
205
- with torch.device("cpu"):
206
- weight_test = getattr(rms_norm.__class__(1), "weight", None)
207
- except Exception:
208
- logger.warning(
209
- "Failed to determine if RMSNorm weight is centered on zero or one. "
210
- "Defaulting to one."
211
- )
212
- weight_test = None
213
- if weight_test is not None and torch.all(weight_test == 0):
214
- return GemmaRMSNorm(**kwargs)
215
- # Otherwise assume it's a regular RMSNorm
216
- kwargs["has_weight"] = getattr(rms_norm, "with_scale", True)
217
- if weight_meta is not None:
218
- kwargs["dtype"] = weight_meta.dtype
219
- else:
220
- # No weight, fall back to weightless RMSNorm
221
- kwargs["has_weight"] = False
222
- return RMSNorm(**kwargs)
223
-
224
-
184
def recursive_replace_linear(
185
model: nn.Module,
186
quant_config: "QuantizationConfig | None",