@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4fc722eca4
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Kernel/Quant] Remove AQLM (#22943)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Michael Goin committed Aug 16, 2025 at 15:38 UTC 4fc722eca4f6ad63edf1936989f4d2171aab3ca2
16 files changed -1534
.buildkite/scripts/hardware_ci/run-amd-test.sh
-1
@@ -121,7 +121,6 @@ fi
121 if [[ $commands == *" kernels/quantization"* ]]; then
122 commands="${commands} \
123 --ignore=kernels/quantization/test_int8_quant.py \
124 - --ignore=kernels/quantization/test_aqlm.py \
124 --ignore=kernels/quantization/test_machete_mm.py \
125 --ignore=kernels/quantization/test_block_fp8.py \
126 --ignore=kernels/quantization/test_block_int8.py \
CMakeLists.txt
-1
@@ -286,7 +286,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
286 FetchContent_MakeAvailable(cutlass)
287
288 list(APPEND VLLM_EXT_SRC
289 - "csrc/quantization/aqlm/gemm_kernels.cu"
289 "csrc/quantization/awq/gemm_kernels.cu"
290 "csrc/permute_cols.cu"
291 "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
benchmarks/kernels/benchmark_aqlm.py deleted
-345
@@ -1,345 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import os
5 -import sys
6 -from typing import Optional
7 -
8 -import torch
9 -import torch.nn.functional as F
10 -
11 -from vllm import _custom_ops as ops
12 -from vllm.model_executor.layers.quantization.aqlm import (
13 - dequantize_weight,
14 - generic_dequantize_gemm,
15 - get_int_dtype,
16 - optimized_dequantize_gemm,
17 -)
18 -from vllm.utils import FlexibleArgumentParser
19 -
20 -os.environ["CUDA_VISIBLE_DEVICES"] = "0"
21 -
22 -
23 -def torch_mult(
24 - # [..., in_features]
25 - input: torch.Tensor,
26 - weights: torch.Tensor,
27 - # [num_out_groups, 1, 1, 1]
28 - scales: torch.Tensor,
29 -) -> torch.Tensor:
30 - output = F.linear(input, weights)
31 - return output
32 -
33 -
34 -def dequant_out_scale(
35 - # [..., in_features]
36 - input: torch.Tensor,
37 - # [num_out_groups, num_in_groups, num_codebooks]
38 - codes: torch.IntTensor,
39 - # [num_codebooks, codebook_size, out_group_size, in_group_size]
40 - codebooks: torch.Tensor,
41 - # [num_out_groups, 1, 1, 1]
42 - scales: torch.Tensor,
43 - output_partition_sizes: torch.IntTensor,
44 - bias: Optional[torch.Tensor],
45 -) -> torch.Tensor:
46 - weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)
47 -
48 - if bias is None:
49 - output = F.linear(input, weights, bias)
50 - orig_shape = output.shape
51 - flattened_output = output.view(-1, output.size(-1))
52 - f_scales = scales.view(-1, scales.shape[0])
53 - b_scales = f_scales.expand(flattened_output.shape[0], -1)
54 - flattened_output *= b_scales
55 - return flattened_output.view(orig_shape)
56 - else:
57 - b_scales = scales.view(scales.shape[:-3] + (-1,)).expand(-1, weights.shape[1])
58 - weights *= b_scales
59 - return F.linear(input, weights, bias)
60 -
61 -
62 -def dequant_weight_scale(
63 - # [..., in_features]
64 - input: torch.Tensor,
65 - # [num_out_groups, num_in_groups, num_codebooks]
66 - codes: torch.IntTensor,
67 - # [num_codebooks, codebook_size, out_group_size, in_group_size]
68 - codebooks: torch.Tensor,
69 - # [num_out_groups, 1, 1, 1]
70 - scales: torch.Tensor,
71 - output_partition_sizes: torch.IntTensor,
72 - bias: Optional[torch.Tensor],
73 -) -> torch.Tensor:
74 - weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)
75 -
76 - b_scales = scales.view(scales.shape[:-3] + (-1,)).expand(-1, weights.shape[1])
77 - weights *= b_scales
78 - return F.linear(input, weights, bias)
79 -
80 -
81 -def dequant_no_scale(
82 - # [..., in_features]
83 - input: torch.Tensor,
84 - # [num_out_groups, num_in_groups, num_codebooks]
85 - codes: torch.IntTensor,
86 - # [num_codebooks, codebook_size, out_group_size, in_group_size]
87 - codebooks: torch.Tensor,
88 - # [num_out_groups, 1, 1, 1]
89 - scales: torch.Tensor,
90 - output_partition_sizes: torch.IntTensor,
91 - bias: Optional[torch.Tensor],
92 -) -> torch.Tensor:
93 - weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)
94 -
95 - return F.linear(input, weights, bias)
96 -
97 -
98 -# Compare the optimized 1x16 and 2x8 cuda decompression/dequant kernels against
99 -# the generic pytorch version.
100 -# Just visual comparison.
101 -def dequant_test(k: int, parts: torch.Tensor, nbooks: int, bits: int) -> None:
102 - n = int(parts.sum().item())
103 -
104 - device = torch.device("cuda:0")
105 -
106 - code_range = (1 << bits) // 2
107 - ingroups = 8
108 -
109 - codes = torch.randint(
110 - -code_range,
111 - code_range,
112 - size=(n, k // ingroups, nbooks),
113 - dtype=get_int_dtype(bits),
114 - device=device,
115 - )
116 -
117 - codebooks = torch.randn(
118 - size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),
119 - dtype=torch.float16,
120 - device=device,
121 - )
122 -
123 - count = 0
124 - for index in range(16):
125 - for i in range(8):
126 - for book in range(nbooks):
127 - codebooks[book, index, 0, i] = count * (10**book)
128 - count += 1
129 -
130 - print("codes shape", codes.shape)
131 -
132 - for i in range(16):
133 - for book in range(nbooks):
134 - codes[0, i, book] = i
135 - codes[0, -i, book] = i
136 -
137 - weights = dequantize_weight(codes, codebooks, None)
138 - weights2 = ops.aqlm_dequant(codes, codebooks, parts)
139 -
140 - print("weights shape:", weights.shape)
141 - print("weights2 shape:", weights2.shape)
142 -
143 - print("weights are:", weights)
144 - print("weights2 are:", weights2)
145 -
146 - print("first 128 weights are", weights[0, 0:128].to(torch.int32))
147 - print("first 128 weights2 are:", weights2[0, 0:128].to(torch.int32))
148 -
149 - print("last 128 weights are", weights[0, -128:])
150 - print("last 128 weights2 are:", weights2[0, -128:])
151 -
152 -
153 -def main():
154 - parser = FlexibleArgumentParser(description="Benchmark aqlm performance.")
155 -
156 - # Add arguments
157 - parser.add_argument(
158 - "--nbooks", type=int, default=1, help="Number of codebooks (default: 1)"
159 - )
160 - parser.add_argument(
161 - "--bits",
162 - type=int,
163 - default=16,
164 - help="Number of bits per code element (default: 16)",
165 - )
166 - parser.add_argument(
167 - "--test",
168 - type=bool,
169 - default=False,
170 - help="Run the decompression/dequant tester rather than benchmarking "
171 - "(default: False)",
172 - )
173 -
174 - # Parse the arguments
175 - args = parser.parse_args()
176 -
177 - # Extract values
178 - nbooks = args.nbooks
179 - bits = args.bits
180 -
181 - if args.test:
182 - dequant_test(4096, torch.tensor((4096,)), nbooks, bits)
183 - return
184 -
185 - # Otherwise, benchmark.
186 - methods = [
187 - ops.aqlm_gemm,
188 - dequant_out_scale,
189 - generic_dequantize_gemm,
190 - optimized_dequantize_gemm,
191 - dequant_weight_scale,
192 - torch_mult,
193 - dequant_no_scale,
194 - ]
195 -
196 - filename = f"./aqlm_benchmark_{nbooks}x{bits}.csv"
197 - print(f"writing benchmarks to file {filename}")
198 - with open(filename, "w") as f:
199 - sys.stdout = f
200 -
201 - print("m | k | n | n parts", end="")
202 - for method in methods:
203 - print(f" | {method.__name__.replace('_', ' ')} (µs)", end="")
204 - print("")
205 -
206 - # These are reasonable prefill sizes.
207 - ksandpartions = (
208 - (4096, (4096, 4096, 4096)),
209 - (4096, (4096,)),
210 - (4096, (11008, 11008)),
211 - (11008, (4096,)),
212 - )
213 -
214 - # reasonable ranges for m.
215 - for m in [
216 - 1,
217 - 2,
218 - 4,
219 - 8,
220 - 10,
221 - 12,
222 - 14,
223 - 16,
224 - 24,
225 - 32,
226 - 48,
227 - 52,
228 - 56,
229 - 64,
230 - 96,
231 - 112,
232 - 128,
233 - 256,
234 - 512,
235 - 1024,
236 - 1536,
237 - 2048,
238 - 3072,
239 - 4096,
240 - ]:
241 - print(f"{m}", file=sys.__stdout__)
242 - for ksp in ksandpartions:
243 - run_grid(m, ksp[0], torch.tensor(ksp[1]), nbooks, bits, methods)
244 -
245 - sys.stdout = sys.__stdout__
246 -
247 -
248 -def run_grid(m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int, methods):
249 - # I didn't see visible improvements from increasing these, but feel free :)
250 - num_warmup_trials = 1
251 - num_trials = 1
252 -
253 - num_calls = 100
254 -
255 - # warmup.
256 - for method in methods:
257 - for _ in range(num_warmup_trials):
258 - run_timing(
259 - num_calls=num_calls,
260 - m=m,
261 - k=k,
262 - parts=parts,
263 - nbooks=nbooks,
264 - bits=bits,
265 - method=method,
266 - )
267 -
268 - n = parts.sum().item()
269 - print(f"{m} | {k} | {n} | {parts.tolist()}", end="")
270 -
271 - for method in methods:
272 - best_time_us = 1e20
273 - for _ in range(num_trials):
274 - kernel_dur_ms = run_timing(
275 - num_calls=num_calls,
276 - m=m,
277 - k=k,
278 - parts=parts,
279 - nbooks=nbooks,
280 - bits=bits,
281 - method=method,
282 - )
283 -
284 - kernel_dur_us = 1000 * kernel_dur_ms
285 -
286 - if kernel_dur_us < best_time_us:
287 - best_time_us = kernel_dur_us
288 -
289 - print(f" | {kernel_dur_us:.0f}", end="")
290 -
291 - print("")
292 -
293 -
294 -def run_timing(
295 - num_calls: int, m: int, k: int, parts: torch.Tensor, nbooks: int, bits: int, method
296 -) -> float:
297 - n = int(parts.sum().item())
298 -
299 - device = torch.device("cuda:0")
300 -
301 - input = torch.randn((1, m, k), dtype=torch.float16, device=device)
302 -
303 - code_range = (1 << bits) // 2
304 - ingroups = 8
305 -
306 - codes = torch.randint(
307 - -code_range,
308 - code_range,
309 - size=(n, k // ingroups, nbooks),
310 - dtype=get_int_dtype(bits),
311 - device=device,
312 - )
313 -
314 - codebooks = torch.randn(
315 - size=(parts.shape[0] * nbooks, 1 << bits, 1, 8),
316 - dtype=torch.float16,
317 - device=device,
318 - )
319 -
320 - scales = torch.randn(size=(n, 1, 1, 1), dtype=torch.float16, device=device)
321 -
322 - # for comparison to just a pytorch mult.
323 - weights = torch.randn((n, k), dtype=torch.float16, device=device)
324 -
325 - start_event = torch.cuda.Event(enable_timing=True)
326 - end_event = torch.cuda.Event(enable_timing=True)
327 -
328 - start_event.record()
329 -
330 - if method is torch_mult:
331 - for i in range(num_calls):
332 - torch_mult(input, weights, scales)
333 - else:
334 - for i in range(num_calls):
335 - method(input, codes, codebooks, scales, parts, None)
336 -
337 - end_event.record()
338 - end_event.synchronize()
339 -
340 - dur_ms = start_event.elapsed_time(end_event) / num_calls
341 - return dur_ms
342 -
343 -
344 -if __name__ == "__main__":
345 - sys.exit(main())
csrc/ops.h
-9
@@ -154,15 +154,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
154 torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
155
156 #ifndef USE_ROCM
157 -torch::Tensor aqlm_gemm(const torch::Tensor& input, const torch::Tensor& codes,
158 - const torch::Tensor& codebooks,
159 - const torch::Tensor& scales,
160 - const std::vector<int64_t>& codebook_partition_sizes,
161 - const std::optional<torch::Tensor>& bias);
162 -
163 -torch::Tensor aqlm_dequant(
164 - const torch::Tensor& codes, const torch::Tensor& codebooks,
165 - const std::vector<int64_t>& codebook_partition_sizes);
157
158 torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
159 torch::Tensor _scaling_factors, torch::Tensor _zeros,
csrc/quantization/aqlm/gemm_kernels.cu deleted
-597
@@ -1,597 +0,0 @@
1 -/*
2 - * Modified by Neural Magic
3 - * Adapted from https://github.com/Vahe1994/AQLM
4 - *
5 - * Licensed under the Apache License, Version 2.0 (the "License");
6 - * you may not use this file except in compliance with the License.
7 - * You may obtain a copy of the License at
8 - *
9 - * http://www.apache.org/licenses/LICENSE-2.0
10 - *
11 - * Unless required by applicable law or agreed to in writing, software
12 - * distributed under the License is distributed on an "AS IS" BASIS,
13 - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 - * See the License for the specific language governing permissions and
15 - * limitations under the License.
16 - */
17 -
18 -#include <cuda.h>
19 -#include <cuda_fp16.h>
20 -#include <cuda_runtime.h>
21 -#include <torch/all.h>
22 -#include <c10/cuda/CUDAStream.h>
23 -#include <c10/cuda/CUDAGuard.h>
24 -
25 -#include <iostream>
26 -#include <cstdlib>
27 -
28 -namespace vllm {
29 -namespace aqlm {
30 -
31 -__global__ void Code1x16MatVec(
32 - const int4* __restrict__ A, const int4* __restrict__ B,
33 - int4* __restrict__ C, const int4* __restrict__ codebook, const int prob_m,
34 - const int prob_k,
35 - const int4 codebook_a_sizes, // cumulative sizes of A spanning each
36 - // codebook, at most 3 long.
37 - const int codebook_stride // as int4.
38 -) {
39 - int a_gl_stride = prob_k / 8 / 8;
40 - int a_gl_rd = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
41 - bool pred = a_gl_rd < prob_m;
42 -
43 - if (pred) {
44 - // advance to the correct codebook, this easy because we only multiply one
45 - // column of the codebook.
46 - auto codebook_size = &codebook_a_sizes.x;
47 - while (a_gl_rd >= *codebook_size) {
48 - codebook += codebook_stride;
49 - ++codebook_size;
50 - }
51 - }
52 -
53 - int b_gl_rd = 0;
54 - int c_gl_wr = a_gl_rd;
55 - a_gl_rd = a_gl_stride * a_gl_rd + threadIdx.x % 32;
56 - int a_gl_end = a_gl_rd + a_gl_stride - threadIdx.x % 32;
57 -
58 - __shared__ int4 sh_b[32 * 9];
59 - float res = 0;
60 -
61 - int iters = (prob_k / 8 + 8 * 32 - 1) / (8 * 32);
62 - while (iters--) {
63 - // We pad shared memory to avoid bank conflicts during reads
64 - __syncthreads();
65 - for (int i = threadIdx.x; i < 32 * 8; i += blockDim.x) {
66 - if (b_gl_rd + i < prob_k / 8) sh_b[9 * (i / 8) + i % 8] = B[b_gl_rd + i];
67 - }
68 - __syncthreads();
69 - b_gl_rd += 32 * 8;
70 -
71 - int b_sh_rd = 9 * (threadIdx.x % 32);
72 - if (pred && a_gl_rd < a_gl_end) {
73 - const uint16_t* enc = reinterpret_cast<const uint16_t*>(&A[a_gl_rd]);
74 -#pragma unroll
75 - for (int i = 0; i < 8; i++) {
76 - uint32_t dec[4];
77 - // We bypass the L1 cache to avoid massive amounts of memory streaming
78 - // that doesn't actually help us; this brings > 2x speedup.
79 - asm volatile("ld.cg.global.v4.u32 {%0, %1, %2, %3}, [%4];"
80 - : "=r"(dec[0]), "=r"(dec[1]), "=r"(dec[2]), "=r"(dec[3])
81 - : "l"((void*)&codebook[enc[i]]));
82 - half2* a = reinterpret_cast<half2*>(&dec);
83 - half2* b = reinterpret_cast<half2*>(&sh_b[b_sh_rd]);
84 - half2 res2 = {};
85 -#pragma unroll
86 - for (int j = 0; j < 4; j++) res2 = __hfma2(a[j], b[j], res2);
87 - res += __half2float(res2.x) + __half2float(res2.y);
88 - b_sh_rd++;
89 - }
90 - a_gl_rd += 32;
91 - }
92 - }
93 -
94 - if (pred) {
95 -#pragma unroll
96 - for (int i = 16; i > 0; i /= 2) res += __shfl_down_sync(0xffffffff, res, i);
97 - if (threadIdx.x % 32 == 0)
98 - reinterpret_cast<__half*>(C)[c_gl_wr] = __float2half(res);
99 - }
100 -}
101 -
102 -__global__ void Code2x8MatVec(
103 - const int4* __restrict__ A, const int4* __restrict__ B,
104 - int4* __restrict__ C, const int4* __restrict__ codebook, int prob_m,
105 - int prob_k,
106 - const int4 codebook_a_sizes, // cumulative sizes of A spanning each
107 - // codebook, at most 3 long.
108 - const int codebook_stride // as int4.
109 -
110 -) {
111 - int a_gl_stride = prob_k / 8 / 8;
112 - int a_gl_rd = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
113 - bool pred = a_gl_rd < prob_m;
114 -
115 - if (pred) {
116 - // advance to the correct codebook, this easy because we only multiply one
117 - // column of the codebook.
118 - auto codebook_size = &codebook_a_sizes.x;
119 - while (a_gl_rd >= *codebook_size) {
120 - codebook += codebook_stride;
121 - ++codebook_size;
122 - }
123 - }
124 -
125 - int b_gl_rd = 0;
126 - int c_gl_wr = a_gl_rd;
127 - a_gl_rd = a_gl_stride * a_gl_rd + threadIdx.x % 32;
128 - int a_gl_end = a_gl_rd + a_gl_stride - threadIdx.x % 32;
129 - int lane = threadIdx.x % 8;
130 -
131 - extern __shared__ int4 sh[];
132 - int4* sh_b = sh;
133 - int4* sh_code = sh_b + 32 * 9;
134 - int4* sh_code0 = sh_code;
135 - int4* sh_code1 = sh_code + 256 * 8;
136 -
137 - for (int i = threadIdx.x; i < 2 * 256; i += blockDim.x) {
138 - int4 dec = codebook[i];
139 -#pragma unroll
140 - for (int j = 0; j < 8; j++) sh_code[8 * i + (j + lane) % 8] = dec;
141 - }
142 - __syncthreads();
143 -
144 - float res = 0;
145 -
146 - int iters = (prob_k / 8 + 8 * 32 - 1) / (8 * 32);
147 - while (iters--) {
148 - // We pad shared memory to avoid bank conflicts during reads
149 - __syncthreads();
150 - for (int i = threadIdx.x; i < 32 * 8; i += blockDim.x) {
151 - if (b_gl_rd + i < prob_k / 8) sh_b[9 * (i / 8) + i % 8] = B[b_gl_rd + i];
152 - }
153 - __syncthreads();
154 - b_gl_rd += 32 * 8;
155 -
156 - int b_sh_rd = 9 * (threadIdx.x % 32);
157 - if (pred && a_gl_rd < a_gl_end) {
158 - const uint8_t* enc = reinterpret_cast<const uint8_t*>(&A[a_gl_rd]);
159 -#pragma unroll
160 - for (int i = 0; i < 8; i++) {
161 - half2* a0 =
162 - reinterpret_cast<half2*>(&sh_code0[8 * enc[2 * i + 0] + lane]);
163 - half2* a1 =
164 - reinterpret_cast<half2*>(&sh_code1[8 * enc[2 * i + 1] + lane]);
165 - half2* b = reinterpret_cast<half2*>(&sh_b[b_sh_rd]);
166 - half2 res2 = {};
167 -#pragma unroll
168 - for (int j = 0; j < 4; j++)
169 - res2 = __hfma2(__hadd2(a0[j], a1[j]), b[j], res2);
170 - res += __half2float(res2.x) + __half2float(res2.y);
171 - b_sh_rd++;
172 - }
173 - a_gl_rd += 32;
174 - }
175 - }
176 -
177 - if (pred) {
178 -#pragma unroll
179 - for (int i = 16; i > 0; i /= 2) res += __shfl_down_sync(0xffffffff, res, i);
180 - if (threadIdx.x % 32 == 0)
181 - reinterpret_cast<__half*>(C)[c_gl_wr] = __float2half(res);
182 - }
183 -}
184 -
185 -__global__ void Code1x16Dequant(
186 - const int4* __restrict__ A, int4* __restrict__ C,
187 - const int4* __restrict__ codebook, int prob_m, int prob_k,
188 - const int4 codebook_a_sizes, // cumulative sizes of A spanning each
189 - // codebook, at most 3 long, sums to m.
190 - const int codebook_stride // as int4
191 -) {
192 - int a_gl_stride = prob_k / 8 / 8;
193 - int a_gl_rd = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
194 - bool pred = a_gl_rd < prob_m;
195 -
196 - if (pred) {
197 - // advance to the correct codebook, this easy because we only multiply one
198 - // column of the codebook.
199 - auto codebook_size = &codebook_a_sizes.x;
200 - while (a_gl_rd >= *codebook_size) {
201 - codebook += codebook_stride;
202 - ++codebook_size;
203 - }
204 - }
205 -
206 - a_gl_rd = a_gl_stride * a_gl_rd + threadIdx.x % 32;
207 - int a_gl_end = a_gl_rd + a_gl_stride - threadIdx.x % 32;
208 -
209 - int c_gl_stride = prob_k / 8;
210 - int c_gl_wr = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
211 - c_gl_wr = c_gl_stride * c_gl_wr + (threadIdx.x % 32) * 8;
212 -
213 - int iters = (prob_k / 8 - 1) / (8 * 32) + 1;
214 - while (iters--) {
215 - if (pred && a_gl_rd < a_gl_end) {
216 - const uint16_t* enc = reinterpret_cast<const uint16_t*>(&A[a_gl_rd]);
217 -#pragma unroll
218 - for (int i = 0; i < 8; i++) {
219 - int4 chunk;
220 - auto dec = reinterpret_cast<uint32_t*>(&chunk);
221 - // We bypass the L1 cache to avoid massive amounts of memory streaming
222 - // that doesn't actually help us; this brings > 2x speedup.
223 - asm volatile("ld.cg.global.v4.u32 {%0, %1, %2, %3}, [%4];"
224 - : "=r"(dec[0]), "=r"(dec[1]), "=r"(dec[2]), "=r"(dec[3])
225 - : "l"((void*)&codebook[enc[i]]));
226 -
227 - C[a_gl_rd * 8 + i] = chunk;
228 - }
229 - }
230 - a_gl_rd += 32;
231 - }
232 -}
233 -
234 -__global__ void Code2x8Dequant(
235 - const int4* __restrict__ A, int4* __restrict__ C,
236 - const int4* __restrict__ codebook, int prob_m, int prob_k,
237 - const int4
238 - codebook_a_sizes, // cumulative sizes of A spanning each codebook, at
239 - // most 3 long, corresponds to cols.
240 - const int codebook_stride // as int4
241 -) {
242 - int a_gl_stride = prob_k / 8 / 8;
243 - int a_gl_rd = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
244 - bool pred = a_gl_rd < prob_m;
245 -
246 - if (pred) {
247 - // advance to the correct codebook, this easy because we only multiply one
248 - // column of the codebook.
249 - auto codebook_size = &codebook_a_sizes.x;
250 - while (a_gl_rd >= *codebook_size) {
251 - codebook += codebook_stride;
252 - ++codebook_size;
253 - }
254 - }
255 -
256 - a_gl_rd = a_gl_stride * a_gl_rd + threadIdx.x % 32;
257 - int a_gl_end = a_gl_rd + a_gl_stride - threadIdx.x % 32;
258 - int lane = threadIdx.x % 8;
259 -
260 - int c_gl_stride = prob_k / 8;
261 - int c_gl_wr = (blockDim.x / 32) * blockIdx.x + (threadIdx.x / 32);
262 - c_gl_wr = c_gl_stride * c_gl_wr + (threadIdx.x % 32) * 8;
263 -
264 - extern __shared__ int4 sh[];
265 - int4* sh_code = sh;
266 - int4* sh_code0 = sh_code;
267 - int4* sh_code1 = sh_code + 256 * 8;
268 -
269 - for (int i = threadIdx.x; i < 2 * 256; i += blockDim.x) {
270 - int4 dec = codebook[i];
271 -#pragma unroll
272 - for (int j = 0; j < 8; j++) sh_code[8 * i + (j + lane) % 8] = dec;
273 - }
274 - __syncthreads();
275 -
276 - int iters = (prob_k / 8 - 1) / (8 * 32) + 1;
277 - while (iters--) {
278 - if (pred && a_gl_rd < a_gl_end) {
279 - const uint8_t* enc = reinterpret_cast<const uint8_t*>(&A[a_gl_rd]);
280 -#pragma unroll
281 - for (int i = 0; i < 8; i++) {
282 - int4 chunk;
283 - half2* a0 =
284 - reinterpret_cast<half2*>(&sh_code0[8 * enc[2 * i + 0] + lane]);
285 - half2* a1 =
286 - reinterpret_cast<half2*>(&sh_code1[8 * enc[2 * i + 1] + lane]);
287 -#pragma unroll
288 - for (int j = 0; j < 4; j++)
289 - reinterpret_cast<half2*>(&chunk)[j] = __hadd2(a0[j], a1[j]);
290 - C[a_gl_rd * 8 + i] = chunk;
291 - }
292 - }
293 - a_gl_rd += 32;
294 - }
295 -}
296 -
297 -inline int ceildiv(int a, int b) { return (a + b - 1) / b; }
298 -
299 -const int THREAD_M = 16;
300 -
301 -void code1x16_matvec_cuda(const void* __restrict__ A,
302 - const void* __restrict__ B, void* __restrict__ C,
303 - const void* __restrict__ codebook, int prob_m,
304 - int prob_k, const int4 codebook_a_sizes,
305 - const int codebook_stride) {
306 - int sms;
307 - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0);
308 - int waves = 0;
309 - int thread_m;
310 - do {
311 - waves++;
312 - thread_m = ceildiv(prob_m, waves * sms);
313 - } while (thread_m > THREAD_M);
314 -
315 - int blocks = ceildiv(prob_m, thread_m);
316 - int threads = 32 * thread_m;
317 - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
318 - Code1x16MatVec<<<blocks, threads, 16 * 32 * 9, stream>>>(
319 - (const int4*)A, (const int4*)B, (int4*)C, (const int4*)codebook, prob_m,
320 - prob_k, codebook_a_sizes, codebook_stride);
321 -}
322 -
323 -void code2x8_matvec_cuda(const void* __restrict__ A, const void* __restrict__ B,
324 - void* __restrict__ C,
325 - const void* __restrict__ codebook, int prob_m,
326 - int prob_k, const int4 codebook_a_sizes,
327 - const int codebook_stride) {
328 - int sms;
329 - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0);
330 - int waves = 0;
331 - int thread_m;
332 - do {
333 - waves++;
334 - thread_m = ceildiv(prob_m, waves * sms);
335 - } while (thread_m > THREAD_M);
336 -
337 - int blocks = ceildiv(prob_m, thread_m);
338 - int threads = 32 * thread_m;
339 - int shared = 16 * (2 * 256 * 8 + 32 * 9);
340 - cudaFuncSetAttribute(Code2x8MatVec,
341 - cudaFuncAttributeMaxDynamicSharedMemorySize, shared);
342 - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
343 - Code2x8MatVec<<<blocks, threads, shared, stream>>>(
344 - (const int4*)A, (const int4*)B, (int4*)C, (const int4*)codebook, prob_m,
345 - prob_k, codebook_a_sizes, codebook_stride);
346 -}
347 -
348 -void code1x16_dequant_cuda(
349 - const void* __restrict__ A, void* __restrict__ C,
350 - const void* __restrict__ codebook, int prob_m, int prob_k,
351 - const int4 codebook_a_sizes, // cumulative sizes of A spanning each
352 - // codebook, at most 3 long.
353 - const int codebook_stride // as int4.
354 -) {
355 - int sms;
356 - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0);
357 - int waves = 0;
358 - int thread_m;
359 - do {
360 - waves++;
361 - thread_m = ceildiv(prob_m, waves * sms);
362 - } while (thread_m > THREAD_M);
363 -
364 - int blocks = ceildiv(prob_m, thread_m);
365 - int threads = 32 * thread_m;
366 - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
367 - Code1x16Dequant<<<blocks, threads, 0, stream>>>(
368 - (const int4*)A, (int4*)C, (const int4*)codebook, prob_m, prob_k,
369 - codebook_a_sizes, // cumulative sizes of A spanning each codebook, at
370 - // most 3 long.
371 - codebook_stride // as int4.
372 - );
373 -}
374 -
375 -// Dequantizes the code and codebook into weights.
376 -void code2x8_dequant_cuda(
377 - const void* __restrict__ A, void* __restrict__ C,
378 - const void* __restrict__ codebook, int prob_m, int prob_k,
379 - const int4
380 - codebook_a_sizes, // cumulative sizes of A spanning each codebook, at
381 - // most 3 long, corresponds to cols.
382 - const int codebook_stride // as int4
383 -) {
384 - int sms;
385 - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, 0);
386 - int waves = 0;
387 - int thread_m;
388 - do {
389 - waves++;
390 - thread_m = ceildiv(prob_m, waves * sms);
391 - } while (thread_m > THREAD_M);
392 -
393 - int blocks = ceildiv(prob_m, thread_m);
394 - int threads = 32 * thread_m;
395 - int shared = 16 * (2 * 256 * 8 + 32 * 9);
396 - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
397 -
398 - cudaFuncSetAttribute(Code2x8Dequant,
399 - cudaFuncAttributeMaxDynamicSharedMemorySize, shared);
400 - Code2x8Dequant<<<blocks, threads, shared, stream>>>(
401 - (const int4*)A, (int4*)C, (const int4*)codebook, prob_m, prob_k,
402 - codebook_a_sizes, codebook_stride);
403 -}
404 -
405 -int codebook_stride(const torch::Tensor& codebooks) {
406 - return codebooks.stride(0) * codebooks.element_size() / sizeof(int4);
407 -}
408 -
409 -void code1x16_matvec(
410 - const torch::Tensor& A, const torch::Tensor& B, torch::Tensor& C,
411 - const torch::Tensor& codebook,
412 - const int4 codebook_a_sizes // cumulative sizes of A spanning each
413 - // codebook, at most 3 long.
414 -) {
415 - const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
416 - int prob_m = C.size(0);
417 - int prob_k = B.size(0);
418 -
419 - code1x16_matvec_cuda(A.data_ptr(), B.data_ptr(), C.data_ptr(),
420 - codebook.data_ptr(), prob_m, prob_k, codebook_a_sizes,
421 - codebook_stride(codebook));
422 -}
423 -
424 -torch::Tensor code1x16_matmat(const torch::Tensor& input,
425 - const torch::Tensor& codes,
426 - const torch::Tensor& codebooks,
427 - const torch::Tensor& scales,
428 - const int4 codebook_a_sizes,
429 - const std::optional<torch::Tensor>& bias) {
430 - auto input_sizes = input.sizes();
431 - auto out_features = codes.size(0) * codebooks.size(2);
432 - auto flat_input = input.reshape({-1, input.size(-1)});
433 - auto flat_output = torch::empty(
434 - {flat_input.size(0), out_features},
435 - torch::TensorOptions().dtype(input.dtype()).device(input.device()));
436 -
437 - for (int i = 0; i < flat_input.size(0); ++i) {
438 - auto input_vec = flat_input.index({i});
439 - auto output_vec = flat_output.index({i});
440 - code1x16_matvec(codes.squeeze(2), input_vec, output_vec, codebooks,
441 - codebook_a_sizes);
442 - }
443 - flat_output *= scales.flatten().unsqueeze(0);
444 -
445 - if (bias.has_value()) {
446 - flat_output += bias->unsqueeze(0);
447 - }
448 -
449 - auto output_sizes = input_sizes.vec();
450 - output_sizes.pop_back();
451 - output_sizes.push_back(-1);
452 - auto output = flat_output.reshape(output_sizes);
453 - return output;
454 -}
455 -
456 -void code2x8_matvec(const torch::Tensor& A, const torch::Tensor& B,
457 - torch::Tensor& C, const torch::Tensor& codebook,
458 - const int4 codebook_a_sizes) {
459 - const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
460 - int prob_m = C.size(0);
461 - int prob_k = B.size(0);
462 - code2x8_matvec_cuda(A.data_ptr(), B.data_ptr(), C.data_ptr(),
463 - codebook.data_ptr(), prob_m, prob_k, codebook_a_sizes,
464 - 2 * codebook_stride(codebook));
465 -}
466 -
467 -torch::Tensor code2x8_matmat(const torch::Tensor& input,
468 - const torch::Tensor& codes,
469 - const torch::Tensor& codebooks,
470 - const torch::Tensor& scales,
471 - const int4 codebook_a_sizes,
472 - const std::optional<torch::Tensor>& bias) {
473 - auto input_sizes = input.sizes();
474 - auto out_features = codes.size(0) * codebooks.size(2);
475 - auto flat_input = input.reshape({-1, input.size(-1)});
476 - auto flat_output = torch::empty(
477 - {flat_input.size(0), out_features},
478 - torch::TensorOptions().dtype(input.dtype()).device(input.device()));
479 -
480 - for (int i = 0; i < flat_input.size(0); ++i) {
481 - auto input_vec = flat_input.index({i});
482 - auto output_vec = flat_output.index({i});
483 - code2x8_matvec(codes.squeeze(2), input_vec, output_vec, codebooks,
484 - codebook_a_sizes);
485 - }
486 - flat_output *= scales.flatten().unsqueeze(0);
487 - if (bias.has_value()) {
488 - flat_output += bias->unsqueeze(0);
489 - }
490 -
491 - auto output_sizes = input_sizes.vec();
492 - output_sizes.pop_back();
493 - output_sizes.push_back(-1);
494 - auto output = flat_output.reshape(output_sizes);
495 - return output;
496 -}
497 -
498 -// Accumulate the partition sizes.
499 -int4 accumulate_sizes(const std::vector<int64_t>& codebook_partition_sizes) {
500 - int4 cumulative_sizes;
501 - auto cumulative_size = &cumulative_sizes.x;
502 - size_t i = 0;
503 - int last = 0;
504 - assert(codebook_partition_sizes.size() <= 4);
505 - for (; i < codebook_partition_sizes.size(); ++i, ++cumulative_size) {
506 - *cumulative_size = codebook_partition_sizes[i] + last;
507 - last = *cumulative_size;
508 - }
509 - // fill in the rest with unreachable.
510 - for (; i < 4; ++i, ++cumulative_size) {
511 - *cumulative_size = last * 10;
512 - }
513 - return cumulative_sizes;
514 -}
515 -
516 -} // namespace aqlm
517 -} // namespace vllm
518 -
519 -torch::Tensor aqlm_gemm(const torch::Tensor& input, const torch::Tensor& codes,
520 - const torch::Tensor& codebooks,
521 - const torch::Tensor& scales,
522 - const std::vector<int64_t>& codebook_partition_sizes,
523 - const std::optional<torch::Tensor>& bias) {
524 - int4 cumulative_sizes =
525 - vllm::aqlm::accumulate_sizes(codebook_partition_sizes);
526 -
527 - int const nbooks = codebooks.size(0) / codebook_partition_sizes.size();
528 - int const entries = codebooks.size(1);
529 -
530 - if (nbooks == 1 && entries == (1 << 16)) {
531 - return vllm::aqlm::code1x16_matmat(input, codes, codebooks, scales,
532 - cumulative_sizes, bias);
533 - }
534 - if (nbooks == 2 && entries == (1 << 8)) {
535 - return vllm::aqlm::code2x8_matmat(input, codes, codebooks, scales,
536 - cumulative_sizes, bias);
537 - }
538 -
539 - TORCH_CHECK(false, "AQLM with ", nbooks, " codebooks and ", entries,
540 - " entries is not currently supported.")
541 - return {};
542 -}
543 -
544 -torch::Tensor aqlm_dequant(
545 - const torch::Tensor& codes, const torch::Tensor& codebooks,
546 - const std::vector<int64_t>& codebook_partition_sizes) {
547 - int4 cumulative_sizes =
548 - vllm::aqlm::accumulate_sizes(codebook_partition_sizes);
549 -
550 - int const nbooks = codebooks.size(0) / codebook_partition_sizes.size();
551 - int const entries = codebooks.size(1);
552 -
553 - const at::cuda::OptionalCUDAGuard device_guard(device_of(codes));
554 - int rows = codes.size(1);
555 - int cols = codes.size(0);
556 -
557 - auto in_features = codes.size(1) * 8;
558 - auto out_features = codes.size(0);
559 -
560 - assert(out_features == std::accumulate(codebook_partition_sizes.begin(),
561 - codebook_partition_sizes.end(), 0));
562 -
563 - auto weights = torch::empty({out_features, in_features},
564 - torch::TensorOptions()
565 - .dtype(codebooks.dtype())
566 - .device(codebooks.device()));
567 -
568 - if (nbooks == 1 && entries == (1 << 16)) {
569 - vllm::aqlm::code1x16_dequant_cuda(codes.data_ptr(), weights.data_ptr(),
570 - codebooks.data_ptr(), out_features,
571 - in_features, cumulative_sizes,
572 - vllm::aqlm::codebook_stride(codebooks));
573 -
574 - // if you wanted to flip to scaling the weights, (though it's 30%-ish slower
575 - // and not consistent with gemv implementation.) weights *=
576 - // scales.index({"...", 0, 0});
577 -
578 - return weights;
579 - }
580 -
581 - if (nbooks == 2 && entries == (1 << 8)) {
582 - vllm::aqlm::code2x8_dequant_cuda(codes.data_ptr(), weights.data_ptr(),
583 - codebooks.data_ptr(), out_features,
584 - in_features, cumulative_sizes,
585 - vllm::aqlm::codebook_stride(codebooks));
586 -
587 - // if you wanted to flip to scaling the weights, (though it's 30%-ish slower
588 - // and not consistent with gemv implementation) weights *=
589 - // scales.index({"...", 0, 0});
590 -
591 - return weights;
592 - }
593 -
594 - TORCH_CHECK(false, "AQLM with ", nbooks, " codebooks and ", entries,
595 - " entries is not currently supported.")
596 - return {};
597 -}
csrc/torch_bindings.cpp
-15
@@ -207,21 +207,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
207
208 // Quantization ops
209 #ifndef USE_ROCM
210 - // Quantized GEMM for AQLM.
211 - ops.def(
212 - "aqlm_gemm(Tensor input, Tensor codes, Tensor codebooks, "
213 - "Tensor scales, int[] codebook_partition_sizes, Tensor? bias) "
214 - "-> Tensor",
215 - {stride_tag});
216 - ops.impl("aqlm_gemm", torch::kCUDA, &aqlm_gemm);
217 -
218 - // Decompression method for AQLM.
219 - ops.def(
220 - "aqlm_dequant(Tensor codes, Tensor codebooks, "
221 - "int[] codebook_partition_sizes) -> Tensor",
222 - {stride_tag});
223 - ops.impl("aqlm_dequant", torch::kCUDA, &aqlm_dequant);
224 -
210 // Quantized GEMM for AWQ.
211 ops.def(
212 "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
docs/features/quantization/supported_hardware.md
-1
@@ -17,7 +17,6 @@ th {
17 | INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ |
18 | FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ✅︎ | ❌ |
19 | BitBLAS (GPTQ) | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
20 -| AQLM | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
20 | bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
21 | DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
22 | GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ | ❌ |
docs/mkdocs/hooks/generate_examples.py
-1
@@ -24,7 +24,6 @@ def fix_case(text: str) -> str:
24 "llm": "LLM",
25 "mae": "MAE",
26 "tpu": "TPU",
27 - "aqlm": "AQLM",
27 "gguf": "GGUF",
28 "lora": "LoRA",
29 "rlhf": "RLHF",
examples/offline_inference/basic/README.md
-14
@@ -52,20 +52,6 @@ Try it yourself with the following argument:
52
53 ### Quantization
54
55 -#### AQLM
56 -
57 -vLLM supports models that are quantized using AQLM.
58 -
59 -Try one yourself by passing one of the following models to the `--model` argument:
60 -
61 -- `ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf`
62 -- `ISTA-DASLab/Llama-2-7b-AQLM-2Bit-2x8-hf`
63 -- `ISTA-DASLab/Llama-2-13b-AQLM-2Bit-1x16-hf`
64 -- `ISTA-DASLab/Mixtral-8x7b-AQLM-2Bit-1x16-hf`
65 -- `BlackSamorez/TinyLlama-1_1B-Chat-v1_0-AQLM-2Bit-1x16-hf`
66 -
67 -> Some of these models are likely to be too large for a single GPU. You can split them across multiple GPUs by setting `--tensor-parallel-size` to the number of required GPUs.
68 -
55 #### GGUF
56
57 vLLM supports models that are quantized using GGUF.
tests/compile/test_full_graph.py
-4
@@ -31,10 +31,6 @@ def models_list(*, all: bool = True, keywords: Optional[list[str]] = None):
31 ]
32
33 if all:
34 - if is_quant_method_supported("aqlm"):
35 - TEST_MODELS.append(("ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf", {
36 - "quantization": "aqlm"
37 - }))
34
35 # TODO: figure out why this fails.
36 if False and is_quant_method_supported("gguf"): # noqa: SIM223
tests/kernels/quantization/test_aqlm.py deleted
-40
@@ -1,40 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -
6 -from tests.kernels.utils import opcheck
7 -from vllm import _custom_ops as ops # noqa: F401
8 -
9 -
10 -def test_aqlm_dequant_opcheck():
11 - codes = torch.randint(-32768,
12 - 32767, (22016, 512, 1),
13 - device='cuda',
14 - dtype=torch.int16)
15 - codebooks = torch.rand((2, 65536, 1, 8),
16 - device='cuda',
17 - dtype=torch.float16)
18 - codebook_partition_sizes = [11008, 11008]
19 -
20 - opcheck(torch.ops._C.aqlm_dequant,
21 - (codes, codebooks, codebook_partition_sizes))
22 -
23 -
24 -def test_aqlm_gemm_opcheck():
25 - input = torch.rand((4, 4096), device='cuda', dtype=torch.float16)
26 - codes = torch.randint(-32768,
27 - 32767, (12288, 512, 1),
28 - device='cuda',
29 - dtype=torch.int16)
30 - codebooks = torch.rand((3, 65536, 1, 8),
31 - device='cuda',
32 - dtype=torch.float16)
33 - scales = torch.rand((12288, 1, 1, 1), device='cuda', dtype=torch.float16)
34 - codebook_partition_sizes = [4096, 4096, 4096]
35 - bias = None
36 -
37 - opcheck(torch.ops._C.aqlm_gemm,
38 - (input, codes, codebooks, scales, codebook_partition_sizes, None))
39 - opcheck(torch.ops._C.aqlm_gemm,
40 - (input, codes, codebooks, scales, codebook_partition_sizes, bias))
tests/models/quantization/test_aqlm.py deleted
-68
@@ -1,68 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import pytest
4 -
5 -from tests.quantization.utils import is_quant_method_supported
6 -from vllm.platforms import current_platform
7 -
8 -# These ground truth generations were generated using `transformers==4.38.1
9 -# aqlm==1.1.0 torch==2.2.0`
10 -# and the below code:
11 -# ```python
12 -# from transformers import AutoTokenizer, AutoModelForCausalLM
13 -# model_id = "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf"
14 -# quantized_model = AutoModelForCausalLM.from_pretrained(model_id,
15 -# torch_dtype="auto", device_map="cuda").cuda()
16 -# tokenizer = AutoTokenizer.from_pretrained(model_id)
17 -# outputs = []
18 -# for prompt in example_prompts:
19 -# input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda")
20 -# hf_outputs = quantized_model.generate(input_ids, max_new_tokens=32)
21 -# outputs.append(tokenizer.decode(hf_outputs[0][input_ids.shape[1]:]))
22 -# print(outputs)
23 -# ```
24 -ground_truth_generations = [
25 - '\n### Features\n\n- **High-throughput**: v',
26 - 'The major milestones in the development of artificial intelligence from '
27 - '195',
28 - 'Compare and contrast artificial intelligence with human intelligence in '
29 - 'terms of processing information. The',
30 - 'Explain the difference between supervised and unsupervised learning.'
31 - '\nExplain',
32 - 'Write a short story about a robot that dreams for the first time. The',
33 - 'Analyze the impact of the COVID-19 pandemic on global economic',
34 - 'The Mona Lisa is a painting by Leonardo da Vinci, and it',
35 - 'The early bird catches the worm.\nThe early bird catches the'
36 -]
37 -
38 -
39 -@pytest.mark.skipif(not is_quant_method_supported("aqlm")
40 - or current_platform.is_rocm()
41 - or not current_platform.is_cuda(),
42 - reason="AQLM is not supported on this GPU type.")
43 -@pytest.mark.parametrize("model", ["ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf"])
44 -@pytest.mark.parametrize("dtype", ["half"])
45 -@pytest.mark.parametrize("max_tokens", [16])
46 -@pytest.mark.parametrize("num_logprobs", [1])
47 -def test_models(
48 - vllm_runner,
49 - example_prompts,
50 - model: str,
51 - dtype: str,
52 - max_tokens: int,
53 - num_logprobs: int,
54 -) -> None:
55 -
56 - with vllm_runner(model, dtype=dtype) as vllm_model:
57 - vllm_outputs = vllm_model.generate_greedy_logprobs(
58 - example_prompts, max_tokens, num_logprobs)
59 -
60 - # loop through the prompts to compare against the ground truth generations
61 - for prompt_idx in range(len(example_prompts)):
62 - vllm_output_ids, vllm_output_str, vllm_logprobs = vllm_outputs[
63 - prompt_idx]
64 -
65 - print("Prompt: ", repr(example_prompts[prompt_idx]))
66 - print("Reference output:", repr(ground_truth_generations[prompt_idx]))
67 - print("Output output: ", repr(vllm_output_str))
68 - assert vllm_output_str == ground_truth_generations[prompt_idx]
vllm/_custom_ops.py
-41
@@ -476,32 +476,6 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
476 dtype=input.dtype,
477 device=input.device).sum(0)
478
479 - @register_fake("_C::aqlm_gemm")
480 - def _aqlm_gemm_fake(input: torch.Tensor, codes: torch.Tensor,
481 - codebooks: torch.Tensor, scales: torch.Tensor,
482 - codebook_partition_sizes: list[int],
483 - bias: Optional[torch.Tensor]) -> torch.Tensor:
484 - out_features = codes.size(0) * codebooks.size(2)
485 - flat_input = input.reshape((-1, input.size(-1)))
486 - flat_output = torch.empty((flat_input.size(0), out_features),
487 - dtype=input.dtype,
488 - device=input.device)
489 -
490 - output_sizes = list(input.shape)
491 - output_sizes.pop()
492 - output_sizes.append(-1)
493 - return flat_output.reshape(tuple(output_sizes))
494 -
495 - @register_fake("_C::aqlm_dequant")
496 - def _aqlm_dequant_fake(
497 - codes: torch.Tensor, codebooks: torch.Tensor,
498 - codebook_partition_sizes: list[int]) -> torch.Tensor:
499 - in_features = codes.size(1) * 8
500 - out_features = codes.size(0)
501 - return torch.empty((out_features, in_features),
502 - dtype=codebooks.dtype,
503 - device=codebooks.device)
504 -
479 @register_fake("_C::machete_mm")
480 def machete_mm_fake(
481 a: torch.Tensor,
@@ -957,21 +931,6 @@ def cutlass_fp4_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
931 sf_offsets)
932
933
960 -# aqlm
961 -def aqlm_gemm(input: torch.Tensor, codes: torch.Tensor,
962 - codebooks: torch.Tensor, scales: torch.Tensor,
963 - codebook_partition_sizes: list[int],
964 - bias: Optional[torch.Tensor]) -> torch.Tensor:
965 - return torch.ops._C.aqlm_gemm(input, codes, codebooks, scales,
966 - codebook_partition_sizes, bias)
967 -
968 -
969 -def aqlm_dequant(codes: torch.Tensor, codebooks: torch.Tensor,
970 - codebook_partition_sizes: list[int]) -> torch.Tensor:
971 - return torch.ops._C.aqlm_dequant(codes, codebooks,
972 - codebook_partition_sizes)
973 -
974 -
934 # gptq_marlin
935 def gptq_marlin_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
936 size_k: int, size_n: int,
vllm/model_executor/layers/linear.py
-18
@@ -692,8 +692,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
692
693 param_data = param.data
694 output_dim = getattr(param, "output_dim", None)
695 - # Special case for AQLM codebooks.
696 - is_metadata = getattr(param, "is_metadata", False)
695 # Special case for per-tensor scale to load scalar into fused array.
696 needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
697
@@ -781,13 +779,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
779 if not is_sharded_weight:
780 loaded_weight = loaded_weight.narrow(output_dim, start_idx,
781 shard_size)
784 - # Special case for AQLM codebooks.
785 - elif is_metadata:
786 - # metadata indicates fixed size concatenated along dim 0
787 - shard_size = loaded_weight.shape[0]
788 - shard_offset = loaded_shard_id * shard_size
789 - param_data = param_data.narrow(0, shard_offset, shard_size)
790 -
782 # Special case for per-tensor scales in fused case.
783 elif needs_scalar_to_array:
784 param_data, loaded_weight = adjust_scalar_to_fused_array(
@@ -1081,8 +1072,6 @@ class QKVParallelLinear(ColumnParallelLinear):
1072
1073 param_data = param.data
1074 output_dim = getattr(param, "output_dim", None)
1084 - # Special case for AQLM codebooks.
1085 - is_metadata = getattr(param, "is_metadata", False)
1075
1076 # Special case for per-tensor scales in fused case.
1077 needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
@@ -1204,13 +1193,6 @@ class QKVParallelLinear(ColumnParallelLinear):
1193 loaded_weight = loaded_weight.narrow(output_dim, start_idx,
1194 shard_size)
1195
1207 - # Special case for for AQLM codebooks.
1208 - elif is_metadata:
1209 - # metadata indicates fixed size concatenated along dim 0
1210 - shard_size = loaded_weight.shape[0]
1211 - shard_index = ["q", "k", "v"].index(loaded_shard_id)
1212 - param_data = param_data.narrow(0, shard_index * shard_size,
1213 - shard_size)
1196 # Special case for per-tensor scales in fused case.
1197 elif needs_scalar_to_array:
1198 param_data, loaded_weight = adjust_scalar_to_fused_array(
vllm/model_executor/layers/quantization/__init__.py
-3
@@ -7,7 +7,6 @@ from vllm.model_executor.layers.quantization.base_config import (
7 QuantizationConfig)
8
9 QuantizationMethods = Literal[
10 - "aqlm",
10 "awq",
11 "deepspeedfp",
12 "tpu_int8",
@@ -88,7 +87,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
87 # lazy import to avoid triggering `torch.compile` too early
88 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
89
91 - from .aqlm import AQLMConfig
90 from .auto_round import AutoRoundConfig
91 from .awq import AWQConfig
92 from .awq_marlin import AWQMarlinConfig
@@ -120,7 +118,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
118 from .tpu_int8 import Int8TpuConfig
119
120 method_to_config: dict[str, type[QuantizationConfig]] = {
123 - "aqlm": AQLMConfig,
121 "awq": AWQConfig,
122 "deepspeedfp": DeepSpeedFPConfig,
123 "tpu_int8": Int8TpuConfig,
vllm/model_executor/layers/quantization/aqlm.py deleted
-376
@@ -1,376 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -# Supports AQLM compression, see https://github.com/Vahe1994/AQLM
5 -# and https://arxiv.org/pdf/2401.06118.pdf
6 -
7 -import math
8 -from typing import Any, Optional
9 -
10 -import torch
11 -import torch.nn.functional as F
12 -from torch.nn.parameter import Parameter
13 -
14 -from vllm import _custom_ops as ops
15 -from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
16 -from vllm.model_executor.layers.quantization import QuantizationMethods
17 -from vllm.model_executor.layers.quantization.base_config import (
18 - QuantizationConfig)
19 -from vllm.model_executor.utils import set_weight_attrs
20 -
21 -
22 -def get_int_dtype(nbits: int) -> torch.dtype:
23 - if nbits <= 8:
24 - return torch.int8
25 - if nbits <= 16:
26 - return torch.int16
27 - if nbits <= 32:
28 - return torch.int32
29 - if nbits <= 64:
30 - return torch.int64
31 - raise ValueError(f"No dtype available for {nbits}-bit codebooks")
32 -
33 -
34 -@torch.inference_mode()
35 -def unpack_int_data(data: torch.IntTensor, nbits: int) -> torch.IntTensor:
36 - return data.to(torch.int64) % (2**nbits)
37 -
38 -
39 -def dequantize_weight(codes: torch.Tensor,
40 - codebooks: torch.Tensor,
41 - scales: Optional[torch.Tensor] = None) -> torch.Tensor:
42 - """
43 - Decode float weights from quantization codes. Differentiable.
44 - :param codes: tensor of integer quantization codes, shape
45 - [*dims, num_out_groups, num_in_groups, num_codebooks]
46 - :param codebooks: tensor of vectors for each quantization code,
47 - [num_codebooks, codebook_size, out_group_size, in_group_size]
48 - :param scales: weight will be multiplied by this factor, must be
49 - broadcastble with
50 - [*dims, out_groups, num_in_groups, out_group_size, in_group_size]
51 - :return: reconstructed weight tensor of shape
52 - [*dims, num_in_groups*group_size]
53 - """
54 - num_out_groups, num_in_groups, num_codebooks = codes.shape[-3:]
55 - num_codebooks, codebook_size, out_group_size, in_group_size = \
56 - codebooks.shape
57 - out_features = num_out_groups * out_group_size
58 - in_features = num_in_groups * in_group_size
59 - codebook_offsets = torch.arange(
60 - 0, num_codebooks * codebook_size, codebook_size,
61 - device=codes.device) # shape: [num_codebooks]
62 - reconstructed_weight_flat = F.embedding_bag(
63 - codes.flatten(0, -2) + codebook_offsets,
64 - codebooks.flatten(0, 1).flatten(-2, -1),
65 - mode="sum"
66 - ) # [prod(dims) * num_out_groups * num_in_groups, out_group_size
67 - # * in_group_size]
68 -
69 - reconstructed_weight_groupwise = reconstructed_weight_flat.view(
70 - list(codes.shape[:-3]) +
71 - [num_out_groups, num_in_groups, out_group_size, in_group_size])
72 - if scales is not None:
73 - reconstructed_weight_groupwise = reconstructed_weight_groupwise.mul(
74 - scales)
75 - return reconstructed_weight_groupwise.swapaxes(
76 - -3, -2).reshape(list(codes.shape[:-3]) + [out_features, in_features])
77 -
78 -
79 -def dequantize_gemm(
80 - input: torch.Tensor, # [..., in_features]
81 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]
82 - codebooks: torch.
83 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]
84 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]
85 - bias: Optional[torch.Tensor],
86 -) -> torch.Tensor:
87 - dequantized_weight = dequantize_weight(
88 - unpack_int_data(codes, codebooks.shape[1].bit_length() - 1),
89 - codebooks,
90 - scales,
91 - )
92 - return F.linear(input, dequantized_weight, bias)
93 -
94 -
95 -# Generic dequantization, slow but flexible.
96 -def generic_dequantize_gemm(
97 - input: torch.Tensor, # [..., in_features]
98 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]
99 - codebooks: torch.
100 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]
101 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]
102 - output_partition_sizes: list[int],
103 - bias: Optional[torch.Tensor],
104 -) -> torch.Tensor:
105 - output_shape = input.shape[:-1] + (scales.shape[0], )
106 - output = torch.empty(output_shape, dtype=input.dtype, device=input.device)
107 - num_outputs = len(output_partition_sizes)
108 -
109 - # break the inputs and codebooks apart then combine the outputs.
110 - # Surprisingly (to me) this is faster than doing 3 de-quants and 1 big
111 - # multiply at the end.
112 - num_codebooks = codebooks.shape[0] // num_outputs
113 - assert (scales.shape[0] == codes.shape[0])
114 - assert (sum(output_partition_sizes) == scales.shape[0])
115 - output_offset = 0
116 - codebooks_offset = 0
117 - for output_size in output_partition_sizes:
118 - shard_output = dequantize_gemm(
119 - input, codes.narrow(0, output_offset, output_size),
120 - codebooks.narrow(0, codebooks_offset, num_codebooks),
121 - scales.narrow(0, output_offset, output_size), None
122 - if bias is None else bias.narrow(0, output_offset, output_size))
123 -
124 - output_slice = output.narrow(-1, output_offset, output_size)
125 - assert (output_slice.shape == shard_output.shape)
126 - output_slice.copy_(shard_output)
127 - output_offset += output_size
128 - codebooks_offset += num_codebooks
129 - return output
130 -
131 -
132 -# Optimized dequnantize/decompression kernels, supports 1x16 and 2x8
133 -# at 6 and 9 times faster than the generic version above, respectively.
134 -def optimized_dequantize_gemm(
135 - input: torch.Tensor, # [..., in_features]
136 - codes: torch.IntTensor, # [num_out_groups, num_in_groups, num_codebooks]
137 - codebooks: torch.
138 - Tensor, # [num_codebooks, codebook_size, out_group_size, in_group_size]
139 - scales: torch.Tensor, # [num_out_groups, 1, 1, 1]
140 - output_partition_sizes: list[int],
141 - bias: Optional[torch.Tensor],
142 -) -> torch.Tensor:
143 - weights = ops.aqlm_dequant(codes, codebooks, output_partition_sizes)
144 -
145 - if bias is None:
146 - # scaling the output is fastest, so we do that when possible.
147 - output = F.linear(input, weights, bias)
148 - orig_shape = output.shape
149 - flattened_output = output.view(-1, output.size(-1))
150 - f_scales = scales.view(-1, scales.shape[0])
151 - b_scales = f_scales.expand(flattened_output.shape[0], -1)
152 - flattened_output *= b_scales
153 - return output.view(orig_shape)
154 - else:
155 - b_scales = scales.view(scales.shape[:-3] + (-1, )).expand(
156 - -1, weights.shape[1])
157 - weights *= b_scales
158 - return F.linear(input, weights, bias)
159 -
160 -
161 -class AQLMConfig(QuantizationConfig):
162 - """Config class for AQLM.
163 -
164 - Reference: https://github.com/Vahe1994/AQLM
165 - """
166 -
167 - def __init__(
168 - self,
169 - in_group_size: int,
170 - nbits_per_codebook: int,
171 - num_codebooks: int,
172 - out_group_size: int,
173 - ) -> None:
174 - super().__init__()
175 - self.in_group_size = in_group_size
176 - self.nbits_per_codebook = nbits_per_codebook
177 - self.num_codebooks = num_codebooks
178 - self.out_group_size = out_group_size
179 -
180 - # out_group_size > 1 is untested, and probably won't work as-is.
181 - assert (self.out_group_size == 1)
182 - self.pack_factor = (self.in_group_size * self.out_group_size)
183 -
184 - def __repr__(self) -> str:
185 - return (f"AQLMConfig(in_group_size={self.in_group_size}, "
186 - f"nbits_per_codebook={self.nbits_per_codebook}, "
187 - f"num_codebooks={self.num_codebooks}, "
188 - f"out_group_size={self.out_group_size})")
189 -
190 - @classmethod
191 - def get_name(cls) -> QuantizationMethods:
192 - return "aqlm"
193 -
194 - @classmethod
195 - def get_supported_act_dtypes(cls) -> list[torch.dtype]:
196 - return [torch.half]
197 -
198 - @classmethod
199 - def get_min_capability(cls) -> int:
200 - return 60
201 -
202 - @classmethod
203 - def get_config_filenames(cls) -> list[str]:
204 - return [] # no extra configs.
205 -
206 - @classmethod
207 - def from_config(cls, config: dict[str, Any]) -> "AQLMConfig":
208 - in_group_size = cls.get_from_keys(config, ["in_group_size"])
209 - nbits_per_codebook = cls.get_from_keys(config, ["nbits_per_codebook"])
210 - num_code_books = cls.get_from_keys(config, ["num_codebooks"])
211 - out_group_size = cls.get_from_keys(config, ["out_group_size"])
212 - return cls(in_group_size, nbits_per_codebook, num_code_books,
213 - out_group_size)
214 -
215 - def get_quant_method(self, layer: torch.nn.Module,
216 - prefix: str) -> Optional["AQLMLinearMethod"]:
217 - if isinstance(layer, LinearBase):
218 - return AQLMLinearMethod(self)
219 - return None
220 -
221 -
222 -class AQLMLinearMethod(LinearMethodBase):
223 - """Linear method for AQLM.
224 -
225 - Args:
226 - quant_config: The AQLM quantization config.
227 - """
228 -
229 - def __init__(self, quant_config: AQLMConfig):
230 - self.quant_config = quant_config
231 -
232 - def create_weights(self, layer: torch.nn.Module,
233 - input_size_per_partition: int,
234 - output_partition_sizes: list[int], input_size: int,
235 - output_size: int, params_dtype: torch.dtype,
236 - **extra_weight_attrs):
237 - del output_size # Unused.
238 - del input_size # Unused.
239 -
240 - if params_dtype != torch.half:
241 - raise ValueError("Only half is currently supported by aqlm")
242 - if input_size_per_partition % self.quant_config.in_group_size != 0:
243 - raise ValueError(
244 - "The input size is not aligned with the quantized "
245 - "weight shape. This can be caused by too large "
246 - "tensor parallel size.")
247 -
248 - output_size_per_partition = sum(output_partition_sizes)
249 - if output_size_per_partition % self.quant_config.out_group_size != 0:
250 - raise ValueError(
251 - "The output size is not aligned with the quantized "
252 - "weight shape. This can be caused by too large "
253 - "tensor parallel size.")
254 -
255 - codes = Parameter(
256 - torch.empty(
257 - # There could actually be two pack factors, one along input and
258 - # one along output, but we don't currently support
259 - # out_group_size, and only the one along output needs to be
260 - # marked with "packed_dim" in order for QKVLinear to work.
261 - output_size_per_partition,
262 - input_size_per_partition // self.quant_config.pack_factor,
263 - self.quant_config.num_codebooks,
264 - dtype=get_int_dtype(self.quant_config.nbits_per_codebook),
265 - ),
266 - requires_grad=False,
267 - )
268 -
269 - set_weight_attrs(
270 - codes,
271 - {
272 - "input_dim": 1,
273 - "output_dim": 0,
274 - "packed_dim": 1,
275 - "pack_factor": self.quant_config.pack_factor,
276 - },
277 - )
278 -
279 - codebooks = Parameter(
280 - torch.empty(
281 - self.quant_config.num_codebooks * len(output_partition_sizes),
282 - 2**self.quant_config.nbits_per_codebook,
283 - self.quant_config.out_group_size,
284 - self.quant_config.in_group_size,
285 - dtype=params_dtype,
286 - ),
287 - requires_grad=False,
288 - )
289 - set_weight_attrs(
290 - codebooks,
291 - {
292 - # metadata indicates fixed size concatenated along dim 0
293 - "is_metadata": True,
294 - "output_partition_sizes": output_partition_sizes
295 - },
296 - )
297 -
298 - scales = Parameter(
299 - torch.empty(
300 - (
301 - output_size_per_partition //
302 - self.quant_config.out_group_size,
303 - 1,
304 - 1,
305 - 1,
306 - ),
307 - dtype=params_dtype,
308 - ),
309 - requires_grad=False,
310 - )
311 - set_weight_attrs(
312 - scales,
313 - {
314 - "output_dim": 0,
315 - "packed_dim": 0,
316 - "pack_factor": self.quant_config.out_group_size
317 - },
318 - )
319 -
320 - layer.register_parameter("codes", codes)
321 - set_weight_attrs(codes, extra_weight_attrs)
322 - layer.register_parameter("codebooks", codebooks)
323 - set_weight_attrs(codebooks, extra_weight_attrs)
324 - layer.register_parameter("scales", scales)
325 - set_weight_attrs(scales, extra_weight_attrs)
326 -
327 - def apply(
328 - self,
329 - layer: torch.nn.Module,
330 - x: torch.Tensor,
331 - bias: Optional[torch.Tensor] = None,
332 - ) -> torch.Tensor:
333 - codebooks = layer.codebooks
334 - codes = layer.codes
335 - scales = layer.scales
336 - output_partition_sizes = getattr(codebooks, "output_partition_sizes",
337 - [])
338 -
339 - nbooks = codes.shape[2]
340 - ingroups = codebooks.shape[3]
341 - outgroups = codebooks.shape[2]
342 - bits = codebooks.shape[1]
343 -
344 - # We support these formats with dedicated gemm and decompression
345 - # kernels.
346 - if ingroups == 8 and outgroups == 1 and (
347 - (bits == 256 and nbooks == 2) or (bits == 65536 and nbooks == 1)):
348 -
349 - # thresholds determined by timings on an A6000, one GPU
350 - use_gemv = math.prod(x.shape[:-1]) <= 6
351 -
352 - return ops.aqlm_gemm(
353 - x,
354 - codes,
355 - codebooks,
356 - scales,
357 - output_partition_sizes,
358 - bias,
359 - ) if use_gemv else optimized_dequantize_gemm(
360 - x,
361 - codes,
362 - codebooks,
363 - scales,
364 - output_partition_sizes,
365 - bias,
366 - )
367 -
368 - # fall back all unoptimized formats
369 - return generic_dequantize_gemm(
370 - x,
371 - codes,
372 - codebooks,
373 - scales,
374 - output_partition_sizes,
375 - bias,
376 - )