| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utility functions used for tests and benchmarks.""" |
| 4 | |
| 5 | import numpy as np |
| 6 | import torch |
| 7 | |
| 8 | from vllm import _custom_ops as ops |
| 9 | from vllm.scalar_type import ScalarType, scalar_types |
| 10 | |
| 11 | from .marlin_utils import GPTQ_MARLIN_TILE, marlin_permute_scales, marlin_zero_points |
| 12 | from .quant_utils import ( |
| 13 | get_pack_factor, |
| 14 | gptq_quantize_weights, |
| 15 | quantize_weights, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | class MarlinWorkspace: |
| 20 | def __init__(self, out_features, min_thread_n, max_parallel): |
| 21 | assert out_features % min_thread_n == 0, ( |
| 22 | "out_features = {} is indivisible by min_thread_n = {}".format( |
| 23 | out_features, min_thread_n |
| 24 | ) |
| 25 | ) |
| 26 | |
| 27 | max_workspace_size = (out_features // min_thread_n) * max_parallel |
| 28 | |
| 29 | self.scratch = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda") |
| 30 | |
| 31 | |
| 32 | def marlin_permute_weights( |
| 33 | q_w, size_k, size_n, perm, tile=GPTQ_MARLIN_TILE, is_a_8bit=False |
| 34 | ): |
| 35 | assert q_w.shape == (size_k, size_n) |
| 36 | assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}" |
| 37 | assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}" |
| 38 | |
| 39 | if is_a_8bit: |
| 40 | # Permute weights to 32x32 marlin tiles |
| 41 | q_w = q_w.reshape((size_k // (tile * 2), tile * 2, size_n // tile, tile)) |
| 42 | else: |
| 43 | # Permute weights to 16x64 marlin tiles |
| 44 | q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile)) |
| 45 | q_w = q_w.permute((0, 2, 1, 3)) |
| 46 | q_w = q_w.reshape((size_k // tile, size_n * tile)) |
| 47 | |
| 48 | q_w = q_w.reshape((-1, perm.numel()))[:, perm].reshape(q_w.shape) |
| 49 | |
| 50 | return q_w |
| 51 | |
| 52 | |
| 53 | def marlin_weights(q_w, size_k, size_n, num_bits, perm, is_a_8bit=False): |
| 54 | # Permute |
| 55 | q_w = marlin_permute_weights(q_w, size_k, size_n, perm, is_a_8bit=is_a_8bit) |
| 56 | |
| 57 | # Pack |
| 58 | pack_factor = get_pack_factor(num_bits) |
| 59 | orig_device = q_w.device |
| 60 | |
| 61 | q_w = q_w.cpu().numpy().astype(np.uint32) |
| 62 | |
| 63 | q_packed = np.zeros((q_w.shape[0], q_w.shape[1] // pack_factor), dtype=np.uint32) |
| 64 | for i in range(pack_factor): |
| 65 | q_packed |= q_w[:, i::pack_factor] << num_bits * i |
| 66 | |
| 67 | q_packed = torch.from_numpy(q_packed.astype(np.int32)).to(orig_device) |
| 68 | |
| 69 | return q_packed |
| 70 | |
| 71 | |
| 72 | def get_weight_perm(num_bits: int, is_a_8bit: bool = False): |
| 73 | perm_list: list[int] = [] |
| 74 | if is_a_8bit: |
| 75 | for i in range(32): |
| 76 | perm1 = [] |
| 77 | col = i // 4 |
| 78 | for block in [0, 1]: |
| 79 | for row in [ |
| 80 | 4 * (i % 4), |
| 81 | 4 * (i % 4) + 1, |
| 82 | 4 * (i % 4) + 2, |
| 83 | 4 * (i % 4) + 3, |
| 84 | 4 * (i % 4 + 4), |
| 85 | 4 * (i % 4 + 4) + 1, |
| 86 | 4 * (i % 4 + 4) + 2, |
| 87 | 4 * (i % 4 + 4) + 3, |
| 88 | ]: |
| 89 | perm1.append(16 * row + col + 8 * block) |
| 90 | for j in range(2): |
| 91 | perm_list.extend([p + 512 * j for p in perm1]) |
| 92 | else: |
| 93 | for i in range(32): |
| 94 | perm1 = [] |
| 95 | col = i // 4 |
| 96 | for block in [0, 1]: |
| 97 | for row in [ |
| 98 | 2 * (i % 4), |
| 99 | 2 * (i % 4) + 1, |
| 100 | 2 * (i % 4 + 4), |
| 101 | 2 * (i % 4 + 4) + 1, |
| 102 | ]: |
| 103 | perm1.append(16 * row + col + 8 * block) |
| 104 | for j in range(4): |
| 105 | perm_list.extend([p + 256 * j for p in perm1]) |
| 106 | |
| 107 | perm = np.array(perm_list) |
| 108 | |
| 109 | if num_bits == 4: |
| 110 | if is_a_8bit: # noqa: SIM108 |
| 111 | interleave = np.array([0, 4, 1, 5, 2, 6, 3, 7]) |
| 112 | else: |
| 113 | interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) |
| 114 | elif num_bits == 8: |
| 115 | if is_a_8bit: # noqa: SIM108 |
| 116 | interleave = np.array([0, 1, 2, 3]) |
| 117 | else: |
| 118 | interleave = np.array([0, 2, 1, 3]) |
| 119 | else: |
| 120 | raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) |
| 121 | |
| 122 | perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel() |
| 123 | perm = torch.from_numpy(perm) |
| 124 | return perm |
| 125 | |
| 126 | |
| 127 | def marlin_quantize( |
| 128 | w: torch.Tensor, |
| 129 | quant_type: ScalarType, |
| 130 | group_size: int, |
| 131 | input_dtype: torch.dtype | None = None, |
| 132 | ): |
| 133 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 134 | |
| 135 | size_k, size_n = w.shape |
| 136 | num_bits = quant_type.size_bits |
| 137 | |
| 138 | # Normalize group_size |
| 139 | if group_size == -1: |
| 140 | group_size = size_k |
| 141 | assert group_size <= size_k |
| 142 | |
| 143 | w_ref, q_w, s = gptq_quantize_weights(w, quant_type, group_size) |
| 144 | |
| 145 | # Reformat to marlin |
| 146 | weight_perm = get_weight_perm(num_bits, is_a_8bit) |
| 147 | marlin_q_w = marlin_weights( |
| 148 | q_w, size_k, size_n, num_bits, weight_perm, is_a_8bit=is_a_8bit |
| 149 | ) |
| 150 | marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, is_a_8bit=is_a_8bit) |
| 151 | |
| 152 | if input_dtype == torch.float8_e4m3fn and quant_type == scalar_types.uint4b8: |
| 153 | ops.marlin_int4_fp8_preprocess(marlin_q_w, inplace=True) |
| 154 | marlin_s = marlin_s * 512 |
| 155 | |
| 156 | # Create result |
| 157 | res_list = [w_ref, marlin_q_w, marlin_s] |
| 158 | for i in range(len(res_list)): |
| 159 | res_list[i] = res_list[i].to(w.device) |
| 160 | |
| 161 | return res_list |
| 162 | |
| 163 | |
| 164 | def awq_marlin_quantize( |
| 165 | w: torch.Tensor, |
| 166 | quant_type: ScalarType, |
| 167 | group_size: int, |
| 168 | input_dtype: torch.dtype | None = None, |
| 169 | ): |
| 170 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 171 | size_k, size_n = w.shape |
| 172 | |
| 173 | # Normalize group_size |
| 174 | if group_size == -1: |
| 175 | group_size = size_k |
| 176 | assert group_size <= size_k |
| 177 | |
| 178 | # Detect num groups |
| 179 | assert size_k % group_size == 0 |
| 180 | num_groups = size_k // group_size |
| 181 | |
| 182 | # Quantize with zp |
| 183 | w_ref, q_w, s, zp = quantize_weights(w, quant_type, group_size, zero_points=True) |
| 184 | |
| 185 | if input_dtype == torch.float8_e4m3fn and quant_type == scalar_types.uint4: |
| 186 | repeated_zp = zp.repeat_interleave(group_size, 0) |
| 187 | q_w_old = q_w |
| 188 | q_w = q_w_old - repeated_zp |
| 189 | q_w[q_w < 0] = 15 - q_w_old[q_w < 0] |
| 190 | s = s * 512 |
| 191 | |
| 192 | # Reformat to marlin |
| 193 | weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit) |
| 194 | marlin_q_w = marlin_weights( |
| 195 | q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit=is_a_8bit |
| 196 | ) |
| 197 | marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, is_a_8bit=is_a_8bit) |
| 198 | marlin_zp = marlin_zero_points( |
| 199 | zp, num_groups, size_n, quant_type.size_bits, is_a_8bit=is_a_8bit |
| 200 | ) |
| 201 | |
| 202 | # Create result |
| 203 | res_list = [w_ref, marlin_q_w, marlin_s, marlin_zp] |
| 204 | for i in range(len(res_list)): |
| 205 | res_list[i] = res_list[i].to(w.device) |
| 206 | |
| 207 | return res_list |