@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 229 lines 7.96 KB
Raw
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()