A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 18:30 UTC
[torch.compile] Refactor Attention Quant Fusion Pass and Remove Boilerplate (#37373)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
BadrBasowid committed
Apr 1, 2026 at 02:15 UTC
077a9a8e3743a8ffe790a043666b01b196ed9854
7 files changed
+278
-215
docs/design/cuda_graphs.md
+1
-1
@@ -225,7 +225,7 @@ outputs = model.generate(
225
226
### Piecewise compilation and full graph custom passes (attention fusion, sequence parallelism)
227
228
-Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs.
228
+Unfortunately, some custom compile passes have to see the whole graph to be effective and hence aren't compatible with piecewise compilation. This includes `AttnQuantFusionPass` and `SequenceParallelismPass`. As a short-term solution, we automatically disable piecewise compilation (by setting `splitting_ops=[]`) when attention fusion is enabled. We use CUDA Graph modes `FULL` or `FULL_DECODE_ONLY` (depending on backend support). However, this leads to another optimization incompatibility and confusing performance tradeoffs.
229
230
Long term, we've added the ability to partition the graph in Inductor instead of right after Dynamo. It can be enabled with `CompilationConfig.use_inductor_graph_partition=True` but is currently experimental and only available with `torch>=2.9`. This also increases compilation time as it has to compile the whole graph and cannot reuse piecewise compilation artifacts. Once vLLM supports 2.9, we plan to make this the default approach as it will also speed up piecewise cudagraph capture.
231
tests/compile/fusions_e2e/conftest.py
+26
-5
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import logging
4
+from collections import defaultdict
5
6
import pytest
7
import regex as re
@@ -52,6 +53,16 @@ def run_model(compile_config: int | CompilationConfig, model: str, **model_kwarg
53
llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints
54
)
55
56
+ # Fetch match table from each worker via RPC and sum across workers.
57
+ worker_tables = llm.llm_engine.engine_core.collective_rpc(
58
+ "get_compilation_match_table"
59
+ )
60
+ combined: defaultdict[str, int] = defaultdict(int)
61
+ for table in worker_tables:
62
+ for k, v in table.items():
63
+ combined[k] += v
64
+ return dict(combined)
65
+
66
67
@pytest.fixture
68
def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
@@ -113,7 +124,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
124
)
125
126
with caplog_mp_spawn(logging.DEBUG) as log_holder:
116
- run_model(full_compilation_config, model_name, **model_kwargs)
127
+ match_table = run_model(full_compilation_config, model_name, **model_kwargs)
128
129
num_compile_ranges = len(full_compilation_config.get_compile_ranges())
130
assert num_compile_ranges in [1, 2, 3]
@@ -155,11 +166,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
166
else:
167
num_ranges_activated = num_compile_ranges
168
169
+ # TODO: Remove log counting in unit tests
170
+ # once all matchers implement VllmFusionPatternMatcherPass
171
n_expected = tp_size * num_ranges_activated
159
- assert len(log_matches) == n_expected, (
160
- f"Could not find {n_expected} {match_name} "
161
- f"(found {len(log_matches)}) in:\n {log_holder.text}"
162
- )
172
+ if match_name != "attn_quant_fusion":
173
+ assert len(log_matches) == n_expected, (
174
+ f"Could not find {n_expected} {match_name} "
175
+ f"(found {len(log_matches)}) in:\n {log_holder.text}"
176
+ )
177
178
expected_matches = getattr(matches, match_name)
179
@@ -215,6 +229,13 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
229
f"{tp_size * (num_ranges_activated - 1)} large-range "
230
f"entries (SP took precedence), found: {log_matches}"
231
)
232
+
233
+ elif match_name == "attn_quant_fusion":
234
+ actual_match = match_table.get(match_name, 0)
235
+ assert actual_match == expected_matches * n_expected, (
236
+ f"Could not find {expected_matches * n_expected} "
237
+ f"{match_name} (found {actual_match})."
238
+ )
239
else:
240
expected_matches_list = [expected_matches] * n_expected
241
assert sorted(log_matches) == expected_matches_list, (
tests/compile/passes/test_fusion_attn.py
+6
-3
@@ -9,7 +9,10 @@ from tests.compile.backend import LazyInitPass, TestBackend
9
from tests.utils import TestFP8Layer, flat_product
10
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
-from vllm.compilation.passes.fusion.attn_quant_fusion import ATTN_OP, AttnFusionPass
12
+from vllm.compilation.passes.fusion.attn_quant_fusion import (
13
+ ATTN_OP,
14
+ AttnQuantFusionPass,
15
+)
16
from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS
17
from vllm.compilation.passes.fx_utils import find_op_nodes
18
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
@@ -384,7 +387,7 @@ def test_attention_quant_pattern(
387
388
# Create test backend with fusion passes enabled
389
noop_pass = NoOpEliminationPass(vllm_config)
387
- attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
390
+ attn_pass = LazyInitPass(AttnQuantFusionPass, vllm_config)
391
cleanup_pass = PostCleanupPass(vllm_config)
392
393
test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
@@ -434,7 +437,7 @@ def test_attention_quant_pattern(
437
# Only output quant ops are fused into attention.
438
test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
439
437
- # access the underlying `AttnFusionPass` on the `LazyInitPass`
440
+ # access the underlying `AttnQuantFusionPass` on the `LazyInitPass`
441
assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
442
443
# Check attention ops in the graph before and after fusion
vllm/compilation/passes/fusion/attn_quant_fusion.py
+96
-197
@@ -1,15 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from abc import ABC, abstractmethod
4
+
5
from collections.abc import Callable
6
-from typing import Any, ParamSpec
6
7
import torch
9
-import torch._inductor.pattern_matcher as pm
10
-from torch import fx
8
from torch._higher_order_ops.auto_functionalize import auto_functionalized
12
-from torch._inductor.pattern_matcher import PatternMatcherPass
9
10
from vllm.config import VllmConfig, get_layers_from_vllm_config
11
from vllm.logger import init_logger
@@ -22,14 +18,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
from vllm.platforms import current_platform
19
from vllm.utils.math_utils import round_up
20
25
-from ..fx_utils import is_func
26
-from ..inductor_pass import enable_fake_mode
27
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21
+from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
22
from .matcher_utils import MatcherQuantFP8
29
-from .rms_quant_fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
23
+from .rms_quant_fusion import QUANT_OPS
24
25
logger = init_logger(__name__)
32
-P = ParamSpec("P")
26
+
27
FP8_DTYPE = current_platform.fp8_dtype()
28
FP4_DTYPE = torch.uint8
29
@@ -37,83 +31,10 @@ ATTN_OP = torch.ops.vllm.unified_attention_with_output.default
31
RESHAPE_OP = torch.ops.aten.reshape.default
32
33
40
-class AttentionQuantPattern(ABC):
41
- """
42
- The base class for Attn+Quant fusions.
43
- Should not be used directly.
44
- """
34
+_FP8_QUANT_KEY = QuantKey(dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=True)
35
46
- def __init__(
47
- self,
48
- layer: Attention,
49
- quant_key: QuantKey,
50
- dtype: torch.dtype,
51
- ) -> None:
52
- self.layer = layer
53
- self.layer_name = layer.layer_name
54
- self.num_heads = layer.num_heads
55
- self.head_size = layer.head_size
56
- self.quant_key = quant_key
57
- self.quant_dtype = quant_key.dtype
58
- self.dtype = dtype
59
-
60
- assert self.quant_key in QUANT_OPS, (
61
- f"unsupported quantization scheme {self.quant_key}"
62
- )
63
- self.QUANT_OP = QUANT_OPS[self.quant_key]
64
-
65
- def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
66
- kwargs = {"dtype": self.dtype, "device": "cuda", **kwargs}
67
- return torch.empty(*args, **kwargs)
68
-
69
- def empty_quant(self, *args: Any, **kwargs: Any) -> torch.Tensor:
70
- kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
71
- return torch.empty(*args, **kwargs)
72
-
73
- @staticmethod
74
- def wrap_trace_fn(
75
- trace_fn: Callable[P, fx.GraphModule],
76
- *process_fx_fns: Callable[[fx.GraphModule], None],
77
- ) -> Callable[P, fx.GraphModule]:
78
- def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
79
- gm = trace_fn(*args, **kwargs)
80
- for process_fx in process_fx_fns:
81
- process_fx(gm)
82
-
83
- return gm
84
-
85
- return wrapped
86
-
87
- @staticmethod
88
- def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
89
- from torch._inductor.fx_passes.post_grad import view_to_reshape
90
-
91
- view_to_reshape(gm)
92
-
93
- @staticmethod
94
- def remove_noop_permutes(gm: torch.fx.GraphModule) -> None:
95
- for node in gm.graph.nodes:
96
- if not is_func(node, torch.ops.aten.permute.default):
97
- continue
98
-
99
- dims = node.args[1]
100
- if any(dim != i for i, dim in enumerate(dims)):
101
- continue
102
-
103
- # this is now an identity op, remove
104
- node.replace_all_uses_with(node.args[0])
105
- gm.graph.erase_node(node)
106
-
107
- def register_if_supported(self, pm_pass: PatternMatcherPass) -> None:
108
- if self.layer.impl.fused_output_quant_supported(self.quant_key):
109
- self._register(pm_pass)
110
-
111
- @abstractmethod
112
- def _register(self, pm_pass: PatternMatcherPass) -> None:
113
- raise NotImplementedError
114
-
115
-
116
-class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
36
+
37
+class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
38
"""
39
Fusion for Attention+Fp8StaticQuant.
40
@@ -123,20 +44,16 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
44
will be passed into Attention op as the `output_scale` argument.
45
"""
46
126
- def __init__(
127
- self,
128
- layer: Attention,
129
- dtype: torch.dtype,
130
- symmetric: bool = True,
131
- ) -> None:
132
- quant_key = QuantKey(
133
- dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=symmetric
134
- )
135
- super().__init__(layer, quant_key, dtype)
136
- self.quant_matcher = MatcherQuantFP8(quant_key)
137
-
138
- def _register(self, pm_pass: PatternMatcherPass) -> None:
139
- def pattern(
47
+ def __init__(self, layer: Attention, dtype: torch.dtype):
48
+ self._layer_name = layer.layer_name
49
+ self._num_heads = layer.num_heads
50
+ self._head_size = layer.head_size
51
+ self._dtype = dtype
52
+ self._quant_matcher = MatcherQuantFP8(_FP8_QUANT_KEY)
53
+
54
+ @property
55
+ def pattern(self) -> Callable[..., torch.Tensor]:
56
+ def _pattern(
57
q: torch.Tensor,
58
k: torch.Tensor,
59
v: torch.Tensor,
@@ -150,18 +67,21 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
67
key=k,
68
value=v,
69
output=output_attn,
153
- layer_name=self.layer_name,
70
+ layer_name=self._layer_name,
71
output_scale=None,
72
output_block_scale=None,
73
kv_cache_dummy_dep=kv_cache_dummy_dep,
74
)
75
attn_out_view = RESHAPE_OP(
159
- at1[1], [q.shape[0], self.num_heads * self.head_size]
76
+ at1[1], [q.shape[0], self._num_heads * self._head_size]
77
)
78
+ return self._quant_matcher(attn_out_view, scale)[0]
79
162
- return self.quant_matcher(attn_out_view, scale)[0]
80
+ return _pattern
81
164
- def replacement(
82
+ @property
83
+ def replacement(self) -> Callable[..., torch.Tensor]:
84
+ def _replacement(
85
q: torch.Tensor,
86
k: torch.Tensor,
87
v: torch.Tensor,
@@ -169,10 +89,9 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
89
scale: torch.Tensor,
90
kv_cache_dummy_dep: torch.Tensor,
91
) -> torch.Tensor:
172
- # attn output in quant_dtype
92
output_attn = torch.empty(
174
- [q.shape[0], self.num_heads, self.head_size],
175
- dtype=self.quant_dtype,
93
+ [q.shape[0], self._num_heads, self._head_size],
94
+ dtype=FP8_DTYPE,
95
device=q.device,
96
)
97
at1 = auto_functionalized(
@@ -181,36 +100,32 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
100
key=k,
101
value=v,
102
output=output_attn,
184
- layer_name=self.layer_name,
103
+ layer_name=self._layer_name,
104
output_scale=scale,
105
output_block_scale=None,
106
kv_cache_dummy_dep=kv_cache_dummy_dep,
107
)
189
- return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
190
-
191
- inputs = [
192
- self.empty(5, self.num_heads, self.head_size), # q
193
- self.empty(5, self.num_heads, self.head_size), # k
194
- self.empty(5, self.num_heads, self.head_size), # v
195
- self.empty(5, self.num_heads, self.head_size), # attn_output
196
- empty_fp32(1, 1), # scale
197
- self.empty(0), # kv_cache_dummy_dep
108
+ return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
109
+
110
+ return _replacement
111
+
112
+ def get_inputs(self):
113
+ dtype = self._dtype
114
+ num_heads = self._num_heads
115
+ head_size = self._head_size
116
+ return [
117
+ self.empty(5, num_heads, head_size, dtype=dtype), # q
118
+ self.empty(5, num_heads, head_size, dtype=dtype), # k
119
+ self.empty(5, num_heads, head_size, dtype=dtype), # v
120
+ self.empty(5, num_heads, head_size, dtype=dtype), # attn_output
121
+ self.empty_fp32(1, 1), # scale
122
+ self.empty(0, dtype=dtype), # kv_cache_dummy_dep
123
]
124
200
- pm.register_replacement(
201
- pattern,
202
- replacement,
203
- inputs,
204
- AttentionQuantPattern.wrap_trace_fn(
205
- pm.fwd_only,
206
- AttentionQuantPattern.fx_view_to_reshape,
207
- AttentionQuantPattern.remove_noop_permutes,
208
- ),
209
- pm_pass,
210
- )
211
-
125
213
-class AttentionNvfp4QuantPattern(AttentionQuantPattern):
126
+class AttnNvfp4QuantPattern(
127
+ VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
128
+):
129
"""
130
Fusion for Attention+Nvfp4Quant.
131
@@ -220,11 +135,16 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
135
will be passed into Attention op as the `output_scale` argument.
136
"""
137
223
- def __init__(self, layer: Attention, dtype: torch.dtype) -> None:
224
- super().__init__(layer, kNvfp4Dynamic, dtype)
138
+ def __init__(self, layer: Attention, dtype: torch.dtype):
139
+ self._layer_name = layer.layer_name
140
+ self._num_heads = layer.num_heads
141
+ self._head_size = layer.head_size
142
+ self._dtype = dtype
143
+ self._QUANT_OP = QUANT_OPS[kNvfp4Dynamic]
144
226
- def _register(self, pm_pass: PatternMatcherPass) -> None:
227
- def pattern(
145
+ @property
146
+ def pattern(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
147
+ def _pattern(
148
q: torch.Tensor,
149
k: torch.Tensor,
150
v: torch.Tensor,
@@ -240,16 +160,16 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
160
key=k,
161
value=v,
162
output=output_attn,
243
- layer_name=self.layer_name,
163
+ layer_name=self._layer_name,
164
output_scale=None,
165
output_block_scale=None,
166
kv_cache_dummy_dep=kv_cache_dummy_dep,
167
)
168
attn_out_view = RESHAPE_OP(
249
- at1[1], [q.shape[0], self.num_heads * self.head_size]
169
+ at1[1], [q.shape[0], self._num_heads * self._head_size]
170
)
171
at2 = auto_functionalized(
252
- self.QUANT_OP,
172
+ self._QUANT_OP,
173
input=attn_out_view,
174
input_scale=input_scale,
175
is_sf_swizzled_layout=True,
@@ -259,23 +179,25 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
179
output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
180
return at2[1], output_scale_view
181
262
- def replacement(
182
+ return _pattern
183
+
184
+ @property
185
+ def replacement(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
186
+ def _replacement(
187
q: torch.Tensor,
188
k: torch.Tensor,
189
v: torch.Tensor,
190
output_attn: torch.Tensor,
267
- output_quant: torch.Tensor,
191
+ _output_quant: torch.Tensor,
192
output_scale: torch.Tensor,
193
input_scale: torch.Tensor,
194
kv_cache_dummy_dep: torch.Tensor,
195
) -> tuple[torch.Tensor, torch.Tensor]:
272
- # attention output in quant_dtype
196
output_attn = torch.empty(
274
- [q.shape[0], self.num_heads, self.head_size // 2],
275
- dtype=self.quant_dtype,
197
+ [q.shape[0], self._num_heads, self._head_size // 2],
198
+ dtype=FP4_DTYPE,
199
device=q.device,
200
)
278
- # attention output block scale
201
output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
202
at2 = auto_functionalized(
203
ATTN_OP,
@@ -283,41 +205,35 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
205
key=k,
206
value=v,
207
output=output_attn,
286
- layer_name=self.layer_name,
208
+ layer_name=self._layer_name,
209
output_scale=input_scale,
210
output_block_scale=output_scale_view,
211
kv_cache_dummy_dep=kv_cache_dummy_dep,
212
)
291
- output = RESHAPE_OP(at2[1], [-1, self.num_heads * self.head_size // 2])
213
+ output = RESHAPE_OP(at2[1], [-1, self._num_heads * self._head_size // 2])
214
return output, at2[2]
215
294
- inputs = [
295
- empty_bf16(5, self.num_heads, self.head_size), # q
296
- empty_bf16(5, self.num_heads, self.head_size), # k
297
- empty_bf16(5, self.num_heads, self.head_size), # v
298
- empty_bf16(5, self.num_heads, self.head_size), # output_attn
299
- self.empty_quant(5, self.num_heads * self.head_size // 2), # output_quant
300
- empty_i32(
301
- 128, round_up(self.num_heads * self.head_size // 16, 4)
216
+ return _replacement
217
+
218
+ def get_inputs(self):
219
+ dtype = self._dtype
220
+ num_heads = self._num_heads
221
+ head_size = self._head_size
222
+ return [
223
+ self.empty_bf16(5, num_heads, head_size), # q
224
+ self.empty_bf16(5, num_heads, head_size), # k
225
+ self.empty_bf16(5, num_heads, head_size), # v
226
+ self.empty_bf16(5, num_heads, head_size), # output_attn
227
+ self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE), # output_quant
228
+ self.empty_i32(
229
+ 128, round_up(num_heads * head_size // 16, 4)
230
), # output_scale
303
- empty_fp32(1, 1), # input_scale
304
- self.empty(0), # kv_cache_dummy_dep
231
+ self.empty_fp32(1, 1), # input_scale
232
+ self.empty(0, dtype=dtype), # kv_cache_dummy_dep
233
]
234
307
- pm.register_replacement(
308
- pattern,
309
- replacement,
310
- inputs,
311
- AttentionQuantPattern.wrap_trace_fn(
312
- pm.fwd_only,
313
- AttentionQuantPattern.fx_view_to_reshape,
314
- AttentionQuantPattern.remove_noop_permutes,
315
- ),
316
- pm_pass,
317
- )
318
-
235
320
-class AttnFusionPass(VllmPatternMatcherPass):
236
+class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
237
"""
238
This pass fuses post-attention quantization onto attention if supported.
239
@@ -330,43 +246,26 @@ class AttnFusionPass(VllmPatternMatcherPass):
246
support are attention kernels, which need to support fusing output quant.
247
"""
248
333
- @enable_fake_mode
249
def __init__(self, config: VllmConfig) -> None:
335
- super().__init__(config)
336
-
337
- self.patterns = PatternMatcherPass(pass_name="attn_fusion_pass")
338
-
339
- attn_layers = get_layers_from_vllm_config(config, Attention)
340
- for layer_name, layer in attn_layers.items():
341
- pattern_fp8 = AttentionFp8StaticQuantPattern(
342
- layer, config.model_config.dtype
343
- )
344
- pattern_fp8.register_if_supported(self.patterns)
250
+ super().__init__(config, "attn_quant_fusion")
251
346
- if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
347
- pattern_nvfp4 = AttentionNvfp4QuantPattern(
348
- layer, config.model_config.dtype
349
- )
350
- pattern_nvfp4.register_if_supported(self.patterns)
252
+ dtype = config.model_config.dtype
253
+ layers = list(get_layers_from_vllm_config(config, Attention).values())
254
352
- if len(attn_layers) == 0:
255
+ if len(layers) == 0:
256
logger.warning(
257
"Attention + quant fusion is enabled, but no attention layers "
258
"were found in CompilationConfig.static_forward_context "
259
"so no fusion patterns were registered."
260
)
261
359
- self.dump_patterns(config, self.patterns)
262
+ for layer in layers:
263
+ if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY):
264
+ self.register(AttnFp8StaticQuantPattern(layer, dtype))
265
361
- @VllmInductorPass.time_and_log
362
- def __call__(self, graph: torch.fx.graph.Graph) -> None:
363
- self.matched_count = self.patterns.apply(graph)
364
- logger.debug("Fused quant onto %s attention nodes", self.matched_count)
266
+ if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
267
+ for layer in layers:
268
+ if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):
269
+ self.register(AttnNvfp4QuantPattern(layer, dtype))
270
366
- def uuid(self) -> str:
367
- return VllmInductorPass.hash_source(
368
- self,
369
- AttentionQuantPattern,
370
- AttentionFp8StaticQuantPattern,
371
- AttentionNvfp4QuantPattern,
372
- )
271
+ self.dump_patterns(config, self.pm_pass)
vllm/compilation/passes/pass_manager.py
+5
-3
@@ -14,7 +14,7 @@ from vllm.logger import init_logger
14
from vllm.platforms import current_platform
15
from vllm.utils.system_utils import set_env_var
16
17
-from .vllm_inductor_pass import VllmInductorPass
17
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
18
19
if rocm_aiter_ops.is_enabled():
20
from .fusion.rocm_aiter_fusion import (
@@ -25,7 +25,7 @@ if rocm_aiter_ops.is_enabled():
25
26
if current_platform.is_cuda_alike():
27
from .fusion.act_quant_fusion import ActivationQuantFusionPass
28
- from .fusion.attn_quant_fusion import AttnFusionPass
28
+ from .fusion.attn_quant_fusion import AttnQuantFusionPass
29
from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
30
from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
31
from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
@@ -108,6 +108,8 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
108
self.fix_functionalization(graph)
109
VllmInductorPass.dump_prefix = None # Cleanup index
110
111
+ VllmPatternMatcherPass.log_match_summary()
112
+
113
def configure(self, config: VllmConfig) -> None:
114
self.pass_config = config.compilation_config.pass_config
115
@@ -144,7 +146,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
146
self.passes += [RopeKVCacheFusionPass(config)]
147
148
if self.pass_config.fuse_attn_quant:
147
- self.passes += [AttnFusionPass(config)]
149
+ self.passes += [AttnQuantFusionPass(config)]
150
151
if self.pass_config.enable_qk_norm_rope_fusion:
152
self.passes += [SplitCoalescingPass(config)]
vllm/compilation/passes/vllm_inductor_pass.py
+139
-6
@@ -3,19 +3,24 @@
3
import functools
4
import operator
5
import time
6
+from abc import ABC, abstractmethod
7
+from collections import defaultdict
8
from collections.abc import Callable
9
from dataclasses import dataclass
8
-from typing import ClassVar
10
+from typing import Any, ClassVar, Generic, ParamSpec, TypeVar
11
12
import regex as re
13
import torch
14
+import torch._inductor.pattern_matcher as pm
15
+from torch import fx
16
from torch._dynamo.utils import lazy_format_graph_code
17
from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter
18
19
from vllm.config import VllmConfig
20
from vllm.logger import init_logger
21
18
-from .inductor_pass import InductorPass
22
+from .fx_utils import is_func
23
+from .inductor_pass import InductorPass, enable_fake_mode
24
25
logger = init_logger(__name__)
26
@@ -79,18 +84,23 @@ class VllmInductorPass(InductorPass):
84
logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)
85
86
87
+def get_match_table() -> dict[str, int]:
88
+ """Return a snapshot of the match table."""
89
+ return dict(VllmPatternMatcherPass.match_table)
90
+
91
+
92
class VllmPatternMatcherPass(VllmInductorPass):
93
"""
94
A VllmInductorPass that uses the Inductor pattern matcher.
85
- Its main use is providing the dump_patterns utility that dumps the
86
- Inductor pattern matcher patterns into a file, which greatly aids debugging.
87
-
88
- TODO(luka) move more utilities to this pass.
95
+ Provides pattern registration with match counting, debug dumping, and logging.
96
"""
97
98
matched_count: int = 0
99
"""The number of matched patterns in the pass."""
100
101
+ match_table: ClassVar[defaultdict[str, int]] = defaultdict(int)
102
+ """Global table mapping pass name to its total match count."""
103
+
104
_OP_OVERLOAD_PATTERN: ClassVar[re.Pattern] = re.compile(
105
r"<OpOverload\(op='([^']*)', overload='([^']*)'\)>"
106
)
@@ -104,6 +114,11 @@ class VllmPatternMatcherPass(VllmInductorPass):
114
)
115
)
116
117
+ @classmethod
118
+ def log_match_summary(cls) -> None:
119
+ if cls.match_table:
120
+ logger.debug("fusion pass matches: %s", dict(cls.match_table))
121
+
122
def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None:
123
"""
124
If debug dumping is enabled, dump the Inductor pattern-matcher patterns
@@ -171,6 +186,124 @@ class VllmPatternMatcherPass(VllmInductorPass):
186
print(f"{pattern_repr}\n", file=f)
187
188
189
+P = ParamSpec("P")
190
+R = TypeVar("R")
191
+
192
+
193
+class VllmPatternReplacement(ABC, Generic[P, R]):
194
+ """
195
+ A pattern/replacement pair for FX graph fusion.
196
+
197
+ Implement the three abstract members below, then pass
198
+ instances to VllmFusionPatternMatcherPass.register(). The pass will
199
+ find every occurrence of `pattern` in the graph and substitute it
200
+ with `replacement`.
201
+ """
202
+
203
+ # TODO(Badr): bound methods work for pattern registration since
204
+ # PyTorch 2.10. Once vLLM requires torch>=2.11, replace these properties
205
+ # with plain methods and drop the closure indirection.
206
+ @property
207
+ @abstractmethod
208
+ def pattern(self) -> Callable[P, R]:
209
+ """Returns a closure defining the FX subgraph to search for."""
210
+ ...
211
+
212
+ @property
213
+ @abstractmethod
214
+ def replacement(self) -> Callable[P, R]:
215
+ """
216
+ Returns a closure defining the FX subgraph to
217
+ substitute in place of each match.
218
+ """
219
+ ...
220
+
221
+ @abstractmethod
222
+ def get_inputs(self) -> list[torch.Tensor]:
223
+ """Example tensors used to trace pattern and replacement."""
224
+ ...
225
+
226
+ # Helpers for get_inputs: uninitialized tensors of common dtypes.
227
+ @staticmethod
228
+ def empty(*args, **kwargs) -> torch.Tensor:
229
+ return torch.empty(*args, device="cuda", **kwargs)
230
+
231
+ @staticmethod
232
+ def empty_bf16(*args, **kwargs) -> torch.Tensor:
233
+ return torch.empty(*args, dtype=torch.bfloat16, device="cuda", **kwargs)
234
+
235
+ @staticmethod
236
+ def empty_fp16(*args, **kwargs) -> torch.Tensor:
237
+ return torch.empty(*args, dtype=torch.float16, device="cuda", **kwargs)
238
+
239
+ @staticmethod
240
+ def empty_fp32(*args, **kwargs) -> torch.Tensor:
241
+ return torch.empty(*args, dtype=torch.float32, device="cuda", **kwargs)
242
+
243
+ @staticmethod
244
+ def empty_i32(*args, **kwargs) -> torch.Tensor:
245
+ return torch.empty(*args, dtype=torch.int32, device="cuda", **kwargs)
246
+
247
+
248
+def _fx_view_to_reshape(gm: fx.GraphModule) -> None:
249
+ from torch._inductor.fx_passes.post_grad import view_to_reshape
250
+
251
+ view_to_reshape(gm)
252
+
253
+
254
+def _remove_noop_permutes(gm: fx.GraphModule) -> None:
255
+ for node in gm.graph.nodes:
256
+ if not is_func(node, torch.ops.aten.permute.default):
257
+ continue
258
+ dims = node.args[1]
259
+ if any(dim != i for i, dim in enumerate(dims)):
260
+ continue
261
+ node.replace_all_uses_with(node.args[0])
262
+ gm.graph.erase_node(node)
263
+
264
+
265
+class VllmFusionPatternMatcherPass(VllmPatternMatcherPass):
266
+ """
267
+ A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects.
268
+ Subclasses register patterns via self.register() in their own __init__.
269
+ """
270
+
271
+ def __init__(self, config: VllmConfig, pass_name: str) -> None:
272
+ super().__init__(config)
273
+ self.pass_name = pass_name
274
+ self.pm_pass = PatternMatcherPass(pass_name=pass_name)
275
+ self._pattern_replacements: list[VllmPatternReplacement] = []
276
+
277
+ @enable_fake_mode
278
+ def register(self, pr: VllmPatternReplacement) -> None:
279
+ pm.register_replacement(
280
+ pr.pattern,
281
+ pr.replacement,
282
+ pr.get_inputs(),
283
+ self._trace_fn,
284
+ self.pm_pass,
285
+ )
286
+ self._pattern_replacements.append(pr)
287
+
288
+ def uuid(self) -> str:
289
+ return VllmInductorPass.hash_source(
290
+ type(self),
291
+ *[type(pr) for pr in self._pattern_replacements],
292
+ )
293
+
294
+ @staticmethod
295
+ def _trace_fn(*args: Any, **kwargs: Any) -> fx.GraphModule:
296
+ gm = pm.fwd_only(*args, **kwargs)
297
+ _fx_view_to_reshape(gm)
298
+ _remove_noop_permutes(gm)
299
+ return gm
300
+
301
+ @VllmInductorPass.time_and_log
302
+ def __call__(self, graph: torch.fx.Graph) -> None:
303
+ self.matched_count = self.pm_pass.apply(graph)
304
+ VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count
305
+
306
+
307
class PrinterInductorPass(VllmInductorPass):
308
def __init__(self, name: str, config: VllmConfig) -> None:
309
super().__init__(config)
vllm/v1/worker/gpu_worker.py
+5
@@ -703,6 +703,11 @@ class Worker(WorkerBase):
703
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
704
return self.model_runner.get_supported_tasks()
705
706
+ def get_compilation_match_table(self) -> dict[str, int]:
707
+ from vllm.compilation.passes.vllm_inductor_pass import get_match_table
708
+
709
+ return get_match_table()
710
+
711
def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
712
"""Get encoder timing stats from model runner."""
713
return self.model_runner.get_encoder_timing_stats()