@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 365 lines 13.2 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 import torch.nn.functional as F
7
8 from vllm.utils.import_utils import has_helion
9
10 if not has_helion():
11 pytest.skip(
12 "Helion is not installed. Install with: pip install vllm[helion]",
13 allow_module_level=True,
14 )
15
16 from vllm.kernels.helion.case_key import CaseKey
17 from vllm.kernels.helion.config_manager import ConfigManager
18 from vllm.kernels.helion.ops.silu_mul_fp8 import (
19 _pick_cache,
20 pick_silu_mul_fp8_config,
21 silu_mul_fp8,
22 silu_mul_fp8_baseline,
23 )
24
25
26 def skip_if_platform_unsupported():
27 try:
28 from vllm.kernels.helion.utils import get_canonical_gpu_name
29
30 if not torch.cuda.is_available():
31 pytest.skip("CUDA not available")
32
33 platform = get_canonical_gpu_name()
34
35 try:
36 config_manager = ConfigManager.get_instance()
37 except RuntimeError:
38 config_manager = ConfigManager()
39
40 configs = config_manager.get_platform_configs("silu_mul_fp8", platform)
41 if len(configs) == 0:
42 pytest.skip("Current GPU platform not supported for silu_mul_fp8 kernel")
43
44 except (ImportError, RuntimeError, KeyError):
45 pytest.skip("Error detecting platform support for silu_mul_fp8 kernel")
46
47
48 @pytest.fixture(autouse=True)
49 def reset_config_manager_singleton():
50 ConfigManager.reset_instance()
51 ConfigManager()
52 yield
53 ConfigManager.reset_instance()
54
55
56 class TestSiluMulFp8ConfigPicker:
57 def setup_method(self):
58 _pick_cache.clear()
59
60 def test_config_picker_exact_match(self):
61 config_keys = [
62 CaseKey({"intermediate": 2048, "numtokens": 256}),
63 CaseKey({"intermediate": 4096, "numtokens": 256}),
64 ]
65
66 input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
67 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
68 args = (input_tensor, scale)
69
70 selected_key = pick_silu_mul_fp8_config(args, config_keys)
71 assert selected_key == CaseKey({"intermediate": 2048, "numtokens": 256})
72
73 def test_config_picker_closest_match(self):
74 config_keys = [
75 CaseKey({"intermediate": 2048, "numtokens": 256}),
76 CaseKey({"intermediate": 4096, "numtokens": 256}),
77 ]
78 input_tensor = torch.randn(32, 7000, dtype=torch.bfloat16, device="cuda")
79 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
80 args = (input_tensor, scale)
81
82 selected_key = pick_silu_mul_fp8_config(args, config_keys)
83 assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 256})
84
85 def test_config_picker_no_configs(self):
86 config_keys: list[CaseKey] = []
87
88 input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
89 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
90 args = (input_tensor, scale)
91
92 selected_key = pick_silu_mul_fp8_config(args, config_keys)
93 assert selected_key is None
94
95 @pytest.mark.parametrize("intermediate_size", [2048, 4096, 5120])
96 def test_config_picker_different_sizes(self, intermediate_size):
97 config_keys = [
98 CaseKey({"intermediate": 2048, "numtokens": 256}),
99 CaseKey({"intermediate": 4096, "numtokens": 256}),
100 CaseKey({"intermediate": 5120, "numtokens": 256}),
101 ]
102
103 input_tensor = torch.randn(
104 32, 2 * intermediate_size, dtype=torch.bfloat16, device="cuda"
105 )
106 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
107 args = (input_tensor, scale)
108
109 selected_key = pick_silu_mul_fp8_config(args, config_keys)
110 assert selected_key == {
111 "intermediate": intermediate_size,
112 "numtokens": 256,
113 }
114
115 def test_config_picker_numtokens_ceiling(self):
116 config_keys = [
117 CaseKey({"intermediate": 4096, "numtokens": 8}),
118 CaseKey({"intermediate": 4096, "numtokens": 32}),
119 CaseKey({"intermediate": 4096, "numtokens": 128}),
120 CaseKey({"intermediate": 4096, "numtokens": 256}),
121 ]
122 input_tensor = torch.randn(20, 8192, dtype=torch.bfloat16, device="cuda")
123 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
124
125 selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
126 assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 32})
127
128 def test_config_picker_numtokens_exact(self):
129 config_keys = [
130 CaseKey({"intermediate": 4096, "numtokens": 8}),
131 CaseKey({"intermediate": 4096, "numtokens": 32}),
132 CaseKey({"intermediate": 4096, "numtokens": 128}),
133 ]
134 input_tensor = torch.randn(32, 8192, dtype=torch.bfloat16, device="cuda")
135 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
136
137 selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
138 assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 32})
139
140 def test_config_picker_numtokens_fallback_to_largest(self):
141 config_keys = [
142 CaseKey({"intermediate": 4096, "numtokens": 8}),
143 CaseKey({"intermediate": 4096, "numtokens": 32}),
144 CaseKey({"intermediate": 4096, "numtokens": 128}),
145 ]
146 input_tensor = torch.randn(512, 8192, dtype=torch.bfloat16, device="cuda")
147 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
148
149 selected_key = pick_silu_mul_fp8_config((input_tensor, scale), config_keys)
150 assert selected_key == CaseKey({"intermediate": 4096, "numtokens": 128})
151
152
153 class TestSiluMulFp8Correctness:
154 @pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
155 @pytest.mark.parametrize("intermediate_size", [2048, 3000, 3500, 4096, 5000])
156 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
157 def test_silu_mul_fp8_correctness(self, batch_size, intermediate_size, dtype):
158 skip_if_platform_unsupported()
159
160 input_size = 2 * intermediate_size
161 input_tensor = torch.randn(batch_size, input_size, dtype=dtype, device="cuda")
162 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
163
164 reference_output = silu_mul_fp8_baseline(input_tensor, scale)
165 helion_output = silu_mul_fp8(input_tensor, scale)
166
167 assert helion_output.shape == reference_output.shape
168 assert helion_output.dtype == torch.float8_e4m3fn
169 assert reference_output.dtype == torch.float8_e4m3fn
170
171 ref_f32 = reference_output.to(torch.float32)
172 helion_f32 = helion_output.to(torch.float32)
173 # FP8 E4M3 has limited precision. Values near quantization boundaries
174 # can round differently due to intermediate precision differences.
175 torch.testing.assert_close(
176 helion_f32,
177 ref_f32,
178 atol=0.05,
179 rtol=0.05,
180 msg=f"Mismatch at batch={batch_size}, size={intermediate_size}",
181 )
182
183 def test_silu_mul_fp8_shape_inference(self):
184 skip_if_platform_unsupported()
185 batch_size, input_size = 32, 8192
186 intermediate_size = input_size // 2
187
188 input_tensor = torch.randn(
189 batch_size, input_size, dtype=torch.bfloat16, device="cuda"
190 )
191 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
192
193 output = silu_mul_fp8(input_tensor, scale)
194
195 expected_shape = (batch_size, intermediate_size)
196 assert output.shape == expected_shape
197 assert output.dtype == torch.float8_e4m3fn
198
199 def test_silu_mul_fp8_scale_variations(self):
200 skip_if_platform_unsupported()
201 batch_size, input_size = 16, 4096
202
203 input_tensor = torch.randn(
204 batch_size, input_size, dtype=torch.bfloat16, device="cuda"
205 )
206
207 scales = [0.1, 0.5, 1.0, 2.0, 10.0]
208
209 for scale_val in scales:
210 scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda")
211
212 reference_output = silu_mul_fp8_baseline(input_tensor, scale)
213 helion_output = silu_mul_fp8(input_tensor, scale)
214 ref_f32 = reference_output.to(torch.float32)
215 helion_f32 = helion_output.to(torch.float32)
216
217 torch.testing.assert_close(
218 helion_f32,
219 ref_f32,
220 atol=0.05,
221 rtol=0.05,
222 msg=f"Mismatch for scale={scale_val}",
223 )
224
225 @pytest.mark.parametrize(
226 "shape",
227 [
228 (1, 4096),
229 (16, 4096),
230 (128, 4096),
231 (1024, 4096),
232 (1, 8192),
233 (16, 8192),
234 (128, 8192),
235 ],
236 )
237 def test_silu_mul_fp8_various_shapes(self, shape):
238 skip_if_platform_unsupported()
239
240 input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
241 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
242
243 reference_output = silu_mul_fp8_baseline(input_tensor, scale)
244 helion_output = silu_mul_fp8(input_tensor, scale)
245
246 assert helion_output.shape == reference_output.shape
247
248 ref_f32 = reference_output.to(torch.float32)
249 helion_f32 = helion_output.to(torch.float32)
250
251 torch.testing.assert_close(
252 helion_f32, ref_f32, atol=0.05, rtol=0.05, msg=f"Mismatch for shape={shape}"
253 )
254
255
256 def silu_mul_fp8_pytorch(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
257 """Pure PyTorch reference using F.silu.
258
259 This matches vLLM's SiluAndMul.forward_native exactly:
260 F.silu(x[..., :d]) * x[..., d:]
261 """
262 d = input.shape[-1] // 2
263 result = F.silu(input[..., :d]) * input[..., d:]
264 return (result.to(torch.float32) / scale).to(torch.float8_e4m3fn)
265
266
267 class TestSiluMulFp8PytorchReference:
268 """Tests comparing Helion kernel against pure PyTorch implementation.
269
270 Uses tighter tolerance since both use PyTorch's FP8 conversion
271 (same rounding mode), unlike the vLLM C++ baseline which uses
272 NVIDIA's hardware FP8 conversion with different rounding.
273 """
274
275 @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 256])
276 @pytest.mark.parametrize("intermediate_size", [1024, 2048, 4096])
277 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
278 def test_silu_mul_fp8_vs_pytorch(self, batch_size, intermediate_size, dtype):
279 skip_if_platform_unsupported()
280
281 input_tensor = torch.randn(
282 batch_size, 2 * intermediate_size, dtype=dtype, device="cuda"
283 )
284 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
285
286 pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
287 helion_output = silu_mul_fp8(input_tensor, scale)
288
289 assert helion_output.shape == pytorch_output.shape
290 assert helion_output.dtype == torch.float8_e4m3fn
291
292 pytorch_f32 = pytorch_output.to(torch.float32)
293 helion_f32 = helion_output.to(torch.float32)
294
295 # Tolerance accounts for FP8 quantization boundary effects
296 torch.testing.assert_close(
297 helion_f32,
298 pytorch_f32,
299 atol=0.05,
300 rtol=0.05,
301 msg=(
302 f"Mismatch at batch={batch_size}, size={intermediate_size}, "
303 f"dtype={dtype}"
304 ),
305 )
306
307 @pytest.mark.parametrize(
308 "shape",
309 [
310 (1, 2, 4096), # 3D input
311 (2, 4, 2048), # 3D input
312 (1, 1, 1, 8192), # 4D input
313 ],
314 )
315 def test_silu_mul_fp8_multidim_vs_pytorch(self, shape):
316 skip_if_platform_unsupported()
317
318 input_tensor = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
319 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
320
321 pytorch_output = silu_mul_fp8_pytorch(input_tensor, scale)
322 helion_output = silu_mul_fp8(input_tensor, scale)
323
324 assert helion_output.shape == pytorch_output.shape
325
326 pytorch_f32 = pytorch_output.to(torch.float32)
327 helion_f32 = helion_output.to(torch.float32)
328
329 torch.testing.assert_close(
330 helion_f32,
331 pytorch_f32,
332 atol=0.05,
333 rtol=0.05,
334 msg=f"Mismatch for shape={shape}",
335 )
336
337
338 class TestSiluMulFp8Integration:
339 def test_kernel_registration_integration(self):
340 from vllm.kernels.helion.register import get_registered_kernels
341
342 registered_kernels = get_registered_kernels()
343 assert "silu_mul_fp8" in registered_kernels
344
345 kernel_wrapper = registered_kernels["silu_mul_fp8"]
346 assert kernel_wrapper.op_name == "silu_mul_fp8"
347 assert kernel_wrapper._config_picker is not None
348
349 def test_fake_impl_functionality(self):
350 skip_if_platform_unsupported()
351 from vllm.kernels.helion.register import get_registered_kernels
352
353 input_tensor = torch.randn(32, 4096, dtype=torch.bfloat16, device="cuda")
354 scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
355 registered_kernels = get_registered_kernels()
356 kernel_wrapper = registered_kernels["silu_mul_fp8"]
357 fake_impl = kernel_wrapper._fake_impl
358 assert fake_impl is not None
359
360 fake_output = fake_impl(input_tensor, scale)
361
362 expected_shape = (32, 2048)
363 assert fake_output.shape == expected_shape
364 assert fake_output.dtype == torch.float8_e4m3fn
365 assert fake_output.device == input_tensor.device