@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 03dc26e639

[Perf][Kernel][Quantization] Fuse ReLU2 with static FP8 activation quantization (#53793)

Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: snordmann <snordmann@nvidia.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Samuel Nordmann committed Sep 14, 2026 at 19:26 UTC 03dc26e639e1927cda51d7b62639a11a354e1c0d
12 files changed +459 -32
tests/compile/passes/test_silu_mul_quant_manual_fusion.py
+6 -5
@@ -34,6 +34,7 @@ from vllm.model_executor.layers.fusion.fused_act_quant import (
34 from vllm.model_executor.layers.fusion.quant_activation import (
35 QuantizedActivation,
36 expose_input_quant_key,
37 + get_input_quant_key,
38 )
39 from vllm.model_executor.layers.quantization.utils.quant_utils import (
40 kFp8Dynamic128Sym,
@@ -45,7 +46,7 @@ from vllm.platforms import current_platform
46
47 # Mock linear layer for testing fusion paths that don't have real kernel support
48 class MockLinearForFusion(torch.nn.Module):
48 - """Mock linear layer that exposes input_quant_key for fusion testing."""
49 + """Mock linear layer that exposes an input key for fusion testing."""
50
51 def __init__(
52 self,
@@ -55,7 +56,7 @@ class MockLinearForFusion(torch.nn.Module):
56 input_global_scale_inv=None,
57 ):
58 super().__init__()
58 - self.input_quant_key = quant_key
59 + self._input_quant_key = quant_key
60 if input_scale is not None:
61 self.input_scale = input_scale
62 if input_global_scale is not None:
@@ -118,7 +119,7 @@ def test_manual_fusion_fp8_static_with_linear(
119 # Enable fusion
120 expose_input_quant_key(fp8_linear, fp8_linear.kernel)
121
121 - if not hasattr(fp8_linear, "input_quant_key"):
122 + if get_input_quant_key(fp8_linear) is None:
123 pytest.skip(
124 f"Kernel {force_kernel.__name__} doesn't support input_quant_key"
125 )
@@ -315,7 +316,7 @@ def test_manual_fusion_nvfp4_dynamic(
316 envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
317 )
318 def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
318 - """Test that maybe_fused_act_quant falls back when no input_quant_key."""
319 + """Test that maybe_fused_act_quant falls back without an input key."""
320 torch.set_default_device("cuda")
321 torch.set_default_dtype(dtype)
322
@@ -327,7 +328,7 @@ def test_manual_fusion_fallback_no_key(dtype: torch.dtype):
328
329 with set_current_vllm_config(config):
330 silu_and_mul = SiluAndMul()
330 - # Linear without input_quant_key attribute
331 + # Linear without an input quantization key
332 mock_linear = torch.nn.Linear(128, 128)
333
334 result = maybe_fused_act_quant(silu_and_mul, x, mock_linear)
tests/fusion/test_quant_activation_contract.py
+3 -2
@@ -42,6 +42,7 @@ from vllm.model_executor.layers.fusion.quant_activation import (
42 QuantizedActivation,
43 as_quantized_activation,
44 expose_input_quant_key,
45 + get_input_quant_key,
46 )
47 from vllm.model_executor.layers.quantization.utils.quant_utils import (
48 kFp8StaticTensorSym,
@@ -118,13 +119,13 @@ def test_bridge_marks_supporting_and_skips_others():
119 supported = _probe(FlashInferCutlassNvFp4LinearKernel)
120 layer = torch.nn.Module()
121 expose_input_quant_key(layer, supported)
121 - assert layer.input_quant_key == kNvfp4Dynamic
122 + assert get_input_quant_key(layer) == kNvfp4Dynamic
123
124 unsupported = _probe(FlashInferTrtllmNvFp4LinearKernel)
125 assert unsupported.input_quant_key() is None
126 layer = torch.nn.Module()
127 expose_input_quant_key(layer, unsupported)
127 - assert not hasattr(layer, "input_quant_key")
128 + assert get_input_quant_key(layer) is None
129
130
131 def test_as_quantized_activation_validates_key():
tests/kernels/test_fused_quant_activation.py
+26 -4
@@ -91,7 +91,7 @@ class MockLinearFp8Static(torch.nn.Module):
91
92 def __init__(self, input_scale: torch.Tensor):
93 super().__init__()
94 - self.input_quant_key = kFp8StaticTensorSym
94 + self._input_quant_key = kFp8StaticTensorSym
95 self.input_scale = input_scale
96
97
@@ -100,15 +100,21 @@ class MockLinearFp8Dynamic128(torch.nn.Module):
100
101 def __init__(self):
102 super().__init__()
103 - self.input_quant_key = kFp8Dynamic128Sym
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)."""
107 + """Mock linear layer with no input quantization key."""
108
109 pass
110
111
112 +class MockLinearRequiresUnquantized(MockLinearFp8Static):
113 + """Mock consumer that needs the original activation for another branch."""
114 +
115 + requires_unquantized_input = True
116 +
117 +
118 @pytest.mark.parametrize("num_tokens", [1, 16, 128])
119 @pytest.mark.parametrize("hidden_size", [128, 512, 1024])
120 @pytest.mark.parametrize("dtype", DTYPES)
@@ -199,7 +205,7 @@ def test_maybe_fused_act_quant_fallback(
205 hidden_size: int,
206 dtype: torch.dtype,
207 ) -> None:
202 - """Test maybe_fused_act_quant falls back when no input_quant_key."""
208 + """Test maybe_fused_act_quant falls back without an input quantization key."""
209 device = "cuda:0"
210 torch.set_default_device(device)
211
@@ -216,3 +222,19 @@ def test_maybe_fused_act_quant_fallback(
222
223 ref_out = act_fn(x)
224 torch.testing.assert_close(result, ref_out)
225 +
226 +
227 +@torch.inference_mode()
228 +def test_maybe_fused_act_quant_preserves_required_unquantized_input(
229 + default_vllm_config,
230 +) -> None:
231 + device = "cuda:0"
232 + act_fn = SiluAndMul()
233 + scale = torch.tensor([0.5], device=device, dtype=torch.float32)
234 + linear = MockLinearRequiresUnquantized(scale)
235 + x = torch.randn(17, 256, dtype=torch.bfloat16, device=device)
236 +
237 + result = maybe_fused_act_quant(act_fn, x, linear)
238 +
239 + assert isinstance(result, torch.Tensor)
240 + torch.testing.assert_close(result, act_fn(x))
tests/kernels/test_relu2_fp8_quant.py new
+229
@@ -0,0 +1,229 @@
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 +from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
8 +from vllm.config.vllm import OptimizationLevel
9 +from vllm.model_executor.layers.activation import ReLUSquaredActivation
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.input_quant_fp8 import QuantFP8
13 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 + GroupShape,
15 + kFp8StaticTensorSym,
16 +)
17 +from vllm.platforms import current_platform
18 +
19 +requires_sm90 = pytest.mark.skipif(
20 + not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
21 + reason="This test requires SM90 or newer",
22 +)
23 +
24 +
25 +class _StaticFp8Linear(torch.nn.Module):
26 + def __init__(self, input_scale: torch.Tensor) -> None:
27 + super().__init__()
28 + self._input_quant_key = kFp8StaticTensorSym
29 + self.input_scale = input_scale
30 +
31 +
32 +def _assert_fp8_bitwise_equal(actual: torch.Tensor, expected: torch.Tensor) -> None:
33 + assert actual.dtype == expected.dtype == current_platform.fp8_dtype()
34 + torch.testing.assert_close(
35 + actual.view(torch.uint8), expected.view(torch.uint8), rtol=0, atol=0
36 + )
37 +
38 +
39 +def _assert_quantized_activation(
40 + result: QuantizedActivation, x: torch.Tensor, scale: torch.Tensor
41 +) -> None:
42 + assert result.scale is scale
43 + assert result.orig_dtype == x.dtype
44 + assert result.orig_shape == x.shape
45 + assert result.quant_key == kFp8StaticTensorSym
46 +
47 +
48 +@pytest.fixture(scope="module")
49 +def o2_relu2_fp8_ops():
50 + config = VllmConfig(
51 + optimization_level=OptimizationLevel.O2,
52 + compilation_config=CompilationConfig(custom_ops=["none"]),
53 + )
54 + with set_current_vllm_config(config):
55 + yield (
56 + ReLUSquaredActivation(),
57 + QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR),
58 + )
59 +
60 +
61 +@requires_sm90
62 +@pytest.mark.parametrize("shape", [(1, 1), (1, 16), (17, 5120), (128, 10240)])
63 +@torch.inference_mode()
64 +def test_relu2_static_fp8_quant(o2_relu2_fp8_ops, shape) -> None:
65 + torch.manual_seed(0)
66 + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
67 + scale = torch.tensor(0.1495535671710968, device="cuda", dtype=torch.float32)
68 +
69 + relu2, quant_fp8 = o2_relu2_fp8_ops
70 + expected, _ = quant_fp8(relu2(x), scale)
71 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
72 +
73 + assert isinstance(result, QuantizedActivation)
74 + _assert_quantized_activation(result, x, scale)
75 + _assert_fp8_bitwise_equal(result.data, expected)
76 +
77 +
78 +@requires_sm90
79 +@pytest.mark.parametrize(
80 + "enabled_ops",
81 + [(), ("relu2",), ("quant_fp8",), ("relu2", "quant_fp8")],
82 +)
83 +@torch.inference_mode()
84 +def test_relu2_static_fp8_quant_ignores_custom_op_dispatch(enabled_ops) -> None:
85 + custom_ops = ["none", *(f"+{op}" for op in enabled_ops)]
86 + config = VllmConfig(
87 + optimization_level=OptimizationLevel.O2,
88 + compilation_config=CompilationConfig(custom_ops=custom_ops),
89 + )
90 + with set_current_vllm_config(config):
91 + torch.manual_seed(0)
92 + relu2 = ReLUSquaredActivation()
93 + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
94 + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
95 + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
96 +
97 + expected, _ = quant_fp8(relu2(x), scale)
98 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
99 +
100 + assert isinstance(result, QuantizedActivation)
101 + _assert_quantized_activation(result, x, scale)
102 + _assert_fp8_bitwise_equal(result.data, expected)
103 +
104 +
105 +@requires_sm90
106 +@torch.inference_mode()
107 +def test_relu2_static_fp8_quant_without_compilation() -> None:
108 + config = VllmConfig(optimization_level=OptimizationLevel.O0)
109 + with set_current_vllm_config(config):
110 + torch.manual_seed(0)
111 + relu2 = ReLUSquaredActivation()
112 + quant_fp8 = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
113 + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
114 + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
115 +
116 + expected, _ = quant_fp8(relu2(x), scale)
117 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
118 +
119 + assert isinstance(result, QuantizedActivation)
120 + _assert_quantized_activation(result, x, scale)
121 + _assert_fp8_bitwise_equal(result.data, expected)
122 +
123 +
124 +@requires_sm90
125 +@pytest.mark.parametrize(
126 + "scale_value",
127 + [
128 + 0.00435965,
129 + 0.007149832788854837,
130 + 0.03125,
131 + 0.0714285746216774,
132 + 0.1495535671710968,
133 + 0.1517857164144516,
134 + 6.03571415,
135 + ],
136 +)
137 +@torch.inference_mode()
138 +def test_relu2_static_fp8_quant_all_bf16_values(o2_relu2_fp8_ops, scale_value) -> None:
139 + bits = torch.arange(1 << 16, dtype=torch.int32).to(torch.uint16)
140 + x = bits.view(torch.bfloat16).reshape(1, -1).to("cuda")
141 + scale = torch.tensor(scale_value, device="cuda", dtype=torch.float32)
142 +
143 + relu2, quant_fp8 = o2_relu2_fp8_ops
144 + expected, _ = quant_fp8(relu2(x), scale)
145 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
146 +
147 + assert isinstance(result, QuantizedActivation)
148 + _assert_quantized_activation(result, x, scale)
149 + _assert_fp8_bitwise_equal(result.data, expected)
150 +
151 +
152 +@requires_sm90
153 +@torch.inference_mode()
154 +def test_relu2_static_fp8_quant_empty(o2_relu2_fp8_ops) -> None:
155 + x = torch.empty(0, device="cuda", dtype=torch.bfloat16)
156 + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
157 +
158 + relu2, _ = o2_relu2_fp8_ops
159 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
160 +
161 + assert isinstance(result, QuantizedActivation)
162 + _assert_quantized_activation(result, x, scale)
163 + assert result.data.shape == x.shape
164 + assert result.data.dtype == current_platform.fp8_dtype()
165 +
166 +
167 +@requires_sm90
168 +@pytest.mark.parametrize("unsupported", ["dtype", "layout", "scale"])
169 +@torch.inference_mode()
170 +def test_relu2_static_fp8_quant_falls_back(o2_relu2_fp8_ops, unsupported: str) -> None:
171 + x = torch.randn((17, 32), device="cuda", dtype=torch.bfloat16)
172 + scale = torch.tensor(0.03125, device="cuda", dtype=torch.float32)
173 + if unsupported == "dtype":
174 + x = x.to(torch.float16)
175 + elif unsupported == "layout":
176 + x = x.T
177 + else:
178 + scale = scale.to(torch.float16)
179 +
180 + relu2, _ = o2_relu2_fp8_ops
181 + expected = relu2(x)
182 + result = maybe_fused_act_quant(relu2, x, _StaticFp8Linear(scale))
183 +
184 + assert isinstance(result, torch.Tensor)
185 + assert not isinstance(result, QuantizedActivation)
186 + torch.testing.assert_close(result, expected, rtol=0, atol=0)
187 +
188 +
189 +@requires_sm90
190 +@torch.inference_mode()
191 +def test_relu2_static_fp8_quant_torch_compile_dynamic(
192 + o2_relu2_fp8_ops,
193 +) -> None:
194 + torch.manual_seed(1)
195 + scale = torch.tensor(0.00435965, device="cuda", dtype=torch.float32)
196 + relu2, quant_fp8 = o2_relu2_fp8_ops
197 + linear = _StaticFp8Linear(scale)
198 +
199 + def fused(x: torch.Tensor) -> torch.Tensor:
200 + result = maybe_fused_act_quant(relu2, x, linear)
201 + assert isinstance(result, QuantizedActivation)
202 + return result.data
203 +
204 + compile_count = 0
205 +
206 + def counting_inductor_backend(gm, example_inputs):
207 + nonlocal compile_count
208 + compile_count += 1
209 + return torch._inductor.compile(gm, example_inputs)
210 +
211 + torch._dynamo.reset()
212 + try:
213 + x = torch.randn((17, 5120), device="cuda", dtype=torch.bfloat16)
214 + torch._dynamo.mark_dynamic(x, 0)
215 + compiled_fused = torch.compile(
216 + fused, backend=counting_inductor_backend, fullgraph=True
217 + )
218 +
219 + for candidate in (
220 + x,
221 + torch.randn((31, 5120), device="cuda", dtype=torch.bfloat16),
222 + torch.randn((63, 5120), device="cuda", dtype=torch.bfloat16),
223 + ):
224 + expected, _ = quant_fp8(relu2(candidate), scale)
225 + _assert_fp8_bitwise_equal(compiled_fused(candidate), expected)
226 +
227 + assert compile_count == 1
228 + finally:
229 + torch._dynamo.reset()
tests/lora/test_layers.py
+10
@@ -31,6 +31,9 @@ from vllm.lora.layers import (
31 )
32 from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
33 from vllm.lora.punica_wrapper import get_punica_wrapper
34 +from vllm.model_executor.layers.fusion.quant_activation import (
35 + get_input_quant_key,
36 +)
37 from vllm.model_executor.layers.linear import (
38 ColumnParallelLinear,
39 MergedColumnParallelLinear,
@@ -56,6 +59,13 @@ TOLERANCES = {
59 torch.bfloat16: (3e-2, 2e-2),
60 }
61
62 +
63 +def test_lora_linear_requires_unquantized_input() -> None:
64 + layer = RowParallelLinearWithLoRA.__new__(RowParallelLinearWithLoRA)
65 + layer._input_quant_key = object()
66 + assert get_input_quant_key(layer) is None
67 +
68 +
69 pytestmark = [
70 pytest.mark.skipif(
71 not (
tests/model_executor/test_nemotron_h_quantization.py
+27
@@ -3,6 +3,8 @@
3
4 from unittest.mock import Mock, patch
5
6 +import torch
7 +
8
9 def test_nemotron_h_lm_head_receives_quant_config():
10 from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
@@ -32,3 +34,28 @@ def test_nemotron_h_lm_head_receives_quant_config():
34 MockLMHead.assert_called_once()
35 call_kwargs = MockLMHead.call_args.kwargs
36 assert call_kwargs["quant_config"] is mock_quant_config
37 +
38 +
39 +def test_relu2_fp8_fusion_uses_registry():
40 + from vllm.model_executor.models.nemotron_h import NemotronHMLP
41 +
42 + projected = torch.empty((1, 1), dtype=torch.bfloat16)
43 + fused = Mock()
44 + act_fn = Mock()
45 + down_proj = Mock(side_effect=lambda x: (x, None))
46 +
47 + mlp = NemotronHMLP.__new__(NemotronHMLP)
48 + torch.nn.Module.__init__(mlp)
49 + mlp.up_proj = Mock(return_value=(projected, None))
50 + mlp.down_proj = down_proj
51 + mlp.act_fn = act_fn
52 +
53 + with patch(
54 + "vllm.model_executor.models.nemotron_h.maybe_fused_act_quant",
55 + return_value=fused,
56 + ) as maybe_fused:
57 + result = mlp(Mock())
58 +
59 + maybe_fused.assert_called_once_with(act_fn, projected, down_proj)
60 + act_fn.assert_not_called()
61 + assert result is fused
vllm/lora/layers/base_linear.py
+4
@@ -68,6 +68,10 @@ if envs.VLLM_LORA_ENABLE_DUAL_STREAM:
68
69
70 class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
71 + # The adapter branch consumes the original activation, independently of
72 + # whether the wrapped base layer can consume a pre-quantized activation.
73 + requires_unquantized_input = True
74 +
75 def __init__(self, base_layer: LinearBase):
76 super().__init__()
77
vllm/model_executor/layers/fusion/fused_act_quant.py
+54 -12
@@ -5,21 +5,27 @@ 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).
8 +linear advertises a consumable input quantization 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.
13 +fires here, the activation and quantization are already consumed, so a compiler
14 +pass cannot fuse the same boundary again.
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
21 +from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul
22 +from vllm.model_executor.layers.fusion.quant_activation import (
23 + QuantizedActivation,
24 + get_input_quant_key,
25 +)
26 +from vllm.model_executor.layers.fusion.relu2_fp8_quant import (
27 + relu_squared_static_fp8_quant,
28 +)
29 from vllm.model_executor.layers.linear import LinearBase
30 from vllm.model_executor.layers.quantization.utils.quant_utils import (
31 QuantKey,
@@ -139,12 +145,46 @@ def _silu_and_mul_nvfp4_dynamic(
145 )
146
147
142 -# (activation module type, consumer input_quant_key) -> fused producer.
143 -# Mirrors ActivationQuantFusionPass.FUSED_OPS; add a row to migrate a scheme.
148 +# (activation module type, consumer input quantization key) -> fused producer.
149 +# Add a row for each supported manual producer.
150 _FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = {
151 (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static,
152 }
153
154 +
155 +def _relu_squared_static_fp8_quant_supported(
156 + act_fn: torch.nn.Module,
157 + x: torch.Tensor,
158 + linear: LinearBase,
159 +) -> bool:
160 + """Return whether the ReLU2 static-FP8 producer can consume this input."""
161 + scale = getattr(linear, "input_scale", None)
162 + return (
163 + isinstance(act_fn, ReLUSquaredActivation)
164 + and x.is_cuda
165 + and x.dtype == torch.bfloat16
166 + and x.is_contiguous()
167 + and isinstance(scale, torch.Tensor)
168 + and scale.dtype == torch.float32
169 + and scale.device == x.device
170 + and scale.numel() == 1
171 + )
172 +
173 +
174 +# Optional per-entry predicates preserve the unconditional fallback contract
175 +# for producers whose supported inputs are narrower than their registry key.
176 +_FUSED_ACT_QUANT_SUPPORT: dict[
177 + tuple[type, QuantKey], Callable[[torch.nn.Module, torch.Tensor, LinearBase], bool]
178 +] = {}
179 +
180 +# This path is validated on SM90 and newer CUDA devices.
181 +if current_platform.is_cuda() and current_platform.has_device_capability(90):
182 + _RELU2_STATIC_FP8_KEY = (ReLUSquaredActivation, kFp8StaticTensorSym)
183 + _FUSED_ACT_QUANT[_RELU2_STATIC_FP8_KEY] = relu_squared_static_fp8_quant
184 + _FUSED_ACT_QUANT_SUPPORT[_RELU2_STATIC_FP8_KEY] = (
185 + _relu_squared_static_fp8_quant_supported
186 + )
187 +
188 # Add CUDA-specific entries for dynamic block quantization
189 if current_platform.is_cuda_alike():
190 _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128
@@ -161,12 +201,14 @@ def maybe_fused_act_quant(
201 ) -> "torch.Tensor | QuantizedActivation":
202 """Apply act_fn, fusing the downstream linear's input quant when possible.
203
164 - Returns a QuantizedActivation when a fused kernel matches
165 - (act_fn, linear.input_quant_key), else the plain activated tensor.
204 + Returns a QuantizedActivation when a fused kernel matches the activation and
205 + the consumer's effective input quantization key, else the plain activation.
206 """
167 - key = getattr(linear, "input_quant_key", None)
207 + key = get_input_quant_key(linear)
208 if key is not None:
169 - producer = _FUSED_ACT_QUANT.get((type(act_fn), key))
170 - if producer is not None:
209 + registry_key = (type(act_fn), key)
210 + producer = _FUSED_ACT_QUANT.get(registry_key)
211 + support = _FUSED_ACT_QUANT_SUPPORT.get(registry_key)
212 + if producer is not None and (support is None or support(act_fn, x, linear)):
213 return producer(x, linear)
214 return act_fn(x)
vllm/model_executor/layers/fusion/quant_activation.py
+15 -7
@@ -5,7 +5,9 @@ A QuantizedActivation is a pre-quantized activation produced by a fused kernel
5 and consumed directly by a linear layer, letting the layer skip its own input
6 quantization. A linear advertises the key its kernel can consume via
7 expose_input_quant_key; the kernel validates and reads the activation via
8 -as_quantized_activation.
8 +as_quantized_activation. Producers query the effective capability through
9 +get_input_quant_key, which hides the key when another consumer branch needs the
10 +original activation.
11 """
12
13 from dataclasses import dataclass, replace
@@ -46,12 +48,11 @@ class QuantizedActivation:
48
49
50 def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
49 - """Advertise the kernel's pre-quantized input key on the layer, if any.
51 + """Store the kernel's pre-quantized input key on the layer, if any.
52
51 - This is the bridge from a kernel's input_quant_key() to the
52 - layer.input_quant_key attribute that fusion call sites read. The attribute
53 - is left unset when the kernel quantizes its own input, so non-supporting
54 - backends never receive a QuantizedActivation.
53 + This is the bridge from a kernel's input_quant_key() to the layer capability
54 + that fusion call sites read through get_input_quant_key. The raw key is left
55 + unset when the kernel quantizes its own input.
56
57 TODO(mgoin): Producers also need the consumer's quantization scales (e.g.
58 static input scale, global scale). Expose those here as well so producers
@@ -59,7 +60,14 @@ def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None:
60 """
61 key = kernel.input_quant_key()
62 if key is not None:
62 - layer.input_quant_key = key
63 + layer._input_quant_key = key
64 +
65 +
66 +def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None:
67 + """Return the pre-quantized input key the complete layer can consume."""
68 + if getattr(layer, "requires_unquantized_input", False):
69 + return None
70 + return getattr(layer, "_input_quant_key", None)
71
72
73 def as_quantized_activation(
vllm/model_executor/layers/fusion/relu2_fp8_quant.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
7 +from vllm.model_executor.layers.linear import LinearBase
8 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
9 + get_fp8_min_max,
10 + kFp8StaticTensorSym,
11 +)
12 +from vllm.platforms import current_platform
13 +from vllm.triton_utils import tl, triton
14 +
15 +_FP8_MIN, _FP8_MAX = get_fp8_min_max()
16 +
17 +
18 +@triton.jit
19 +def _relu_squared_static_fp8_quant_kernel(
20 + x_ptr,
21 + scale_ptr,
22 + output_ptr,
23 + num_elements,
24 + FP8_MIN: tl.constexpr,
25 + FP8_MAX: tl.constexpr,
26 + BLOCK_SIZE: tl.constexpr,
27 +):
28 + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
29 + mask = offsets < num_elements
30 + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
31 + input_is_nan = x != x
32 + relu = tl.maximum(x, 0.0)
33 + # Preserve the standalone ReLU2 kernel's BF16 output boundary.
34 + activated = (relu * relu).to(tl.bfloat16).to(tl.float32)
35 + # Match O2's Inductor-compiled QuantFP8 reciprocal-then-multiply.
36 + quantized = activated * (1.0 / tl.load(scale_ptr))
37 + quantized = tl.maximum(tl.minimum(quantized, FP8_MAX), FP8_MIN)
38 + # Match the O2 native chain by preserving NaN through the FP8 conversion.
39 + quantized = tl.where(input_is_nan, x, quantized)
40 + tl.store(output_ptr + offsets, quantized, mask=mask)
41 +
42 +
43 +def relu_squared_static_fp8_quant(
44 + x: torch.Tensor, linear: LinearBase
45 +) -> QuantizedActivation:
46 + """BF16 ReLU2 followed by static per-tensor FP8 quantization."""
47 + assert x.dtype == torch.bfloat16
48 + assert x.is_contiguous()
49 + scale = linear.input_scale
50 + assert scale.dtype == torch.float32
51 + assert x.device == scale.device
52 + assert scale.numel() == 1
53 +
54 + output = torch.empty_like(x, dtype=current_platform.fp8_dtype())
55 + if x.numel() != 0:
56 + block_size = min(triton.next_power_of_2(x.shape[-1]), 2048)
57 + num_warps = min(max(block_size // 256, 1), 4)
58 + grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
59 + _relu_squared_static_fp8_quant_kernel[grid](
60 + x,
61 + scale,
62 + output,
63 + x.numel(),
64 + FP8_MIN=_FP8_MIN,
65 + FP8_MAX=_FP8_MAX,
66 + BLOCK_SIZE=block_size,
67 + num_warps=num_warps,
68 + )
69 +
70 + return QuantizedActivation(
71 + data=output,
72 + scale=scale,
73 + orig_dtype=x.dtype,
74 + orig_shape=x.shape,
75 + quant_key=kFp8StaticTensorSym,
76 + )
vllm/model_executor/layers/linear.py
+2
@@ -252,6 +252,8 @@ class LinearBase(PluggableLayer):
252 disable_tp: If true, tensor parallelism will be disabled for this layer.
253 """
254
255 + requires_unquantized_input = False
256 +
257 def __init__(
258 self,
259 input_size: int,
vllm/model_executor/models/nemotron_h.py
+7 -2
@@ -25,7 +25,11 @@ import torch
25 from torch import nn
26
27 from vllm.compilation.decorators import support_torch_compile
28 -from vllm.config import CacheConfig, ModelConfig, VllmConfig
28 +from vllm.config import (
29 + CacheConfig,
30 + ModelConfig,
31 + VllmConfig,
32 +)
33 from vllm.config.parallel import ParallelConfig
34 from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
35 from vllm.distributed.communication_op import tensor_model_parallel_all_gather
@@ -38,6 +42,7 @@ from vllm.model_executor.layers.fused_moe import (
42 activation_without_mul,
43 fused_moe_make_expert_params_mapping,
44 )
45 +from vllm.model_executor.layers.fusion.fused_act_quant import maybe_fused_act_quant
46 from vllm.model_executor.layers.layernorm import RMSNorm
47 from vllm.model_executor.layers.linear import (
48 ColumnParallelLinear,
@@ -118,7 +123,7 @@ class NemotronHMLP(nn.Module):
123
124 def forward(self, x: torch.Tensor):
125 x, _ = self.up_proj(x)
121 - x = self.act_fn(x)
126 + x = maybe_fused_act_quant(self.act_fn, x, self.down_proj)
127 x, _ = self.down_proj(x)
128 return x
129