@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 177 lines 5.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the AWQ Triton kernel.
4
5 Run `pytest tests/kernels/quantization/test_awq_triton.py`.
6 """
7
8 import pytest
9 import torch
10
11 from vllm.model_executor.layers.quantization.awq_triton import (
12 AWQ_TRITON_SUPPORTED_GROUP_SIZES,
13 awq_dequantize_triton,
14 awq_gemm_triton,
15 )
16 from vllm.platforms import current_platform
17 from vllm.utils.torch_utils import set_random_seed
18
19 pytestmark = pytest.mark.skipif(
20 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
21 reason="AWQ Triton kernels require CUDA/ROCm or XPU.",
22 )
23
24 device = current_platform.device_type
25
26
27 def reverse_awq_order(t: torch.Tensor):
28 bits = 4
29 AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
30 reverse_order_tensor = torch.arange(
31 t.shape[-1],
32 dtype=torch.int32,
33 device=t.device,
34 )
35 reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
36 reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
37 reverse_order_tensor = reverse_order_tensor.view(-1)
38
39 t = t[:, reverse_order_tensor] & 0xF
40 return t
41
42
43 # qweights - [R , C // 8], int32
44 # scales - [R // G, C ], float16
45 # zeros - [R // G, C // 8], int32
46 def awq_dequantize_torch(
47 qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
48 ) -> torch.Tensor:
49 if group_size == -1:
50 group_size = qweight.shape[0]
51
52 bits = 4
53 shifts = torch.arange(0, 32, bits, device=qzeros.device)
54
55 iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
56 torch.int8
57 )
58
59 iweights = iweights.view(iweights.shape[0], -1)
60
61 zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
62 torch.int8
63 )
64 zeros = zeros.view(qzeros.shape[0], -1)
65 zeros = reverse_awq_order(zeros)
66
67 iweights = reverse_awq_order(iweights)
68
69 iweights = torch.bitwise_and(iweights, (2**bits) - 1)
70 zeros = torch.bitwise_and(zeros, (2**bits) - 1)
71
72 scales = scales.repeat_interleave(group_size, dim=0)
73 zeros = zeros.repeat_interleave(group_size, dim=0)
74 return (iweights - zeros) * scales
75
76
77 # qweights - [R , C // 8], int32
78 # scales - [R // G, C ], float16
79 # zeros - [R // G, C // 8], int32
80 @pytest.mark.parametrize("qweight_rows", [3584, 18944, 128, 256, 512, 1024])
81 @pytest.mark.parametrize("qweight_cols", [448, 576, 4736, 16, 32, 64, 128])
82 @pytest.mark.parametrize("group_size", AWQ_TRITON_SUPPORTED_GROUP_SIZES)
83 def test_dequantize(qweight_rows, qweight_cols, group_size):
84 if group_size == -1:
85 group_size = qweight_rows
86
87 qweight_dtype = torch.int32
88 scales_rows = qweight_rows // group_size
89 scales_cols = qweight_cols * 8
90 scales_dtype = torch.float16
91 zeros_rows = scales_rows
92 zeros_cols = qweight_cols
93 zeros_dtype = torch.int32
94
95 set_random_seed(0)
96
97 qweight = torch.randint(
98 0,
99 torch.iinfo(torch.int32).max,
100 (qweight_rows, qweight_cols),
101 dtype=qweight_dtype,
102 device=device,
103 )
104 scales = torch.rand(scales_rows, scales_cols, dtype=scales_dtype, device=device)
105 zeros = torch.randint(
106 0,
107 torch.iinfo(torch.int32).max,
108 (zeros_rows, zeros_cols),
109 dtype=zeros_dtype,
110 device=device,
111 )
112
113 iweights_triton = awq_dequantize_triton(qweight, scales, zeros)
114
115 assert not torch.any(torch.isinf(iweights_triton)) and not torch.any(
116 torch.isnan(iweights_triton)
117 )
118
119 iweights_torch = awq_dequantize_torch(qweight, scales, zeros, group_size)
120
121 torch.testing.assert_close(iweights_triton, iweights_torch)
122
123
124 # input - [N, K]
125 # qweight - [K, M // 8]
126 # qzeros - [K // G, M // 8]
127 # scales - [K // G, M]
128 @pytest.mark.parametrize("N", [1, 2, 4, 8, 14, 17, 23, 32])
129 @pytest.mark.parametrize("K", [128])
130 @pytest.mark.parametrize("M", [16, 24, 32])
131 @pytest.mark.parametrize("group_size", AWQ_TRITON_SUPPORTED_GROUP_SIZES)
132 @pytest.mark.parametrize("splitK", [1, 8])
133 def test_gemm(N, K, M, splitK, group_size):
134 if group_size == -1:
135 group_size = K
136
137 split_k_iters = splitK
138
139 input_rows = N
140 input_cols = K
141 input_dtype = torch.float32
142 qweight_rows = input_cols
143 qweight_cols = M // 8
144 scales_rows = qweight_rows // group_size
145 scales_cols = M
146 scales_dtype = torch.float32
147 qzeros_rows = scales_rows
148 qzeros_cols = qweight_cols
149
150 set_random_seed(0)
151
152 input = torch.rand((input_rows, input_cols), dtype=input_dtype, device=device)
153 qweight = torch.randint(
154 0, torch.iinfo(torch.int32).max, (qweight_rows, qweight_cols), device=device
155 )
156 qzeros = torch.randint(
157 0, torch.iinfo(torch.int32).max, (qzeros_rows, qzeros_cols), device=device
158 )
159 scales = torch.rand((scales_rows, scales_cols), dtype=scales_dtype, device=device)
160
161 output_triton = awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
162
163 assert not torch.any(torch.isinf(output_triton)) and not torch.any(
164 torch.isnan(output_triton)
165 )
166
167 dequantized_weights = awq_dequantize_triton(qweight, scales, qzeros)
168
169 output_torch = torch.matmul(input, dequantized_weights)
170
171 assert not torch.any(torch.isinf(output_torch)) and not torch.any(
172 torch.isnan(output_torch)
173 )
174
175 torch.testing.assert_close(
176 output_triton.cpu(), output_torch.cpu(), atol=1e-1, rtol=1e-1
177 )