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 14, 05:30 UTC
[torch.compile] Cleanup compilation tests and custom passes, add debug utils, fix DCE bug (#23091), fix test (#24376), and prep for custom op matching (#24604) (#24542)
Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: luka <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Luka Govedič committed
Sep 22, 2025 at 15:30 UTC
d5e0fca264f7d277ed372f7c075827ed9a0c5e7e
24 files changed
+407
-464
tests/compile/backend.py
+27
-1
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import weakref
5
from collections.abc import Sequence
6
from copy import deepcopy
7
from typing import Callable, Union
@@ -10,7 +11,26 @@ from torch._ops import OpOverload
11
12
from vllm.compilation.fx_utils import find_op_nodes
13
from vllm.compilation.inductor_pass import InductorPass
13
-from vllm.config import get_current_vllm_config
14
+from vllm.compilation.pass_manager import with_pattern_match_debug
15
+from vllm.compilation.vllm_inductor_pass import VllmInductorPass
16
+from vllm.config import VllmConfig, get_current_vllm_config
17
+
18
+
19
+class LazyInitPass(InductorPass):
20
+ """
21
+ If there's a pass that we want to initialize lazily in a test,
22
+ we can wrap it in LazyInitPass, which will initialize the pass when invoked
23
+ and then immediately invoke it.
24
+ """
25
+
26
+ def __init__(self, pass_cls: type[VllmInductorPass],
27
+ vllm_config: VllmConfig):
28
+ self.pass_cls = pass_cls
29
+ self.vllm_config = weakref.proxy(vllm_config) # avoid cycle
30
+
31
+ def __call__(self, graph: fx.Graph) -> None:
32
+ self.pass_ = self.pass_cls(self.vllm_config)
33
+ self.pass_(graph)
34
35
36
class TestBackend:
@@ -40,10 +60,16 @@ class TestBackend:
60
example_inputs,
61
config_patches=self.inductor_config)
62
63
+ @with_pattern_match_debug
64
def post_pass(self, graph: fx.Graph):
65
self.graph_pre_pass = deepcopy(graph)
66
+
67
+ VllmInductorPass.dump_prefix = 0
68
for pass_ in self.custom_passes:
69
pass_(graph)
70
+ VllmInductorPass.dump_prefix += 1
71
+
72
+ VllmInductorPass.dump_prefix = None
73
74
self.graph_post_pass = deepcopy(graph)
75
# assign by reference, will reflect the final state of the graph
tests/compile/test_async_tp.py
+2
@@ -294,6 +294,8 @@ def async_tp_pass_on_test_model(local_rank: int, world_size: int,
294
compiled_model = torch.compile(model, backend=backend)
295
compiled_model(hidden_states)
296
297
+ assert async_tp_pass.matched_count == 1
298
+
299
# In pre-nodes, all gather or reduce scatter should exist,
300
# fused_matmul_reduce_scatter or fused_all_gather_matmul should not
301
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
tests/compile/test_config.py
+9
-1
@@ -4,7 +4,7 @@ import pytest
4
5
import vllm
6
from vllm.compilation.counter import compilation_counter
7
-from vllm.config import VllmConfig
7
+from vllm.config import CompilationConfig, VllmConfig
8
from vllm.utils import _is_torch_equal_or_newer
9
10
@@ -26,6 +26,14 @@ def test_use_cudagraphs_dynamic(monkeypatch):
26
assert not vllm_config.compilation_config.use_cudagraph
27
28
29
+def test_custom_op():
30
+ # proper syntax
31
+ _ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"])
32
+
33
+ with pytest.raises(ValueError, match="Invalid syntax '"):
34
+ _ = CompilationConfig(custom_ops=["quant_fp8"])
35
+
36
+
37
# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
38
@pytest.mark.forked
39
# NB: We don't test VLLM_DISABLE_COMPILE_CACHE=0 because that depends
tests/compile/test_functionalization.py
+6
-4
@@ -8,9 +8,10 @@ import vllm.envs as envs
8
from vllm import LLM, SamplingParams
9
from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
10
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
11
-from vllm.compilation.fusion import FUSED_OPS, FusionPass
11
+from vllm.compilation.fusion import FUSED_OPS, RMSNormQuantFusionPass
12
from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
13
from vllm.compilation.noop_elimination import NoOpEliminationPass
14
+from vllm.compilation.post_cleanup import PostCleanupPass
15
from vllm.config import CompilationConfig, PassConfig, VllmConfig
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
QuantKey, kFp8DynamicTokenSym, kFp8StaticTensorSym)
@@ -58,11 +59,12 @@ def test_fix_functionalization(model: str, quant_key: QuantKey,
59
vllm_config.compilation_config = CompilationConfig(
60
pass_config=PassConfig(enable_fusion=do_fusion, enable_noop=True))
61
noop_pass = NoOpEliminationPass(vllm_config)
61
- fusion_pass = FusionPass.instance(vllm_config)
62
+ fusion_pass = RMSNormQuantFusionPass(vllm_config)
63
+ cleanup_pass = PostCleanupPass(vllm_config)
64
act_quant_fusion_pass = ActivationQuantFusionPass(vllm_config)
65
64
- passes = [noop_pass, fusion_pass, act_quant_fusion_pass
65
- ] if do_fusion else [noop_pass]
66
+ passes = [noop_pass, fusion_pass, act_quant_fusion_pass, cleanup_pass
67
+ ] if do_fusion else [noop_pass, cleanup_pass]
68
func_pass = FixFunctionalizationPass(vllm_config)
69
backend_func = TestBackend(*passes, func_pass)
70
backend_no_func = TestBackend(*passes)
tests/compile/test_fusion.py
+10
-7
@@ -4,11 +4,11 @@
4
import pytest
5
import torch
6
7
-import vllm.envs as envs
7
import vllm.plugins
8
from vllm.compilation.fusion import (FUSED_OPS, QUANT_OPS, FusedRMSQuantKey,
10
- FusionPass)
9
+ RMSNormQuantFusionPass)
10
from vllm.compilation.noop_elimination import NoOpEliminationPass
11
+from vllm.compilation.post_cleanup import PostCleanupPass
12
from vllm.config import (CompilationConfig, CompilationLevel, PassConfig,
13
VllmConfig)
14
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -79,15 +79,15 @@ class TestModel(torch.nn.Module):
79
80
81
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
82
-@pytest.mark.parametrize("hidden_size", [64, 3392, 4096])
83
-@pytest.mark.parametrize("num_tokens", [7, 256, 533, 2048, 2049])
82
+@pytest.mark.parametrize("hidden_size", [64])
83
+@pytest.mark.parametrize("num_tokens", [257])
84
@pytest.mark.parametrize("eps", [1e-5, 1e-6])
85
@pytest.mark.parametrize("static", [True, False])
86
# cuda_force_torch used to test torch code path on platforms that
87
# cutlass_fp8_supported() == True.
88
@pytest.mark.parametrize("cuda_force_torch",
89
[True, False] if cutlass_fp8_supported() else [True])
90
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
90
+@pytest.mark.skipif(not current_platform.is_cuda_alike(),
91
reason="Only test on CUDA and ROCm")
92
def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
93
cuda_force_torch):
@@ -104,9 +104,10 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
104
with vllm.config.set_current_vllm_config(vllm_config):
105
# Reshape pass is needed for the fusion pass to work
106
noop_pass = NoOpEliminationPass(vllm_config)
107
- fusion_pass = FusionPass.instance(vllm_config)
107
+ fusion_pass = RMSNormQuantFusionPass(vllm_config)
108
+ cleanup_pass = PostCleanupPass(vllm_config)
109
109
- backend = TestBackend(noop_pass, fusion_pass)
110
+ backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
111
model = TestModel(hidden_size, eps, static, cuda_force_torch)
112
113
# First dimension dynamic
@@ -128,6 +129,8 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
129
130
torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
131
132
+ assert fusion_pass.matched_count == 2
133
+
134
# In pre-nodes, fp8 quant should be there and fused kernels should not
135
backend.check_before_ops(model.ops_in_model_before())
136
tests/compile/test_fusion_all_reduce.py
+5
-1
@@ -9,6 +9,7 @@ import vllm.envs as envs
9
from vllm.compilation.collective_fusion import AllReduceFusionPass
10
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
11
from vllm.compilation.noop_elimination import NoOpEliminationPass
12
+from vllm.compilation.post_cleanup import PostCleanupPass
13
from vllm.config import (CompilationConfig, CompilationLevel, DeviceConfig,
14
ModelConfig, PassConfig, VllmConfig)
15
from vllm.distributed import tensor_model_parallel_all_reduce
@@ -215,8 +216,10 @@ def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
216
all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)
217
noop_pass = NoOpEliminationPass(vllm_config)
218
func_pass = FixFunctionalizationPass(vllm_config)
219
+ cleanup_pass = PostCleanupPass(vllm_config)
220
219
- backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass)
221
+ backend = TestBackend(all_reduce_fusion_pass, noop_pass, func_pass,
222
+ cleanup_pass)
223
224
token_num = batch_size * seq_len
225
model = test_model_cls(hidden_size, token_num)
@@ -227,6 +230,7 @@ def all_reduce_fusion_pass_on_test_model(local_rank: int, world_size: int,
230
compiled_model = torch.compile(model, backend=backend)
231
compiled_model(hidden_states, residual)
232
233
+ assert all_reduce_fusion_pass.matched_count == 1
234
backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
235
backend.check_after_ops(model.ops_in_model_after())
236
del all_reduce_fusion_pass
tests/compile/test_fusion_attn.py
+13
-7
@@ -6,18 +6,19 @@ from typing import Optional
6
import pytest
7
import torch._dynamo
8
9
-from tests.compile.backend import TestBackend
9
+from tests.compile.backend import LazyInitPass, TestBackend
10
from tests.models.utils import check_outputs_equal
11
from tests.v1.attention.utils import (BatchSpec, _Backend,
12
create_common_attn_metadata)
13
from vllm import LLM, SamplingParams
14
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
15
-from vllm.attention import Attention
15
+from vllm.attention import Attention, AttentionMetadata
16
from vllm.attention.selector import global_force_attn_backend_context_manager
17
from vllm.compilation.fusion import QUANT_OPS
18
from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
19
from vllm.compilation.fx_utils import find_op_nodes
20
from vllm.compilation.noop_elimination import NoOpEliminationPass
21
+from vllm.compilation.post_cleanup import PostCleanupPass
22
from vllm.config import (CacheConfig, CompilationConfig, CompilationLevel,
23
ModelConfig, PassConfig, SchedulerConfig, VllmConfig,
24
set_current_vllm_config)
@@ -104,7 +105,7 @@ def test_attention_fusion_v0(example_prompts, monkeypatch, model: str,
105
106
# AttnFusionPass needs attention layers to be registered in config upon init
107
# so we initialize it during compilation.
107
- attn_pass = lambda *args, **kw: AttnFusionPass(vllm_config)(*args, **kw)
108
+ attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
109
backend = TestBackend(NoOpEliminationPass(vllm_config), attn_pass)
110
llm2 = LLM(model,
111
enforce_eager=True,
@@ -197,7 +198,8 @@ class AttentionQuantPatternModel(torch.nn.Module):
198
device=self.device,
199
)
200
200
- def build_attn_metadata(self, batch_size: int, use_hnd: bool):
201
+ def build_attn_metadata(self, batch_size: int, use_hnd: bool) \
202
+ -> AttentionMetadata:
203
"""Initialize attention metadata."""
204
205
# Create common attn metadata
@@ -447,9 +449,10 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
449
450
# Create test backend with fusion passes enabled
451
noop_pass = NoOpEliminationPass(vllm_config)
450
- attn_pass = lambda *args, **kw: AttnFusionPass(vllm_config)(*args, **kw
451
- )
452
- test_backend = TestBackend(noop_pass, attn_pass)
452
+ attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
453
+ cleanup_pass = PostCleanupPass(vllm_config)
454
+
455
+ test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)
456
457
# Compile model with fusion enabled
458
model_compiled = torch.compile(model_fused,
@@ -485,6 +488,9 @@ def test_attention_quant_pattern(num_qo_heads: int, num_kv_heads: int,
488
test_backend.check_before_ops([QUANT_OPS[quant_key]],
489
fully_replaced=True)
490
491
+ # access the underlying `AttnFusionPass` on the `LazyInitPass`
492
+ assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
493
+
494
# Check attention ops in the graph before and after fusion
495
attn_nodes_pre = list(find_op_nodes(ATTN_OP, test_backend.graph_pre_pass))
496
attn_nodes_post = list(find_op_nodes(ATTN_OP,
tests/compile/test_sequence_parallelism.py
+14
-7
@@ -6,10 +6,12 @@ import torch
6
7
import vllm.envs as envs
8
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
9
-from vllm.compilation.fusion import FusionPass
9
+from vllm.compilation.fusion import RMSNormQuantFusionPass
10
from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe, is_func
11
from vllm.compilation.noop_elimination import NoOpEliminationPass
12
+from vllm.compilation.post_cleanup import PostCleanupPass
13
from vllm.compilation.sequence_parallelism import SequenceParallelismPass
14
+from vllm.compilation.vllm_inductor_pass import VllmInductorPass
15
from vllm.config import (CompilationConfig, DeviceConfig, ModelConfig,
16
PassConfig, VllmConfig)
17
from vllm.distributed import tensor_model_parallel_all_reduce
@@ -104,7 +106,7 @@ class TestQuantModel(torch.nn.Module):
106
# Initialize weights
107
torch.nn.init.normal_(self.gate_proj, std=0.02)
108
107
- self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=False)
109
+ self.fp8_linear = Fp8LinearOp(act_quant_static=True)
110
111
self.scale = torch.rand(1, dtype=torch.float32)
112
# Create a weight that is compatible with torch._scaled_mm,
@@ -137,8 +139,7 @@ class TestQuantModel(torch.nn.Module):
139
# layer normalization
140
norm_output, residual_output = self.norm(all_reduce, residual)
141
140
- # for static input quantization
141
- # self.fp8_linear is initialized with use_per_token_if_dynamic=False
142
+ # scaled_mm with static input quantization
143
fp8_linear_result = self.fp8_linear.apply(norm_output,
144
self.w,
145
self.wscale,
@@ -253,16 +254,20 @@ def sequence_parallelism_pass_on_test_model(
254
dtype=dtype,
255
seed=42)
256
256
- sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
257
noop_pass = NoOpEliminationPass(vllm_config)
258
+ sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
259
func_pass = FixFunctionalizationPass(vllm_config)
260
+ cleanup_pass = PostCleanupPass(vllm_config)
261
260
- passes_for_backend = [noop_pass, sequence_parallelism_pass]
262
+ passes_for_backend: list[VllmInductorPass] = \
263
+ [noop_pass, sequence_parallelism_pass]
264
265
if enable_fusion:
263
- fusion_pass = FusionPass.instance(vllm_config)
266
+ fusion_pass = RMSNormQuantFusionPass(vllm_config)
267
passes_for_backend.append(fusion_pass)
268
269
+ passes_for_backend.append(cleanup_pass)
270
+
271
backend_no_func = TestBackend(*passes_for_backend)
272
backend_func = TestBackend(*passes_for_backend, func_pass)
273
@@ -279,6 +284,8 @@ def sequence_parallelism_pass_on_test_model(
284
compiled_model_func = torch.compile(model, backend=backend_func)
285
compiled_model_func(hidden_states, residual)
286
287
+ assert sequence_parallelism_pass.matched_count == 1
288
+
289
# In pre-nodes, all reduce should be there,
290
# reduce scatter and all gather should not
291
backend_no_func.check_before_ops(model.ops_in_model_before())
tests/compile/test_silu_mul_quant_fusion.py
+12
-1
@@ -15,6 +15,7 @@ from vllm.compilation.activation_quant_fusion import (
15
# yapf: enable
16
from vllm.compilation.fusion import QUANT_OPS
17
from vllm.compilation.noop_elimination import NoOpEliminationPass
18
+from vllm.compilation.post_cleanup import PostCleanupPass
19
from vllm.config import CompilationConfig, PassConfig, VllmConfig
20
from vllm.model_executor.layers.activation import SiluAndMul
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -69,6 +70,10 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
70
71
def __init__(self, hidden_size: int, x: torch.Tensor, **kwargs):
72
super().__init__()
73
+ from vllm.compilation.activation_quant_fusion import (
74
+ silu_and_mul_nvfp4_quant_supported)
75
+ assert silu_and_mul_nvfp4_quant_supported
76
+
77
self.silu_and_mul = SiluAndMul()
78
79
# create nvfp4 weight
@@ -127,7 +132,11 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, dtype, model_class,
132
pass_config=PassConfig(enable_fusion=True, enable_noop=True))
133
fusion_pass = ActivationQuantFusionPass(config)
134
130
- backend = TestBackend(NoOpEliminationPass(config), fusion_pass)
135
+ passes = [
136
+ NoOpEliminationPass(config), fusion_pass,
137
+ PostCleanupPass(config)
138
+ ]
139
+ backend = TestBackend(*passes)
140
model = model_class(hidden_size=hidden_size,
141
cuda_force_torch=cuda_force_torch,
142
x=x)
@@ -151,6 +160,8 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, dtype, model_class,
160
atol=atol,
161
rtol=rtol)
162
163
+ assert fusion_pass.matched_count == 1
164
+
165
# In pre-nodes, quant op should be present and fused kernels should not
166
backend.check_before_ops(model.ops_in_model_before())
167
vllm/compilation/activation_quant_fusion.py
+7
-11
@@ -17,7 +17,7 @@ from vllm.platforms import current_platform
17
18
from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
19
from .inductor_pass import enable_fake_mode
20
-from .vllm_inductor_pass import VllmInductorPass
20
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21
22
logger = init_logger(__name__)
23
@@ -152,7 +152,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
152
register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
153
154
155
-class ActivationQuantFusionPass(VllmInductorPass):
155
+class ActivationQuantFusionPass(VllmPatternMatcherPass):
156
"""
157
This pass fuses a pre-defined set of custom ops into fused ops.
158
It uses the torch pattern matcher to find the patterns and replace them.
@@ -176,16 +176,12 @@ class ActivationQuantFusionPass(VllmInductorPass):
176
pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern()
177
pattern_silu_mul_nvfp4.register(self.patterns)
178
179
- def __call__(self, graph: torch.fx.Graph):
180
- self.begin()
181
- self.dump_graph(graph, "before_act_quant_fusion")
182
-
183
- count = self.patterns.apply(graph)
184
- logger.debug("Replaced %s patterns in ActivationQuantFusionPass",
185
- count)
179
+ self.dump_patterns(config, self.patterns)
180
187
- self.dump_graph(graph, "after_act_quant_fusion")
188
- self.end_and_log()
181
+ @VllmInductorPass.time_and_log
182
+ def __call__(self, graph: torch.fx.Graph):
183
+ self.matched_count = self.patterns.apply(graph)
184
+ logger.debug("Replaced %s patterns", self.matched_count)
185
186
def uuid(self):
187
return VllmInductorPass.hash_source(self, ActivationQuantPattern,
vllm/compilation/collective_fusion.py
+14
-15
@@ -20,7 +20,7 @@ from vllm.platforms import current_platform
20
from vllm.utils import direct_register_custom_op
21
22
from .inductor_pass import enable_fake_mode
23
-from .vllm_inductor_pass import VllmInductorPass
23
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
24
25
FP8_DTYPE = current_platform.fp8_dtype()
26
@@ -348,7 +348,7 @@ class AllGatherCutlassScaledMMPattern(BasePattern):
348
pm.fwd_only, pm_pass)
349
350
351
-class AsyncTPPass(VllmInductorPass):
351
+class AsyncTPPass(VllmPatternMatcherPass):
352
353
@enable_fake_mode
354
def __init__(self, config: VllmConfig):
@@ -378,18 +378,17 @@ class AsyncTPPass(VllmInductorPass):
378
AllGatherCutlassScaledMMPattern(
379
self.model_dtype, self.device).register(self.patterns)
380
381
+ self.dump_patterns(config, self.patterns)
382
+
383
def is_applicable_for_shape(self, shape: Optional[int]) -> bool:
384
# only do replace for specific shapes
385
tp_size = get_tensor_model_parallel_world_size()
386
return shape is not None and shape % tp_size == 0
387
388
+ @VllmInductorPass.time_and_log
389
def __call__(self, graph: fx.Graph):
387
- self.begin()
388
- self.dump_graph(graph, "before_async_tp_pass")
389
- count = self.patterns.apply(graph)
390
- logger.debug("Replaced %s patterns with async TP pass.", count)
391
- self.dump_graph(graph, "after_async_tp_pass")
392
- self.end_and_log()
390
+ self.matched_count = self.patterns.apply(graph)
391
+ logger.debug("Replaced %s patterns", self.matched_count)
392
393
394
if flashinfer_comm is not None:
@@ -1068,7 +1067,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
1067
pm.fwd_only, pm_pass)
1068
1069
1071
-class AllReduceFusionPass(VllmInductorPass):
1070
+class AllReduceFusionPass(VllmPatternMatcherPass):
1071
1072
def __init__(self, config: VllmConfig):
1073
super().__init__(config)
@@ -1124,6 +1123,7 @@ class AllReduceFusionPass(VllmInductorPass):
1123
fuse_rms_quant=config.compilation_config.pass_config.enable_fusion)
1124
1125
self.register_patterns()
1126
+ self.dump_patterns(config, self.patterns)
1127
1128
@enable_fake_mode
1129
def register_patterns(self):
@@ -1172,15 +1172,14 @@ class AllReduceFusionPass(VllmInductorPass):
1172
1173
self.disabled = False
1174
1175
+ @VllmInductorPass.time_and_log
1176
def __call__(self, graph: fx.Graph):
1177
if self.disabled:
1178
+ logger.debug("AllReduceFusionPass disabled")
1179
return
1178
- self.begin()
1179
- self.dump_graph(graph, "before_all_reduce_fusion_pass")
1180
- count = self.patterns.apply(graph)
1181
- logger.debug("Replaced %s patterns", count)
1182
- self.dump_graph(graph, "after_all_reduce_fusion_pass")
1183
- self.end_and_log()
1180
+
1181
+ self.matched_count = self.patterns.apply(graph)
1182
+ logger.debug("Replaced %s patterns", self.matched_count)
1183
1184
def __del__(self):
1185
if getattr(self, "disabled", True):
vllm/compilation/fix_functionalization.py
+3
-6
@@ -26,6 +26,7 @@ class FixFunctionalizationPass(VllmInductorPass):
26
To add new nodes to defunctionalize, add to the if-elif chain in __call__.
27
"""
28
29
+ @VllmInductorPass.time_and_log
30
def __call__(self, graph: torch.fx.Graph):
31
# XPU does not support auto-functionalization yet.
32
# Will enable this when switch to vllm-xpu-kernels.
@@ -34,9 +35,6 @@ class FixFunctionalizationPass(VllmInductorPass):
35
"pass currently.")
36
return
37
37
- self.begin()
38
- self.dump_graph(graph, "before_fix_functionalization")
39
-
38
self.nodes_to_remove: list[torch.fx.Node] = []
39
count = 0
40
for node in graph.nodes:
@@ -111,7 +109,7 @@ class FixFunctionalizationPass(VllmInductorPass):
109
110
count += 1
111
114
- self.dump_graph(graph, "before_fix_functionalization_cleanup")
112
+ self.dump_graph(graph, "before_cleanup")
113
114
# Remove the nodes all at once
115
count_removed = len(self.nodes_to_remove)
@@ -120,8 +118,7 @@ class FixFunctionalizationPass(VllmInductorPass):
118
119
logger.debug("De-functionalized %s nodes, removed %s nodes", count,
120
count_removed)
123
- self.dump_graph(graph, "after_fix_functionalization")
124
- self.end_and_log()
121
+ self.nodes_to_remove.clear()
122
123
def _remove(self, node_or_nodes: Union[torch.fx.Node,
124
Iterable[torch.fx.Node]]):
vllm/compilation/fusion.py
+31
-248
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Callable, NamedTuple, Optional
3
+from typing import Any, NamedTuple
4
5
import torch
6
import torch._inductor.pattern_matcher as pm
@@ -16,10 +16,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
kFp8StaticTensorSym, kNvfp4Quant, kStaticTensorScale)
17
from vllm.platforms import current_platform
18
19
-from .fx_utils import find_getitem_maybe
19
from .inductor_pass import enable_fake_mode
21
-from .multi_output_match import MultiOutputMatch
22
-from .vllm_inductor_pass import VllmInductorPass
20
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21
22
logger = init_logger(__name__)
23
FP8_DTYPE = current_platform.fp8_dtype()
@@ -50,8 +48,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
48
torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
49
}
50
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
53
- QUANT_OPS[
54
- kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default # noqa: E501
51
+ QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default
52
53
54
class FusedRMSQuantKey(NamedTuple):
@@ -80,68 +77,6 @@ FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
77
}
78
79
83
-class QuantMultiOutputMatch(MultiOutputMatch):
84
-
85
- def __init__(self, match: pm.Match, quant_op, fused_op):
86
- super().__init__(match)
87
- assert isinstance(quant_op, OpOverload)
88
- assert isinstance(fused_op, OpOverload)
89
- self.QUANT_OP = quant_op # in-place quant op
90
- self.FUSED_OP = fused_op # in-place fused quant op
91
-
92
- def insert_fused_node(self, fused_return_mapping: dict[int, tuple[fx.Node,
93
- int]],
94
- **kwargs):
95
- """
96
- This utility function inserts an auto-functionalized node for FUSED_OP.
97
- It also correctly sets its meta value and rebinds the users of the
98
- unfused nodes to use the fused node instead.
99
-
100
- :param fused_return_mapping: A dictionary, mapping from getitem indices
101
- of the fused node result to a tuple of the old node and a getitem index.
102
- :param kwargs: kwargs that get directly forwarded to the auto_fn node
103
-
104
- Example:
105
- If we want to replace this graph:
106
- _, x1, x2 = auto_fn(op1)
107
- _, y1, y2 = auto_fn(op2)
108
-
109
- with
110
- _, x1, y2, x2 = auto_fn(FUSED_OP)
111
-
112
- we would call:
113
- insert_fused_node({1: (op1_node, 1), 2: (op2_node, 2), 3: (op1_node, 2)}
114
-
115
- Note that the 0th element is None for auto-functionalized in-place ops.
116
- Hence, others appear 1-indexed.
117
- """
118
- fused_node = self.insert_auto_fn(self.FUSED_OP, kwargs)
119
- indices = fused_return_mapping.keys()
120
- getitem_nodes = self.insert_getitems(fused_node, indices)
121
-
122
- # Prepare the meta value, use a list so it's mutable
123
- meta_val = [None] * (max(indices) + 1)
124
-
125
- # Iterate through elements of the tuple produced by fused_node
126
- for idx, getitem_node in zip(indices, getitem_nodes):
127
- old_node, old_idx = fused_return_mapping[idx]
128
-
129
- # If the old value was never used, the old_getitem might not exist
130
- old_getitem = find_getitem_maybe(old_node, old_idx)
131
- if old_getitem is not None:
132
- # Rebind the users of match getitem nodes to use the new nodes.
133
- # The old nodes will be removed by DCE at the end of the pass.
134
- old_getitem.replace_all_uses_with(getitem_node)
135
- getitem_node.meta["val"] = old_getitem.meta["val"]
136
-
137
- # Extract the appropriate meta value
138
- # It is present even if the getitem node does not exist
139
- meta_val[idx] = old_node.meta["val"][old_idx]
140
-
141
- # Fix the meta value on the new fused node
142
- fused_node.meta["val"] = tuple(meta_val)
143
-
144
-
80
class RMSNormQuantPattern:
81
82
def __init__(self, epsilon: float, key: FusedRMSQuantKey):
@@ -224,8 +159,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
159
symmetric=symmetric))
160
super().__init__(epsilon, key)
161
227
- def register(self, pm_pass: PatternMatcherPass,
228
- record_match: Callable[[MultiOutputMatch], bool]):
162
+ def register(self, pm_pass: PatternMatcherPass):
163
164
def pattern(result: torch.Tensor, input: torch.Tensor,
165
residual: torch.Tensor, weight: torch.Tensor,
@@ -271,36 +205,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
205
inputs,
206
pm.fwd_only,
207
pm_pass,
274
- extra_check=lambda m: record_match(
275
- self.Match(m, self.QUANT_OP, self.FUSED_OP)))
276
-
277
- class Match(QuantMultiOutputMatch):
278
-
279
- def process(self):
280
- # Find the nodes in the match that we need to rebind
281
- rms_node = self.find_auto_fn(RMS_ADD_OP)
282
- quant_node = self.find_auto_fn(self.QUANT_OP)
283
-
284
- assert len(rms_node.users) == 2
285
- assert len(quant_node.users) == 1
286
-
287
- # First, insert a new auto_functionalized node for the fused op,
288
- # as well as getitem nodes to extract the result and residual.
289
- # The auto_fn node returns a tuple of (None, result, residual).
290
- #
291
- # The resulting graph looks like this:
292
- # at = auto_functionalized(torch.ops._C.fused_add_rms_norm_static_fp8_quant.default, ...) # noqa
293
- # result_node_new = at[1]
294
- # residual_node_new = at[2]
295
- with self.inserting_after_match():
296
- # Missing epsilon, scalars cannot be inputs to the pattern
297
- kwargs = self.match.kwargs.copy()
298
-
299
- # 0 is always None
300
- fused_return_mapping = {1: (quant_node, 1), 2: (rms_node, 2)}
301
- self.insert_fused_node(fused_return_mapping,
302
- **kwargs,
303
- epsilon=rms_node.kwargs["epsilon"])
208
+ )
209
210
211
class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
@@ -317,8 +222,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
222
symmetric=symmetric))
223
super().__init__(epsilon, key)
224
320
- def register(self, pm_pass: PatternMatcherPass,
321
- record_match: Callable[[MultiOutputMatch], bool]):
225
+ def register(self, pm_pass: PatternMatcherPass):
226
227
def pattern(result: torch.Tensor, result_rms: torch.Tensor,
228
input: torch.Tensor, weight: torch.Tensor,
@@ -366,39 +270,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
270
inputs,
271
pm.fwd_only,
272
pm_pass,
369
- extra_check=lambda m: record_match(
370
- self.Match(m, self.QUANT_OP, self.FUSED_OP)))
371
-
372
- class Match(QuantMultiOutputMatch):
373
-
374
- def process(self):
375
- # Find the nodes in the match that we need to rebind
376
- rms_node = self.find_auto_fn(RMS_OP)
377
- quant_node = self.find_auto_fn(self.QUANT_OP)
378
-
379
- assert len(rms_node.users) == 1
380
- assert len(quant_node.users) == 2
381
-
382
- # First, insert a new auto_functionalized node for the fused op,
383
- # as well as getitem nodes to extract the result and scale.
384
- # The auto_fn node returns a tuple of (None, result, scale).
385
- #
386
- # The resulting graph looks like this:
387
- # at = auto_functionalized(torch.ops._C.rms_norm_dynamic_per_token_quant.default, ...) # noqa
388
- # result_node_new = at[1]
389
- # scale_node_new = at[2]
390
- with self.inserting_after_match():
391
- # Missing epsilon, scalars cannot be inputs to the pattern
392
- kwargs = self.match.kwargs.copy()
393
- del kwargs["result_rms"] # not used in the fused op
394
-
395
- fused_return_mapping = {1: (quant_node, 1), 2: (quant_node, 2)}
396
- self.insert_fused_node(
397
- fused_return_mapping,
398
- epsilon=rms_node.kwargs["epsilon"],
399
- scale_ub=None, # not used but required
400
- residual=None, # not used but required
401
- **kwargs)
273
+ )
274
275
276
class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
@@ -415,8 +287,7 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
287
symmetric=symmetric))
288
super().__init__(epsilon, key)
289
418
- def register(self, pm_pass: PatternMatcherPass,
419
- record_match: Callable[[MultiOutputMatch], bool]):
290
+ def register(self, pm_pass: PatternMatcherPass):
291
292
def pattern(result: torch.Tensor, input: torch.Tensor,
293
residual: torch.Tensor, weight: torch.Tensor,
@@ -464,137 +335,49 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
335
inputs,
336
pm.fwd_only,
337
pm_pass,
467
- extra_check=lambda m: record_match(
468
- self.Match(m, self.QUANT_OP, self.FUSED_OP)))
469
-
470
- class Match(QuantMultiOutputMatch):
471
-
472
- def process(self):
473
- # Find the nodes in the match that we need to rebind
474
- rms_node = self.find_auto_fn(RMS_ADD_OP)
475
- quant_node = self.find_auto_fn(self.QUANT_OP)
476
-
477
- assert len(rms_node.users) == 2
478
- assert len(quant_node.users) == 2
479
-
480
- # First, insert a new auto_functionalized node for the fused op,
481
- # as well as getitem nodes to extract result, scale, and residual.
482
- # The auto_fn node returns a tuple (None, result, scale, residual).
483
- #
484
- # The resulting graph looks like this:
485
- # at = auto_functionalized(torch.ops._C.rms_norm_dynamic_per_token_quant.default, ...) # noqa
486
- # result_node_new = at[1]
487
- # scale_node_new = at[2]
488
- # residual_node_new = at[3]
489
- with self.inserting_after_match():
490
- # Missing epsilon, scalars cannot be inputs to the pattern
491
- kwargs = self.match.kwargs.copy()
492
-
493
- fused_return_mapping = {
494
- 1: (quant_node, 1), # result
495
- 2: (quant_node, 2), # scale
496
- 3: (rms_node, 2), # residual
497
- }
498
- self.insert_fused_node(
499
- fused_return_mapping,
500
- epsilon=rms_node.kwargs["epsilon"],
501
- scale_ub=None, # not used but required
502
- **kwargs)
503
-
504
-
505
-class FusionPass(VllmInductorPass):
338
+ )
339
+
340
+
341
+class RMSNormQuantFusionPass(VllmPatternMatcherPass):
342
"""
507
- This pass fuses a pre-defined set of custom ops into fused ops.
508
- It uses the torch pattern matcher to find the patterns and replace them.
509
- It also manually processes multi-output matches, as those are broken in
510
- the torch pattern matcher.
511
-
512
- Because patterns can only be registered once, the pass is a singleton.
513
- This will be addressed in a future version of PyTorch:
514
- https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
343
+ This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
344
+ It also supports fused_add_rms_norm.
345
"""
346
517
- _instance: 'Optional[FusionPass]' = None
518
-
519
- @classmethod
520
- def instance(cls, config: VllmConfig):
521
- """
522
- Get the singleton instance of the FusionPass.
523
- If the instance exists, the config is updated but
524
- initialization is not repeated.
525
- """
526
- if cls._instance is None:
527
- cls._instance = FusionPass(config)
528
- else:
529
- cls._instance.pass_config = config.compilation_config.pass_config
530
- return cls._instance
531
-
347
@enable_fake_mode
348
def __init__(self, config: VllmConfig):
534
- assert self.__class__._instance is None, \
535
- "FusionPass singleton instance already exists"
349
super().__init__(config)
350
538
- self.matches: list[MultiOutputMatch] = []
351
self.patterns: PatternMatcherPass = PatternMatcherPass(
540
- pass_name="fusion_pass")
352
+ pass_name="rmsnorm_quant_fusion_pass")
353
354
for epsilon in [1e-5, 1e-6]:
355
# Fuse rms_norm + static fp8 quant
356
RMSNormStaticQuantPattern(epsilon,
357
FP8_DTYPE).register(self.patterns)
358
547
- # Matches for patterns below have 2 or more outputs,
548
- # so we need to process them manually (see process_matches)
549
-
550
- # Fuse rms_norm + static fp8 quant
359
+ # Fuse fused_add_rms_norm + static fp8 quant
360
FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(
552
- self.patterns, self.record_match)
361
+ self.patterns)
362
363
# Fuse rms_norm + dynamic per-token fp8 quant
555
- RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
556
- self.patterns, self.record_match)
364
+ RMSNormDynamicQuantPattern(epsilon,
365
+ FP8_DTYPE).register(self.patterns)
366
367
# Fuse fused_add_rms_norm + dynamic per-token fp8 quant
368
FusedAddRMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
560
- self.patterns, self.record_match)
561
-
562
- # WARNING: This is a hack to clear the pattern matcher cache
563
- # and allow multiple values of epsilon.
564
- torch._inductor.pattern_matcher._seen_patterns.clear()
565
-
566
- def record_match(self, match: MultiOutputMatch) -> bool:
567
- # Hijack the extra_check to record the match and
568
- # save it for post-processing.
569
- self.matches.append(match)
570
-
571
- # Return False to prevent automatic replacement.
572
- return False
573
-
574
- def process_matches(self, graph: fx.Graph):
575
- """
576
- Manually process multi-output matches and replace them with fused nodes.
577
- See MultiOutputMatch for more details.
578
- """
579
- for match in self.matches:
580
- match.process()
369
+ self.patterns)
370
582
- # Finally, remove matched nodes
583
- graph.eliminate_dead_code()
584
- assert all(node not in graph.nodes for match in self.matches
585
- for node in match.match.nodes)
371
+ self.dump_patterns(config, self.patterns)
372
373
+ @VllmInductorPass.time_and_log
374
def __call__(self, graph: fx.Graph):
588
- self.begin()
589
- self.dump_graph(graph, "before_fusion")
590
-
591
- count = self.patterns.apply(graph)
592
- logger.debug("Replaced %s patterns", count)
593
- self.dump_graph(graph, "after_pattern_match")
594
-
595
- # Manually process multi-output matches (and run DCE)
596
- self.process_matches(graph)
597
- logger.debug("Post-processed %s matches", len(self.matches))
598
- self.dump_graph(graph, "after_fusion")
599
- self.matches.clear()
600
- self.end_and_log()
375
+ self.matched_count = self.patterns.apply(graph)
376
+ logger.debug("Replaced %s patterns", self.matched_count)
377
+
378
+ def uuid(self) -> Any:
379
+ return self.hash_source(self, RMSNormQuantPattern,
380
+ RMSNormStaticQuantPattern,
381
+ RMSNormDynamicQuantPattern,
382
+ FusedAddRMSNormStaticQuantPattern,
383
+ FusedAddRMSNormDynamicQuantPattern)
vllm/compilation/fusion_attn.py
+7
-15
@@ -18,7 +18,7 @@ from vllm.utils import round_up
18
19
from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
20
from .inductor_pass import enable_fake_mode
21
-from .vllm_inductor_pass import VllmInductorPass
21
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
22
23
logger = init_logger(__name__)
24
@@ -245,7 +245,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
245
pm_pass)
246
247
248
-class AttnFusionPass(VllmInductorPass):
248
+class AttnFusionPass(VllmPatternMatcherPass):
249
"""
250
This pass fuses post-attention quantization onto attention if supported.
251
@@ -282,20 +282,12 @@ class AttnFusionPass(VllmInductorPass):
282
"were found in CompilationConfig.static_forward_context "
283
"so no fusion patterns were registered.")
284
285
- def __call__(self, graph: torch.fx.graph.Graph) -> None:
286
- self.begin()
287
- self.dump_graph(graph, "before_attn_fusion")
288
-
289
- count = self.patterns.apply(graph)
285
+ self.dump_patterns(config, self.patterns)
286
291
- # TODO: Move this to pass_manager.py after the fx graph broken issue
292
- # has been resolved.
293
- # see https://github.com/vllm-project/vllm/issues/23091
294
- graph.eliminate_dead_code()
295
-
296
- logger.debug("Fused quantization onto %s attention nodes", count)
297
- self.dump_graph(graph, "after_attn_fusion")
298
- self.end_and_log()
287
+ @VllmInductorPass.time_and_log
288
+ def __call__(self, graph: torch.fx.graph.Graph) -> None:
289
+ self.matched_count = self.patterns.apply(graph)
290
+ logger.debug("Fused quant onto %s attention nodes", self.matched_count)
291
292
def uuid(self):
293
return VllmInductorPass.hash_source(self, AttentionQuantPattern,
vllm/compilation/multi_output_match.py
deleted
-109
@@ -1,109 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import abc
5
-import operator
6
-from abc import abstractmethod
7
-from collections.abc import Iterable
8
-
9
-from torch import fx
10
-from torch._higher_order_ops.auto_functionalize import auto_functionalized
11
-from torch._inductor import pattern_matcher as pm
12
-from torch._ops import OpOverload
13
-from torch.fx import Node
14
-
15
-from vllm.compilation.fx_utils import find_auto_fn
16
-
17
-
18
-class MultiOutputMatch(abc.ABC):
19
- """
20
- This class provides utilities to process multi-output matches and
21
- manually insert replacements.
22
-
23
- This is necessary because the automatic replacement for multi-output
24
- matches is broken: https://github.com/pytorch/pytorch/issues/137280
25
- """
26
-
27
- def __init__(self, match: pm.Match):
28
- self.match = match
29
-
30
- @abstractmethod
31
- def process(self):
32
- """
33
- Process a multi-output match and manually insert the replacement.
34
-
35
- This method should:
36
- 1. Insert the replacement nodes after the last node in the match.
37
- 2. Rebind the users of nodes in the match to use the new nodes.
38
- 3. Set meta["val"] for de-functionalization.
39
-
40
- The result of an auto-functionalized node is a tuple of tensors.
41
- The first element is the return value of the function, usually None.
42
- The remaining elements are the mutated args of the function.
43
-
44
- All auto-functionalized nodes must contain a proper meta["val"],
45
- as it is used by de-functionalization. meta["val"] has to contain the
46
- value of the node (tuple of tensors) that would be returned by the
47
- functionalized node during tracing.
48
-
49
- Existing nodes in the graph all have this property set, but we have
50
- to set it manually for new nodes we insert.
51
-
52
- Example:
53
- # op schema: foo(a: Tensor!, b: Tensor, c: Tensor!) -> None
54
- at = auto_functionalized(torch.ops._C.foo.default, a, b, c)
55
- # at.meta["val"] = (None, a, c)
56
- """
57
- raise NotImplementedError
58
-
59
- @property
60
- def nodes(self) -> list[fx.Node]:
61
- return self.match.nodes
62
-
63
- @property
64
- def graph(self) -> fx.Graph:
65
- return self.match.graph
66
-
67
- def find_auto_fn(self, op) -> fx.Node:
68
- """
69
- Find the first auto_functionalized node with the given op in the match.
70
- """
71
- return find_auto_fn(self.nodes, op)
72
-
73
- def inserting_after_match(self):
74
- """
75
- Insert nodes after the last node in the match.
76
- This is done to avoid use-before-definition errors after inserting
77
- replacement nodes.
78
- """
79
-
80
- # match.nodes is not guaranteed to be sorted.
81
- # Find the last node in the match.
82
- for last_node_in_match in reversed(self.graph.nodes):
83
- if last_node_in_match in self.match.nodes:
84
- break
85
- else:
86
- raise ValueError("No nodes in graph")
87
-
88
- return self.graph.inserting_after(last_node_in_match)
89
-
90
- def insert_getitems(self, tuple_node: fx.Node,
91
- indices: Iterable[int]) -> tuple[fx.Node, ...]:
92
- """
93
- Insert operator.getitem nodes to extract elements from a tuple node.
94
-
95
- :param tuple_node: The tuple node to extract elements from.
96
- :param indices: The indices of the elements to extract.
97
- :return: Tuple of the new getitem nodes, corresponding to the indices.
98
- """
99
- with self.graph.inserting_after(tuple_node):
100
- return tuple(
101
- self.graph.call_function(operator.getitem, (tuple_node, idx))
102
- for idx in indices)
103
-
104
- def insert_auto_fn(self, op: OpOverload, kwargs) -> Node:
105
- """
106
- Insert an auto_functionalized node with the given op and kwargs.
107
- """
108
- return self.graph.call_function(auto_functionalized, (op, ),
109
- kwargs=kwargs)
vllm/compilation/noop_elimination.py
+1
-4
@@ -64,9 +64,8 @@ class NoOpEliminationPass(VllmInductorPass):
64
out: "f16[s0, 4096]" = at[1]
65
"""
66
67
+ @VllmInductorPass.time_and_log
68
def __call__(self, graph: torch.fx.Graph):
68
- self.begin()
69
- self.dump_graph(graph, "before_noop_elimination")
69
count = 0
70
# Remove no-op reshapes/views:
71
for node in graph.nodes:
@@ -121,8 +120,6 @@ class NoOpEliminationPass(VllmInductorPass):
120
count += 1
121
122
logger.debug("Removed %s no-op reshapes and slices", count)
124
- self.dump_graph(graph, "after_noop_elimination")
125
- self.end_and_log()
123
124
# ---------------------- Reshape helpers ----------------------
125
def reshape_dims_equivalent(self, dim: Union[int, torch.fx.Node],
vllm/compilation/pass_manager.py
+43
-6
@@ -1,15 +1,21 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import functools
4
5
from torch import fx as fx
6
7
+from vllm import envs
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
+from vllm.utils import set_env_var
12
+
13
+from .post_cleanup import PostCleanupPass
14
+from .vllm_inductor_pass import VllmInductorPass
15
16
if current_platform.is_cuda_alike():
17
from .activation_quant_fusion import ActivationQuantFusionPass
12
- from .fusion import FusionPass
18
+ from .fusion import RMSNormQuantFusionPass
19
from .fusion_attn import AttnFusionPass
20
21
if current_platform.is_cuda():
@@ -19,11 +25,28 @@ from .fix_functionalization import FixFunctionalizationPass
25
from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
26
from .noop_elimination import NoOpEliminationPass
27
from .sequence_parallelism import SequenceParallelismPass
22
-from .vllm_inductor_pass import VllmInductorPass
28
29
logger = init_logger(__name__)
30
31
32
+def with_pattern_match_debug(fn):
33
+ """
34
+ Function decorator that turns on inductor pattern match debug
35
+ for the duration of the call.
36
+ Used to avoid logging builtin Inductor pattern matching.
37
+ """
38
+
39
+ @functools.wraps(fn)
40
+ def wrapper(*args, **kwargs):
41
+ if (debug_val := envs.VLLM_PATTERN_MATCH_DEBUG) is not None:
42
+ # optionally check rank here
43
+ with set_env_var("TORCHINDUCTOR_PATTERN_MATCH_DEBUG", debug_val):
44
+ return fn(*args, **kwargs)
45
+ return fn(*args, **kwargs)
46
+
47
+ return wrapper
48
+
49
+
50
class PostGradPassManager(CustomGraphPass):
51
"""
52
The pass manager for post-grad passes.
@@ -40,16 +63,26 @@ class PostGradPassManager(CustomGraphPass):
63
"""
64
65
def __init__(self):
43
- self.passes: list[VllmInductorPass] = []
66
+ self.passes: list[InductorPass] = []
67
68
+ @with_pattern_match_debug
69
def __call__(self, graph: fx.Graph):
70
+ VllmInductorPass.dump_prefix = 0 # reset dump index
71
+
72
shape = get_pass_context().runtime_shape
73
for pass_ in self.passes:
74
if pass_.is_applicable_for_shape(shape):
75
pass_(graph)
76
+ VllmInductorPass.dump_prefix += 1
77
+
78
+ # post-cleanup goes before fix_functionalization
79
+ # because it requires a functional graph
80
+ self.post_cleanup(graph)
81
+ VllmInductorPass.dump_prefix += 1
82
83
# always run fix_functionalization last
84
self.fix_functionalization(graph)
85
+ VllmInductorPass.dump_prefix = None # Cleanup index
86
87
def configure(self, config: VllmConfig):
88
self.pass_config = config.compilation_config.pass_config
@@ -61,14 +94,18 @@ class PostGradPassManager(CustomGraphPass):
94
if self.pass_config.enable_async_tp:
95
self.passes += [AsyncTPPass(config)]
96
97
+ if self.pass_config.enable_fi_allreduce_fusion:
98
+ self.passes += [AllReduceFusionPass(config)]
99
+
100
if self.pass_config.enable_fusion:
65
- self.passes += [FusionPass.instance(config)]
101
+ self.passes += [RMSNormQuantFusionPass(config)]
102
self.passes += [ActivationQuantFusionPass(config)]
103
104
if self.pass_config.enable_attn_fusion:
105
self.passes += [AttnFusionPass(config)]
70
- if self.pass_config.enable_fi_allreduce_fusion:
71
- self.passes += [AllReduceFusionPass(config)]
106
+
107
+ # needs a functional graph
108
+ self.post_cleanup = PostCleanupPass(config)
109
self.fix_functionalization = FixFunctionalizationPass(config)
110
111
def add(self, pass_: InductorPass):
vllm/compilation/post_cleanup.py
new
+20
@@ -0,0 +1,20 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from torch import fx
4
+
5
+from vllm.compilation.vllm_inductor_pass import VllmInductorPass
6
+
7
+
8
+class PostCleanupPass(VllmInductorPass):
9
+ """
10
+ This pass performs cleanup after custom passes.
11
+ It topologically sorts the graph and removes unused nodes.
12
+ This is needed because the pattern matcher does not guarantee producing
13
+ a topologically sorted graph, and there may be unused nodes left around.
14
+ """
15
+
16
+ @VllmInductorPass.time_and_log
17
+ def __call__(self, graph: fx.Graph) -> None:
18
+ from torch._inductor.pattern_matcher import stable_topological_sort
19
+ stable_topological_sort(graph)
20
+ graph.eliminate_dead_code()
vllm/compilation/sequence_parallelism.py
+6
-12
@@ -15,7 +15,7 @@ from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
17
from .inductor_pass import enable_fake_mode
18
-from .vllm_inductor_pass import VllmInductorPass
18
+from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
19
20
logger = init_logger(__name__)
21
@@ -417,7 +417,7 @@ class LastAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
417
pm.fwd_only, pm_pass)
418
419
420
-class SequenceParallelismPass(VllmInductorPass):
420
+class SequenceParallelismPass(VllmPatternMatcherPass):
421
"""
422
This pass enables sequence parallelism for models.
423
It identifies patterns where an AllReduce operation is followed by
@@ -466,19 +466,13 @@ class SequenceParallelismPass(VllmInductorPass):
466
467
LastAllReduceRMSNormPattern(epsilon, self.model_dtype,
468
self.device).register(self.patterns)
469
-
470
- # WARNING: This is a hack to clear the pattern matcher cache
471
- # and allow multiple values of epsilon.
472
- torch._inductor.pattern_matcher._seen_patterns.clear()
469
+ self.dump_patterns(config, self.patterns)
470
471
def is_applicable_for_shape(self, shape: Optional[int]) -> bool:
472
tp_size = get_tensor_model_parallel_world_size()
473
return shape is not None and shape % tp_size == 0
474
475
+ @VllmInductorPass.time_and_log
476
def __call__(self, graph: fx.Graph):
479
- self.begin()
480
- self.dump_graph(graph, "before_sequence_parallelism_pass")
481
- count = self.patterns.apply(graph)
482
- logger.debug("Replaced %s patterns with sequence parallelism", count)
483
- self.dump_graph(graph, "after_sequence_parallelism_pass")
484
- self.end_and_log()
477
+ self.matched_count = self.patterns.apply(graph)
478
+ logger.debug("Replaced %s patterns", self.matched_count)
vllm/compilation/vllm_inductor_pass.py
+108
-2
@@ -1,10 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import functools
4
+import operator
5
import time
6
+from pathlib import Path
7
+from typing import ClassVar, Optional
8
9
+import regex as re
10
import torch
11
from torch._dynamo.utils import lazy_format_graph_code
12
+from torch._inductor.pattern_matcher import (PatternMatcherPass,
13
+ PatternPrettyPrinter)
14
15
from vllm.config import VllmConfig
16
from vllm.logger import init_logger
@@ -19,6 +25,8 @@ class VllmInductorPass(InductorPass):
25
An inductor pass with access to vLLM PassConfig.
26
It provides timing, logging, and dumping utilities.
27
"""
28
+ dump_prefix: ClassVar[Optional[int]] = None
29
+ """Keep track of pass index for debug dump ordering."""
30
31
def __init__(self, config: VllmConfig):
32
self.pass_config = config.compilation_config.pass_config
@@ -28,8 +36,24 @@ class VllmInductorPass(InductorPass):
36
else None
37
self.pass_name = self.__class__.__name__
38
39
+ @staticmethod
40
+ def time_and_log(call_fn):
41
+
42
+ @functools.wraps(call_fn)
43
+ def wrapped(self: VllmInductorPass, graph: torch.fx.Graph):
44
+ self.begin()
45
+ self.dump_graph(graph, "before")
46
+ call_fn(self, graph)
47
+ self.dump_graph(graph, "after")
48
+ self.end_and_log()
49
+
50
+ return wrapped
51
+
52
def dump_graph(self, graph: torch.fx.Graph, stage: str):
32
- lazy_format_graph_code(stage, graph.owning_module)
53
+ i = VllmInductorPass.dump_prefix
54
+ i_str = "" if i is None else f".{i}"
55
+ lazy_format_graph_code(f"post_grad{i_str}.{self.pass_name}.{stage}",
56
+ graph.owning_module)
57
58
def begin(self):
59
self._start_time = time.perf_counter_ns()
@@ -40,6 +64,88 @@ class VllmInductorPass(InductorPass):
64
logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)
65
66
67
+class VllmPatternMatcherPass(VllmInductorPass):
68
+ """
69
+ A VllmInductorPass that uses the Inductor pattern matcher.
70
+ Its main use is providing the dump_patterns utility that dumps the
71
+ Inductor pattern matcher patterns into a file, which greatly aids debugging.
72
+
73
+ TODO(luka) move more utilities to this pass.
74
+ """
75
+ matched_count: int = 0
76
+ """The number of matched patterns in the pass."""
77
+
78
+ _OP_OVERLOAD_PATTERN: ClassVar[re.Pattern] = re.compile(
79
+ r"<OpOverload\(op='([^']*)', overload='([^']*)'\)>")
80
+
81
+ def _replace_op_overloads(self, string: str) -> str:
82
+ """Replace <OpOverload(..., ...)> with nicer formulations"""
83
+ return self._OP_OVERLOAD_PATTERN.sub(
84
+ lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}",
85
+ string,
86
+ )
87
+
88
+ def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass):
89
+ """
90
+ If debug dumping is enabled, dump the Inductor pattern-matcher patterns
91
+ into the debug_dump_path folder next to the dumped fx graphs.
92
+
93
+ This method does its best to print something that looks like Python code
94
+ for easier debugging and potentially navigation. If any errors appear in
95
+ the output, please add to this method.
96
+
97
+ TODO(luka): use pattern object to manually produce pattern graph
98
+ """
99
+ debug_dump_path = config.compilation_config.debug_dump_path
100
+ if not debug_dump_path:
101
+ return
102
+
103
+ rank = config.parallel_config.rank
104
+ debug_dump_path = Path(debug_dump_path) / f"rank_{rank}"
105
+ debug_dump_path.mkdir(parents=True, exist_ok=True)
106
+
107
+ from vllm.utils import unique_filepath
108
+ file_path = unique_filepath(
109
+ lambda i: debug_dump_path / f"patterns.{self.pass_name}.{i}.py")
110
+
111
+ with file_path.open("w") as f:
112
+ print(
113
+ f'# This file was produced by VllmPatternMatcherPass.'
114
+ f'dump_patterns for {self.pass_name}.\n'
115
+ f'# It does its best to produce valid-Python-looking code but'
116
+ f' please add to dump_patterns if there are any errors.\n\n'
117
+ f'from torch._higher_order_ops.auto_functionalize import '
118
+ f'auto_functionalized as auto_functionalized\n'
119
+ f'from torch._inductor.pattern_matcher import *',
120
+ file=f)
121
+
122
+ for node, patterns in pm_pass.patterns.items():
123
+ # fix the operator.getitem repr
124
+ if node[1] == operator.getitem:
125
+ node_repr = f"({repr(node[0])}, operator.getitem)"
126
+ else:
127
+ node_repr = repr(node)
128
+
129
+ node_repr = self._replace_op_overloads(node_repr)
130
+
131
+ print(f"\n\n# Patterns for op: {node_repr}", file=f)
132
+ for i, pattern in enumerate(patterns):
133
+ # reserve auto_functionalized ahead of time
134
+ pp = PatternPrettyPrinter()
135
+ pp.namespace.create_name("auto_functionalized", None)
136
+
137
+ # Assemble pattern
138
+ out_node = pp.pretty_print(pattern.pattern)
139
+ pattern_repr = "\n".join([f"def pattern_{i}():"] + [
140
+ f"{pp.memoized_objs_names[key]} = "
141
+ f"{pp.memoized_objs_pp[key]}"
142
+ for key in pp.memoized_objs_names
143
+ ] + [f"return {out_node}"]).replace("\n", "\n ")
144
+
145
+ pattern_repr = self._replace_op_overloads(pattern_repr)
146
+ print(f"{pattern_repr}\n", file=f)
147
+
148
+
149
class PrinterInductorPass(VllmInductorPass):
150
151
def __init__(self, name: str, config: VllmConfig):
vllm/config/__init__.py
+3
-4
@@ -905,10 +905,9 @@ def set_current_vllm_config(vllm_config: VllmConfig,
905
except Exception:
906
raise
907
else:
908
- logger.debug("enabled custom ops: %s",
909
- vllm_config.compilation_config.enabled_custom_ops)
910
- logger.debug("disabled custom ops: %s",
911
- vllm_config.compilation_config.disabled_custom_ops)
908
+ if check_compile:
909
+ vllm_config.compilation_config.custom_op_log_check()
910
+
911
if check_compile and \
912
vllm_config.compilation_config.level == CompilationLevel.PIECEWISE \
913
and compilation_counter.num_models_seen == num_models_seen:
vllm/config/compilation.py
+46
-2
@@ -487,6 +487,12 @@ class CompilationConfig:
487
"supported with torch>=2.9.0.dev. Set "
488
"use_inductor_graph_partition=False instead.")
489
490
+ for op in self.custom_ops:
491
+ if op[0] not in {'+', '-'} and op not in {'all', 'none'}:
492
+ raise ValueError(f"Invalid syntax '{op}' for custom op, "
493
+ "must be 'all', 'none', '+op' or '-op' "
494
+ "(where 'op' is the registered op name)")
495
+
496
def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:
497
if self.level == CompilationLevel.NO_COMPILATION:
498
raise ValueError("No compilation level is set.")
@@ -532,8 +538,8 @@ class CompilationConfig:
538
for x in self.compile_sizes:
539
if isinstance(x, str):
540
assert x == "cudagraph_capture_sizes", \
535
- "Unrecognized size type in compile_sizes, " \
536
- f"expect 'cudagraph_capture_sizes', got {x}"
541
+ "Unrecognized size type in compile_sizes, " \
542
+ f"expect 'cudagraph_capture_sizes', got {x}"
543
computed_compile_sizes.extend(self.cudagraph_capture_sizes)
544
else:
545
assert isinstance(x, int)
@@ -628,3 +634,41 @@ class CompilationConfig:
634
635
return use_fx_graph_piecewise_compilation or \
636
use_inductor_piecewise_compilation
637
+
638
+ def custom_op_log_check(self):
639
+ """
640
+ This method logs the enabled/disabled custom ops and checks that the
641
+ passed custom_ops field only contains relevant ops.
642
+ It is called at the end of set_current_vllm_config,
643
+ after the custom ops have been instantiated.
644
+ """
645
+
646
+ if len(self.enabled_custom_ops) + len(self.disabled_custom_ops) == 0:
647
+ logger.debug("No custom ops found in model.")
648
+ return
649
+
650
+ logger.debug("enabled custom ops: %s", self.enabled_custom_ops)
651
+ logger.debug("disabled custom ops: %s", self.disabled_custom_ops)
652
+
653
+ all_ops_in_model = (self.enabled_custom_ops | self.disabled_custom_ops)
654
+ for op in self.custom_ops:
655
+ if op in {"all", "none"}:
656
+ continue
657
+
658
+ assert op[0] in {'+', '-'}, "Invalid custom op syntax " \
659
+ "(should be checked during init)"
660
+
661
+ # check if op name exists in model
662
+ op_name = op[1:]
663
+ if op_name not in all_ops_in_model:
664
+ from vllm.model_executor.custom_op import CustomOp
665
+
666
+ # Does op exist at all or is it just not present in this model?
667
+ # Note: Only imported op classes appear in the registry.
668
+ missing_str = "doesn't exist (or wasn't imported/registered)" \
669
+ if op_name not in CustomOp.op_registry \
670
+ else "not present in model"
671
+
672
+ enable_str = "enabling" if op[0] == '+' else "disabling"
673
+ logger.warning_once("Op '%s' %s, %s with '%s' has no effect",
674
+ op_name, missing_str, enable_str, op)
vllm/envs.py
+6
@@ -190,6 +190,7 @@ if TYPE_CHECKING:
190
VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False
191
VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True
192
VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER"
193
+ VLLM_PATTERN_MATCH_DEBUG: Optional[str] = None
194
195
196
def get_default_cache_root():
@@ -442,6 +443,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
443
"VLLM_USE_STANDALONE_COMPILE":
444
lambda: os.environ.get("VLLM_USE_STANDALONE_COMPILE", "0") == "1",
445
446
+ # Debug pattern matching inside custom passes.
447
+ # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3').
448
+ "VLLM_PATTERN_MATCH_DEBUG":
449
+ lambda: os.environ.get("VLLM_PATTERN_MATCH_DEBUG", None),
450
+
451
# local rank of the process in the distributed setting, used to determine
452
# the GPU device id
453
"LOCAL_RANK":
vllm/utils/__init__.py
+14
-1
@@ -3392,7 +3392,7 @@ def length_from_prompt_token_ids_or_embeds(
3392
prompt_token_ids: Optional[list[int]],
3393
prompt_embeds: Optional[torch.Tensor],
3394
) -> int:
3395
- """Calculate the request length (in number of tokens) give either
3395
+ """Calculate the request length (in number of tokens) give either
3396
prompt_token_ids or prompt_embeds.
3397
"""
3398
prompt_token_len = None if prompt_token_ids is None else len(
@@ -3413,3 +3413,16 @@ def length_from_prompt_token_ids_or_embeds(
3413
f" prompt_token_ids={prompt_token_len}"
3414
f" prompt_embeds={prompt_embeds_len}")
3415
return prompt_token_len
3416
+
3417
+
3418
+@contextlib.contextmanager
3419
+def set_env_var(key, value):
3420
+ old = os.environ.get(key)
3421
+ os.environ[key] = value
3422
+ try:
3423
+ yield
3424
+ finally:
3425
+ if old is None:
3426
+ del os.environ[key]
3427
+ else:
3428
+ os.environ[key] = old