@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 215 lines 6.01 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 tests.kernels.utils import to_int8
8 from vllm import _custom_ops as ops
9 from vllm.platforms import current_platform
10
11 if not current_platform.is_cpu():
12 pytest.skip("skipping CPU-only tests", allow_module_level=True)
13
14 NK_FACTORS = [
15 (256, 128),
16 (4096, 4096),
17 (16384, 4096),
18 (64128, 4096),
19 (1023, 491),
20 (1001, 15),
21 ]
22 M_FACTORS = [
23 (16, 1, 32, 128, 64),
24 (1, 17, 1, 31, 17),
25 ]
26 CACHE_SIZES = [2]
27 DTYPE = [torch.bfloat16]
28
29
30 def rand_int8(shape: tuple, device: str = "cpu"):
31 return to_int8(torch.rand(shape, device=device) * 255 - 128)
32
33
34 def ref_int8_scaled_mm(
35 a: torch.Tensor,
36 b: torch.Tensor,
37 scale_a: torch.Tensor,
38 scale_b: torch.Tensor,
39 azp: torch.Tensor | None,
40 bias: torch.Tensor | None,
41 output_type: torch.dtype,
42 ):
43 if azp is not None:
44 a = a.to(dtype=torch.float32) - azp.to(dtype=torch.float32)
45 output = torch.mm(
46 (scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
47 )
48 if bias is not None:
49 output += bias.float()
50
51 return output.to(dtype=output_type)
52
53
54 def onednn_int8_gemm_test_helper(
55 primitive_cache_size: int,
56 m: int,
57 n: int,
58 k: int,
59 per_tensor_a_quant: bool,
60 per_tensor_b_quant: bool,
61 use_azp: bool,
62 use_bias: bool,
63 out_dtype: torch.dtype = torch.bfloat16,
64 device: str = "cpu",
65 ):
66 # Test for a oneDNN kernel with per-tensor / per-token activation
67 # quantization and per-tensor / per-output channel weight quantization.
68 a = to_int8(torch.randn((m, k), device=device) * 5)
69 b = to_int8(torch.randn((n, k), device=device).t() * 5)
70
71 a_scales_shape = (1, 1) if per_tensor_a_quant else (m, 1)
72 b_scales_shape = (1, 1) if per_tensor_b_quant else (1, n)
73
74 scale_a = torch.randn(a_scales_shape, device=device, dtype=torch.float32)
75 scale_b = torch.randn(b_scales_shape, device=device, dtype=torch.float32)
76
77 if use_azp:
78 azp = torch.rand(a_scales_shape, dtype=torch.float32) * 10 + 1.5
79 azp = (azp / scale_a).round().to(dtype=torch.int32)
80 azp_adj = scale_b * b.sum(dim=0, keepdim=True, dtype=torch.float32)
81 else:
82 azp = None
83 azp_adj = None
84
85 bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None
86
87 handler = ops.create_onednn_scaled_mm(
88 b,
89 scale_b,
90 out_dtype,
91 not per_tensor_a_quant,
92 use_azp,
93 primitive_cache_size,
94 )
95
96 out = torch.zeros((m, n), dtype=out_dtype)
97 ops.onednn_scaled_mm(handler, a, out, scale_a, azp, azp_adj, bias)
98 baseline = ref_int8_scaled_mm(a, b, scale_a, scale_b, azp, bias, out_dtype)
99
100 torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
101
102 if use_bias:
103 # To test runtime bias setting
104 out = torch.zeros((m, n), dtype=out_dtype)
105 ops.onednn_scaled_mm(handler, a, out, scale_a, azp, azp_adj, None)
106 baseline = ref_int8_scaled_mm(a, b, scale_a, scale_b, azp, None, out_dtype)
107
108 torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
109
110
111 def onednn_gemm_test_helper(
112 primitive_cache_size: int,
113 m: int,
114 n: int,
115 k: int,
116 use_bias: bool,
117 use_stride: bool,
118 dtype: torch.dtype = torch.bfloat16,
119 device: str = "cpu",
120 ):
121 if use_stride:
122 a = torch.rand((m, 2 * k), dtype=dtype, device=device) * 1.5
123 a = a[:, :k]
124 else:
125 a = torch.rand((m, k), dtype=dtype, device=device) * 1.5
126
127 b = torch.rand((n, k), dtype=dtype, device=device) * 1.5
128
129 if use_bias:
130 bias = torch.rand((n,), device=device, dtype=dtype) * 5
131 bias_f32 = bias.float()
132 else:
133 bias = None
134 bias_f32 = None
135
136 handler = ops.create_onednn_mm(
137 b.t(),
138 primitive_cache_size,
139 )
140
141 out = ops.onednn_mm(handler, a, bias)
142 baseline = torch.nn.functional.linear(a.float(), b.float(), bias_f32).to(
143 dtype=a.dtype
144 )
145
146 torch.testing.assert_close(out, baseline)
147
148 if use_bias:
149 # To test runtime bias setting
150 out = ops.onednn_mm(handler, a, None)
151 baseline = torch.nn.functional.linear(a.float(), b.float(), None).to(
152 dtype=a.dtype
153 )
154
155 torch.testing.assert_close(out, baseline)
156
157
158 @pytest.mark.parametrize("n,k", NK_FACTORS)
159 @pytest.mark.parametrize("m_list", M_FACTORS)
160 @pytest.mark.parametrize("per_tensor_a_scale", [True, False])
161 @pytest.mark.parametrize("per_tensor_b_scale", [True, False])
162 @pytest.mark.parametrize("use_bias", [True, False])
163 @pytest.mark.parametrize("use_azp", [True, False])
164 @pytest.mark.parametrize("output_type", DTYPE)
165 @pytest.mark.parametrize("primitive_cache_size", CACHE_SIZES)
166 def test_onednn_int8_scaled_gemm(
167 n: int,
168 k: int,
169 m_list: tuple[int, ...],
170 per_tensor_a_scale: bool,
171 per_tensor_b_scale: bool,
172 use_bias: bool,
173 use_azp: bool,
174 output_type: torch.dtype,
175 primitive_cache_size: int,
176 ):
177 for m in m_list:
178 onednn_int8_gemm_test_helper(
179 primitive_cache_size=primitive_cache_size,
180 m=m,
181 n=n,
182 k=k,
183 per_tensor_a_quant=per_tensor_a_scale,
184 per_tensor_b_quant=per_tensor_b_scale,
185 use_bias=use_bias,
186 use_azp=use_azp,
187 out_dtype=output_type,
188 )
189
190
191 @pytest.mark.parametrize("n,k", NK_FACTORS)
192 @pytest.mark.parametrize("m_list", M_FACTORS)
193 @pytest.mark.parametrize("use_bias", [True, False])
194 @pytest.mark.parametrize("use_stride", [True, False])
195 @pytest.mark.parametrize("dtype", DTYPE)
196 @pytest.mark.parametrize("primitive_cache_size", CACHE_SIZES)
197 def test_onednn_gemm(
198 n: int,
199 k: int,
200 m_list: tuple[int, ...],
201 use_bias: bool,
202 use_stride: bool,
203 dtype: torch.dtype,
204 primitive_cache_size: int,
205 ):
206 for m in m_list:
207 onednn_gemm_test_helper(
208 primitive_cache_size=primitive_cache_size,
209 m=m,
210 n=n,
211 k=k,
212 use_bias=use_bias,
213 use_stride=use_stride,
214 dtype=dtype,
215 )