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] Enable silu_mul_fp8_quant fusion without custom ops enabled (#27146)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Oct 22, 2025 at 12:22 UTC
ab3e80042eac24dd362408e6d63ad98768046359
4 files changed
+128
-70
tests/compile/test_silu_mul_quant_fusion.py
+75
-43
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import cast
3
+import itertools
4
5
import pytest
6
import torch
@@ -16,7 +16,13 @@ from vllm.compilation.activation_quant_fusion import (
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
19
+from vllm.config import (
20
+ CompilationConfig,
21
+ CompilationMode,
22
+ PassConfig,
23
+ VllmConfig,
24
+ set_current_vllm_config,
25
+)
26
from vllm.model_executor.layers.activation import SiluAndMul
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
GroupShape,
@@ -25,7 +31,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
)
32
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
33
Fp8LinearOp,
28
- cutlass_fp8_supported,
34
+ maybe_create_device_identity,
35
)
36
from vllm.platforms import current_platform
37
@@ -54,6 +60,8 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
60
act_quant_static=True,
61
act_quant_group_shape=GroupShape.PER_TENSOR,
62
)
63
+ self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
64
+ self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
65
66
def forward(self, x):
67
y = self.silu_and_mul(x)
@@ -61,7 +69,14 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
69
return x2
70
71
def ops_in_model_before(self):
64
- return [SILU_MUL_OP, QUANT_OPS[kFp8StaticTensorSym]]
72
+ return [
73
+ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,
74
+ (
75
+ QUANT_OPS[kFp8StaticTensorSym]
76
+ if self.enable_quant_fp8_custom_op
77
+ else torch.ops.aten.reciprocal
78
+ ),
79
+ ]
80
81
def ops_in_model_after(self):
82
return [FUSED_OPS[kFp8StaticTensorSym]]
@@ -77,6 +92,7 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
92
assert silu_and_mul_nvfp4_quant_supported
93
94
self.silu_and_mul = SiluAndMul()
95
+ self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
96
97
# create nvfp4 weight
98
w = torch.rand((hidden_size, hidden_size))
@@ -101,7 +117,10 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
117
return out
118
119
def ops_in_model_before(self):
104
- return [SILU_MUL_OP, QUANT_OPS[kNvfp4Quant]]
120
+ return [
121
+ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,
122
+ QUANT_OPS[kNvfp4Quant],
123
+ ]
124
125
def ops_in_model_after(self):
126
return [FUSED_OPS[kNvfp4Quant]]
@@ -110,67 +129,80 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
129
@pytest.mark.parametrize("num_tokens", [32, 64])
130
@pytest.mark.parametrize("hidden_size", [128, 256])
131
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
132
+@pytest.mark.parametrize("enable_silu_mul_custom_op", [True, False])
133
@pytest.mark.parametrize(
114
- "model_class",
115
- cast(
116
- list[type],
117
- [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
118
- if is_nvfp4_supported()
119
- else [TestSiluMulFp8QuantModel],
120
- ),
134
+ "model_class, enable_quant_fp8_custom_op, cuda_force_torch",
135
+ list(itertools.product([TestSiluMulFp8QuantModel], [True, False], [True, False]))
136
+ + [(TestSiluMulNvfp4QuantModel, False, False)],
137
)
138
# cuda_force_torch used to test torch code path on platforms that
139
# cutlass_fp8_supported() == True.
124
-@pytest.mark.parametrize(
125
- "cuda_force_torch", [True, False] if cutlass_fp8_supported() else [True]
126
-)
140
@pytest.mark.skipif(
141
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
142
)
143
def test_fusion_silu_and_mul_quant(
131
- num_tokens, hidden_size, dtype, model_class, cuda_force_torch
144
+ num_tokens: int,
145
+ hidden_size: int,
146
+ dtype: torch.dtype,
147
+ model_class: type[TestSiluMulFp8QuantModel | TestSiluMulNvfp4QuantModel],
148
+ enable_silu_mul_custom_op: bool,
149
+ enable_quant_fp8_custom_op: bool,
150
+ cuda_force_torch: bool,
151
):
133
- if model_class == TestSiluMulNvfp4QuantModel and cuda_force_torch:
134
- pytest.skip("Duplicate tests for NVFP4")
152
+ if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():
153
+ pytest.skip("NVFP4 is not supported on this GPU.")
154
155
torch.set_default_device("cuda")
156
torch.set_default_dtype(dtype)
157
+ maybe_create_device_identity()
158
159
x = torch.rand(num_tokens, hidden_size * 2)
160
161
# Reshape pass is needed for the fusion pass to work
142
- config = VllmConfig()
143
- config.compilation_config = CompilationConfig(
144
- pass_config=PassConfig(enable_fusion=True, enable_noop=True)
162
+ custom_ops = []
163
+ if enable_silu_mul_custom_op:
164
+ custom_ops.append("+silu_and_mul")
165
+ if enable_quant_fp8_custom_op:
166
+ custom_ops.append("+quant_fp8")
167
+ config = VllmConfig(
168
+ compilation_config=CompilationConfig(
169
+ mode=CompilationMode.VLLM_COMPILE,
170
+ custom_ops=custom_ops,
171
+ pass_config=PassConfig(enable_fusion=True, enable_noop=True),
172
+ ),
173
)
146
- fusion_pass = ActivationQuantFusionPass(config)
174
148
- passes = [NoOpEliminationPass(config), fusion_pass, PostCleanupPass(config)]
149
- backend = TestBackend(*passes)
150
- model = model_class(hidden_size=hidden_size, cuda_force_torch=cuda_force_torch, x=x)
175
+ with set_current_vllm_config(config):
176
+ fusion_pass = ActivationQuantFusionPass(config)
177
152
- # First dimension dynamic
153
- torch._dynamo.mark_dynamic(x, 0)
178
+ passes = [NoOpEliminationPass(config), fusion_pass, PostCleanupPass(config)]
179
+ backend = TestBackend(*passes)
180
+ model = model_class(
181
+ hidden_size=hidden_size, cuda_force_torch=cuda_force_torch, x=x
182
+ )
183
155
- result = model(x)
184
+ # First dimension dynamic
185
+ torch._dynamo.mark_dynamic(x, 0)
186
157
- model2 = torch.compile(model, backend=backend)
158
- result2 = model2(x)
187
+ result = model(x)
188
160
- # Check that it gives the same answer
161
- if model_class == TestSiluMulFp8QuantModel:
162
- atol, rtol = 1e-3, 1e-3
163
- elif model_class == TestSiluMulNvfp4QuantModel:
164
- atol, rtol = 1e-1, 1e-1
189
+ model2 = torch.compile(model, backend=backend)
190
+ result2 = model2(x)
191
166
- torch.testing.assert_close(
167
- result[0].to(dtype=dtype), result2[0].to(dtype=dtype), atol=atol, rtol=rtol
168
- )
192
+ # Check that it gives the same answer
193
+ if model_class == TestSiluMulFp8QuantModel:
194
+ atol, rtol = 1e-3, 1e-3
195
+ elif model_class == TestSiluMulNvfp4QuantModel:
196
+ atol, rtol = 1e-1, 1e-1
197
+
198
+ torch.testing.assert_close(
199
+ result[0].to(dtype=dtype), result2[0].to(dtype=dtype), atol=atol, rtol=rtol
200
+ )
201
170
- assert fusion_pass.matched_count == 1
202
+ assert fusion_pass.matched_count == 1
203
172
- # In pre-nodes, quant op should be present and fused kernels should not
173
- backend.check_before_ops(model.ops_in_model_before())
204
+ # In pre-nodes, quant op should be present and fused kernels should not
205
+ backend.check_before_ops(model.ops_in_model_before())
206
175
- # In post-nodes, fused kernels should be present and quant op should not
176
- backend.check_after_ops(model.ops_in_model_after())
207
+ # In post-nodes, fused kernels should be present and quant op should not
208
+ backend.check_after_ops(model.ops_in_model_after())
vllm/compilation/activation_quant_fusion.py
+21
-26
@@ -18,12 +18,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
QuantKey,
19
kFp8StaticTensorSym,
20
kNvfp4Quant,
21
- kStaticTensorScale,
21
)
22
from vllm.platforms import current_platform
23
24
from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
25
from .inductor_pass import enable_fake_mode
26
+from .matcher_utils import MatcherQuantFP8, MatcherSiluAndMul
27
from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
28
29
logger = init_logger(__name__)
@@ -66,6 +66,8 @@ class ActivationQuantPattern(ABC):
66
)
67
self.FUSED_OP = FUSED_OPS[self.quant_key]
68
69
+ self.silu_and_mul_matcher = MatcherSiluAndMul()
70
+
71
def empty_quant(self, *args, **kwargs):
72
kwargs = {"dtype": self.quant_dtype, "device": "cuda", **kwargs}
73
return torch.empty(*args, **kwargs)
@@ -80,42 +82,38 @@ class SiluMulFp8StaticQuantPattern(ActivationQuantPattern):
82
Fusion for SiluMul+Fp8StaticQuant Pattern
83
"""
84
83
- def __init__(self, symmetric: bool = True):
84
- quant_key = QuantKey(
85
- dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=symmetric
86
- )
87
- super().__init__(quant_key)
85
+ def __init__(self):
86
+ super().__init__(kFp8StaticTensorSym)
87
+ self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
88
89
def register(self, pm_pass: PatternMatcherPass):
90
def pattern(
91
- result: torch.Tensor,
92
- result_silu_mul: torch.Tensor,
91
input: torch.Tensor,
92
scale: torch.Tensor,
93
):
96
- at1 = auto_functionalized(SILU_MUL_OP, result=result_silu_mul, input=input)
97
- at2 = auto_functionalized(
98
- self.QUANT_OP, result=result, input=at1[1], scale=scale
99
- )
100
- return at2[1]
94
+ result_silu_mul = self.silu_and_mul_matcher(input)
95
+ result_quant = self.quant_matcher(result_silu_mul, scale)
96
+ return result_quant[0]
97
98
def replacement(
103
- result: torch.Tensor,
104
- result_silu_mul: torch.Tensor,
99
input: torch.Tensor,
100
scale: torch.Tensor,
101
):
102
+ d = input.shape[-1] // 2
103
+ output_shape = input.shape[:-1] + (d,)
104
+ result = torch.empty(
105
+ output_shape, device=input.device, dtype=self.quant_dtype
106
+ )
107
at = auto_functionalized(
108
self.FUSED_OP, result=result, input=input, scale=scale
109
)
110
return at[1]
111
112
inputs = [
114
- self.empty_quant(5, 4), # result
115
- empty_bf16(5, 4), # result_silu_mul
116
- empty_bf16(5, 4), # input
117
- empty_fp32(1, 1), # scale
113
+ *self.silu_and_mul_matcher.inputs(), # input
114
+ self.quant_matcher.inputs()[1], # scale
115
]
116
+ pattern(*inputs)
117
118
register_replacement(pattern, replacement, inputs, fwd_only, pm_pass)
119
@@ -132,24 +130,22 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
130
def pattern(
131
result: torch.Tensor,
132
output_scale: torch.Tensor,
135
- result_silu_mul: torch.Tensor,
133
input: torch.Tensor,
134
scale: torch.Tensor,
135
):
139
- at1 = auto_functionalized(SILU_MUL_OP, result=result_silu_mul, input=input)
140
- at2 = auto_functionalized(
136
+ result_silu_mul = self.silu_and_mul_matcher(input)
137
+ at = auto_functionalized(
138
self.QUANT_OP,
139
output=result,
143
- input=at1[1],
140
+ input=result_silu_mul,
141
output_scale=output_scale,
142
input_scale=scale,
143
)
147
- return at2[1], at2[2]
144
+ return at[1], at[2]
145
146
def replacement(
147
result: torch.Tensor,
148
output_scale: torch.Tensor,
152
- result_silu_mul: torch.Tensor,
149
input: torch.Tensor,
150
scale: torch.Tensor,
151
):
@@ -165,7 +161,6 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
161
inputs = [
162
self.empty_quant(5, 32), # result
163
empty_i32(128, 4), # output_scale
168
- empty_bf16(5, 64), # result_silu_mul
164
empty_bf16(5, 64), # input
165
empty_fp32(1, 1), # scale
166
]
vllm/compilation/matcher_utils.py
+30
@@ -7,6 +7,7 @@ from torch._higher_order_ops import auto_functionalized
7
from torch._ops import OpOverload
8
9
from vllm.config import get_current_vllm_config
10
+from vllm.model_executor.layers.activation import SiluAndMul
11
from vllm.model_executor.layers.layernorm import RMSNorm
12
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -31,6 +32,8 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
32
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
33
QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default # noqa: E501
34
35
+SILU_MUL_OP = torch.ops._C.silu_and_mul.default
36
+
37
38
class MatcherCustomOp(ABC):
39
def __init__(self, enabled: bool):
@@ -206,3 +209,30 @@ class MatcherQuantFP8(MatcherCustomOp):
209
return [input, self.empty_f32(1, 1)]
210
211
return [input]
212
+
213
+
214
+class MatcherSiluAndMul(MatcherCustomOp):
215
+ def __init__(self, enabled: bool | None = None):
216
+ if enabled is None:
217
+ enabled = SiluAndMul.enabled()
218
+ super().__init__(enabled)
219
+
220
+ def inputs(self) -> list[torch.Tensor]:
221
+ input = self.empty(5, 4)
222
+ return [input]
223
+
224
+ def forward_custom(
225
+ self,
226
+ x: torch.Tensor,
227
+ ) -> torch.Tensor:
228
+ d = x.shape[-1] // 2
229
+ output_shape = x.shape[:-1] + (d,)
230
+ out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
231
+ result = auto_functionalized(SILU_MUL_OP, result=out, input=x)
232
+ return result[1]
233
+
234
+ def forward_native(
235
+ self,
236
+ x: torch.Tensor,
237
+ ) -> torch.Tensor:
238
+ return SiluAndMul.forward_native(x)
vllm/model_executor/layers/activation.py
+2
-1
@@ -80,7 +80,8 @@ class SiluAndMul(CustomOp):
80
elif current_platform.is_cpu():
81
self._forward_method = self.forward_native
82
83
- def forward_native(self, x: torch.Tensor) -> torch.Tensor:
83
+ @staticmethod
84
+ def forward_native(x: torch.Tensor) -> torch.Tensor:
85
"""PyTorch-native implementation equivalent to forward()."""
86
d = x.shape[-1] // 2
87
return F.silu(x[..., :d]) * x[..., d:]