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, 01:00 UTC
[Compilation] Add Unit Tests for VllmFusionPatternMatcherPass (#39692)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
BadrBasowid committed
Apr 17, 2026 at 06:57 UTC
29057d3bee1e4a5f84a41eb0cbd2f67b9fa35816
1 file changed
+128
tests/compile/passes/test_vllm_fusion_pattern_matcher_pass.py
new
+128
@@ -0,0 +1,128 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm.config
8
+from tests.compile.backend import TestBackend
9
+from vllm.platforms import current_platform
10
+from vllm.compilation.passes.vllm_inductor_pass import (
11
+ VllmFusionPatternMatcherPass,
12
+ VllmPatternMatcherPass,
13
+ VllmPatternReplacement,
14
+)
15
+from vllm.config import CompilationConfig, CompilationMode, VllmConfig
16
+
17
+
18
+
19
+class ReluToAbsPattern(VllmPatternReplacement):
20
+ """Replaces relu(x) with abs(x) — a minimal test fixture."""
21
+
22
+ @property
23
+ def pattern(self):
24
+ def _pattern(x: torch.Tensor) -> torch.Tensor:
25
+ return torch.ops.aten.relu.default(x)
26
+
27
+ return _pattern
28
+
29
+ @property
30
+ def replacement(self):
31
+ def _replacement(x: torch.Tensor) -> torch.Tensor:
32
+ return torch.ops.aten.abs.default(x)
33
+
34
+ return _replacement
35
+
36
+ def get_inputs(self) -> list[torch.Tensor]:
37
+ return [self.empty_fp32(4)]
38
+
39
+
40
+class ExpToSqrtPattern(VllmPatternReplacement):
41
+ """A second distinct pattern type — used to test uuid differentiation."""
42
+
43
+ @property
44
+ def pattern(self):
45
+ def _pattern(x: torch.Tensor) -> torch.Tensor:
46
+ return torch.ops.aten.exp.default(x)
47
+
48
+ return _pattern
49
+
50
+ @property
51
+ def replacement(self):
52
+ def _replacement(x: torch.Tensor) -> torch.Tensor:
53
+ return torch.ops.aten.sqrt.default(x)
54
+
55
+ return _replacement
56
+
57
+ def get_inputs(self) -> list[torch.Tensor]:
58
+ return [self.empty_fp32(4)]
59
+
60
+
61
+
62
+class ReluFusionPass(VllmFusionPatternMatcherPass):
63
+ def __init__(self, config: VllmConfig) -> None:
64
+ super().__init__(config, "test_relu_fusion")
65
+ self.register(ReluToAbsPattern())
66
+
67
+
68
+class TwoPatternFusionPass(VllmFusionPatternMatcherPass):
69
+ def __init__(self, config: VllmConfig) -> None:
70
+ super().__init__(config, "test_two_pattern_fusion")
71
+ self.register(ReluToAbsPattern())
72
+ self.register(ExpToSqrtPattern())
73
+
74
+
75
+
76
+@pytest.fixture
77
+def vllm_config():
78
+ return VllmConfig(
79
+ compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE),
80
+ )
81
+
82
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
83
+def test_register_tracks_patterns(vllm_config):
84
+ """register() appends each VllmPatternReplacement to _pattern_replacements."""
85
+ with vllm.config.set_current_vllm_config(vllm_config):
86
+ single = ReluFusionPass(vllm_config)
87
+ two = TwoPatternFusionPass(vllm_config)
88
+
89
+ assert len(single._pattern_replacements) == 1
90
+ assert len(two._pattern_replacements) == 2
91
+
92
+
93
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
94
+def test_uuid_stable(vllm_config):
95
+ """Two instances of the same pass class produce identical uuids."""
96
+ with vllm.config.set_current_vllm_config(vllm_config):
97
+ p1 = ReluFusionPass(vllm_config)
98
+ p2 = ReluFusionPass(vllm_config)
99
+ p3= TwoPatternFusionPass(vllm_config)
100
+
101
+ assert p1.uuid() == p2.uuid()
102
+ assert p1.uuid() != p3.uuid()
103
+ assert p2.uuid() != p3.uuid()
104
+
105
+
106
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Requires CUDA")
107
+@pytest.mark.parametrize("N", [1, 2, 4])
108
+def test_matched_count_and_match_table(vllm_config, N):
109
+ """matched_count and match_table reflect the number of matched patterns."""
110
+
111
+ class Model(torch.nn.Module):
112
+ def forward(self, *inputs):
113
+ # N independent relus
114
+ return sum(torch.relu(x) for x in inputs)
115
+
116
+ with vllm.config.set_current_vllm_config(vllm_config):
117
+ torch.set_default_device("cuda")
118
+ torch.set_default_dtype(torch.float32)
119
+
120
+ fusion_pass = ReluFusionPass(vllm_config)
121
+ backend = TestBackend(fusion_pass)
122
+ model = torch.compile(Model(), backend=backend)
123
+
124
+ inputs = [torch.rand(8) for _ in range(N)]
125
+ model(*inputs)
126
+
127
+ assert fusion_pass.matched_count == N
128
+ assert VllmPatternMatcherPass.match_table["test_relu_fusion"] >= N