@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b7624069ff

[Fusion] Manual `ActivationQuantFusionPass` initial application (#51415)

Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: ElizaWszola <ewszola@redhat.com>

Misha Goin committed Sep 3, 2026 at 13:37 UTC b7624069ff023043b39ba4ea67ae5f6045556819
10 files changed +740 -19
tests/compile/fusions_e2e/common.py
+39
@@ -43,6 +43,45 @@ is_blackwell = lambda: current_platform.is_device_capability_family(100)
43 """Are we running on Blackwell, a lot of tests depend on it"""
44
45
46 +def nvfp4_kernel_exposes_input_quant_key() -> bool:
47 + """Check if the NVFP4 kernel selected on this platform exposes input_quant_key.
48 +
49 + FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to
50 + layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant)
51 + output format and the cutedsl backend's expected input format.
52 +
53 + FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports
54 + manual fusion.
55 + """
56 + if not current_platform.is_cuda():
57 + return False
58 +
59 + try:
60 + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
61 + FlashInferCuteDslNvFp4LinearKernel,
62 + )
63 +
64 + is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported()
65 + if is_supported:
66 + return False
67 + except ImportError:
68 + pass
69 +
70 + try:
71 + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
72 + FlashInferCutlassNvFp4LinearKernel,
73 + )
74 +
75 + is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported()
76 + if is_supported:
77 + return True
78 + except ImportError:
79 + pass
80 +
81 + # Fallback: assume no manual fusion support
82 + return False
83 +
84 +
85 def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
86 """Generate all combinations of custom ops for parametrization."""
87 custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]
tests/compile/fusions_e2e/test_tp1_quant.py
+10 -5
@@ -14,6 +14,7 @@ from .common import (
14 Matches,
15 custom_ops_combos,
16 is_blackwell,
17 + nvfp4_kernel_exposes_input_quant_key,
18 )
19 from .models import (
20 FLASHINFER_ATTN,
@@ -77,7 +78,6 @@ def test_tp1_fp8_fusions(
78 inductor_graph_partition: bool,
79 use_deepgemm: bool,
80 run_e2e_fusion_test,
80 - monkeypatch,
81 ):
82 if use_deepgemm and not current_platform.is_cuda():
83 pytest.skip("DeepGemm only supported on CUDA")
@@ -109,7 +109,7 @@ def test_tp1_fp8_fusions(
109 custom_ops=custom_ops.split(","),
110 pass_config=PassConfig(
111 fuse_norm_quant=True,
112 - fuse_act_quant=True,
112 + fuse_act_quant=False,
113 fuse_attn_quant=True,
114 enable_qk_norm_rope_fusion=True,
115 ),
@@ -119,7 +119,6 @@ def test_tp1_fp8_fusions(
119
120 matches_check = [
121 "rms_quant_fusion",
122 - "act_quant_fusion",
122 "norm_rope_fusion",
123 "attn_quant_fusion",
124 ]
@@ -168,6 +167,12 @@ def test_tp1_fp4_fusions(
167 inductor_graph_partition: bool,
168 run_e2e_fusion_test,
169 ):
170 + if nvfp4_kernel_exposes_input_quant_key():
171 + pytest.skip(
172 + "NVFP4 kernel exposes input_quant_key; manual fusion fires "
173 + "instead of compiler pass-based fusion"
174 + )
175 +
176 matches = matches_fn(n_layers)
177
178 # Reduce size of model and skip weight loading time
@@ -181,13 +186,13 @@ def test_tp1_fp4_fusions(
186 custom_ops=custom_ops.split(","),
187 pass_config=PassConfig(
188 fuse_norm_quant=True,
184 - fuse_act_quant=True,
189 + fuse_act_quant=False,
190 fuse_attn_quant=True,
191 enable_qk_norm_rope_fusion=True,
192 ),
193 )
194
190 - matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"]
195 + matches_check = ["attn_quant_fusion", "norm_rope_fusion"]
196
197 run_e2e_fusion_test(
198 model_name,
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+9 -6
@@ -14,6 +14,7 @@ from .common import (
14 Matches,
15 custom_ops_combos,
16 is_blackwell,
17 + nvfp4_kernel_exposes_input_quant_key,
18 )
19 from .models import (
20 FLASHINFER_ATTN,
@@ -68,7 +69,6 @@ def test_tp2_ar_rms_fp8_fusions(
69 custom_ops: str,
70 inductor_graph_partition: bool,
71 run_e2e_fusion_test,
71 - monkeypatch,
72 ):
73 matches = matches_fn(n_layers)
74
@@ -89,7 +89,7 @@ def test_tp2_ar_rms_fp8_fusions(
89 custom_ops=custom_ops.split(","),
90 pass_config=PassConfig(
91 fuse_norm_quant=True,
92 - fuse_act_quant=True,
92 + fuse_act_quant=False,
93 fuse_attn_quant=True,
94 enable_qk_norm_rope_fusion=True,
95 fuse_allreduce_rms=True,
@@ -98,7 +98,6 @@ def test_tp2_ar_rms_fp8_fusions(
98
99 matches_check = [
100 "rms_quant_fusion",
101 - "act_quant_fusion",
101 "norm_rope_fusion",
102 "attn_quant_fusion",
103 "ar_rms_fusion",
@@ -139,8 +138,13 @@ def test_tp2_ar_rms_fp4_fusions(
138 custom_ops: str,
139 inductor_graph_partition: bool,
140 run_e2e_fusion_test,
142 - monkeypatch,
141 ):
142 + if nvfp4_kernel_exposes_input_quant_key():
143 + pytest.skip(
144 + "NVFP4 kernel exposes input_quant_key; manual fusion fires "
145 + "instead of compiler pass-based fusion"
146 + )
147 +
148 matches = matches_fn(n_layers)
149
150 # Reduce size of model and skip weight loading time
@@ -154,14 +158,13 @@ def test_tp2_ar_rms_fp4_fusions(
158 use_inductor_graph_partition=inductor_graph_partition,
159 custom_ops=custom_ops.split(","),
160 pass_config=PassConfig(
157 - fuse_act_quant=True,
161 + fuse_act_quant=False,
162 fuse_attn_quant=True,
163 fuse_allreduce_rms=True,
164 ),
165 )
166
167 matches_check = [
164 - "act_quant_fusion",
168 "attn_quant_fusion",
169 "ar_rms_fusion",
170 ]
tests/compile/fusions_e2e/test_tp2_async_tp.py
+10 -6
@@ -14,6 +14,7 @@ from .common import (
14 Matches,
15 custom_ops_combos,
16 is_blackwell,
17 + nvfp4_kernel_exposes_input_quant_key,
18 )
19 from .models import (
20 FLASHINFER_ATTN,
@@ -61,7 +62,7 @@ def test_tp2_async_tp_fp8_fusions(
62 custom_ops=custom_ops.split(","),
63 pass_config=PassConfig(
64 fuse_norm_quant=True,
64 - fuse_act_quant=True,
65 + fuse_act_quant=False,
66 fuse_attn_quant=True,
67 enable_qk_norm_rope_fusion=True,
68 enable_sp=True,
@@ -74,7 +75,6 @@ def test_tp2_async_tp_fp8_fusions(
75
76 matches_check = [
77 "rms_quant_fusion",
77 - "act_quant_fusion",
78 "norm_rope_fusion",
79 "attn_quant_fusion",
80 "sequence_parallel",
@@ -114,6 +114,12 @@ def test_tp2_async_tp_nvfp4_fusions(
114 inductor_graph_partition: bool,
115 run_e2e_fusion_test,
116 ):
117 + if nvfp4_kernel_exposes_input_quant_key():
118 + pytest.skip(
119 + "NVFP4 kernel exposes input_quant_key; manual fusion fires "
120 + "instead of compiler pass-based fusion"
121 + )
122 +
123 # NVFP4 currently wires the all-gather + GEMM path only.
124 matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2)
125
@@ -127,7 +133,7 @@ def test_tp2_async_tp_nvfp4_fusions(
133 use_inductor_graph_partition=inductor_graph_partition,
134 custom_ops=custom_ops.split(","),
135 pass_config=PassConfig(
130 - fuse_act_quant=True,
136 + fuse_act_quant=False,
137 fuse_attn_quant=True,
138 enable_sp=True,
139 fuse_gemm_comms=True,
@@ -138,7 +144,6 @@ def test_tp2_async_tp_nvfp4_fusions(
144 )
145
146 matches_check = [
141 - "act_quant_fusion",
147 "attn_quant_fusion",
148 "sequence_parallel",
149 "async_tp",
@@ -245,7 +250,7 @@ def test_tp2_sp_ar_rms_fp8_fusions(
250 custom_ops=custom_ops.split(","),
251 pass_config=PassConfig(
252 fuse_norm_quant=True,
248 - fuse_act_quant=True,
253 + fuse_act_quant=False,
254 fuse_attn_quant=True,
255 enable_qk_norm_rope_fusion=True,
256 enable_sp=True,
@@ -258,7 +263,6 @@ def test_tp2_sp_ar_rms_fp8_fusions(
263
264 matches_check = [
265 "rms_quant_fusion",
261 - "act_quant_fusion",
266 "norm_rope_fusion",
267 "attn_quant_fusion",
268 "ar_rms_fusion",
tests/compile/passes/test_silu_mul_quant_manual_fusion.py new
+348
@@ -0,0 +1,348 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests for manual fusion via maybe_fused_act_quant.
5 +
6 +Tests all fusion paths in _FUSED_ACT_QUANT:
7 +- kFp8StaticTensorSym: all platforms
8 +- kFp8Dynamic128Sym: CUDA only
9 +- kNvfp4Dynamic: CUDA SM100+ only
10 +"""
11 +
12 +import pytest
13 +import torch
14 +
15 +import vllm.envs as envs
16 +from tests.utils import TestFP8Layer
17 +from vllm.config import (
18 + CompilationConfig,
19 + VllmConfig,
20 + set_current_vllm_config,
21 +)
22 +from vllm.model_executor.kernels.linear import (
23 + CutlassFP8ScaledMMLinearKernel,
24 + FlashInferFP8ScaledMMLinearKernel,
25 + FP8ScaledMMLinearKernel,
26 + PerTensorTorchFP8ScaledMMLinearKernel,
27 + ROCmFP8ScaledMMLinearKernel,
28 +)
29 +from vllm.model_executor.layers.activation import SiluAndMul
30 +from vllm.model_executor.layers.fusion.fused_act_quant import (
31 + _FUSED_ACT_QUANT,
32 + maybe_fused_act_quant,
33 +)
34 +from vllm.model_executor.layers.fusion.quant_activation import (
35 + QuantizedActivation,
36 + expose_input_quant_key,
37 +)
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 + kFp8Dynamic128Sym,
40 + kFp8StaticTensorSym,
41 + kNvfp4Dynamic,
42 +)
43 +from vllm.platforms import current_platform
44 +
45 +
46 +# Mock linear layer for testing fusion paths that don't have real kernel support
47 +class MockLinearForFusion(torch.nn.Module):
48 + """Mock linear layer that exposes input_quant_key for fusion testing."""
49 +
50 + def __init__(self, quant_key, input_scale=None, input_global_scale=None):
51 + super().__init__()
52 + self.input_quant_key = quant_key
53 + if input_scale is not None:
54 + self.input_scale = input_scale
55 + if input_global_scale is not None:
56 + self.input_global_scale = input_global_scale
57 +
58 +
59 +ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel]
60 +CUDA_KERNELS = [
61 + FlashInferFP8ScaledMMLinearKernel,
62 + CutlassFP8ScaledMMLinearKernel,
63 + PerTensorTorchFP8ScaledMMLinearKernel,
64 +]
65 +TEST_KERNELS = ROCM_KERNELS if current_platform.is_rocm() else CUDA_KERNELS
66 +
67 +
68 +@pytest.mark.parametrize("num_tokens", [32, 64])
69 +@pytest.mark.parametrize("hidden_size", [128, 256])
70 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
71 +@pytest.mark.parametrize("force_kernel", TEST_KERNELS)
72 +@pytest.mark.skipif(
73 + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
74 +)
75 +def test_manual_fusion_fp8_static_with_linear(
76 + num_tokens: int,
77 + hidden_size: int,
78 + dtype: torch.dtype,
79 + force_kernel: FP8ScaledMMLinearKernel,
80 +):
81 + """Test manual fusion with real FP8 linear layer (kFp8StaticTensorSym).
82 +
83 + This is an end-to-end test that verifies the full flow:
84 + unfused (silu_and_mul -> in-kernel quant) vs fused (silu_and_mul_quant).
85 + """
86 + torch.set_default_device("cuda")
87 + torch.set_default_dtype(dtype)
88 +
89 + x = torch.rand(num_tokens, hidden_size * 2)
90 +
91 + config = VllmConfig(
92 + compilation_config=CompilationConfig(custom_ops=["none"]),
93 + )
94 +
95 + with set_current_vllm_config(config):
96 + silu_and_mul = SiluAndMul()
97 + fp8_linear = TestFP8Layer(
98 + weight_shape=(hidden_size, hidden_size),
99 + activation_quant_key=kFp8StaticTensorSym,
100 + weight_quant_key=kFp8StaticTensorSym,
101 + force_kernel=force_kernel,
102 + input_dtype=dtype,
103 + )
104 +
105 + # Run without fusion: silu_and_mul returns plain tensor
106 + y_unfused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear)
107 + assert isinstance(y_unfused, torch.Tensor)
108 + result_unfused = fp8_linear(y_unfused)
109 +
110 + # Enable fusion
111 + expose_input_quant_key(fp8_linear, fp8_linear.kernel)
112 +
113 + if not hasattr(fp8_linear, "input_quant_key"):
114 + pytest.skip(
115 + f"Kernel {force_kernel.__name__} doesn't support input_quant_key"
116 + )
117 +
118 + # Run with fusion: silu_and_mul returns QuantizedActivation
119 + y_fused = maybe_fused_act_quant(silu_and_mul, x, fp8_linear)
120 + assert isinstance(y_fused, QuantizedActivation)
121 + assert y_fused.quant_key == kFp8StaticTensorSym
122 + assert y_fused.data.dtype == current_platform.fp8_dtype()
123 + assert y_fused.data.shape == (num_tokens, hidden_size)
124 + assert y_fused.orig_dtype == dtype
125 + assert y_fused.orig_shape == (num_tokens, hidden_size)
126 +
127 + result_fused = fp8_linear(y_fused)
128 +
129 + torch.testing.assert_close(
130 + result_fused.to(dtype=dtype),
131 + result_unfused.to(dtype=dtype),
132 + atol=5e-2,
133 + rtol=5e-2,
134 + )
135 +
136 +
137 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
138 +@pytest.mark.skipif(
139 + not current_platform.is_cuda_alike(), reason="Dynamic block quant CUDA only"
140 +)
141 +@pytest.mark.skipif(
142 + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
143 +)
144 +def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype):
145 + """Test kFp8Dynamic128Sym fusion path (group_size=128).
146 +
147 + Compares fused (silu_and_mul_per_block_quant) vs unfused (silu_and_mul)
148 + by dequantizing the fused result and comparing with unfused.
149 + """
150 + if (SiluAndMul, kFp8Dynamic128Sym) not in _FUSED_ACT_QUANT:
151 + pytest.skip("kFp8Dynamic128Sym fusion not available")
152 +
153 + torch.set_default_device("cuda")
154 + torch.set_default_dtype(dtype)
155 +
156 + # hidden_size must be divisible by group_size (128)
157 + num_tokens, hidden_size = 32, 256
158 + group_size = 128
159 + x = torch.rand(num_tokens, hidden_size * 2)
160 +
161 + config = VllmConfig(
162 + compilation_config=CompilationConfig(custom_ops=["none"]),
163 + )
164 +
165 + with set_current_vllm_config(config):
166 + silu_and_mul = SiluAndMul()
167 +
168 + # Unfused path: just apply silu_and_mul
169 + mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size)
170 + result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key)
171 + assert isinstance(result_unfused, torch.Tensor)
172 +
173 + # Fused path: apply silu_and_mul + per-block quantization
174 + mock_linear_with_key = MockLinearForFusion(kFp8Dynamic128Sym)
175 + result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key)
176 +
177 + # Verify fused result structure
178 + assert isinstance(result_fused, QuantizedActivation)
179 + assert result_fused.quant_key == kFp8Dynamic128Sym
180 + assert result_fused.data.dtype == current_platform.fp8_dtype()
181 + assert result_fused.data.shape == (num_tokens, hidden_size)
182 + assert result_fused.orig_dtype == dtype
183 +
184 + # Check scale shape
185 + expected_num_groups = hidden_size // group_size
186 + assert result_fused.scale.shape == (num_tokens, expected_num_groups)
187 +
188 + # Dequantize fused result and compare with unfused
189 + # Per-block dequant: data * scale (broadcast scale across group)
190 + dequant_data = result_fused.data.to(dtype).view(
191 + num_tokens, expected_num_groups, group_size
192 + )
193 + scales_expanded = result_fused.scale.unsqueeze(
194 + -1
195 + ) # (num_tokens, num_groups, 1)
196 + dequant_result = (
197 + (dequant_data * scales_expanded).view(num_tokens, hidden_size).to(dtype)
198 + )
199 +
200 + torch.testing.assert_close(
201 + dequant_result,
202 + result_unfused,
203 + atol=5e-2,
204 + rtol=5e-2,
205 + )
206 +
207 +
208 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
209 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="NVFP4 CUDA only")
210 +@pytest.mark.skipif(
211 + not current_platform.has_device_capability(100), reason="NVFP4 requires SM100+"
212 +)
213 +@pytest.mark.skipif(
214 + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
215 +)
216 +def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
217 + """Test kNvfp4Dynamic fusion path.
218 +
219 + Compares fused (silu_and_mul_nvfp4_quant) vs unfused (silu_and_mul)
220 + by dequantizing the fused result and comparing with unfused.
221 + """
222 + if (SiluAndMul, kNvfp4Dynamic) not in _FUSED_ACT_QUANT:
223 + pytest.skip("kNvfp4Dynamic fusion not available")
224 +
225 + from tests.kernels.quantization.nvfp4_utils import dequantize_nvfp4_to_dtype
226 +
227 + torch.set_default_device("cuda")
228 + torch.set_default_dtype(dtype)
229 +
230 + # NVFP4 requires hidden_size divisible by 16 (block size) and by 2 (packing).
231 + # M (num_tokens) must be >= 128 for the 128x4 swizzled scale layout.
232 + num_tokens, hidden_size = 128, 256
233 + x = torch.rand(num_tokens, hidden_size * 2)
234 + input_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
235 +
236 + config = VllmConfig(
237 + compilation_config=CompilationConfig(custom_ops=["none"]),
238 + )
239 +
240 + with set_current_vllm_config(config):
241 + silu_and_mul = SiluAndMul()
242 +
243 + # Unfused path: just apply silu_and_mul
244 + mock_linear_no_key = torch.nn.Linear(hidden_size, hidden_size)
245 + result_unfused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_no_key)
246 + assert isinstance(result_unfused, torch.Tensor)
247 +
248 + # Fused path: apply silu_and_mul + NVFP4 quantization
249 + mock_linear_with_key = MockLinearForFusion(
250 + kNvfp4Dynamic, input_global_scale=input_global_scale
251 + )
252 + result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key)
253 +
254 + # Verify fused result structure
255 + assert isinstance(result_fused, QuantizedActivation)
256 + assert result_fused.quant_key == kNvfp4Dynamic
257 + # NVFP4 packs 2 values into 1 byte
258 + assert result_fused.data.dtype == torch.uint8
259 + assert result_fused.data.shape == (num_tokens, hidden_size // 2)
260 + assert result_fused.orig_dtype == dtype
261 + assert result_fused.orig_shape == (num_tokens, hidden_size)
262 +
263 + # Dequantize fused result and compare with unfused
264 + dequant_result = dequantize_nvfp4_to_dtype(
265 + tensor_fp4=result_fused.data,
266 + tensor_sf=result_fused.scale,
267 + global_scale=input_global_scale,
268 + dtype=dtype,
269 + device="cuda",
270 + block_size=16,
271 + is_sf_128x4_layout=True,
272 + )
273 +
274 + torch.testing.assert_close(
275 + dequant_result,
276 + result_unfused,
277 + atol=3e-1,
278 + rtol=3e-1,
279 + )
280 +
281 +
282 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
283 +@pytest.mark.skipif(
284 + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
285 +)
286 +def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
287 + """Test that maybe_fused_act_quant falls back when no input_quant_key."""
288 + torch.set_default_device("cuda")
289 + torch.set_default_dtype(dtype)
290 +
291 + x = torch.rand(32, 256)
292 +
293 + config = VllmConfig(
294 + compilation_config=CompilationConfig(custom_ops=["none"]),
295 + )
296 +
297 + with set_current_vllm_config(config):
298 + silu_and_mul = SiluAndMul()
299 + # Linear without input_quant_key attribute
300 + mock_linear = torch.nn.Linear(128, 128)
301 +
302 + result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
303 +
304 + # Should fall back to plain silu_and_mul
305 + assert isinstance(result, torch.Tensor)
306 + assert result.shape == (32, 128)
307 + assert result.dtype == dtype
308 +
309 +
310 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
311 +@pytest.mark.skipif(
312 + envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
313 +)
314 +def test_manual_fusion_fallback_unsupported_key(dtype: torch.dtype):
315 + """Test that maybe_fused_act_quant falls back for unsupported quant keys."""
316 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
317 + GroupShape,
318 + QuantKey,
319 + ScaleDesc,
320 + )
321 +
322 + torch.set_default_device("cuda")
323 + torch.set_default_dtype(dtype)
324 +
325 + x = torch.rand(32, 256)
326 +
327 + config = VllmConfig(
328 + compilation_config=CompilationConfig(custom_ops=["none"]),
329 + )
330 +
331 + with set_current_vllm_config(config):
332 + # Create an unsupported quant key
333 + unsupported_key = QuantKey(
334 + dtype=torch.int8,
335 + scale=ScaleDesc(
336 + dtype=torch.float32, static=False, group_shape=GroupShape(1, 1)
337 + ),
338 + )
339 +
340 + silu_and_mul = SiluAndMul()
341 + mock_linear = MockLinearForFusion(unsupported_key)
342 +
343 + result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
344 +
345 + # Should fall back to plain silu_and_mul since key not in _FUSED_ACT_QUANT
346 + assert isinstance(result, torch.Tensor)
347 + assert result.shape == (32, 128)
348 + assert result.dtype == dtype
tests/fusion/test_quant_activation_contract.py
+4
@@ -25,6 +25,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
25 from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
26 FlashInferFP8ScaledMMLinearKernel,
27 )
28 +from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
29 + PerTensorTorchFP8ScaledMMLinearKernel,
30 +)
31 from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
32 FP8ScaledMMLinearLayerConfig,
33 Int8ScaledMMLinearKernel,
@@ -46,6 +49,7 @@ SUPPORTING = {
49 CutlassFP8ScaledMMLinearKernel,
50 FlashInferFP8ScaledMMLinearKernel,
51 FlashInferCutlassNvFp4LinearKernel,
52 + PerTensorTorchFP8ScaledMMLinearKernel,
53 }
54
55
tests/kernels/test_fused_quant_activation.py
+147
@@ -2,10 +2,20 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 +import torch.nn.functional as F
6
7 import vllm._custom_ops as ops
8 from tests.kernels.utils import opcheck
9 from vllm.model_executor.layers.activation import SiluAndMul
10 +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
11 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
12 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13 + per_token_group_quant_fp8,
14 +)
15 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 + kFp8Dynamic128Sym,
17 + kFp8StaticTensorSym,
18 +)
19 from vllm.platforms import current_platform
20 from vllm.utils.torch_utils import set_random_seed
21
@@ -69,3 +79,140 @@ def test_silu_and_mul(
79 ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
80 )
81 opcheck(torch.ops._C.silu_and_mul_quant, (ops_out, x, scale))
82 +
83 +
84 +# ---------------------------------------------------------------------------
85 +# Tests for maybe_fused_act_quant interface
86 +# ---------------------------------------------------------------------------
87 +
88 +
89 +class MockLinearFp8Static(torch.nn.Module):
90 + """Mock linear layer advertising kFp8StaticTensorSym."""
91 +
92 + def __init__(self, input_scale: torch.Tensor):
93 + super().__init__()
94 + self.input_quant_key = kFp8StaticTensorSym
95 + self.input_scale = input_scale
96 +
97 +
98 +class MockLinearFp8Dynamic128(torch.nn.Module):
99 + """Mock linear layer advertising kFp8Dynamic128Sym."""
100 +
101 + def __init__(self):
102 + super().__init__()
103 + self.input_quant_key = kFp8Dynamic128Sym
104 +
105 +
106 +class MockLinearNoQuant(torch.nn.Module):
107 + """Mock linear layer with no input_quant_key (no fusion)."""
108 +
109 + pass
110 +
111 +
112 +@pytest.mark.parametrize("num_tokens", [1, 16, 128])
113 +@pytest.mark.parametrize("hidden_size", [128, 512, 1024])
114 +@pytest.mark.parametrize("dtype", DTYPES)
115 +@torch.inference_mode()
116 +def test_maybe_fused_act_quant_fp8_static(
117 + default_vllm_config,
118 + num_tokens: int,
119 + hidden_size: int,
120 + dtype: torch.dtype,
121 +) -> None:
122 + """Test maybe_fused_act_quant with FP8 static per-tensor quantization."""
123 + device = "cuda:0"
124 + torch.set_default_device(device)
125 +
126 + act_fn = SiluAndMul()
127 + scale = torch.tensor([0.5], device=device, dtype=torch.float32)
128 + linear = MockLinearFp8Static(scale)
129 +
130 + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
131 + result = maybe_fused_act_quant(act_fn, x, linear)
132 +
133 + assert isinstance(result, QuantizedActivation)
134 + assert result.quant_key == kFp8StaticTensorSym
135 + assert result.data.dtype == current_platform.fp8_dtype()
136 + assert result.orig_dtype == dtype
137 + assert result.orig_shape == (num_tokens, hidden_size)
138 +
139 + ref_out = ref_impl(act_fn, x, scale)
140 + torch.testing.assert_close(result.data.to(torch.float32), ref_out.to(torch.float32))
141 +
142 +
143 +@pytest.mark.parametrize("num_tokens", [1, 16, 128])
144 +@pytest.mark.parametrize("hidden_size", [128, 512, 1024])
145 +@pytest.mark.parametrize("dtype", DTYPES)
146 +@torch.inference_mode()
147 +def test_maybe_fused_act_quant_fp8_dynamic_block(
148 + default_vllm_config,
149 + num_tokens: int,
150 + hidden_size: int,
151 + dtype: torch.dtype,
152 +) -> None:
153 + """Test maybe_fused_act_quant with FP8 dynamic per-block quantization."""
154 + group_size = 128 # We only support 128 for now
155 +
156 + device = "cuda:0"
157 + torch.set_default_device(device)
158 +
159 + act_fn = SiluAndMul()
160 + linear = MockLinearFp8Dynamic128()
161 +
162 + scale = 1 / hidden_size
163 + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device) * scale
164 + result = maybe_fused_act_quant(act_fn, x, linear)
165 +
166 + assert isinstance(result, QuantizedActivation)
167 + assert result.quant_key == kFp8Dynamic128Sym
168 + assert result.data.dtype == current_platform.fp8_dtype()
169 + assert result.orig_dtype == dtype
170 + assert result.orig_shape == (num_tokens, hidden_size)
171 +
172 + num_groups = hidden_size // group_size
173 + assert result.scale.shape == (num_tokens, num_groups)
174 +
175 + gate, up = x.split(hidden_size, dim=-1)
176 + silu_out = F.silu(gate) * up
177 + ref_out, ref_scales = per_token_group_quant_fp8(
178 + silu_out, group_size=group_size, use_ue8m0=False
179 + )
180 +
181 + torch.testing.assert_close(result.scale, ref_scales, rtol=1e-5, atol=1e-5)
182 +
183 + ref_deq = ref_out.to(torch.float32) * ref_scales.repeat_interleave(
184 + group_size, dim=1
185 + )
186 + result_deq = result.data.to(torch.float32) * result.scale.repeat_interleave(
187 + group_size, dim=1
188 + )
189 + torch.testing.assert_close(ref_deq, result_deq, atol=5e-2, rtol=5e-2)
190 +
191 +
192 +@pytest.mark.parametrize("num_tokens", [1, 16, 128])
193 +@pytest.mark.parametrize("hidden_size", [128, 512])
194 +@pytest.mark.parametrize("dtype", DTYPES)
195 +@torch.inference_mode()
196 +def test_maybe_fused_act_quant_fallback(
197 + default_vllm_config,
198 + num_tokens: int,
199 + hidden_size: int,
200 + dtype: torch.dtype,
201 +) -> None:
202 + """Test maybe_fused_act_quant falls back when no input_quant_key."""
203 + device = "cuda:0"
204 + torch.set_default_device(device)
205 +
206 + act_fn = SiluAndMul()
207 + linear = MockLinearNoQuant()
208 + x = torch.randn(num_tokens, hidden_size * 2, dtype=dtype, device=device)
209 +
210 + result = maybe_fused_act_quant(act_fn, x, linear)
211 +
212 + assert isinstance(result, torch.Tensor)
213 + assert not isinstance(result, QuantizedActivation)
214 + assert result.dtype == dtype
215 + assert result.shape == (num_tokens, hidden_size)
216 +
217 + ref_out = act_fn(x)
218 + torch.testing.assert_close(result, ref_out)
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+10 -1
@@ -7,7 +7,11 @@ import torch
7
8 from vllm.config import CompilationMode, get_current_vllm_config
9 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
10 -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
10 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 + GroupShape,
12 + QuantKey,
13 + kFp8StaticTensorSym,
14 +)
15 from vllm.platforms import current_platform
16
17 from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
@@ -88,6 +92,11 @@ class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
92 return False, "requires per tensor activation and weight scales."
93 return True, None
94
95 + def input_quant_key(self) -> QuantKey | None:
96 + if self.config.activation_quant_key == kFp8StaticTensorSym:
97 + return kFp8StaticTensorSym
98 + return None
99 +
100 def apply_scaled_mm(
101 self,
102 *,
vllm/model_executor/layers/fusion/fused_act_quant.py new
+161
@@ -0,0 +1,161 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Producer side of the QuantizedActivation contract for activation layers.
5 +
6 +Given an activation module and the downstream linear it feeds, fuse the
7 +activation with that linear's input quantization into a single kernel when the
8 +linear advertises a consumable input_quant_key (see quant_activation.py).
9 +Falls back to the plain activation when nothing matches, so a model forward can
10 +always call maybe_fused_act_quant unconditionally.
11 +
12 +This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion
13 +fires here the silu_and_mul pattern is already consumed, so the compiler pass
14 +finds nothing to rewrite and the two never double-fuse.
15 +"""
16 +
17 +from collections.abc import Callable
18 +
19 +import torch
20 +
21 +from vllm.model_executor.layers.activation import SiluAndMul
22 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
23 +from vllm.model_executor.layers.linear import LinearBase
24 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
25 + QuantKey,
26 + kFp8Dynamic128Sym,
27 + kFp8StaticTensorSym,
28 + kNvfp4Dynamic,
29 +)
30 +from vllm.platforms import current_platform
31 +
32 +FP8_DTYPE = current_platform.fp8_dtype()
33 +FP4_DTYPE = torch.uint8
34 +
35 +
36 +def _silu_and_mul_fp8_static(
37 + x: torch.Tensor, linear: LinearBase
38 +) -> QuantizedActivation:
39 + """SiluAndMul + FP8 static per-tensor quantization."""
40 + d = x.shape[-1] // 2
41 + out_shape = x.shape[:-1] + (d,)
42 + result = torch.empty(out_shape, dtype=FP8_DTYPE, device=x.device)
43 + # TODO(mgoin): read the consumer scale via the contract instead of reaching
44 + # into the kernel-specific input_scale attribute.
45 + scale = linear.input_scale
46 + torch.ops._C.silu_and_mul_quant(result, x, scale)
47 + return QuantizedActivation(
48 + data=result,
49 + scale=scale,
50 + orig_dtype=x.dtype,
51 + orig_shape=out_shape,
52 + quant_key=kFp8StaticTensorSym,
53 + )
54 +
55 +
56 +def _silu_and_mul_fp8_dynamic_block(
57 + x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey
58 +) -> QuantizedActivation:
59 + """SiluAndMul + FP8 dynamic per-block quantization."""
60 + assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}"
61 +
62 + d = x.shape[-1] // 2
63 + out_shape = x.shape[:-1] + (d,)
64 + num_tokens = x.shape[0]
65 + num_groups = d // group_size
66 +
67 + result = torch.empty((num_tokens, d), dtype=FP8_DTYPE, device=x.device)
68 + scales = torch.empty((num_tokens, num_groups), dtype=torch.float32, device=x.device)
69 +
70 + torch.ops._C.silu_and_mul_per_block_quant(
71 + out=result,
72 + input=x,
73 + scales=scales,
74 + group_size=group_size,
75 + scale_ub=None,
76 + is_scale_transposed=False,
77 + )
78 +
79 + return QuantizedActivation(
80 + data=result.view(out_shape),
81 + scale=scales.view(out_shape[:-1] + (num_groups,)),
82 + orig_dtype=x.dtype,
83 + orig_shape=out_shape,
84 + quant_key=quant_key,
85 + )
86 +
87 +
88 +def _silu_and_mul_fp8_dynamic_128(
89 + x: torch.Tensor, linear: LinearBase
90 +) -> QuantizedActivation:
91 + """SiluAndMul + FP8 dynamic per-group (group=128) quantization."""
92 + return _silu_and_mul_fp8_dynamic_block(x, linear, 128, kFp8Dynamic128Sym)
93 +
94 +
95 +def _silu_and_mul_nvfp4_dynamic(
96 + x: torch.Tensor, linear: LinearBase
97 +) -> QuantizedActivation:
98 + """SiluAndMul + NVFP4 dynamic quantization."""
99 + assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}"
100 +
101 + d = x.shape[-1] // 2
102 + out_shape = x.shape[:-1] + (d,)
103 + num_tokens = x.shape[0]
104 +
105 + # NVFP4 packs 2 values into 1 byte
106 + result = torch.empty((num_tokens, d // 2), dtype=FP4_DTYPE, device=x.device)
107 +
108 + # Block scale output shape: swizzled layout for tensor cores
109 + # Each group of 16 elements shares one FP8 scale
110 + num_k_tiles = (d + 63) // 64
111 + block_scale = torch.empty(
112 + (num_tokens, num_k_tiles * 4), dtype=FP8_DTYPE, device=x.device
113 + )
114 +
115 + input_global_scale = getattr(linear, "input_global_scale", None)
116 + assert input_global_scale is not None, (
117 + "input_global_scale is required for NVFP4 quantization"
118 + )
119 +
120 + torch.ops._C.silu_and_mul_nvfp4_quant(result, block_scale, x, input_global_scale)
121 +
122 + return QuantizedActivation(
123 + data=result.view(out_shape[:-1] + (d // 2,)),
124 + scale=block_scale,
125 + orig_dtype=x.dtype,
126 + orig_shape=out_shape,
127 + quant_key=kNvfp4Dynamic,
128 + )
129 +
130 +
131 +# (activation module type, consumer input_quant_key) -> fused producer.
132 +# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme.
133 +_FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = {
134 + (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static,
135 +}
136 +
137 +# Add CUDA-specific entries for dynamic block quantization
138 +if current_platform.is_cuda_alike():
139 + _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128
140 +
141 +# Add NVFP4 if supported (requires SM100+)
142 +if current_platform.is_cuda() and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant"):
143 + _FUSED_ACT_QUANT[(SiluAndMul, kNvfp4Dynamic)] = _silu_and_mul_nvfp4_dynamic
144 +
145 +
146 +def maybe_fused_act_quant(
147 + act_fn: torch.nn.Module,
148 + x: torch.Tensor,
149 + linear: LinearBase,
150 +) -> "torch.Tensor | QuantizedActivation":
151 + """Apply act_fn, fusing the downstream linear's input quant when possible.
152 +
153 + Returns a QuantizedActivation when a fused kernel matches
154 + (act_fn, linear.input_quant_key), else the plain activated tensor.
155 + """
156 + key = getattr(linear, "input_quant_key", None)
157 + if key is not None:
158 + producer = _FUSED_ACT_QUANT.get((type(act_fn), key))
159 + if producer is not None:
160 + return producer(x, linear)
161 + return act_fn(x)
vllm/model_executor/models/llama.py
+2 -1
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.attention import (
39 Attention,
40 EncoderOnlyAttention,
41 )
42 +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
43 from vllm.model_executor.layers.layernorm import RMSNorm
44 from vllm.model_executor.layers.linear import (
45 MergedColumnParallelLinear,
@@ -114,7 +115,7 @@ class LlamaMLP(nn.Module):
115
116 def forward(self, x):
117 x, _ = self.gate_up_proj(x)
117 - x = self.act_fn(x)
118 + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
119 x, _ = self.down_proj(x)
120 return x
121