@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 154 lines 5.48 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the dynamic_per_token_scaled_fp8_quant helion kernel.
4
5 Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`.
6 """
7
8 from typing import Any
9
10 import pytest
11 import torch
12 from torch._subclasses.fake_tensor import FakeTensorMode
13
14 from tests.kernels.helion.utils import skip_if_platform_unsupported
15 from tests.kernels.quant_utils import FP8_DTYPE
16 from vllm.kernels.helion.case_key import CaseKey
17 from vllm.kernels.helion.config_manager import ConfigManager
18 from vllm.kernels.helion.ops.dynamic_per_token_scaled_fp8_quant import (
19 _pick_cache,
20 baseline,
21 dynamic_per_token_scaled_fp8_quant,
22 pick_config,
23 )
24 from vllm.platforms import current_platform
25 from vllm.utils.import_utils import has_helion
26 from vllm.utils.torch_utils import set_random_seed
27
28 if not has_helion():
29 pytest.skip(
30 "Helion is not installed. Install with: pip install vllm[helion]",
31 allow_module_level=True,
32 )
33
34
35 def _generate_fake_input(num_tokens: int, hidden_size: int) -> tuple[Any, ...]:
36 with FakeTensorMode():
37 input = torch.randn(
38 num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
39 )
40 result = torch.empty(
41 input.shape, device=input.device, dtype=current_platform.fp8_dtype()
42 )
43 scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
44 scale_ub = torch.mean(input).to(torch.float32)
45 args = (result, input, scale, scale_ub)
46 return args
47
48
49 @pytest.fixture(autouse=True)
50 def reset_config_manager_singleton():
51 ConfigManager.reset_instance()
52 ConfigManager()
53 yield
54 ConfigManager.reset_instance()
55
56
57 class TestDynamicPerTokenScaledFp8QuantConfigPicker:
58 def setup_method(self):
59 _pick_cache.clear()
60
61 def test_config_picker_exact_match(self):
62 config_keys = [
63 CaseKey({"hidden_size": 2048, "num_tokens": 16}),
64 CaseKey({"hidden_size": 4096, "num_tokens": 16}),
65 ]
66
67 args = _generate_fake_input(16, 4096)
68 selected_key = pick_config(args, config_keys)
69 assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
70
71 def test_config_picker_closest_match(self):
72 config_keys = [
73 CaseKey({"hidden_size": 2048, "num_tokens": 16}),
74 CaseKey({"hidden_size": 2048, "num_tokens": 32}),
75 CaseKey({"hidden_size": 4096, "num_tokens": 16}),
76 CaseKey({"hidden_size": 4096, "num_tokens": 32}),
77 ]
78
79 args = _generate_fake_input(20, 3000)
80 selected_key = pick_config(args, config_keys)
81 assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
82
83 def test_config_picker_no_configs(self):
84 config_keys: list[CaseKey] = []
85
86 args = _generate_fake_input(16, 4096)
87 selected_key = pick_config(args, config_keys)
88 assert selected_key is None
89
90 def test_config_picker_fallback_to_largest(self):
91 config_keys = [
92 CaseKey({"hidden_size": 2048, "num_tokens": 16}),
93 CaseKey({"hidden_size": 4096, "num_tokens": 16}),
94 ]
95
96 args = _generate_fake_input(32, 8192)
97 selected_key = pick_config(args, config_keys)
98 assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
99
100
101 class TestDynamicPerTokenScaledFp8QuantCorrectness:
102 @pytest.mark.parametrize("num_tokens", [1, 7, 4096])
103 @pytest.mark.parametrize("hidden_size", [17, 1024, 1025, 1026, 5137, 8193])
104 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float])
105 @pytest.mark.parametrize("has_scale_ub", [True, False])
106 @pytest.mark.parametrize("seed", [0])
107 def test_dynamic_per_token_fp8_quant(
108 self,
109 num_tokens: int,
110 hidden_size: int,
111 dtype: torch.dtype,
112 has_scale_ub: bool,
113 seed: int,
114 ) -> None:
115 skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
116 set_random_seed(seed)
117
118 x = (
119 torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
120 ) # avoid nans
121
122 scale_ub = (
123 torch.mean(x).to(dtype=torch.float32, device="cuda")
124 if has_scale_ub
125 else None
126 )
127
128 ref_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
129 ref_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
130 baseline(ref_out, x, ref_scales, scale_ub)
131
132 ops_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
133 ops_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
134 dynamic_per_token_scaled_fp8_quant(ops_out, x, ops_scales, scale_ub)
135
136 torch.testing.assert_close(ref_scales, ops_scales)
137 # allow 1 ULP difference
138 assert (
139 ref_out.view(torch.uint8).to(torch.int16)
140 - ops_out.view(torch.uint8).to(torch.int16)
141 ).abs().max() <= 1
142
143
144 class TestDynamicPerTokenScaledFp8QuantIntegration:
145 def test_kernel_registration_integration(self):
146 from vllm.kernels.helion.register import get_registered_kernels
147
148 registered_kernels = get_registered_kernels()
149 assert "dynamic_per_token_scaled_fp8_quant" in registered_kernels
150
151 kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
152 assert kernel_wrapper.op_name == "dynamic_per_token_scaled_fp8_quant"
153 assert kernel_wrapper._config_picker is not None
154 assert kernel_wrapper._mutates_args == ["result", "scale"]