@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / af9b69f977
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 9, 11:15 UTC

[Quantization][Deprecation] Remove Marlin 24 (#32688)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Robert Shaw committed Jan 28, 2026 at 07:54 UTC af9b69f977bd1166ed63c46f9ccbd3a02344ae4f
20 files changed +159 -3161
CMakeLists.txt
-1
@@ -458,7 +458,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
458 endif()
459
460 set(MARLIN_SRCS
461 - "csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
461 "csrc/quantization/marlin/marlin.cu"
462 "csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
463 "csrc/quantization/marlin/gptq_marlin_repack.cu"
benchmarks/kernels/benchmark_marlin.py
+1 -49
@@ -6,12 +6,6 @@ import torch.utils.benchmark as benchmark
6 from benchmark_shapes import WEIGHT_SHAPES
7
8 from vllm import _custom_ops as ops
9 -from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
10 - GPTQ_MARLIN_24_MAX_PARALLEL,
11 - GPTQ_MARLIN_24_MIN_THREAD_N,
12 - GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES,
13 - GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES,
14 -)
9 from vllm.model_executor.layers.quantization.utils.allspark_utils import (
10 ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
11 ALLSPARK_SUPPORTED_QUANT_TYPES,
@@ -34,9 +28,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
28 awq_marlin_quantize,
29 marlin_quantize,
30 )
37 -from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
38 - marlin_24_quantize,
39 -)
31 from vllm.model_executor.layers.quantization.utils.quant_utils import (
32 gptq_pack,
33 gptq_quantize_weights,
@@ -78,14 +69,7 @@ def bench_run(
69 if size_k % group_size != 0:
70 return
71
81 - marlin_24_supported = (
82 - quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES
83 - and group_size in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES
84 - )
85 - repack_supported = (
86 - quant_type in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES
87 - and group_size in MARLIN_SUPPORTED_GROUP_SIZES
88 - )
72 + repack_supported = group_size in MARLIN_SUPPORTED_GROUP_SIZES
73 allspark_supported = (
74 quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
75 and group_size == -1
@@ -126,14 +110,6 @@ def bench_run(
110 marlin_sort_indices,
111 )
112
129 - def gen_marlin_24_params():
130 - marlin_24_w_ref = marlin_24_q_w_comp = marlin_24_meta = marlin_24_s = None
131 - if marlin_24_supported:
132 - (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s) = (
133 - marlin_24_quantize(b, quant_type, group_size)
134 - )
135 - return (marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s)
136 -
113 def gen_repack_params():
114 q_w_gptq = None
115 repack_sort_indices = None
@@ -188,9 +164,6 @@ def bench_run(
164 marlin_g_idx,
165 marlin_sort_indices,
166 ) = gen_marlin_params()
191 - marlin_24_w_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s = (
192 - gen_marlin_24_params()
193 - )
167 q_w_gptq, repack_sort_indices = gen_repack_params()
168 qw_reorder, s_reorder, zp_reorder, sm_count, sm_version, CUBLAS_M_THRESHOLD = (
169 gen_allspark_params()
@@ -200,9 +173,6 @@ def bench_run(
173 marlin_workspace = MarlinWorkspace(
174 size_n, GPTQ_MARLIN_MIN_THREAD_N, GPTQ_MARLIN_MAX_PARALLEL
175 )
203 - marlin_24_workspace = MarlinWorkspace(
204 - size_n, GPTQ_MARLIN_24_MIN_THREAD_N, GPTQ_MARLIN_24_MAX_PARALLEL
205 - )
176
177 globals = {
178 # Gen params
@@ -222,12 +192,6 @@ def bench_run(
192 "marlin_sort_indices": marlin_sort_indices,
193 "marlin_workspace": marlin_workspace,
194 "is_k_full": is_k_full,
225 - # Marlin_24 params
226 - "marlin_24_w_ref": marlin_24_w_ref,
227 - "marlin_24_q_w_comp": marlin_24_q_w_comp,
228 - "marlin_24_meta": marlin_24_meta,
229 - "marlin_24_s": marlin_24_s,
230 - "marlin_24_workspace": marlin_24_workspace,
195 # GPTQ params
196 "q_w_gptq": q_w_gptq,
197 "repack_sort_indices": repack_sort_indices,
@@ -240,7 +204,6 @@ def bench_run(
204 "CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD,
205 # Kernels
206 "marlin_gemm": ops.marlin_gemm,
243 - "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
207 "gptq_marlin_repack": ops.gptq_marlin_repack,
208 "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
209 }
@@ -281,17 +244,6 @@ def bench_run(
244 ).blocked_autorange(min_run_time=min_run_time)
245 )
246
284 - if marlin_24_supported:
285 - results.append(
286 - benchmark.Timer(
287 - stmt="output = gptq_marlin_24_gemm(a, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s, marlin_24_workspace.scratch, quant_type, size_m, size_n, size_k)", # noqa: E501
288 - globals=globals,
289 - label=label,
290 - sub_label=sub_label,
291 - description="gptq_marlin_24_gemm",
292 - ).blocked_autorange(min_run_time=min_run_time)
293 - )
294 -
247 if repack_supported:
248 results.append(
249 benchmark.Timer(
csrc/quantization/marlin/sparse/LICENSE deleted
-203
@@ -1,203 +0,0 @@
1 -Contains code from https://github.com/IST-DASLab/Sparse-Marlin/
2 -
3 - Apache License
4 - Version 2.0, January 2004
5 - http://www.apache.org/licenses/
6 -
7 - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
8 -
9 - 1. Definitions.
10 -
11 - "License" shall mean the terms and conditions for use, reproduction,
12 - and distribution as defined by Sections 1 through 9 of this document.
13 -
14 - "Licensor" shall mean the copyright owner or entity authorized by
15 - the copyright owner that is granting the License.
16 -
17 - "Legal Entity" shall mean the union of the acting entity and all
18 - other entities that control, are controlled by, or are under common
19 - control with that entity. For the purposes of this definition,
20 - "control" means (i) the power, direct or indirect, to cause the
21 - direction or management of such entity, whether by contract or
22 - otherwise, or (ii) ownership of fifty percent (50%) or more of the
23 - outstanding shares, or (iii) beneficial ownership of such entity.
24 -
25 - "You" (or "Your") shall mean an individual or Legal Entity
26 - exercising permissions granted by this License.
27 -
28 - "Source" form shall mean the preferred form for making modifications,
29 - including but not limited to software source code, documentation
30 - source, and configuration files.
31 -
32 - "Object" form shall mean any form resulting from mechanical
33 - transformation or translation of a Source form, including but
34 - not limited to compiled object code, generated documentation,
35 - and conversions to other media types.
36 -
37 - "Work" shall mean the work of authorship, whether in Source or
38 - Object form, made available under the License, as indicated by a
39 - copyright notice that is included in or attached to the work
40 - (an example is provided in the Appendix below).
41 -
42 - "Derivative Works" shall mean any work, whether in Source or Object
43 - form, that is based on (or derived from) the Work and for which the
44 - editorial revisions, annotations, elaborations, or other modifications
45 - represent, as a whole, an original work of authorship. For the purposes
46 - of this License, Derivative Works shall not include works that remain
47 - separable from, or merely link (or bind by name) to the interfaces of,
48 - the Work and Derivative Works thereof.
49 -
50 - "Contribution" shall mean any work of authorship, including
51 - the original version of the Work and any modifications or additions
52 - to that Work or Derivative Works thereof, that is intentionally
53 - submitted to Licensor for inclusion in the Work by the copyright owner
54 - or by an individual or Legal Entity authorized to submit on behalf of
55 - the copyright owner. For the purposes of this definition, "submitted"
56 - means any form of electronic, verbal, or written communication sent
57 - to the Licensor or its representatives, including but not limited to
58 - communication on electronic mailing lists, source code control systems,
59 - and issue tracking systems that are managed by, or on behalf of, the
60 - Licensor for the purpose of discussing and improving the Work, but
61 - excluding communication that is conspicuously marked or otherwise
62 - designated in writing by the copyright owner as "Not a Contribution."
63 -
64 - "Contributor" shall mean Licensor and any individual or Legal Entity
65 - on behalf of whom a Contribution has been received by Licensor and
66 - subsequently incorporated within the Work.
67 -
68 - 2. Grant of Copyright License. Subject to the terms and conditions of
69 - this License, each Contributor hereby grants to You a perpetual,
70 - worldwide, non-exclusive, no-charge, royalty-free, irrevocable
71 - copyright license to reproduce, prepare Derivative Works of,
72 - publicly display, publicly perform, sublicense, and distribute the
73 - Work and such Derivative Works in Source or Object form.
74 -
75 - 3. Grant of Patent License. Subject to the terms and conditions of
76 - this License, each Contributor hereby grants to You a perpetual,
77 - worldwide, non-exclusive, no-charge, royalty-free, irrevocable
78 - (except as stated in this section) patent license to make, have made,
79 - use, offer to sell, sell, import, and otherwise transfer the Work,
80 - where such license applies only to those patent claims licensable
81 - by such Contributor that are necessarily infringed by their
82 - Contribution(s) alone or by combination of their Contribution(s)
83 - with the Work to which such Contribution(s) was submitted. If You
84 - institute patent litigation against any entity (including a
85 - cross-claim or counterclaim in a lawsuit) alleging that the Work
86 - or a Contribution incorporated within the Work constitutes direct
87 - or contributory patent infringement, then any patent licenses
88 - granted to You under this License for that Work shall terminate
89 - as of the date such litigation is filed.
90 -
91 - 4. Redistribution. You may reproduce and distribute copies of the
92 - Work or Derivative Works thereof in any medium, with or without
93 - modifications, and in Source or Object form, provided that You
94 - meet the following conditions:
95 -
96 - (a) You must give any other recipients of the Work or
97 - Derivative Works a copy of this License; and
98 -
99 - (b) You must cause any modified files to carry prominent notices
100 - stating that You changed the files; and
101 -
102 - (c) You must retain, in the Source form of any Derivative Works
103 - that You distribute, all copyright, patent, trademark, and
104 - attribution notices from the Source form of the Work,
105 - excluding those notices that do not pertain to any part of
106 - the Derivative Works; and
107 -
108 - (d) If the Work includes a "NOTICE" text file as part of its
109 - distribution, then any Derivative Works that You distribute must
110 - include a readable copy of the attribution notices contained
111 - within such NOTICE file, excluding those notices that do not
112 - pertain to any part of the Derivative Works, in at least one
113 - of the following places: within a NOTICE text file distributed
114 - as part of the Derivative Works; within the Source form or
115 - documentation, if provided along with the Derivative Works; or,
116 - within a display generated by the Derivative Works, if and
117 - wherever such third-party notices normally appear. The contents
118 - of the NOTICE file are for informational purposes only and
119 - do not modify the License. You may add Your own attribution
120 - notices within Derivative Works that You distribute, alongside
121 - or as an addendum to the NOTICE text from the Work, provided
122 - that such additional attribution notices cannot be construed
123 - as modifying the License.
124 -
125 - You may add Your own copyright statement to Your modifications and
126 - may provide additional or different license terms and conditions
127 - for use, reproduction, or distribution of Your modifications, or
128 - for any such Derivative Works as a whole, provided Your use,
129 - reproduction, and distribution of the Work otherwise complies with
130 - the conditions stated in this License.
131 -
132 - 5. Submission of Contributions. Unless You explicitly state otherwise,
133 - any Contribution intentionally submitted for inclusion in the Work
134 - by You to the Licensor shall be under the terms and conditions of
135 - this License, without any additional terms or conditions.
136 - Notwithstanding the above, nothing herein shall supersede or modify
137 - the terms of any separate license agreement you may have executed
138 - with Licensor regarding such Contributions.
139 -
140 - 6. Trademarks. This License does not grant permission to use the trade
141 - names, trademarks, service marks, or product names of the Licensor,
142 - except as required for reasonable and customary use in describing the
143 - origin of the Work and reproducing the content of the NOTICE file.
144 -
145 - 7. Disclaimer of Warranty. Unless required by applicable law or
146 - agreed to in writing, Licensor provides the Work (and each
147 - Contributor provides its Contributions) on an "AS IS" BASIS,
148 - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
149 - implied, including, without limitation, any warranties or conditions
150 - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
151 - PARTICULAR PURPOSE. You are solely responsible for determining the
152 - appropriateness of using or redistributing the Work and assume any
153 - risks associated with Your exercise of permissions under this License.
154 -
155 - 8. Limitation of Liability. In no event and under no legal theory,
156 - whether in tort (including negligence), contract, or otherwise,
157 - unless required by applicable law (such as deliberate and grossly
158 - negligent acts) or agreed to in writing, shall any Contributor be
159 - liable to You for damages, including any direct, indirect, special,
160 - incidental, or consequential damages of any character arising as a
161 - result of this License or out of the use or inability to use the
162 - Work (including but not limited to damages for loss of goodwill,
163 - work stoppage, computer failure or malfunction, or any and all
164 - other commercial damages or losses), even if such Contributor
165 - has been advised of the possibility of such damages.
166 -
167 - 9. Accepting Warranty or Additional Liability. While redistributing
168 - the Work or Derivative Works thereof, You may choose to offer,
169 - and charge a fee for, acceptance of support, warranty, indemnity,
170 - or other liability obligations and/or rights consistent with this
171 - License. However, in accepting such obligations, You may act only
172 - on Your own behalf and on Your sole responsibility, not on behalf
173 - of any other Contributor, and only if You agree to indemnify,
174 - defend, and hold each Contributor harmless for any liability
175 - incurred by, or claims asserted against, such Contributor by reason
176 - of your accepting any such warranty or additional liability.
177 -
178 - END OF TERMS AND CONDITIONS
179 -
180 - APPENDIX: How to apply the Apache License to your work.
181 -
182 - To apply the Apache License to your work, attach the following
183 - boilerplate notice, with the fields enclosed by brackets "[]"
184 - replaced with your own identifying information. (Don't include
185 - the brackets!) The text should be enclosed in the appropriate
186 - comment syntax for the file format. We also recommend that a
187 - file or class name and description of purpose be included on the
188 - same "printed page" as the copyright notice for easier
189 - identification within third-party archives.
190 -
191 - Copyright [yyyy] [name of copyright owner]
192 -
193 - Licensed under the Apache License, Version 2.0 (the "License");
194 - you may not use this file except in compliance with the License.
195 - You may obtain a copy of the License at
196 -
197 - http://www.apache.org/licenses/LICENSE-2.0
198 -
199 - Unless required by applicable law or agreed to in writing, software
200 - distributed under the License is distributed on an "AS IS" BASIS,
201 - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
202 - See the License for the specific language governing permissions and
203 - limitations under the License.
\ No newline at end of file
csrc/quantization/marlin/sparse/common/base.h deleted
-51
@@ -1,51 +0,0 @@
1 -/*
2 - * Copyright (C) 2024 Roberto Lopez Castro (roberto.lopez.castro@udc.es). All
3 - * Rights Reserved.
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 -#pragma once
19 -
20 -namespace marlin_24 {
21 -
22 -constexpr int ceildiv(int a, int b) { return (a + b - 1) / b; }
23 -
24 -// Instances of `Vec` are used to organize groups of >>registers<<, as needed
25 -// for instance as inputs to tensor core operations. Consequently, all
26 -// corresponding index accesses must be compile-time constants, which is why we
27 -// extensively use `#pragma unroll` throughout the kernel code to guarantee
28 -// this.
29 -template <typename T, int n>
30 -struct Vec {
31 - T elems[n];
32 - __device__ T& operator[](int i) { return elems[i]; }
33 -};
34 -
35 -template <int M_, int N_, int K_>
36 -struct ShapeBase {
37 - static constexpr int M = M_, N = N_, K = K_;
38 -};
39 -
40 -using I4 = Vec<int, 4>;
41 -
42 -// Matrix fragments for tensor core instructions; their precise layout is
43 -// documented here:
44 -// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
45 -using FragA = Vec<half2, 4>;
46 -using FragB = Vec<half2, 2>;
47 -using FragM = Vec<uint, 1>;
48 -using FragC = Vec<float, 4>;
49 -using FragS = Vec<half2, 1>; // quantization scales
50 -
51 -} // namespace marlin_24
csrc/quantization/marlin/sparse/common/mem.h deleted
-136
@@ -1,136 +0,0 @@
1 -/*
2 - * Copyright (C) 2024 Roberto Lopez Castro (roberto.lopez.castro@udc.es). All
3 - * Rights Reserved.
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 -#pragma once
19 -#include "base.h"
20 -
21 -namespace marlin_24 {
22 -// Predicated asynchronous global->shared copy; used for inputs A where we apply
23 -// predication to handle batchsizes that are not multiples of 16.
24 -__device__ inline void cp_async4_pred_zfill(void* smem_ptr,
25 - const void* glob_ptr,
26 - bool pred = true,
27 - const bool zfill = false) {
28 - const int BYTES = 16;
29 - int src_in_bytes = (zfill ? 0 : BYTES);
30 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
31 - asm volatile(
32 - "{\n"
33 - " .reg .pred p;\n"
34 - " setp.ne.b32 p, %0, 0;\n"
35 - " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
36 - "}\n" ::"r"((int)pred),
37 - "r"(smem), "l"(glob_ptr), "n"(BYTES), "r"(src_in_bytes));
38 -}
39 -
40 -__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
41 - bool pred = true) {
42 - const int BYTES = 16;
43 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
44 - asm volatile(
45 - "{\n"
46 - " .reg .pred p;\n"
47 - " setp.ne.b32 p, %0, 0;\n"
48 - " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
49 - "}\n" ::"r"((int)pred),
50 - "r"(smem), "l"(glob_ptr), "n"(BYTES));
51 -}
52 -
53 -// Asynchronous global->shared copy
54 -__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
55 - const int BYTES = 16;
56 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
57 - asm volatile(
58 - "{\n"
59 - " cp.async.cg.shared.global [%0], [%1], %2;\n"
60 - "}\n" ::"r"(smem),
61 - "l"(glob_ptr), "n"(BYTES));
62 -}
63 -
64 -// Async copy fence.
65 -__device__ inline void cp_async_fence() {
66 - asm volatile("cp.async.commit_group;\n" ::);
67 -}
68 -
69 -// Wait until at most `n` async copy stages are still pending.
70 -template <int n>
71 -__device__ inline void cp_async_wait() {
72 - asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
73 -}
74 -
75 -// Instruction for loading a full 16x16 matrix fragment of operand A from shared
76 -// memory, directly in tensor core layout.
77 -__device__ inline void ldsm4(FragA& frag_a, const void* smem_ptr) {
78 - uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
79 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
80 - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
81 - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
82 - : "r"(smem));
83 -}
84 -
85 -__device__ inline void ldsm4_m(FragM& frag_m, const void* smem_ptr) {
86 - uint32_t* a = reinterpret_cast<uint32_t*>(&frag_m);
87 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
88 - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
89 - : "=r"(a[0]), "=r"(a[1])
90 - : "r"(smem));
91 -}
92 -
93 -// Instruction for loading a full 16x16 matrix fragment of operand A from shared
94 -// memory, directly in tensor core layout.
95 -__device__ inline void ldsm4_t(FragA& frag_a, const void* smem_ptr) {
96 - uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
97 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
98 - asm volatile(
99 - "ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n"
100 - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
101 - : "r"(smem));
102 -}
103 -
104 -// Wait until barrier reaches `count`, then lock for current threadblock.
105 -__device__ inline void barrier_acquire(int* lock, int count) {
106 - if (threadIdx.x == 0) {
107 - int state = -1;
108 - do
109 - // Guarantee that subsequent writes by this threadblock will be visible
110 - // globally.
111 - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
112 - : "=r"(state)
113 - : "l"(lock));
114 - while (state != count);
115 - }
116 - __syncthreads();
117 -}
118 -
119 -// Release barrier and increment visitation count.
120 -__device__ inline void barrier_release(int* lock, bool reset = false) {
121 - __syncthreads();
122 - if (threadIdx.x == 0) {
123 - if (reset) {
124 - lock[0] = 0;
125 - return;
126 - }
127 - int val = 1;
128 - // Make sure that all writes since acquiring this barrier are visible
129 - // globally, while releasing the barrier.
130 - asm volatile("fence.acq_rel.gpu;\n");
131 - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
132 - :
133 - : "l"(lock), "r"(val));
134 - }
135 -}
136 -} // namespace marlin_24
csrc/quantization/marlin/sparse/common/mma.h deleted
-191
@@ -1,191 +0,0 @@
1 -/*
2 - * Copyright (C) 2024 Roberto Lopez Castro (roberto.lopez.castro@udc.es). All
3 - * Rights Reserved.
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 -#pragma once
19 -#include "base.h"
20 -#include <cudaTypedefs.h>
21 -
22 -namespace marlin_24 {
23 -
24 -// On CUDA earlier than 12.5, the ordered_metadata version of this instruction
25 -// is not supported. On later versions of CUDA the version without ordered
26 -// metadata results in the following warning:
27 -// | Advisory: Modifier ‘.sp::ordered_metadata’ should be used on instruction
28 -// | ‘mma’ instead of modifier ‘.sp’ as it is expected to have substantially
29 -// | reduced performance on some future architectures
30 -#if defined CUDA_VERSION && CUDA_VERSION >= 12050
31 - #define MMA_SP_INST \
32 - "mma.sp::ordered_metadata.sync.aligned.m16n8k32.row.col.f32.f16.f16.f32 "
33 -#else
34 - #define MMA_SP_INST "mma.sp.sync.aligned.m16n8k32.row.col.f32.f16.f16.f32 "
35 -#endif
36 -
37 -// m16n8k32 sparse tensor core mma instruction with fp16 inputs and fp32
38 -// output/accumulation.
39 -__device__ inline void mma_sp(const FragB& a_frag0, const FragB& a_frag1,
40 - const FragA& frag_b, FragC& frag_c, FragM& frag_m,
41 - const int psel) {
42 - const uint32_t* a0 = reinterpret_cast<const uint32_t*>(&a_frag0);
43 - const uint32_t* a1 = reinterpret_cast<const uint32_t*>(&a_frag1);
44 - const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
45 - const uint32_t* e = reinterpret_cast<const uint32_t*>(&frag_m);
46 -
47 - float* c = reinterpret_cast<float*>(&frag_c);
48 - if (psel == 0) {
49 - asm volatile(MMA_SP_INST
50 - "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9, %10,%11}, "
51 - "{%12,%13,%14,%15}, %16, 0x0;\n"
52 - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
53 - : "r"(a0[0]), "r"(a1[0]), "r"(a0[1]), "r"(a1[1]), "r"(b[0]),
54 - "r"(b[2]), "r"(b[4]), "r"(b[6]), "f"(c[0]), "f"(c[1]),
55 - "f"(c[2]), "f"(c[3]), "r"(e[0]));
56 - asm volatile(MMA_SP_INST
57 - "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9, %10,%11}, "
58 - "{%12,%13,%14,%15}, %16, 0x0;\n"
59 - : "=f"(c[4]), "=f"(c[5]), "=f"(c[6]), "=f"(c[7])
60 - : "r"(a0[0]), "r"(a1[0]), "r"(a0[1]), "r"(a1[1]), "r"(b[1]),
61 - "r"(b[3]), "r"(b[5]), "r"(b[7]), "f"(c[4]), "f"(c[5]),
62 - "f"(c[6]), "f"(c[7]), "r"(e[0]));
63 - } else {
64 - asm volatile(MMA_SP_INST
65 - "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9, %10,%11}, "
66 - "{%12,%13,%14,%15}, %16, 0x1;\n"
67 - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
68 - : "r"(a0[0]), "r"(a1[0]), "r"(a0[1]), "r"(a1[1]), "r"(b[0]),
69 - "r"(b[2]), "r"(b[4]), "r"(b[6]), "f"(c[0]), "f"(c[1]),
70 - "f"(c[2]), "f"(c[3]), "r"(e[0]));
71 - asm volatile(MMA_SP_INST
72 - "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9, %10,%11}, "
73 - "{%12,%13,%14,%15}, %16, 0x1;\n"
74 - : "=f"(c[4]), "=f"(c[5]), "=f"(c[6]), "=f"(c[7])
75 - : "r"(a0[0]), "r"(a1[0]), "r"(a0[1]), "r"(a1[1]), "r"(b[1]),
76 - "r"(b[3]), "r"(b[5]), "r"(b[7]), "f"(c[4]), "f"(c[5]),
77 - "f"(c[6]), "f"(c[7]), "r"(e[0]));
78 - }
79 -}
80 -
81 -// Lookup-table based 3-input logical operation; explicitly used for
82 -// dequantization as the compiler does not seem to automatically recognize it in
83 -// all cases.
84 -template <int lut>
85 -__device__ inline int lop3(int a, int b, int c) {
86 - int res;
87 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
88 - : "=r"(res)
89 - : "r"(a), "r"(b), "r"(c), "n"(lut));
90 - return res;
91 -}
92 -
93 -__device__ __forceinline__ uint2 to_half4(float c0, float c1, float c2,
94 - float c3) {
95 - uint2 r;
96 - asm("{\n\t"
97 - ".reg .f16 a, b, c, d; \n\t"
98 - "cvt.rn.f16.f32 a, %2; \n\t"
99 - "cvt.rn.f16.f32 b, %3; \n\t"
100 - "cvt.rn.f16.f32 c, %4; \n\t"
101 - "cvt.rn.f16.f32 d, %5; \n\t"
102 - "mov.b32 %0, {a, b}; \n\t"
103 - "mov.b32 %1, {c, d}; \n\t"
104 - "}"
105 - : "=r"(r.x), "=r"(r.y)
106 - : "f"(c0), "f"(c1), "f"(c2), "f"(c3));
107 - return r;
108 -}
109 -
110 -// Constructs destination register by taking bytes from 2 sources (based on
111 -// mask)
112 -template <int start_byte, int mask>
113 -__device__ inline uint32_t prmt(uint32_t a) {
114 - uint32_t res;
115 - asm volatile("prmt.b32 %0, %1, %2, %3;\n"
116 - : "=r"(res)
117 - : "r"(a), "n"(start_byte), "n"(mask));
118 - return res;
119 -}
120 -
121 -// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16
122 -// values. We mostly follow the strategy in the link below, with some small
123 -// changes:
124 -// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
125 -__device__ inline FragB dequant_4bit(int q) {
126 - const int LO = 0x000f000f;
127 - const int HI = 0x00f000f0;
128 - const int EX = 0x64006400;
129 - // Guarantee that the `(a & b) | c` operations are LOP3s.
130 - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
131 - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
132 - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
133 - // directly into `SUB` and `ADD`.
134 - const int SUB = 0x64086408;
135 - const int MUL = 0x2c002c00;
136 - const int ADD = 0xd480d480;
137 -
138 - FragB frag_b;
139 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
140 - *reinterpret_cast<const half2*>(&SUB));
141 - frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
142 - *reinterpret_cast<const half2*>(&MUL),
143 - *reinterpret_cast<const half2*>(&ADD));
144 - return frag_b;
145 -}
146 -
147 -// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16
148 -// values. We mostly follow the strategy in the link below, with some small
149 -// changes:
150 -// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
151 -__device__ inline FragB dequant_8bit(int q) {
152 - static constexpr uint32_t mask_for_elt_01 = 0x5250;
153 - static constexpr uint32_t mask_for_elt_23 = 0x5351;
154 - static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
155 -
156 - uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
157 - uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
158 -
159 - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
160 -
161 - FragB frag_b;
162 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
163 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
164 - frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
165 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
166 - return frag_b;
167 -}
168 -
169 -// Multiply dequantized values by the corresponding quantization scale; used
170 -// only for grouped quantization.
171 -__device__ inline void scale(FragB& frag_b, FragS& frag_s, int i) {
172 - half2 s = __half2half2(reinterpret_cast<__half*>(&frag_s)[i]);
173 - frag_b[0] = __hmul2(frag_b[0], s);
174 - frag_b[1] = __hmul2(frag_b[1], s);
175 -}
176 -
177 -__device__ inline void scale_floats(float* c0, float* c1, float* c2, float* c3,
178 - FragS& s0, float* c4, float* c5, float* c6,
179 - float* c7, FragS& s1) {
180 - *c0 = __fmul_rn(*c0, __half2float(s0[0].x));
181 - *c1 = __fmul_rn(*c1, __half2float(s0[0].y));
182 - *c2 = __fmul_rn(*c2, __half2float(s0[1].x));
183 - *c3 = __fmul_rn(*c3, __half2float(s0[1].y));
184 -
185 - *c4 = __fmul_rn(*c4, __half2float(s1[0].x));
186 - *c5 = __fmul_rn(*c5, __half2float(s1[0].y));
187 - *c6 = __fmul_rn(*c6, __half2float(s1[1].x));
188 - *c7 = __fmul_rn(*c7, __half2float(s1[1].y));
189 -}
190 -
191 -} // namespace marlin_24
csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu deleted
-1145
@@ -1,1145 +0,0 @@
1 -/*
2 - * Notice: This file was modified by Neuralmagic inc to include 8-bit support
3 - *
4 - * Copyright (C) 2024 Roberto Lopez Castro (roberto.lopez.castro@udc.es). All
5 - * Rights Reserved.
6 - *
7 - * Licensed under the Apache License, Version 2.0 (the "License");
8 - * you may not use this file except in compliance with the License.
9 - * You may obtain a copy of the License at
10 - *
11 - * http://www.apache.org/licenses/LICENSE-2.0
12 - *
13 - * Unless required by applicable law or agreed to in writing, software
14 - * distributed under the License is distributed on an "AS IS" BASIS,
15 - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 - * See the License for the specific language governing permissions and
17 - * limitations under the License.
18 - */
19 -#include <torch/all.h>
20 -
21 -#include <ATen/cuda/CUDAContext.h>
22 -#include <c10/cuda/CUDAGuard.h>
23 -#include <cuda.h>
24 -#include <cuda_fp16.h>
25 -#include <cuda_runtime.h>
26 -
27 -#include <iostream>
28 -
29 -#include "common/base.h"
30 -#include "core/scalar_type.hpp"
31 -#include "core/registration.h"
32 -
33 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
34 -
35 -#else
36 -
37 - #include "common/mem.h"
38 - #include "common/mma.h"
39 -
40 -#endif
41 -
42 -template <typename T>
43 -inline std::string str(T x) {
44 - return std::to_string(x);
45 -}
46 -
47 -namespace marlin_24 {
48 -
49 -// 8 warps are a good choice since every SM has 4 schedulers and having more
50 -// than 1 warp per schedule allows some more latency hiding. At the same time,
51 -// we want relatively few warps to have many registers per warp and small tiles.
52 -static constexpr int THREADS = 256;
53 -static constexpr int STAGES = 4;
54 -
55 -static constexpr int min_thread_n = 128;
56 -
57 -static constexpr int tile_size = 16;
58 -static constexpr int max_par = 64;
59 -
60 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
61 -
62 -template <const int num_bits, // weight bits
63 - const int threads, // number of threads in a threadblock
64 - const int thread_m_blocks, // number of 16x16 blocks in the m
65 - // dimension (batchsize) of the
66 - // threadblock
67 - const int thread_n_blocks, // same for n dimension (output)
68 - const int thread_k_blocks, // same for k dimension (reduction)
69 - const int stages, // number of stages for the async global->shared
70 - // fetch pipeline
71 - const int group_blocks = -1 // number of consecutive 16x16 blocks
72 - // with a separate quantization scale
73 - >
74 -__global__ void Marlin_24(
75 - const int4* __restrict__ A, // fp16 input matrix of shape mxk
76 - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
77 - const int4* __restrict__ meta, // 2bit metadata information about 2:4
78 - // format on B
79 - int4* __restrict__ C, // fp16 output buffer of shape mxn
80 - const int4* __restrict__ s, // fp16 quantization scales of shape
81 - // (k/groupsize)xn
82 - int prob_m, // batch dimension m
83 - int prob_n, // output dimension n
84 - int prob_k, // reduction dimension k
85 - int* locks // extra global storage for barrier synchronization
86 -) {}
87 -
88 -torch::Tensor gptq_marlin_24_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
89 - torch::Tensor& b_meta,
90 - torch::Tensor& b_scales,
91 - torch::Tensor& workspace,
92 - vllm::ScalarTypeId const b_q_type_id,
93 - int64_t size_m, int64_t size_n,
94 - int64_t size_k) {
95 - TORCH_CHECK_NOT_IMPLEMENTED(
96 - false, "gptq_marlin_24_gemm(..) requires CUDA_ARCH >= 8.0");
97 - return torch::empty({1, 1});
98 -}
99 -
100 -#else
101 -
102 -template <const int num_bits, // weight bits
103 - const int threads, // number of threads in a threadblock
104 - const int thread_m_blocks, // number of 16x16 blocks in the m
105 - // dimension (batchsize) of the
106 - // threadblock
107 - const int thread_n_blocks, // same for n dimension (output)
108 - const int thread_k_blocks, // same for k dimension (reduction)
109 - const int stages, // number of stages for the async global->shared
110 - // fetch pipeline
111 - const int group_blocks = -1 // number of consecutive 16x16 blocks
112 - // with a separate quantization scale
113 - >
114 -__global__ void Marlin_24(
115 - const int4* __restrict__ A, // fp16 input matrix of shape mxk
116 - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
117 - const int4* __restrict__ meta, // 2bit metadata information about 2:4
118 - // format on B
119 - int4* __restrict__ C, // fp16 output buffer of shape mxn
120 - const int4* __restrict__ s, // fp16 quantization scales of shape
121 - // (k/groupsize)xn
122 - int prob_m, // batch dimension m
123 - int prob_n, // output dimension n
124 - int prob_k, // reduction dimension k
125 - int* locks // extra global storage for barrier synchronization
126 -) {
127 - // Each threadblock processes one "stripe" of the B matrix with (roughly) the
128 - // same size, which might involve multiple column "slices" (of width 16 *
129 - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
130 - // example:
131 - // 0 1 3
132 - // 0 2 3
133 - // 1 2 4
134 - // While this kind of partitioning makes things somewhat more complicated, it
135 - // ensures good utilization of all SMs for many kinds of shape and GPU
136 - // configurations, while requiring as few slow global cross-threadblock
137 - // reductions as possible.
138 -
139 - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a
140 - // better partitioning with less reductions
141 - int parallel = 1;
142 - if (prob_m > 16 * thread_m_blocks) {
143 - parallel = prob_m / (16 * thread_m_blocks);
144 - prob_m = 16 * thread_m_blocks;
145 - }
146 -
147 - // number of thread_k_blocks in k-dim
148 - int k_tiles = prob_k / 32 / thread_k_blocks;
149 - // number of thread_n_blocks in n-dim
150 - int n_tiles = prob_n / 16 / thread_n_blocks;
151 - // iters needed to cover all slices
152 - int iters = ceildiv(k_tiles * n_tiles * parallel, gridDim.x);
153 -
154 - // Ensure that the number of tiles in each stripe is a multiple of the
155 - // groupsize; this avoids an annoying special case where a stripe starts in
156 - // the middle of group.
157 - if (group_blocks != -1)
158 - iters = (group_blocks / thread_k_blocks) *
159 - ceildiv(iters, (group_blocks / thread_k_blocks));
160 -
161 - int slice_row = (iters * blockIdx.x) % k_tiles;
162 - int slice_col_par = (iters * blockIdx.x) / k_tiles;
163 - int slice_col = slice_col_par;
164 - // number of threadblock tiles in the current slice
165 - int slice_iters;
166 - // total number of active threadblocks in the current slice
167 - int slice_count = 0;
168 - // index of threadblock in current slice; numbered bottom to top
169 - int slice_idx;
170 -
171 - // We can easily implement parallel problem execution by just remapping
172 - // indices and advancing global pointers
173 - if (slice_col_par >= n_tiles) {
174 - A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8;
175 - C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
176 - locks += (slice_col_par / n_tiles) * n_tiles;
177 - slice_col = slice_col_par % n_tiles;
178 - }
179 -
180 - // Compute all information about the current slice which is required for
181 - // synchronization.
182 - auto init_slice = [&]() {
183 - slice_iters =
184 - iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
185 - if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
186 - if (slice_iters == 0) return;
187 - if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
188 - slice_count = 1;
189 - slice_idx = 0;
190 - int col_first = iters * ceildiv(k_tiles * slice_col_par, iters);
191 - if (col_first <= k_tiles * (slice_col_par + 1)) {
192 - int col_off = col_first - k_tiles * slice_col_par;
193 - slice_count = ceildiv(k_tiles - col_off, iters);
194 - if (col_off > 0) slice_count++;
195 - int delta_first = iters * blockIdx.x - col_first;
196 - if (delta_first < 0 || (col_off == 0 && delta_first == 0))
197 - slice_idx = slice_count - 1;
198 - else {
199 - slice_idx = slice_count - 1 - delta_first / iters;
200 - if (col_off > 0) slice_idx--;
201 - }
202 - }
203 - if (slice_col == n_tiles) {
204 - A += 16 * thread_m_blocks * prob_k / 8;
205 - C += 16 * thread_m_blocks * prob_n / 8;
206 - locks += n_tiles;
207 - slice_col = 0;
208 - }
209 - };
210 - init_slice();
211 -
212 - // RLC: 8 is vec_size -> 128-bit instructions, 8 fp16 elements
213 - int a_gl_stride = prob_k / 8; // stride of the A matrix in global memory
214 -
215 - // stride of an A matrix tile in shared memory
216 - constexpr int a_sh_stride = 32 * thread_k_blocks / 8;
217 - // delta between subsequent A tiles in global memory
218 - constexpr int a_gl_rd_delta_o = 32 * thread_k_blocks / 8;
219 - // between subsequent accesses within a tile
220 - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
221 - // between shared memory writes
222 - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
223 - // between shared memory tile reads //RLC: 2 * #warps k-dim
224 - constexpr int a_sh_rd_delta_o = 4 * ((threads / 32) / (thread_n_blocks / 4));
225 - // within a shared memory tile
226 - constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
227 - // overall size of a tile
228 - constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);
229 - // number of shared write iterations for a tile
230 - constexpr int a_sh_wr_iters = ceildiv(a_sh_stage, a_sh_wr_delta);
231 -
232 - constexpr int pack_factor = 32 / num_bits;
233 -
234 - int b_gl_stride = 16 * prob_n / (pack_factor * 4);
235 - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
236 - constexpr int b_thread_vecs = num_bits == 4 ? 1 : 2;
237 - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
238 - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
239 - int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
240 - constexpr int b_sh_wr_delta = threads * b_thread_vecs;
241 - constexpr int b_sh_rd_delta = threads * b_thread_vecs;
242 - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
243 - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
244 -
245 - int m_gl_stride = 2 * prob_n / 8; // (16*2*4 / 8) = 16
246 - constexpr int m_sh_stride =
247 - (16 * thread_n_blocks) / 4; // #warps n-dim * threads/warp
248 - int m_gl_rd_delta_o = m_gl_stride * thread_k_blocks;
249 - int m_gl_rd_delta_i = m_gl_stride * (threads / m_sh_stride);
250 - constexpr int m_sh_wr_delta = threads / 2;
251 - constexpr int m_sh_rd_delta = threads / 2;
252 - constexpr int m_sh_stage = m_sh_stride * thread_k_blocks;
253 - constexpr int m_sh_iters = ceildiv(m_sh_stage, m_sh_wr_delta);
254 -
255 - int s_gl_stride = prob_n / 8;
256 - constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
257 - constexpr int s_sh_stage = s_sh_stride;
258 - int s_gl_rd_delta = s_gl_stride;
259 -
260 - // Global A read index of current thread.
261 - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
262 - (threadIdx.x % a_gl_rd_delta_o);
263 - a_gl_rd += a_gl_rd_delta_o * slice_row;
264 - // Shared write index of current thread.
265 - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
266 - (threadIdx.x % a_gl_rd_delta_o);
267 - // Shared read index.
268 - int a_sh_rd =
269 - a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16;
270 - a_sh_rd += 4 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
271 -
272 - int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
273 - (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
274 - b_gl_rd += b_sh_stride * slice_col;
275 - b_gl_rd += b_gl_rd_delta_o * slice_row;
276 - auto b_sh_wr = threadIdx.x * b_thread_vecs;
277 - auto b_sh_rd = threadIdx.x * b_thread_vecs;
278 -
279 - int m_gl_rd = m_gl_stride * (threadIdx.x / (m_sh_stride)) +
280 - (threadIdx.x % (m_sh_stride));
281 - m_gl_rd += (m_sh_stride)*slice_col;
282 - m_gl_rd += m_gl_rd_delta_o * slice_row;
283 - auto m_sh_wr = threadIdx.x;
284 - auto m_sh_rd = threadIdx.x % 16 + (threadIdx.x / 32) * 16;
285 -
286 - int s_gl_rd;
287 - if constexpr (group_blocks == -1) {
288 - s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
289 - } else {
290 - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
291 - s_sh_stride * slice_col + threadIdx.x;
292 - }
293 -
294 - auto s_sh_wr = threadIdx.x;
295 - int s_sh_rd;
296 - // We use a different scale layout for grouped and column-wise quantization as
297 - // we scale a `half2` tile in column-major layout in the former and in
298 - // row-major in the latter case.
299 - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
300 - (threadIdx.x % 32) / 4; // Note that in the original Marlin kernel
301 - // this is (threadIdx.x % 32) / 4
302 -
303 - // Precompute which thread should not read memory in which iterations; this is
304 - // needed if there are more threads than required for a certain tilesize or
305 - // when the batchsize is not a multiple of 16.
306 - bool a_sh_wr_pred[a_sh_wr_iters];
307 - #pragma unroll
308 - for (int i = 0; i < a_sh_wr_iters; i++) {
309 - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;
310 - }
311 - bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
312 -
313 - // To ensure that writing and reading A tiles to/from shared memory, the
314 - // latter in fragment format, is fully bank conflict free, we need to use a
315 - // rather fancy XOR-based layout. The key here is that neither reads nor
316 - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
317 - // same shared memory banks. Further, it seems (based on NSight-Compute) that
318 - // each warp must also write a consecutive memory segment?
319 - auto transform_a = [&](int i) {
320 - int row = i / a_gl_rd_delta_o;
321 - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
322 - };
323 - // Since the computation of this remapping is non-trivial and, due to our main
324 - // loop unrolls, all shared memory accesses are static, we simply precompute
325 - // both transformed reads and writes.
326 - int a_sh_wr_trans[a_sh_wr_iters];
327 - #pragma unroll
328 - for (int i = 0; i < a_sh_wr_iters; i++)
329 - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
330 - int a_sh_rd_trans[2][b_sh_wr_iters][thread_m_blocks];
331 - #pragma unroll
332 - for (int i = 0; i < b_sh_wr_iters; i++) {
333 - #pragma unroll
334 - for (int j = 0; j < thread_m_blocks; j++) {
335 - a_sh_rd_trans[0][i][j] =
336 - transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
337 - a_sh_rd_trans[1][i][j] =
338 - transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd + 2);
339 - }
340 - }
341 -
342 - // Since B-accesses have non-constant stride they have to be computed at
343 - // runtime; we break dependencies between subsequent accesses with a tile by
344 - // maintining multiple pointers (we have enough registers), a tiny
345 - // optimization.
346 - const int4* B_ptr[b_sh_wr_iters];
347 - #pragma unroll
348 - for (int i = 0; i < b_sh_wr_iters; i++)
349 - B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
350 -
351 - bool m_sh_wr_pred = threadIdx.x < m_sh_wr_delta;
352 - const int4* meta_ptr[m_sh_iters];
353 - #pragma unroll
354 - for (int i = 0; i < m_sh_iters; i++)
355 - meta_ptr[i] = meta + m_gl_rd_delta_i * i + m_gl_rd;
356 -
357 - extern __shared__ int4 sh[];
358 - // Shared memory storage for global fetch pipelines.
359 - int4* sh_a = sh;
360 - int4* sh_b = sh_a + (stages * a_sh_stage);
361 - int4* sh_s = sh_b + (stages * b_sh_stage);
362 - int4* sh_m = sh_s + (stages * s_sh_stage);
363 - // Register storage for double buffer of shared memory reads.
364 - FragA frag_a[2][thread_m_blocks][2];
365 - I4 frag_b_quant[2][b_thread_vecs];
366 - FragM frag_m[2][2];
367 - FragC frag_c[thread_m_blocks][4][2];
368 - FragS frag_s[2][4];
369 -
370 - // Zero accumulators.
371 - auto zero_accums = [&]() {
372 - #pragma unroll
373 - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
374 - reinterpret_cast<float*>(frag_c)[i] = 0;
375 - };
376 -
377 - // Asynchronously fetch the next A, B and s tile from global to the next
378 - // shared memory pipeline location.
379 - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
380 - if (pred) {
381 - int4* sh_a_stage = sh_a + a_sh_stage * pipe;
382 - #pragma unroll
383 - for (int i = 0; i < a_sh_wr_iters; i++) {
384 - cp_async4_pred(
385 - &sh_a_stage[a_sh_wr_trans[i]],
386 - &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
387 - a_sh_wr_pred[i]);
388 - }
389 - int4* sh_b_stage = sh_b + b_sh_stage * pipe;
390 - #pragma unroll
391 - for (int i = 0; i < b_sh_wr_iters; i++) {
392 - #pragma unroll
393 - for (int j = 0; j < b_thread_vecs; j++) {
394 - cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j);
395 - }
396 - B_ptr[i] += b_gl_rd_delta_o;
397 - }
398 - int4* sh_meta_stage = sh_m + m_sh_stage * pipe;
399 - #pragma unroll
400 - for (int i = 0; i < m_sh_iters; i++) {
401 - if (m_sh_wr_pred)
402 - cp_async4(&sh_meta_stage[m_sh_wr_delta * i + m_sh_wr], meta_ptr[i]);
403 - meta_ptr[i] += m_gl_rd_delta_o;
404 - }
405 - // Only fetch scales if this tile starts a new group
406 - if constexpr (group_blocks != -1) {
407 - // This assumes group_blocks >= thread_k_blocks
408 - // and would need to be modified to support smaller groups.
409 - static_assert(group_blocks >= thread_k_blocks);
410 - if (pipe % (group_blocks / thread_k_blocks) == 0) {
411 - int4* sh_s_stage = sh_s + s_sh_stage * pipe;
412 - if (s_sh_wr_pred) cp_async4(&sh_s_stage[s_sh_wr], &s[s_gl_rd]);
413 - s_gl_rd += s_gl_rd_delta;
414 - }
415 - }
416 - }
417 - // Insert a fence even when we are winding down the pipeline to ensure that
418 - // waiting is also correct at this point.
419 - cp_async_fence();
420 - };
421 -
422 - // Wait until the next thread tile has been loaded to shared memory.
423 - auto wait_for_stage = [&]() {
424 - // We only have `stages - 2` active fetches since we are double buffering
425 - // and can only issue the next fetch when it is guaranteed that the previous
426 - // shared memory load is fully complete (as it may otherwise be
427 - // overwritten).
428 - cp_async_wait<stages - 2>();
429 - __syncthreads();
430 - };
431 -
432 - // Load the next sub-tile from the current location in the shared memory pipe
433 - // into the current register buffer.
434 - auto fetch_to_registers = [&](int k, int pipe) {
435 - // It may seem inefficient that we reload the groups for every sub-tile;
436 - // however, this does not seem to be a significant bottleneck, while some
437 - // theoretically better attempts have lead to bad instruction ordering by
438 - // the compiler and correspondingly a noticeable drop in performance.
439 - if constexpr (group_blocks != -1) {
440 - // This assumes group_blocks >= thread_k_blocks
441 - // and would need to be modified to support smaller groups.
442 - static_assert(group_blocks >= thread_k_blocks);
443 - int4* sh_s_stage =
444 - sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
445 - (pipe / (group_blocks / thread_k_blocks)));
446 - reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
447 - }
448 - int4* sh_a_stage = sh_a + a_sh_stage * pipe;
449 - #pragma unroll
450 - for (int i = 0; i < thread_m_blocks; i++) {
451 - ldsm4(frag_a[k % 2][i][0],
452 - &sh_a_stage[a_sh_rd_trans[0][k % b_sh_wr_iters][i]]);
453 - ldsm4(frag_a[k % 2][i][1],
454 - &sh_a_stage[a_sh_rd_trans[1][k % b_sh_wr_iters][i]]);
455 - }
456 -
457 - int4* sh_b_stage = sh_b + b_sh_stage * pipe;
458 - #pragma unroll
459 - for (int i = 0; i < b_thread_vecs; i++) {
460 - frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
461 - &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
462 - }
463 -
464 - // Load meta with ldsm4
465 - int4* sh_m_stage = sh_m + m_sh_stage * pipe;
466 - ldsm4_m(frag_m[k % 2][0],
467 - &sh_m_stage[m_sh_rd_delta * (k % m_sh_iters) + m_sh_rd]);
468 - };
469 -
470 - // Execute the actual tensor core matmul of a sub-tile.
471 - auto matmul = [&](int k) {
472 - // We have the m dimension as the inner loop in order to encourage overlapping
473 - // dequantization and matmul operations.
474 - #pragma unroll
475 - for (int j = 0; j < 4; j++) {
476 - FragB frag_b0;
477 - FragB frag_b1;
478 -
479 - if constexpr (num_bits == 4) {
480 - int b_quant = frag_b_quant[k % 2][0][j];
481 - int b_quant_shift = b_quant >> 8;
482 -
483 - frag_b0 = dequant_4bit(b_quant);
484 - frag_b1 = dequant_4bit(b_quant_shift);
485 -
486 - } else {
487 - int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k % 2]);
488 - int b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
489 - int b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
490 -
491 - frag_b0 = dequant_8bit(b_quant_0);
492 - frag_b1 = dequant_8bit(b_quant_1);
493 - }
494 -
495 - // If there are no groups, we can just scale the final output once and can
496 - // avoid doing so for each weight.
497 - if constexpr (group_blocks != -1) {
498 - scale(frag_b0, frag_s[k % 2][j], 0);
499 - }
500 - if constexpr (group_blocks != -1) {
501 - scale(frag_b1, frag_s[k % 2][j], 1);
502 - }
503 -
504 - #pragma unroll
505 - for (int i = 0; i < thread_m_blocks; i++) {
506 - mma_sp(frag_b0, frag_b1, frag_a[k % 2][i][0], frag_c[i][j][0],
507 - frag_m[k % 2][j / 2], j % 2);
508 - }
509 - }
510 - };
511 -
512 - // Since we slice across the k dimension of a tile in order to increase the
513 - // number of warps while keeping the n dimension of a tile reasonable, we have
514 - // multiple warps that accumulate their partial sums of the same output
515 - // location; which we have to reduce over in the end. We do in shared memory.
516 - auto thread_block_reduce = [&]() {
517 - constexpr int red_off = threads / b_sh_stride_threads / 2;
518 - if (red_off >= 1) {
519 - auto red_idx = threadIdx.x / b_sh_stride_threads;
520 - constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
521 - constexpr int red_sh_delta = b_sh_stride_threads;
522 - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
523 - (threadIdx.x % b_sh_stride_threads);
524 -
525 - // Parallel logarithmic shared memory reduction. We make sure to avoid any
526 - // unnecessary read or write iterations, e.g., for two warps we write only
527 - // once by warp 1 and read only once by warp 0.
528 - #pragma unroll
529 - for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
530 - #pragma unroll
531 - for (int i = red_off; i > 0; i /= 2) {
532 - if (i <= red_idx && red_idx < 2 * i) {
533 - #pragma unroll
534 - for (int j = 0; j < 4 * 2; j++) {
535 - int red_sh_wr =
536 - red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
537 - if (i < red_off) {
538 - float* c_rd =
539 - reinterpret_cast<float*>(&sh[red_sh_delta * j + red_sh_rd]);
540 - float* c_wr = reinterpret_cast<float*>(&sh[red_sh_wr]);
541 - #pragma unroll
542 - for (int k = 0; k < 4; k++)
543 - reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
544 - c_rd[k] + c_wr[k];
545 - }
546 - sh[red_sh_wr] =
547 - reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
548 - }
549 - }
550 - __syncthreads();
551 - }
552 - if (red_idx == 0) {
553 - #pragma unroll
554 - for (int i = 0; i < 4 * 2; i++) {
555 - float* c_rd =
556 - reinterpret_cast<float*>(&sh[red_sh_delta * i + red_sh_rd]);
557 - #pragma unroll
558 - for (int j = 0; j < 4; j++)
559 - reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
560 - c_rd[j];
561 - }
562 - }
563 - __syncthreads();
564 - }
565 - }
566 - };
567 -
568 - // Since multiple threadblocks may process parts of the same column slice, we
569 - // finally have to globally reduce over the results. As the striped
570 - // partitioning minimizes the number of such reductions and our outputs are
571 - // usually rather small, we perform this reduction serially in L2 cache.
572 - auto global_reduce = [&](bool first = false, bool last = false) {
573 - // We are very careful here to reduce directly in the output buffer to
574 - // maximize L2 cache utilization in this step. To do this, we write out
575 - // results in FP16 (but still reduce with FP32 compute).
576 - constexpr int active_threads = 32 * thread_n_blocks / 4;
577 - if (threadIdx.x < active_threads) {
578 - int c_gl_stride = prob_n / 8;
579 - int c_gl_wr_delta_o = 2 * 4 * c_gl_stride;
580 - int c_gl_wr_delta_i =
581 - c_gl_stride; // 8 threads (e.g., 0,4,8,12,16,20,24,28)
582 - int c_gl_wr = 2 * c_gl_stride * (threadIdx.x % 4) +
583 - 8 * (threadIdx.x / 32) + (threadIdx.x % 32) / 4;
584 - c_gl_wr += (2 * thread_n_blocks) * slice_col;
585 - constexpr int c_sh_wr_delta = active_threads;
586 - auto c_sh_wr = threadIdx.x;
587 -
588 - int col = 2 * ((threadIdx.x % 32) % 4);
589 -
590 - if (!first) {
591 - // Interestingly, doing direct global accesses here really seems to mess up
592 - // the compiler and lead to slowdowns, hence we also use async-copies even
593 - // though these fetches are not actually asynchronous.
594 - #pragma unroll
595 - for (int i = 0; i < thread_m_blocks * 4; i++) {
596 - cp_async4_pred(&sh[c_sh_wr + c_sh_wr_delta * i],
597 - &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
598 - c_gl_wr_delta_i * (i % 2)],
599 - i < (thread_m_blocks - 1) * 4 ||
600 - 8 * (i / 2) + col + (i % 2) < prob_m);
601 - }
602 - cp_async_fence();
603 - cp_async_wait<0>();
604 - }
605 -
606 - #pragma unroll
607 - for (int i = 0; i < thread_m_blocks * 4; i++) {
608 - if (i < (thread_m_blocks - 1) * 4 ||
609 - 8 * (i / 2) + col + (i % 2) < prob_m) {
610 - if (!first) {
611 - int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta];
612 - #pragma unroll
613 - for (int j2 = 0; j2 < 2; j2++) {
614 - #pragma unroll
615 - for (int j1 = 0; j1 < 4; j1++) {
616 - reinterpret_cast<float*>(
617 - &frag_c)[4 * 2 * 4 * (i / 4) + 8 * j1 + 2 * j2 +
618 - 4 * ((i % 4) / 2) + i % 2] +=
619 - __half2float(
620 - reinterpret_cast<__half*>(&c_red)[(j2 * 4 + j1)]);
621 - }
622 - }
623 - }
624 - if (!last) {
625 - int4 c;
626 - #pragma unroll
627 - for (int j2 = 0; j2 < 2; j2++) {
628 - #pragma unroll
629 - for (int j1 = 0; j1 < 4; j1++) {
630 - reinterpret_cast<__half*>(&c)[(j2 * 4 + j1)] =
631 - __float2half(reinterpret_cast<float*>(
632 - &frag_c)[4 * 2 * 4 * (i / 4) + 8 * j1 + 2 * j2 +
633 - 4 * ((i % 4) / 2) + i % 2]);
634 - }
635 - }
636 - C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] =
637 - c;
638 - }
639 - }
640 - }
641 - }
642 - };
643 -
644 - // Write out the reduce final result in the correct layout. We only actually
645 - // reshuffle matrix fragments in this step, the reduction above is performed
646 - // in fragment layout.
647 - auto write_result = [&]() {
648 - int c_gl_stride = prob_n / 8;
649 -
650 - constexpr int c_sh_stride = 2 * thread_n_blocks; // RLC:
651 - constexpr int c_sh_stride_2 = 2 * c_sh_stride + 2; // RLC:
652 - constexpr int c_sh_stride_3 = 2 * (2 * thread_n_blocks) + 2; // RLC:
653 -
654 - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
655 -
656 - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
657 - (threadIdx.x % (2 * thread_n_blocks));
658 - c_gl_wr += (2 * thread_n_blocks) * slice_col;
659 -
660 - int c_sh_wr = c_sh_stride_2 * ((threadIdx.x % 32) % 4) +
661 - ((threadIdx.x % 32) / 4); // RLC:
662 - c_sh_wr += 8 * (threadIdx.x / 32); // 128/4(half4)
663 -
664 - constexpr int c_sh_rd_delta =
665 - c_sh_stride_3 * (threads / (2 * 2 * thread_n_blocks)); // RLC:
666 - int c_sh_rd = c_sh_stride_3 * (threadIdx.x / (2 * 2 * thread_n_blocks)) +
667 - (threadIdx.x % (2 * 2 * thread_n_blocks));
668 -
669 - int c_gl_wr_end = c_gl_stride * prob_m;
670 -
671 - auto write = [&](int idx, float c0, float c1, float c2, float c3, FragS& s0,
672 - float c4, float c5, float c6, float c7, FragS& s1) {
673 - uint2 res[2];
674 - res[0] = to_half4(c0, c1, c2, c3);
675 - res[1] = to_half4(c4, c5, c6, c7);
676 - half2* tmp = (half2*)&res;
677 - // for per-column quantization we finally apply the scale here
678 - if constexpr (group_blocks == -1 && num_bits == 4) {
679 - tmp[0] = __hmul2(tmp[0], s0[0]);
680 - tmp[1] = __hmul2(tmp[1], s0[1]);
681 - tmp[2] = __hmul2(tmp[2], s1[0]);
682 - tmp[3] = __hmul2(tmp[3], s1[1]);
683 - }
684 - ((int4*)sh)[idx] = *((int4*)&res[0]);
685 - };
686 -
687 - // RLC: only warp 0 and 1 baseline example
688 - if (threadIdx.x / 32 < thread_n_blocks / 4) {
689 - #pragma unroll
690 - for (int i = 0; i < thread_m_blocks; i++) {
691 - int wr = c_sh_wr;
692 - write(wr, frag_c[i][0][0][0], frag_c[i][1][0][0], frag_c[i][2][0][0],
693 - frag_c[i][3][0][0], frag_s[0][0], frag_c[i][0][0][2],
694 - frag_c[i][1][0][2], frag_c[i][2][0][2], frag_c[i][3][0][2],
695 - frag_s[0][2]);
696 - write(wr + c_sh_stride, frag_c[i][0][0][1], frag_c[i][1][0][1],
697 - frag_c[i][2][0][1], frag_c[i][3][0][1], frag_s[0][0],
698 - frag_c[i][0][0][3], frag_c[i][1][0][3], frag_c[i][2][0][3],
699 - frag_c[i][3][0][3], frag_s[0][2]);
700 - write(wr + 4 * c_sh_stride_2, frag_c[i][0][1][0], frag_c[i][1][1][0],
701 - frag_c[i][2][1][0], frag_c[i][3][1][0], frag_s[0][0],
702 - frag_c[i][0][1][2], frag_c[i][1][1][2], frag_c[i][2][1][2],
703 - frag_c[i][3][1][2], frag_s[0][2]);
704 - write(wr + 4 * c_sh_stride_2 + c_sh_stride, frag_c[i][0][1][1],
705 - frag_c[i][1][1][1], frag_c[i][2][1][1], frag_c[i][3][1][1],
706 - frag_s[0][0], frag_c[i][0][1][3], frag_c[i][1][1][3],
707 - frag_c[i][2][1][3], frag_c[i][3][1][3], frag_s[0][2]);
708 -
709 - c_sh_wr += 8 * c_sh_stride_2;
710 - }
711 - }
712 - __syncthreads();
713 -
714 - #pragma unroll
715 - for (int i = 0;
716 - i < ceildiv(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
717 - i++) {
718 - if (c_gl_wr < c_gl_wr_end) {
719 - C[c_gl_wr] = sh[c_sh_rd];
720 - c_gl_wr += c_gl_wr_delta;
721 - c_sh_rd += c_sh_rd_delta;
722 - }
723 - }
724 - };
725 -
726 - // Start global fetch and register load pipelines.
727 - auto start_pipes = [&]() {
728 - #pragma unroll
729 - for (int i = 0; i < stages - 1; i++) fetch_to_shared(i, i, i < slice_iters);
730 - zero_accums();
731 - wait_for_stage();
732 - fetch_to_registers(0, 0);
733 - a_gl_rd += a_gl_rd_delta_o * (stages - 1);
734 - };
735 - start_pipes();
736 -
737 - // Main loop.
738 - while (slice_iters) {
739 - // We unroll over both the global fetch and the register load pipeline to
740 - // ensure all shared memory accesses are static. Note that both pipelines have
741 - // even length meaning that the next iteration will always start at index 0.
742 - #pragma unroll
743 - for (int pipe = 0; pipe < stages;) {
744 - fetch_to_shared((pipe + stages - 1) % stages, pipe,
745 - slice_iters >= stages);
746 - matmul(pipe);
747 - wait_for_stage();
748 -
749 - fetch_to_registers(pipe + 1, (pipe + 1) % stages);
750 -
751 - pipe++;
752 - slice_iters--;
753 - if (slice_iters == 0) break;
754 - }
755 - a_gl_rd += a_gl_rd_delta_o * stages;
756 -
757 - // Process results and, if necessary, proceed to the next column slice.
758 - // While this pattern may not be the most readable, other ways of writing
759 - // the loop seemed to noticeably worse performance after compilation.
760 - if (slice_iters == 0) {
761 - cp_async_wait<0>();
762 - bool last = slice_idx == slice_count - 1;
763 - // For per-column scales, we only fetch them here in the final step before
764 - // write-out
765 - if constexpr (group_blocks == -1) {
766 - if constexpr (num_bits == 8) {
767 - if (s_sh_wr_pred) cp_async4(&sh_s[s_sh_wr], &s[s_gl_rd]);
768 - cp_async_fence();
769 - } else {
770 - if (last) {
771 - if (s_sh_wr_pred) cp_async4(&sh_s[s_sh_wr], &s[s_gl_rd]);
772 - cp_async_fence();
773 - }
774 - }
775 - }
776 - thread_block_reduce();
777 -
778 - if constexpr (group_blocks == -1) {
779 - if constexpr (num_bits == 8) {
780 - cp_async_wait<0>();
781 - __syncthreads();
782 - if (threadIdx.x / 32 < thread_n_blocks / 4) {
783 - *(float4*)(frag_s) = *(float4*)(&sh_s[s_sh_rd]);
784 - }
785 - } else {
786 - if (last) {
787 - cp_async_wait<0>();
788 - __syncthreads();
789 - if (threadIdx.x / 32 < thread_n_blocks / 4) {
790 - *(float4*)(frag_s) = *(float4*)(&sh_s[s_sh_rd]);
791 - }
792 - }
793 - }
794 - }
795 -
796 - // For 8-bit channelwise, we apply the scale before the global reduction
797 - // that converts the fp32 results to fp16 (so that we avoid possible
798 - // overflow in fp16)
799 - if constexpr (group_blocks == -1 && num_bits == 8) {
800 - if (threadIdx.x / 32 < thread_n_blocks / 4) {
801 - #pragma unroll
802 - for (int i = 0; i < thread_m_blocks; i++) {
803 - scale_floats(&frag_c[i][0][0][0], &frag_c[i][1][0][0],
804 - &frag_c[i][2][0][0], &frag_c[i][3][0][0], frag_s[0][0],
805 - &frag_c[i][0][0][2], &frag_c[i][1][0][2],
806 - &frag_c[i][2][0][2], &frag_c[i][3][0][2],
807 - frag_s[0][2]);
808 -
809 - scale_floats(&frag_c[i][0][0][1], &frag_c[i][1][0][1],
810 - &frag_c[i][2][0][1], &frag_c[i][3][0][1], frag_s[0][0],
811 - &frag_c[i][0][0][3], &frag_c[i][1][0][3],
812 - &frag_c[i][2][0][3], &frag_c[i][3][0][3],
813 - frag_s[0][2]);
814 -
815 - scale_floats(&frag_c[i][0][1][0], &frag_c[i][1][1][0],
816 - &frag_c[i][2][1][0], &frag_c[i][3][1][0], frag_s[0][0],
817 - &frag_c[i][0][1][2], &frag_c[i][1][1][2],
818 - &frag_c[i][2][1][2], &frag_c[i][3][1][2],
819 - frag_s[0][2]);
820 -
821 - scale_floats(&frag_c[i][0][1][1], &frag_c[i][1][1][1],
822 - &frag_c[i][2][1][1], &frag_c[i][3][1][1], frag_s[0][0],
823 - &frag_c[i][0][1][3], &frag_c[i][1][1][3],
824 - &frag_c[i][2][1][3], &frag_c[i][3][1][3],
825 - frag_s[0][2]);
826 - }
827 - }
828 - }
829 -
830 - if (slice_count > 1) { // only globally reduce if there is more than one
831 - // block in a slice
832 - barrier_acquire(&locks[slice_col], slice_idx);
833 - global_reduce(slice_idx == 0, last);
834 - barrier_release(&locks[slice_col], last);
835 - }
836 - if (last) // only the last block in a slice actually writes the result
837 - write_result();
838 -
839 - slice_row = 0;
840 - slice_col_par++;
841 - slice_col++;
842 - init_slice();
843 - if (slice_iters) {
844 - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
845 - (threadIdx.x % a_gl_rd_delta_o);
846 - #pragma unroll
847 - for (int i = 0; i < b_sh_wr_iters; i++)
848 - B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
849 - #pragma unroll
850 - for (int i = 0; i < m_sh_iters; i++)
851 - meta_ptr[i] += (m_sh_stride)-m_gl_rd_delta_o * k_tiles;
852 - if (slice_col == 0) {
853 - #pragma unroll
854 - for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
855 - #pragma unroll
856 - for (int i = 0; i < m_sh_iters; i++) meta_ptr[i] -= m_gl_stride;
857 - }
858 - s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
859 - start_pipes();
860 - }
861 - }
862 - }
863 -}
864 -
865 -#endif
866 -
867 -#define CALL_IF_2_4(NUM_BITS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, \
868 - THREAD_K_BLOCKS, GROUP_BLOCKS) \
869 - else if (num_bits == NUM_BITS && thread_m_blocks == THREAD_M_BLOCKS && \
870 - thread_n_blocks == THREAD_N_BLOCKS && \
871 - thread_k_blocks == THREAD_K_BLOCKS && \
872 - group_blocks == GROUP_BLOCKS) { \
873 - cudaFuncSetAttribute( \
874 - Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS, \
875 - THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>, \
876 - cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
877 - Marlin_24<NUM_BITS, THREADS, THREAD_N_BLOCKS, THREAD_M_BLOCKS, \
878 - THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS> \
879 - <<<blocks, THREADS, max_shared_mem, stream>>>(A_ptr, B_ptr, meta_ptr, \
880 - C_ptr, s_ptr, prob_n, \
881 - prob_m, prob_k, locks); \
882 - }
883 -
884 -void marlin_cuda_2_4(const void* A, const void* B, const void* meta, void* C,
885 - void* s, int prob_m, int prob_n, int prob_k,
886 - void* workspace, int num_bits, int groupsize = -1,
887 - int dev = 0, cudaStream_t stream = 0, int thread_k = -1,
888 - int thread_m = -1, int sms = -1, int max_par = 16) {
889 - int tot_n = prob_n;
890 - int tot_n_blocks = ceildiv(tot_n, 16);
891 - int pad = 16 * tot_n_blocks - tot_n;
892 -
893 - if (sms == -1) {
894 - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
895 - }
896 - TORCH_CHECK(sms > 0);
897 -
898 - int max_shared_mem = 0;
899 - cudaDeviceGetAttribute(&max_shared_mem,
900 - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
901 - TORCH_CHECK(max_shared_mem > 0);
902 -
903 - if (thread_k == -1 || thread_m == -1) {
904 - if (prob_n <= 16) {
905 - // For small batchizes, better partitioningif is slightly more important
906 - // than better compute utilization
907 - thread_k = 128;
908 - thread_m = 128;
909 - } else {
910 - thread_k = 64;
911 - thread_m = 256;
912 - }
913 - // Also had
914 - // if prob_n > 256
915 - // thread_k = 32;
916 - // thread_m = 512;
917 - // but this is broken,
918 - // TODO(Lucas, Alex M): figure out why
919 - }
920 -
921 - int thread_k_blocks = thread_k / 32; // 2:4 version with m16n8k32 instruction
922 - int thread_m_blocks = thread_m / 16;
923 - int group_blocks = (groupsize == -1) ? -1 : groupsize / 16;
924 - int blocks = sms;
925 -
926 - TORCH_CHECK(prob_m % thread_m == 0, "prob_m = ", prob_m,
927 - " is not divisible by thread_m = ", thread_m);
928 - TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
929 - " is not divisible by thread_k = ", thread_k);
930 - if (group_blocks != -1) {
931 - TORCH_CHECK((prob_k / 2) % group_blocks == 0, "prob_k/2 = ", prob_k / 2,
932 - " is not divisible by group_blocks = ", group_blocks);
933 - }
934 -
935 - TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
936 - ", ", prob_n, ", ", prob_k, "]");
937 -
938 - const int4* A_ptr = (const int4*)A;
939 - const int4* B_ptr = (const int4*)B;
940 - const int4* meta_ptr = (const int4*)meta;
941 - int4* C_ptr = (int4*)C;
942 - const int4* s_ptr = (const int4*)s;
943 -
944 - constexpr int max_m_blocks = 4;
945 -
946 - int* locks = (int*)workspace;
947 - for (int i = 0; i < tot_n_blocks; i += max_m_blocks) {
948 - int thread_n_blocks = tot_n_blocks - i;
949 - prob_n = tot_n - 16 * i;
950 - int par = 1;
951 - if (thread_n_blocks > max_m_blocks) {
952 - // Note that parallel > 1 currently only works for inputs without any
953 - // padding
954 - par = (16 * thread_n_blocks - pad) / (max_m_blocks * 16);
955 - if (par > max_par) par = max_par;
956 - prob_n = (max_m_blocks * 16) * par;
957 - i += max_m_blocks * (par - 1);
958 - thread_n_blocks = max_m_blocks;
959 - }
960 -
961 - // For compilation speed, we only define the kernel configurations that have
962 - // seemed useful (in terms of performance) in our testing, however many more
963 - // are, in principle, possible.
964 -
965 - // the false is start of the CALL_IF macros
966 - if (false) {
967 - } // BMxBNxBK, group
968 - // 4-bit
969 - CALL_IF_2_4(4, 8, 1, 4, -1) // e.g., 16x128x128
970 - CALL_IF_2_4(4, 8, 1, 4, 4) // e.g., 16x128x128, 64
971 -
972 - CALL_IF_2_4(4, 16, 1, 2, -1) // e.g., 16x256x64
973 - CALL_IF_2_4(4, 16, 1, 2, 4) // e.g., 16x256x64, 64
974 - CALL_IF_2_4(4, 16, 2, 2, -1) // e.g.. 32x256x64
975 - CALL_IF_2_4(4, 16, 2, 2, 4)
976 - CALL_IF_2_4(4, 16, 3, 2, -1)
977 - CALL_IF_2_4(4, 16, 3, 2, 4)
978 - CALL_IF_2_4(4, 16, 4, 2, -1)
979 - CALL_IF_2_4(4, 16, 4, 2, 4)
980 -
981 - CALL_IF_2_4(4, 32, 1, 1, -1) // e.g., 16x256x64
982 - CALL_IF_2_4(4, 32, 1, 1, 4) // e.g., 16x256x64, 64
983 - CALL_IF_2_4(4, 32, 2, 1, -1) // e.g.. 32x256x64
984 - CALL_IF_2_4(4, 32, 2, 1, 4)
985 - CALL_IF_2_4(4, 32, 3, 1, -1)
986 - CALL_IF_2_4(4, 32, 3, 1, 4)
987 - CALL_IF_2_4(4, 32, 4, 1, -1)
988 - CALL_IF_2_4(4, 32, 4, 1, 4)
989 -
990 - // 8-bit
991 - CALL_IF_2_4(8, 8, 1, 4, -1) // e.g., 16x128x128
992 - CALL_IF_2_4(8, 8, 1, 4, 4) // e.g., 16x128x128, 64
993 -
994 - CALL_IF_2_4(8, 16, 1, 2, -1) // e.g., 16x256x64
995 - CALL_IF_2_4(8, 16, 1, 2, 4) // e.g., 16x256x64, 64
996 - CALL_IF_2_4(8, 16, 2, 2, -1) // e.g.. 32x256x64
997 - CALL_IF_2_4(8, 16, 2, 2, 4)
998 - CALL_IF_2_4(8, 16, 3, 2, -1)
999 - CALL_IF_2_4(8, 16, 3, 2, 4)
1000 - CALL_IF_2_4(8, 16, 4, 2, -1)
1001 - CALL_IF_2_4(8, 16, 4, 2, 4)
1002 -
1003 - CALL_IF_2_4(8, 32, 1, 1, -1) // e.g., 16x256x64
1004 - CALL_IF_2_4(8, 32, 1, 1, 4) // e.g., 16x256x64, 64
1005 - CALL_IF_2_4(8, 32, 2, 1, -1) // e.g.. 32x256x64
1006 - CALL_IF_2_4(8, 32, 2, 1, 4)
1007 - CALL_IF_2_4(8, 32, 3, 1, -1)
1008 - CALL_IF_2_4(8, 32, 3, 1, 4)
1009 - CALL_IF_2_4(8, 32, 4, 1, -1)
1010 - CALL_IF_2_4(8, 32, 4, 1, 4)
1011 - else {
1012 - throw std::runtime_error("Unsupported shapes: MKN = [" + str(prob_m) +
1013 - ", " + str(prob_k) + ", " + str(prob_n) + "]" +
1014 - ", groupsize = " + str(groupsize) +
1015 - ", thread_m_blocks = " + str(thread_m_blocks) +
1016 - ", thread_n_blocks = " + str(thread_n_blocks) +
1017 - ", thread_k_blocks = " + str(thread_k_blocks));
1018 - }
1019 -
1020 - A_ptr += 16 * thread_n_blocks * (prob_k / 8) * par;
1021 - C_ptr += 16 * thread_n_blocks * (prob_m / 8) * par;
1022 - }
1023 -}
1024 -
1025 -} // namespace marlin_24
1026 -
1027 -torch::Tensor gptq_marlin_24_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
1028 - torch::Tensor& b_meta,
1029 - torch::Tensor& b_scales,
1030 - torch::Tensor& workspace,
1031 - vllm::ScalarTypeId const b_q_type_id,
1032 - int64_t size_m, int64_t size_n,
1033 - int64_t size_k) {
1034 - vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
1035 - // Verify num_bits
1036 - TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128,
1037 - "num_bits must be uint4b8 or uint8b128. Got = ", b_q_type.str());
1038 - int pack_factor = 32 / b_q_type.size_bits();
1039 -
1040 - // Verify M
1041 - TORCH_CHECK(size_m == a.size(0),
1042 - "Shape mismatch: a.size(0) = " + str(a.size(0)) +
1043 - ", size_m = " + str(size_m));
1044 -
1045 - // Verify K
1046 - TORCH_CHECK(size_k == a.size(1),
1047 - "Shape mismatch: a.size(1) = " + str(a.size(1)) +
1048 - ", size_k = " + str(size_k));
1049 - TORCH_CHECK(size_k % marlin_24::tile_size == 0,
1050 - "size_k = " + str(size_k) + " is not divisible by tile_size = " +
1051 - str(marlin_24::tile_size));
1052 - TORCH_CHECK((size_k / marlin_24::tile_size / 2) == b_q_weight.size(0),
1053 - "Shape mismatch: b_q_weight.size(0) = " +
1054 - str(b_q_weight.size(0)) + ", size_k = " + str(size_k) +
1055 - ", tile_size = " + str(marlin_24::tile_size));
1056 -
1057 - // Verify N
1058 - TORCH_CHECK(b_scales.size(1) == size_n,
1059 - "b_scales.size(1) = " + str(b_scales.size(1)) +
1060 - ", size_n = " + str(size_n));
1061 - TORCH_CHECK(
1062 - b_q_weight.size(1) % marlin_24::tile_size == 0,
1063 - "b_q_weight.size(1) = " + str(b_q_weight.size(1)) +
1064 - " is not divisible by tile_size = " + str(marlin_24::tile_size));
1065 -
1066 - int actual_size_n = (b_q_weight.size(1) / marlin_24::tile_size) * pack_factor;
1067 - TORCH_CHECK(
1068 - size_n == actual_size_n,
1069 - "size_n = " + str(size_n) + ", actual_size_n = " + str(actual_size_n));
1070 -
1071 - // Verify meta
1072 - TORCH_CHECK(b_meta.size(0) == size_k / 8 / 2 / 2,
1073 - "b_meta.size(0) = ", b_meta.size(0),
1074 - " is not size_k / 8 / 2 / 2 = ", size_k / 8 / 2 / 2);
1075 - TORCH_CHECK(b_meta.size(1) == size_n * 2, "b_meta.size(1) = ", b_meta.size(1),
1076 - " is not size_n * 2 = ", size_n * 2);
1077 -
1078 - // Verify A device and strides
1079 - TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
1080 - TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
1081 - TORCH_CHECK(a.dtype() == torch::kFloat16,
1082 - "A is not float16, currently only float16 is supported");
1083 -
1084 - // Verify B device and strides
1085 - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
1086 - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
1087 -
1088 - // Verify b_meta device and strides
1089 - TORCH_CHECK(b_meta.device().is_cuda(), "b_meta is not on GPU");
1090 - TORCH_CHECK(b_meta.is_contiguous(), "b_meta is not contiguous");
1091 -
1092 - // Verify scales device and strides
1093 - TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
1094 - TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
1095 - TORCH_CHECK(b_scales.dtype() == torch::kFloat16,
1096 - "A is not float16, currently only float16 is supported");
1097 -
1098 - // Alloc C matrix
1099 - const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
1100 - auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
1101 - torch::Tensor c = torch::empty({size_m, size_n}, options);
1102 -
1103 - int thread_k = -1;
1104 - int thread_m = -1;
1105 - int sms = -1;
1106 - int max_par = marlin_24::max_par;
1107 -
1108 - int groupsize = -1;
1109 - if (b_scales.size(0) > 1) {
1110 - TORCH_CHECK(size_k % b_scales.size(0) == 0,
1111 - "size_k = " + str(size_k) +
1112 - ", is not divisible by b_scales.size(0) = " +
1113 - str(b_scales.size(0)));
1114 - groupsize = size_k / b_scales.size(0);
1115 - groupsize /= 2; // Because of 24
1116 - }
1117 -
1118 - // Verify groupsize
1119 - TORCH_CHECK(groupsize == -1 || groupsize == 64,
1120 - "Unexpected groupsize = " + str(groupsize));
1121 -
1122 - // Verify workspace size
1123 - TORCH_CHECK(size_n % marlin_24::min_thread_n == 0,
1124 - "size_n = " + str(size_n) +
1125 - ", is not divisible by min_thread_n = " +
1126 - str(marlin_24::min_thread_n));
1127 - int min_workspace_size =
1128 - (size_n / marlin_24::min_thread_n) * marlin_24::max_par;
1129 - TORCH_CHECK(workspace.numel() >= min_workspace_size,
1130 - "workspace.numel = " + str(workspace.numel()) +
1131 - " is below min_workspace_size = " + str(min_workspace_size));
1132 -
1133 - int dev = a.get_device();
1134 - marlin_24::marlin_cuda_2_4(
1135 - a.data_ptr(), b_q_weight.data_ptr(), b_meta.data_ptr(), c.data_ptr(),
1136 - b_scales.data_ptr(), size_n, size_m, size_k, workspace.data_ptr(),
1137 - b_q_type.size_bits(), groupsize, dev, at::cuda::getCurrentCUDAStream(dev),
1138 - thread_k, thread_m, sms, max_par);
1139 -
1140 - return c;
1141 -}
1142 -
1143 -TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1144 - m.impl("gptq_marlin_24_gemm", &gptq_marlin_24_gemm);
1145 -}
csrc/torch_bindings.cpp
-8
@@ -259,14 +259,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
259 // custom types:
260 // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
261
262 - // Marlin_24 (Sparse) Optimized Quantized GEMM for GPTQ.
263 - ops.def(
264 - "gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
265 - "Tensor b_scales, Tensor workspace, "
266 - "int b_q_type, "
267 - "SymInt size_m, SymInt size_n, SymInt size_k) -> Tensor");
268 - // conditionally compiled so impl in source file
269 -
262 // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
263 ops.def(
264 "machete_supported_schedules("
tests/compile/fullgraph/test_full_graph.py
-11
@@ -58,17 +58,6 @@ def models_list(*, all: bool = True, keywords: list[str] | None = None):
58 )
59 )
60
61 - if is_quant_method_supported("gptq_marlin_24"):
62 - TEST_MODELS.append(
63 - (
64 - "alexm-nm/tinyllama-24-marlin24-4bit-g128",
65 - {
66 - "quantization": "gptq_marlin_24",
67 - "allow_deprecated_quantization": True,
68 - },
69 - )
70 - )
71 -
61 if not current_platform.is_rocm() and is_quant_method_supported("awq"):
62 TEST_MODELS.append(
63 ("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {"quantization": "AWQ"})
tests/kernels/quantization/test_marlin_gemm.py
+2 -107
@@ -10,15 +10,9 @@ import itertools
10 import pytest
11 import torch
12
13 -from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
13 +from tests.kernels.utils import opcheck
14 from tests.quantization.utils import is_quant_method_supported
15 from vllm import _custom_ops as ops
16 -from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
17 - GPTQ_MARLIN_24_MAX_PARALLEL,
18 - GPTQ_MARLIN_24_MIN_THREAD_N,
19 - GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES,
20 - GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES,
21 -)
16 from vllm.model_executor.layers.quantization.utils.int8_utils import (
17 per_token_quant_int8,
18 )
@@ -36,15 +30,11 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
30 marlin_quant_fp8_torch,
31 )
32 from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
39 - MarlinWorkspace,
33 awq_marlin_quantize,
34 get_weight_perm,
35 marlin_quantize,
36 marlin_weights,
37 )
45 -from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
46 - marlin_24_quantize,
47 -)
38 from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 awq_pack,
40 gptq_pack,
@@ -57,9 +47,7 @@ from vllm.scalar_type import scalar_types
47
48 if current_platform.is_rocm():
49 pytest.skip(
60 - "These tests require gptq_marlin_repack,"
61 - "marlin_int4_fp8_preprocess, gptq_marlin_24_gemm,"
62 - "or marlin_gemm which are not supported on ROCm.",
50 + "These tests require marlin, which is not supported on ROCm.",
51 allow_module_level=True,
52 )
53
@@ -71,9 +59,6 @@ USE_FP32_REDUCE_OPTS = [True]
59 MARLIN_K_CHUNKS = [128]
60 MARLIN_N_CHUNKS = [64, 256]
61
74 -MARLIN_24_K_CHUNKS = [128]
75 -MARLIN_24_N_CHUNKS = [512]
76 -
62 MARLIN_REPACK_NK_FACTORS = [
63 (4, 8),
64 (7, 5),
@@ -538,96 +523,6 @@ def test_marlin_gemm(
523 assert max_diff < 0.04
524
525
541 -# TODO: find better way to test this?
542 -@torch.compile(fullgraph=True)
543 -def marlin_24_gemm_tester(
544 - a_input,
545 - marlin_24_q_w_comp,
546 - marlin_24_meta,
547 - marlin_24_s,
548 - scratch,
549 - quant_type,
550 - size_m,
551 - size_n,
552 - size_k,
553 -):
554 - return ops.gptq_marlin_24_gemm(
555 - a_input,
556 - marlin_24_q_w_comp,
557 - marlin_24_meta,
558 - marlin_24_s,
559 - scratch,
560 - quant_type,
561 - size_m,
562 - size_n,
563 - size_k,
564 - )
565 -
566 -
567 -@pytest.mark.skipif(
568 - not is_quant_method_supported("gptq_marlin"),
569 - reason="Marlin is not supported on this GPU type.",
570 -)
571 -@pytest.mark.parametrize("k_chunk", MARLIN_24_K_CHUNKS)
572 -@pytest.mark.parametrize("n_chunk", MARLIN_24_N_CHUNKS)
573 -@pytest.mark.parametrize("quant_type", GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)
574 -@pytest.mark.parametrize("group_size", GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES)
575 -@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
576 -def test_gptq_marlin_24_gemm(k_chunk, n_chunk, quant_type, group_size, mnk_factors):
577 - m_factor, n_factor, k_factor = mnk_factors
578 -
579 - size_m = m_factor
580 - size_k = k_chunk * k_factor
581 - size_n = n_chunk * n_factor
582 -
583 - a_input = rand_data((size_m, size_k))
584 - b_weight = rand_data((size_k, size_n))
585 -
586 - (w_24_ref, marlin_24_q_w_comp, marlin_24_meta, marlin_24_s) = marlin_24_quantize(
587 - b_weight, quant_type, group_size
588 - )
589 -
590 - workspace_24 = MarlinWorkspace(
591 - size_n, GPTQ_MARLIN_24_MIN_THREAD_N, GPTQ_MARLIN_24_MAX_PARALLEL
592 - )
593 -
594 - output_ref = torch.matmul(a_input, w_24_ref)
595 -
596 - opcheck(
597 - torch.ops._C.gptq_marlin_24_gemm,
598 - (
599 - a_input,
600 - marlin_24_q_w_comp,
601 - marlin_24_meta,
602 - marlin_24_s,
603 - workspace_24.scratch,
604 - quant_type.id,
605 - a_input.shape[0],
606 - b_weight.shape[1],
607 - a_input.shape[1],
608 - ),
609 - test_utils=DEFAULT_OPCHECK_TEST_UTILS,
610 - )
611 -
612 - output = marlin_24_gemm_tester(
613 - a_input,
614 - marlin_24_q_w_comp,
615 - marlin_24_meta,
616 - marlin_24_s,
617 - workspace_24.scratch,
618 - quant_type,
619 - a_input.shape[0],
620 - b_weight.shape[1],
621 - a_input.shape[1],
622 - )
623 -
624 - torch.cuda.synchronize()
625 -
626 - max_diff = compute_max_diff(output, output_ref)
627 -
628 - assert max_diff < 0.04
629 -
630 -
526 def test_marlin_gemm_subset_input():
527 quant_type = scalar_types.uint4b8
528 group_size = 128
tests/models/quantization/test_gptq_marlin_24.py deleted
-87
@@ -1,87 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Compare the outputs of a GPTQ model to a Marlin_24 model.
4 -
5 -Note: GPTQ and Marlin_24 do not have bitwise correctness.
6 -As a result, in this test, we just confirm that the top selected tokens of the
7 -Marlin/GPTQ models are in the top 3 selections of each other.
8 -"""
9 -
10 -from dataclasses import dataclass
11 -
12 -import pytest
13 -
14 -from tests.quantization.utils import is_quant_method_supported
15 -from vllm.platforms import current_platform
16 -
17 -from ..utils import check_logprobs_close
18 -
19 -
20 -@dataclass
21 -class ModelPair:
22 - model_marlin: str
23 - model_gptq: str
24 -
25 -
26 -model_pairs = [
27 - # 4-bit, group_size == 128
28 - ModelPair(
29 - model_marlin="alexm-nm/tinyllama-24-marlin24-4bit-g128",
30 - model_gptq="alexm-nm/tinyllama-24-gptq-4bit-g128",
31 - ),
32 - # # 4-bit, group_size == channelwise
33 - # ModelPair(model_marlin="alexm-nm/tinyllama-24-marlin24-4bit-channelwise",
34 - # model_gptq="alexm-nm/tinyllama-24-gptq-4bit-channelwise"),
35 - # 8-bit, group_size == 128
36 - ModelPair(
37 - model_marlin="alexm-nm/tinyllama-24-marlin24-8bit-g128",
38 - model_gptq="alexm-nm/tinyllama-24-gptq-8bit-g128",
39 - ),
40 - # # 8-bit, group_size == channelwise
41 - # ModelPair(model_marlin="alexm-nm/tinyllama-24-marlin24-8bit-channelwise",
42 - # model_gptq="alexm-nm/tinyllama-24-gptq-8bit-channelwise"),
43 -]
44 -
45 -
46 -@pytest.mark.flaky(reruns=2)
47 -@pytest.mark.skipif(
48 - not is_quant_method_supported("gptq_marlin_24")
49 - or current_platform.is_rocm()
50 - or not current_platform.is_cuda(),
51 - reason="Marlin24 is not supported on this GPU type.",
52 -)
53 -@pytest.mark.parametrize("model_pair", model_pairs)
54 -@pytest.mark.parametrize("dtype", ["half"])
55 -@pytest.mark.parametrize("max_tokens", [8])
56 -@pytest.mark.parametrize("num_logprobs", [5])
57 -def test_models(
58 - vllm_runner,
59 - example_prompts,
60 - model_pair: ModelPair,
61 - dtype: str,
62 - max_tokens: int,
63 - num_logprobs: int,
64 -) -> None:
65 - with vllm_runner(
66 - model_pair.model_marlin,
67 - dtype=dtype,
68 - quantization="gptq_marlin_24",
69 - allow_deprecated_quantization=True,
70 - ) as marlin_24_model:
71 - marlin_24_outputs = marlin_24_model.generate_greedy_logprobs(
72 - example_prompts, max_tokens, num_logprobs
73 - )
74 -
75 - with vllm_runner(
76 - model_pair.model_gptq, dtype=dtype, quantization="gptq"
77 - ) as gptq_model:
78 - gptq_outputs = gptq_model.generate_greedy_logprobs(
79 - example_prompts, max_tokens, num_logprobs
80 - )
81 -
82 - check_logprobs_close(
83 - outputs_0_lst=gptq_outputs,
84 - outputs_1_lst=marlin_24_outputs,
85 - name_0="gptq",
86 - name_1="marlin_24",
87 - )
tests/quantization/test_compressed_tensors.py
-23
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
17 CompressedTensorsW4A4Fp4,
18 CompressedTensorsW4A8Fp8,
19 CompressedTensorsW4A16Fp4,
20 - CompressedTensorsW4A16Sparse24,
20 CompressedTensorsW8A8Fp8,
21 CompressedTensorsW8A8Int8,
22 CompressedTensorsW8A16Fp8,
@@ -307,28 +306,6 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
306 assert output
307
308
310 -@pytest.mark.skipif(
311 - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
312 -)
313 -def test_compressed_tensors_w4a16_marlin24(vllm_runner):
314 - model_path = "nm-testing/llama7b-one-shot-2_4-w4a16-marlin24-t"
315 - with vllm_runner(model_path, enforce_eager=True) as llm:
316 -
317 - def check_model(model):
318 - layer = model.model.layers[0]
319 -
320 - qkv_proj = layer.self_attn.qkv_proj
321 -
322 - assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
323 - assert isinstance(qkv_proj.scheme, CompressedTensorsW4A16Sparse24)
324 - assert qkv_proj.weight_packed.dtype is torch.int32
325 -
326 - llm.apply_model(check_model)
327 -
328 - output = llm.generate_greedy("Hello my name is", max_tokens=4)
329 - assert output
330 -
331 -
309 def test_compressed_tensors_fp8(vllm_runner):
310 model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test"
311 with vllm_runner(model_path, enforce_eager=True) as llm:
vllm/_custom_ops.py
+143 -150
@@ -499,6 +499,23 @@ def awq_dequantize(
499 return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy)
500
501
502 +if hasattr(torch.ops._C, "awq_dequantize"):
503 +
504 + @register_fake("_C::awq_dequantize")
505 + def _awq_dequantize_fake(
506 + qweight: torch.Tensor,
507 + scales: torch.Tensor,
508 + zeros: torch.Tensor,
509 + split_k_iters: torch.SymInt,
510 + thx: int,
511 + thy: int,
512 + ) -> torch.Tensor:
513 + in_c = qweight.size(0)
514 + qout_c = qweight.size(1)
515 + out_c = qout_c * 8
516 + return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device)
517 +
518 +
519 def awq_gemm(
520 input: torch.Tensor,
521 qweight: torch.Tensor,
@@ -513,6 +530,24 @@ def awq_gemm(
530 return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters)
531
532
533 +if hasattr(torch.ops._C, "awq_gemm"):
534 +
535 + @register_fake("_C::awq_gemm")
536 + def _awq_gemm_fake(
537 + input: torch.Tensor,
538 + qweight: torch.Tensor,
539 + scales: torch.Tensor,
540 + qzeros: torch.Tensor,
541 + split_k_iters: torch.SymInt,
542 + ) -> torch.Tensor:
543 + num_in_feats = input.size(0)
544 + return torch.empty(
545 + (split_k_iters, num_in_feats, qweight.size(1) * 8),
546 + dtype=input.dtype,
547 + device=input.device,
548 + ).sum(0)
549 +
550 +
551 # gptq
552 def gptq_gemm(
553 a: torch.Tensor,
@@ -558,152 +593,6 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor, bit: int) -> None
593 torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
594
595
561 -# marlin_24
562 -def gptq_marlin_24_gemm(
563 - a: torch.Tensor,
564 - b_q_weight: torch.Tensor,
565 - b_meta: torch.Tensor,
566 - b_scales: torch.Tensor,
567 - workspace: torch.Tensor,
568 - b_q_type: ScalarType,
569 - size_m: int,
570 - size_n: int,
571 - size_k: int,
572 -) -> torch.Tensor:
573 - return torch.ops._C.gptq_marlin_24_gemm(
574 - a, b_q_weight, b_meta, b_scales, workspace, b_q_type.id, size_m, size_n, size_k
575 - )
576 -
577 -
578 -if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
579 -
580 - @register_fake("_C::gptq_marlin_24_gemm")
581 - def _gptq_marlin_24_gemm_fake(
582 - a: torch.Tensor,
583 - b_q_weight: torch.Tensor,
584 - b_meta: torch.Tensor,
585 - b_scales: torch.Tensor,
586 - workspace: torch.Tensor,
587 - b_q_type: ScalarType,
588 - size_m: torch.SymInt,
589 - size_n: torch.SymInt,
590 - size_k: torch.SymInt,
591 - ) -> torch.Tensor:
592 - return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
593 -
594 - @register_fake("_C::marlin_gemm")
595 - def _marlin_gemm_fake(
596 - a: torch.Tensor,
597 - c: torch.Tensor | None,
598 - b_q_weight: torch.Tensor,
599 - b_bias: torch.Tensor | None,
600 - b_scales: torch.Tensor,
601 - a_scales: torch.Tensor | None,
602 - global_scale: torch.Tensor | None,
603 - b_zeros: torch.Tensor | None,
604 - g_idx: torch.Tensor | None,
605 - perm: torch.Tensor | None,
606 - workspace: torch.Tensor,
607 - b_q_type_id: int,
608 - size_m: torch.SymInt,
609 - size_n: torch.SymInt,
610 - size_k: torch.SymInt,
611 - is_k_full: bool = True,
612 - use_atomic_add: bool = False,
613 - use_fp32_reduce: bool = False,
614 - is_zp_float: bool = False,
615 - ) -> torch.Tensor:
616 - dtype = a.dtype
617 - if dtype not in [torch.half, torch.bfloat16]:
618 - dtype = b_scales.dtype
619 - return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
620 -
621 - @register_fake("_C::awq_dequantize")
622 - def _awq_dequantize_fake(
623 - qweight: torch.Tensor,
624 - scales: torch.Tensor,
625 - zeros: torch.Tensor,
626 - split_k_iters: torch.SymInt,
627 - thx: int,
628 - thy: int,
629 - ) -> torch.Tensor:
630 - in_c = qweight.size(0)
631 - qout_c = qweight.size(1)
632 - out_c = qout_c * 8
633 - return torch.empty((in_c, out_c), dtype=scales.dtype, device=scales.device)
634 -
635 - @register_fake("_C::awq_gemm")
636 - def _awq_gemm_fake(
637 - input: torch.Tensor,
638 - qweight: torch.Tensor,
639 - scales: torch.Tensor,
640 - qzeros: torch.Tensor,
641 - split_k_iters: torch.SymInt,
642 - ) -> torch.Tensor:
643 - num_in_feats = input.size(0)
644 - return torch.empty(
645 - (split_k_iters, num_in_feats, qweight.size(1) * 8),
646 - dtype=input.dtype,
647 - device=input.device,
648 - ).sum(0)
649 -
650 - @register_fake("_C::machete_mm")
651 - def machete_mm_fake(
652 - a: torch.Tensor,
653 - # b_q Should be the tensor returned by machete_prepack_B
654 - b_q: torch.Tensor,
655 - b_type: ScalarType,
656 - out_type: torch.dtype | None = None,
657 - b_group_scales: torch.Tensor | None = None,
658 - b_group_zeros: torch.Tensor | None = None,
659 - b_group_size: int | None = None,
660 - b_channel_scales: torch.Tensor | None = None,
661 - a_token_scales: torch.Tensor | None = None,
662 - schedule: str | None = None,
663 - ) -> torch.Tensor:
664 - m = a.size(0)
665 - n = b_q.size(1)
666 - return torch.empty((m, n), device=a.device, dtype=a.dtype)
667 -
668 - @register_fake("_C::machete_prepack_B")
669 - def machete_prepack_B_fake(
670 - b_q_weight: torch.Tensor,
671 - a_type: torch.dtype,
672 - b_type: ScalarType,
673 - group_scales_type: torch.dtype | None,
674 - ) -> torch.Tensor:
675 - return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format)
676 -
677 - @register_fake("_C::cutlass_w4a8_mm")
678 - def cutlass_w4a8_mm_fake(
679 - a: torch.Tensor,
680 - # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
681 - b_q: torch.Tensor,
682 - b_group_scales: torch.Tensor,
683 - b_group_size: int,
684 - b_channel_scales: torch.Tensor,
685 - a_token_scales: torch.Tensor,
686 - out_type: torch.dtype | None = None,
687 - maybe_schedule: str | None = None,
688 - ) -> torch.Tensor:
689 - m = a.size(0)
690 - n = b_q.size(1)
691 - out_dtype = out_type if out_type is not None else torch.bfloat16
692 - return torch.empty((m, n), device=a.device, dtype=out_dtype)
693 -
694 - @register_fake("_C::cutlass_pack_scale_fp8")
695 - def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor:
696 - return torch.empty_like(scales, memory_format=torch.contiguous_format)
697 -
698 - @register_fake("_C::cutlass_encode_and_reorder_int4b")
699 - def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
700 - return torch.empty_like(b, memory_format=torch.contiguous_format)
701 -
702 - @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped")
703 - def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor:
704 - return torch.empty_like(b, memory_format=torch.contiguous_format)
705 -
706 -
596 if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
597
598 @register_fake("_C::allspark_w8a16_gemm")
@@ -1356,6 +1245,36 @@ def marlin_gemm(
1245 )
1246
1247
1248 +if hasattr(torch.ops._C, "marlin_gemm"):
1249 +
1250 + @register_fake("_C::marlin_gemm")
1251 + def _marlin_gemm_fake(
1252 + a: torch.Tensor,
1253 + c: torch.Tensor | None,
1254 + b_q_weight: torch.Tensor,
1255 + b_bias: torch.Tensor | None,
1256 + b_scales: torch.Tensor,
1257 + a_scales: torch.Tensor | None,
1258 + global_scale: torch.Tensor | None,
1259 + b_zeros: torch.Tensor | None,
1260 + g_idx: torch.Tensor | None,
1261 + perm: torch.Tensor | None,
1262 + workspace: torch.Tensor,
1263 + b_q_type_id: int,
1264 + size_m: torch.SymInt,
1265 + size_n: torch.SymInt,
1266 + size_k: torch.SymInt,
1267 + is_k_full: bool = True,
1268 + use_atomic_add: bool = False,
1269 + use_fp32_reduce: bool = False,
1270 + is_zp_float: bool = False,
1271 + ) -> torch.Tensor:
1272 + dtype = a.dtype
1273 + if dtype not in [torch.half, torch.bfloat16]:
1274 + dtype = b_scales.dtype
1275 + return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
1276 +
1277 +
1278 # machete
1279 def machete_supported_schedules(
1280 a_type: torch.dtype,
@@ -1404,6 +1323,27 @@ def machete_mm(
1323 )
1324
1325
1326 +if hasattr(torch.ops._C, "machete_mm"):
1327 +
1328 + @register_fake("_C::machete_mm")
1329 + def machete_mm_fake(
1330 + a: torch.Tensor,
1331 + # b_q Should be the tensor returned by machete_prepack_B
1332 + b_q: torch.Tensor,
1333 + b_type: ScalarType,
1334 + out_type: torch.dtype | None = None,
1335 + b_group_scales: torch.Tensor | None = None,
1336 + b_group_zeros: torch.Tensor | None = None,
1337 + b_group_size: int | None = None,
1338 + b_channel_scales: torch.Tensor | None = None,
1339 + a_token_scales: torch.Tensor | None = None,
1340 + schedule: str | None = None,
1341 + ) -> torch.Tensor:
1342 + m = a.size(0)
1343 + n = b_q.size(1)
1344 + return torch.empty((m, n), device=a.device, dtype=a.dtype)
1345 +
1346 +
1347 def machete_prepack_B(
1348 b_q_weight: torch.Tensor,
1349 a_type: torch.dtype,
@@ -1415,6 +1355,18 @@ def machete_prepack_B(
1355 )
1356
1357
1358 +if hasattr(torch.ops._C, "machete_prepack_B"):
1359 +
1360 + @register_fake("_C::machete_prepack_B")
1361 + def machete_prepack_B_fake(
1362 + b_q_weight: torch.Tensor,
1363 + a_type: torch.dtype,
1364 + b_type: ScalarType,
1365 + group_scales_type: torch.dtype | None,
1366 + ) -> torch.Tensor:
1367 + return torch.empty_like(b_q_weight, memory_format=torch.contiguous_format)
1368 +
1369 +
1370 # CUTLASS W4A8
1371 def cutlass_w4a8_mm(
1372 a: torch.Tensor,
@@ -1439,14 +1391,48 @@ def cutlass_w4a8_mm(
1391 )
1392
1393
1394 +if hasattr(torch.ops._C, "cutlass_w4a8_mm"):
1395 +
1396 + @register_fake("_C::cutlass_w4a8_mm")
1397 + def cutlass_w4a8_mm_fake(
1398 + a: torch.Tensor,
1399 + # b_q Should be the tensor returned by cutlass_encode_and_reorder_int4b
1400 + b_q: torch.Tensor,
1401 + b_group_scales: torch.Tensor,
1402 + b_group_size: int,
1403 + b_channel_scales: torch.Tensor,
1404 + a_token_scales: torch.Tensor,
1405 + out_type: torch.dtype | None = None,
1406 + maybe_schedule: str | None = None,
1407 + ) -> torch.Tensor:
1408 + m = a.size(0)
1409 + n = b_q.size(1)
1410 + out_dtype = out_type if out_type is not None else torch.bfloat16
1411 + return torch.empty((m, n), device=a.device, dtype=out_dtype)
1412 +
1413 +
1414 def cutlass_pack_scale_fp8(scales: torch.Tensor) -> torch.Tensor:
1415 return torch.ops._C.cutlass_pack_scale_fp8(scales)
1416
1417
1418 +if hasattr(torch.ops._C, "cutlass_pack_scale_fp8"):
1419 +
1420 + @register_fake("_C::cutlass_pack_scale_fp8")
1421 + def cutlass_pack_scale_fp8_fake(scales: torch.Tensor) -> torch.Tensor:
1422 + return torch.empty_like(scales, memory_format=torch.contiguous_format)
1423 +
1424 +
1425 def cutlass_encode_and_reorder_int4b(b: torch.Tensor) -> torch.Tensor:
1426 return torch.ops._C.cutlass_encode_and_reorder_int4b(b)
1427
1428
1429 +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b"):
1430 +
1431 + @register_fake("_C::cutlass_encode_and_reorder_int4b")
1432 + def cutlass_encode_and_reorder_int4b_fake(b: torch.Tensor) -> torch.Tensor:
1433 + return torch.empty_like(b, memory_format=torch.contiguous_format)
1434 +
1435 +
1436 def cutlass_w4a8_moe_mm(
1437 out_tensors: torch.Tensor,
1438 a_tensors: torch.Tensor,
@@ -1519,17 +1505,24 @@ def cutlass_encode_and_reorder_int4b_grouped(
1505 return torch.ops._C.cutlass_encode_and_reorder_int4b_grouped(b_tensors)
1506
1507
1522 -if hasattr(torch.ops._C, "permute_cols"):
1508 +if hasattr(torch.ops._C, "cutlass_encode_and_reorder_int4b_grouped"):
1509
1524 - @register_fake("_C::permute_cols")
1525 - def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1526 - return torch.empty_like(a)
1510 + @register_fake("_C::cutlass_encode_and_reorder_int4b_grouped")
1511 + def cutlass_encode_and_reorder_int4b_grouped_fake(b: torch.Tensor) -> torch.Tensor:
1512 + return torch.empty_like(b, memory_format=torch.contiguous_format)
1513
1514
1515 def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1516 return torch.ops._C.permute_cols(a, perm)
1517
1518
1519 +if hasattr(torch.ops._C, "permute_cols"):
1520 +
1521 + @register_fake("_C::permute_cols")
1522 + def _permute_cols_fake(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
1523 + return torch.empty_like(a)
1524 +
1525 +
1526 # fp4
1527 def scaled_fp4_quant(
1528 input: torch.Tensor,
vllm/config/model.py
-1
@@ -890,7 +890,6 @@ class ModelConfig:
890 # `override_quantization_method` method) must be checked in order
891 # of preference (this is particularly important for GPTQ).
892 overrides = [
893 - "gptq_marlin_24",
893 "gptq_marlin",
894 "awq_marlin",
895 "ipex",
vllm/model_executor/layers/quantization/__init__.py
-4
@@ -18,7 +18,6 @@ QuantizationMethods = Literal[
18 "modelopt",
19 "modelopt_fp4",
20 "gguf",
21 - "gptq_marlin_24",
21 "gptq_marlin",
22 "awq_marlin",
23 "gptq",
@@ -41,7 +40,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
40 "ptpc_fp8",
41 "fbgemm_fp8",
42 "fp_quant",
44 - "gptq_marlin_24",
43 "experts_int8",
44 "ipex",
45 "petit_nvfp4",
@@ -122,7 +120,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
120 from .gguf import GGUFConfig
121 from .gptq import GPTQConfig
122 from .gptq_marlin import GPTQMarlinConfig
125 - from .gptq_marlin_24 import GPTQMarlin24Config
123 from .inc import INCConfig
124 from .ipex_quant import IPEXConfig
125 from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config
@@ -140,7 +137,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
137 "modelopt": ModelOptFp8Config,
138 "modelopt_fp4": ModelOptNvFp4Config,
139 "gguf": GGUFConfig,
143 - "gptq_marlin_24": GPTQMarlin24Config,
140 "gptq_marlin": GPTQMarlinConfig,
141 "awq_marlin": AWQMarlinConfig,
142 "gptq": GPTQConfig,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+13 -25
@@ -40,7 +40,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
40 CompressedTensorsMoEMethod,
41 )
42 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
43 - W4A16SPARSE24_SUPPORTED_BITS,
43 WNA16_SUPPORTED_BITS,
44 CompressedTensors24,
45 CompressedTensorsScheme,
@@ -49,7 +48,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
48 CompressedTensorsW4A8Int,
49 CompressedTensorsW4A16Fp4,
50 CompressedTensorsW4A16Mxfp4,
52 - CompressedTensorsW4A16Sparse24,
51 CompressedTensorsW8A8Fp8,
52 CompressedTensorsW8A8Int8,
53 CompressedTensorsW8A16Fp8,
@@ -610,29 +608,19 @@ class CompressedTensorsConfig(QuantizationConfig):
608 actorder=weight_quant.actorder,
609 )
610
613 - if self._is_wNa16_group_channel(weight_quant, input_quant):
614 - if (
615 - format == CompressionFormat.marlin_24.value
616 - and weight_quant.num_bits in W4A16SPARSE24_SUPPORTED_BITS
617 - ):
618 - assert weight_quant.symmetric
619 - return CompressedTensorsW4A16Sparse24(
620 - strategy=weight_quant.strategy,
621 - num_bits=weight_quant.num_bits,
622 - group_size=weight_quant.group_size,
623 - )
624 - if (
625 - format == CompressionFormat.pack_quantized.value
626 - and weight_quant.num_bits in WNA16_SUPPORTED_BITS
627 - ):
628 - return CompressedTensorsWNA16(
629 - num_bits=weight_quant.num_bits,
630 - strategy=weight_quant.strategy,
631 - symmetric=weight_quant.symmetric,
632 - group_size=weight_quant.group_size,
633 - actorder=weight_quant.actorder,
634 - layer_name=layer_name,
635 - )
611 + if (
612 + self._is_wNa16_group_channel(weight_quant, input_quant)
613 + and (format == CompressionFormat.pack_quantized.value)
614 + and (weight_quant.num_bits in WNA16_SUPPORTED_BITS)
615 + ):
616 + return CompressedTensorsWNA16(
617 + num_bits=weight_quant.num_bits,
618 + strategy=weight_quant.strategy,
619 + symmetric=weight_quant.symmetric,
620 + group_size=weight_quant.group_size,
621 + actorder=weight_quant.actorder,
622 + layer_name=layer_name,
623 + )
624
625 act_quant_format = is_activation_quantization_format(format)
626 if act_quant_format:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
-6
@@ -5,10 +5,6 @@ from .compressed_tensors_scheme import CompressedTensorsScheme
5 from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
6 from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
7 from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
8 -from .compressed_tensors_w4a16_24 import (
9 - W4A16SPARSE24_SUPPORTED_BITS,
10 - CompressedTensorsW4A16Sparse24,
11 -)
8 from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4
9 from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
10 from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
@@ -23,11 +19,9 @@ __all__ = [
19 "CompressedTensorsScheme",
20 "CompressedTensorsWNA16",
21 "CompressedTensorsW8A16Fp8",
26 - "CompressedTensorsW4A16Sparse24",
22 "CompressedTensorsW8A8Int8",
23 "CompressedTensorsW8A8Fp8",
24 "WNA16_SUPPORTED_BITS",
30 - "W4A16SPARSE24_SUPPORTED_BITS",
25 "CompressedTensors24",
26 "CompressedTensorsW4A16Fp4",
27 "CompressedTensorsW4A16Mxfp4",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_24.py deleted
-176
@@ -1,176 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from collections.abc import Callable
5 -
6 -import torch
7 -from torch.nn import Parameter
8 -
9 -from vllm import _custom_ops as ops
10 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11 - CompressedTensorsScheme,
12 -)
13 -from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
14 - GPTQ_MARLIN_24_MAX_PARALLEL,
15 - GPTQ_MARLIN_24_MIN_THREAD_N,
16 -)
17 -from vllm.model_executor.parameter import (
18 - BasevLLMParameter,
19 - ChannelQuantScaleParameter,
20 - GroupQuantScaleParameter,
21 - PackedvLLMParameter,
22 -)
23 -from vllm.scalar_type import scalar_types
24 -
25 -__all__ = ["CompressedTensorsW4A16Sparse24"]
26 -W4A16SPARSE24_SUPPORTED_TYPES_MAP = {
27 - 4: scalar_types.uint4b8,
28 -}
29 -W4A16SPARSE24_SUPPORTED_BITS = list(W4A16SPARSE24_SUPPORTED_TYPES_MAP.keys())
30 -
31 -
32 -class CompressedTensorsW4A16Sparse24(CompressedTensorsScheme):
33 - def __init__(self, strategy: str, num_bits: int, group_size: int | None = None):
34 - self.strategy = strategy
35 - self.group_size = group_size
36 - self.tile_size = 16
37 -
38 - if num_bits not in W4A16SPARSE24_SUPPORTED_TYPES_MAP:
39 - raise ValueError(
40 - f"Unsupported num_bits = {num_bits}. "
41 - f"Supported num_bits = {W4A16SPARSE24_SUPPORTED_BITS}"
42 - )
43 -
44 - self.quant_type = W4A16SPARSE24_SUPPORTED_TYPES_MAP[num_bits]
45 -
46 - if self.strategy == "group" and self.group_size is None:
47 - raise ValueError("group_size must be given when using strategy group")
48 -
49 - @classmethod
50 - def get_min_capability(cls) -> int:
51 - # ampere + up
52 - return 80
53 -
54 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
55 - # required by torch.compile to be torch.nn.Parameter
56 - layer.weight_packed = Parameter(layer.weight_packed.data, requires_grad=False)
57 - layer.scale_packed = Parameter(layer.scale_packed.data, requires_grad=False)
58 - layer.meta = Parameter(layer.meta.data, requires_grad=False)
59 -
60 - def create_weights(
61 - self,
62 - layer: torch.nn.Module,
63 - input_size: int,
64 - output_partition_sizes: list[int],
65 - input_size_per_partition: int,
66 - params_dtype: torch.dtype,
67 - weight_loader: Callable,
68 - **kwargs,
69 - ):
70 - assert params_dtype == torch.float16, (
71 - "float16 is required for marlin24 compressed models. Set dtype=torch.float16" # noqa: E501
72 - )
73 -
74 - pack_factor = 32 // self.quant_type.size_bits
75 - output_size_per_partition = sum(output_partition_sizes)
76 -
77 - qweight = PackedvLLMParameter(
78 - data=torch.empty(
79 - input_size_per_partition // self.tile_size // 2,
80 - output_size_per_partition * self.tile_size // pack_factor,
81 - dtype=torch.int32,
82 - ),
83 - input_dim=0,
84 - output_dim=1,
85 - packed_dim=1,
86 - packed_factor=pack_factor,
87 - marlin_tile_size=self.tile_size,
88 - weight_loader=weight_loader,
89 - )
90 -
91 - input_groups = (
92 - 1
93 - if self.group_size is None
94 - else input_size_per_partition // self.group_size
95 - )
96 -
97 - weight_scale_args = {
98 - "data": torch.empty(
99 - input_groups,
100 - output_size_per_partition,
101 - dtype=params_dtype,
102 - ),
103 - "weight_loader": weight_loader,
104 - }
105 -
106 - if self.group_size is not None:
107 - scales = GroupQuantScaleParameter(
108 - output_dim=1, input_dim=0, **weight_scale_args
109 - )
110 - else:
111 - scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
112 -
113 - weight_shape = BasevLLMParameter(
114 - data=torch.empty(2, dtype=torch.int64), weight_loader=weight_loader
115 - )
116 -
117 - meta = PackedvLLMParameter(
118 - data=torch.empty(
119 - input_size_per_partition // 8 // 2 // 2,
120 - output_size_per_partition * 2,
121 - dtype=torch.int16,
122 - ),
123 - input_dim=0,
124 - output_dim=1,
125 - packed_dim=1,
126 - packed_factor=1,
127 - marlin_tile_size=2,
128 - weight_loader=weight_loader,
129 - )
130 -
131 - layer.register_parameter("weight_packed", qweight)
132 - layer.register_parameter("weight_shape", weight_shape)
133 - layer.register_parameter("scale_packed", scales)
134 - layer.register_parameter("meta", meta)
135 -
136 - max_workspace_size = (
137 - output_size_per_partition // GPTQ_MARLIN_24_MIN_THREAD_N
138 - ) * GPTQ_MARLIN_24_MAX_PARALLEL
139 -
140 - workspace = Parameter(
141 - torch.zeros(max_workspace_size, dtype=torch.int), requires_grad=False
142 - )
143 - layer.workspace = workspace
144 -
145 - def apply_weights(
146 - self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None
147 - ) -> torch.Tensor:
148 - qweight = layer.weight_packed
149 - meta = layer.meta
150 - scales = layer.scale_packed
151 - workspace = layer.workspace
152 -
153 - x_2d = x.view(-1, x.shape[-1])
154 -
155 - size_m = x_2d.shape[0]
156 - size_k = x_2d.shape[1]
157 - size_n = scales.shape[1]
158 -
159 - output_2d = ops.gptq_marlin_24_gemm(
160 - x_2d,
161 - qweight,
162 - meta,
163 - scales,
164 - workspace,
165 - self.quant_type,
166 - size_m,
167 - size_n,
168 - size_k,
169 - )
170 -
171 - output = output_2d.view(x.shape[:-1] + (output_2d.shape[1],))
172 -
173 - if bias is not None:
174 - output.add_(bias) # In-place add
175 -
176 - return output
vllm/model_executor/layers/quantization/gptq_marlin_24.py deleted
-320
@@ -1,320 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import Any, Optional
5 -
6 -import torch
7 -from torch.nn.parameter import Parameter
8 -
9 -from vllm import _custom_ops as ops
10 -from vllm.logger import init_logger
11 -from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
12 -from vllm.model_executor.layers.quantization import (
13 - QuantizationConfig,
14 - QuantizationMethods,
15 -)
16 -from vllm.model_executor.parameter import (
17 - BasevLLMParameter,
18 - ChannelQuantScaleParameter,
19 - GroupQuantScaleParameter,
20 - PackedvLLMParameter,
21 -)
22 -from vllm.scalar_type import scalar_types
23 -
24 -logger = init_logger(__name__)
25 -
26 -GPTQ_MARLIN_24_TILE = 16
27 -GPTQ_MARLIN_24_MIN_THREAD_N = 128
28 -GPTQ_MARLIN_24_MIN_THREAD_K = 128
29 -GPTQ_MARLIN_24_MAX_PARALLEL = 64
30 -
31 -GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
32 -GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES = [-1, 128]
33 -
34 -
35 -class GPTQMarlin24Config(QuantizationConfig):
36 - """Config class for Marlin24."""
37 -
38 - def __init__(
39 - self,
40 - weight_bits: int,
41 - group_size: int,
42 - ) -> None:
43 - super().__init__()
44 - quant_type = {
45 - 4: scalar_types.uint4b8,
46 - 8: scalar_types.uint8b128,
47 - }.get(weight_bits)
48 -
49 - self.group_size = group_size
50 -
51 - # Verify
52 - if quant_type is None or quant_type not in GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES:
53 - raise ValueError(
54 - f"Marlin_24 does not support quant_type = {quant_type}. "
55 - f"Only weight_bits = {GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES} "
56 - "are supported."
57 - )
58 - if self.group_size not in GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES:
59 - raise ValueError(
60 - f"Marlin_24 does not support group_size = {self.group_size}. "
61 - f"Only group_sizes = {GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES} "
62 - "are supported."
63 - )
64 -
65 - self.quant_type = quant_type
66 -
67 - # 4 Bits packed into 32 bit datatype.
68 - self.pack_factor = 32 // self.quant_type.size_bits
69 -
70 - # Tile size used by marlin kernels.
71 - self.tile_size = 16
72 -
73 - # Min out_features dim
74 - self.min_n_threads = GPTQ_MARLIN_24_MIN_THREAD_N
75 -
76 - # Min in_features dim
77 - self.min_k_threads = GPTQ_MARLIN_24_MIN_THREAD_K
78 -
79 - # Max parallel problems to solve at once (improves large
80 - # batch performance)
81 - self.max_parallel = GPTQ_MARLIN_24_MAX_PARALLEL
82 -
83 - # Permutation length used by the marlin kernels.
84 - self.perm_len = 1024
85 -
86 - def __repr__(self) -> str:
87 - return "Marlin24Config(quant_type={}, group_size={})".format(
88 - self.quant_type, self.group_size
89 - )
90 -
91 - @classmethod
92 - def get_name(cls) -> QuantizationMethods:
93 - return "gptq_marlin_24"
94 -
95 - @classmethod
96 - def get_supported_act_dtypes(cls) -> list[torch.dtype]:
97 - return [torch.half]
98 -
99 - @classmethod
100 - # Need to figure it out
101 - def get_min_capability(cls) -> int:
102 - return 80
103 -
104 - @classmethod
105 - def get_config_filenames(cls) -> list[str]:
106 - return ["quantize_config.json"]
107 -
108 - @classmethod
109 - def from_config(cls, config: dict[str, Any]) -> "GPTQMarlin24Config":
110 - weight_bits = cls.get_from_keys(config, ["bits"])
111 - group_size = cls.get_from_keys(config, ["group_size"])
112 - return cls(weight_bits, group_size)
113 -
114 - @classmethod
115 - def override_quantization_method(
116 - cls, hf_quant_cfg, user_quant
117 - ) -> QuantizationMethods | None:
118 - is_marlin_24_format = hf_quant_cfg.get("checkpoint_format") == "marlin_24"
119 -
120 - is_valid_user_quant = (
121 - user_quant is None or user_quant == "gptq" or user_quant == "gptq_marlin_24"
122 - )
123 -
124 - if is_marlin_24_format and is_valid_user_quant:
125 - msg = "The model is serialized in {} format. Using {} kernel.".format(
126 - cls.get_name(), cls.get_name()
127 - )
128 - logger.info(msg)
129 - return cls.get_name()
130 -
131 - return None
132 -
133 - def get_quant_method(
134 - self, layer: torch.nn.Module, prefix: str
135 - ) -> Optional["GPTQMarlin24LinearMethod"]:
136 - if isinstance(layer, LinearBase):
137 - return GPTQMarlin24LinearMethod(self)
138 - return None
139 -
140 -
141 -class GPTQMarlin24LinearMethod(LinearMethodBase):
142 - """Linear method for Marlin24.
143 -
144 - Args:
145 - quant_config: The Marlin24 quantization config.
146 - """
147 -
148 - def __init__(self, quant_config: GPTQMarlin24Config):
149 - self.quant_config = quant_config
150 -
151 - def create_weights(
152 - self,
153 - layer: torch.nn.Module,
154 - input_size_per_partition: int,
155 - output_partition_sizes: list[int],
156 - input_size: int,
157 - output_size: int,
158 - params_dtype: torch.dtype,
159 - **extra_weight_attrs,
160 - ):
161 - del output_size # Unused.
162 - weight_loader = extra_weight_attrs["weight_loader"]
163 - if params_dtype != torch.float16:
164 - raise ValueError(
165 - f"The params dtype must be float16, but got {params_dtype}"
166 - )
167 -
168 - # Validate output_size_per_partition
169 - output_size_per_partition = sum(output_partition_sizes)
170 - if output_size_per_partition % self.quant_config.min_n_threads != 0:
171 - raise ValueError(
172 - f"Weight output_size_per_partition = "
173 - f"{output_size_per_partition} is not divisible by "
174 - f"min_n_threads = {self.quant_config.min_n_threads}."
175 - )
176 - if output_size_per_partition % self.quant_config.pack_factor != 0:
177 - raise ValueError(
178 - f"Weight output_size_per_partition = "
179 - f"{output_size_per_partition} is not divisible by "
180 - f"pack_factor = {self.quant_config.pack_factor}."
181 - )
182 -
183 - # Validate input_size_per_partition
184 - if input_size_per_partition % self.quant_config.min_k_threads != 0:
185 - raise ValueError(
186 - f"Weight input_size_per_partition = "
187 - f"{input_size_per_partition} is not divisible by "
188 - f"min_k_threads = {self.quant_config.min_k_threads}."
189 - )
190 - if (
191 - self.quant_config.group_size != -1
192 - and input_size_per_partition % self.quant_config.group_size != 0
193 - ):
194 - raise ValueError(
195 - f"Weight input_size_per_partition = "
196 - f"{input_size_per_partition} is not divisible by "
197 - f"group_size = {self.quant_config.group_size}."
198 - )
199 -
200 - # Check that we have at least 4 tiles horizontally in the shard
201 - num_tiles_per_perm = self.quant_config.perm_len // (
202 - self.quant_config.tile_size**2
203 - )
204 - if output_size_per_partition % num_tiles_per_perm != 0:
205 - raise ValueError("Each permutation group must reside on the same gpu")
206 -
207 - # Quantized 4Bit weights packed into Int32.
208 - qweight = PackedvLLMParameter(
209 - data=torch.empty(
210 - input_size_per_partition // self.quant_config.tile_size // 2,
211 - output_size_per_partition
212 - * self.quant_config.tile_size
213 - // self.quant_config.pack_factor,
214 - device="cuda",
215 - dtype=torch.int32,
216 - ),
217 - input_dim=0,
218 - output_dim=1,
219 - packed_dim=1,
220 - packed_factor=self.quant_config.pack_factor,
221 - marlin_tile_size=self.quant_config.tile_size,
222 - weight_loader=weight_loader,
223 - )
224 -
225 - # Meta
226 - meta = PackedvLLMParameter(
227 - data=torch.empty(
228 - input_size_per_partition // 8 // 2 // 2,
229 - output_size_per_partition * 2,
230 - device="cuda",
231 - dtype=torch.int16,
232 - ),
233 - input_dim=0,
234 - output_dim=1,
235 - packed_dim=1,
236 - packed_factor=1,
237 - marlin_tile_size=2,
238 - weight_loader=weight_loader,
239 - )
240 -
241 - # Determine if channelwise or not
242 - input_groups = (
243 - 1
244 - if self.quant_config.group_size == -1
245 - else input_size_per_partition // self.quant_config.group_size
246 - )
247 -
248 - weight_scale_args = {
249 - "data": torch.empty(
250 - input_groups,
251 - output_size_per_partition,
252 - device="cuda",
253 - dtype=params_dtype,
254 - ),
255 - "weight_loader": weight_loader,
256 - }
257 - if input_groups == 1:
258 - scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
259 - else:
260 - scales = GroupQuantScaleParameter(
261 - output_dim=1, input_dim=0, **weight_scale_args
262 - )
263 -
264 - # Allocate workspace (Used for internal locking mechanism)
265 - max_workspace_size = (
266 - output_size_per_partition // self.quant_config.min_n_threads
267 - ) * self.quant_config.max_parallel
268 -
269 - workspace = BasevLLMParameter(
270 - data=torch.zeros(max_workspace_size, device="cuda", dtype=torch.int),
271 - weight_loader=weight_loader,
272 - )
273 -
274 - layer.register_parameter("B_24", qweight)
275 - layer.register_parameter("B_meta", meta)
276 - layer.register_parameter("s", scales)
277 - layer.register_parameter("workspace", workspace)
278 -
279 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
280 - # required by torch.compile
281 - layer.B_24 = Parameter(layer.B_24.data, requires_grad=False)
282 - layer.s = Parameter(layer.s.data, requires_grad=False)
283 - layer.B_meta = Parameter(layer.B_meta.data, requires_grad=False)
284 - layer.workspace = Parameter(layer.workspace.data, requires_grad=False)
285 -
286 - def apply(
287 - self,
288 - layer: torch.nn.Module,
289 - x: torch.Tensor,
290 - bias: torch.Tensor | None = None,
291 - ) -> torch.Tensor:
292 - qweight = layer.B_24
293 - meta = layer.B_meta
294 - scales = layer.s
295 - workspace = layer.workspace
296 -
297 - x_2d = x.view(-1, x.shape[-1])
298 -
299 - size_m = x_2d.shape[0]
300 - size_k = x_2d.shape[1]
301 - size_n = scales.shape[1]
302 -
303 - output_2d = ops.gptq_marlin_24_gemm(
304 - x_2d,
305 - qweight,
306 - meta,
307 - scales,
308 - workspace,
309 - self.quant_config.quant_type,
310 - size_m,
311 - size_n,
312 - size_k,
313 - )
314 -
315 - output = output_2d.view(x.shape[:-1] + (output_2d.shape[1],))
316 -
317 - if bias is not None:
318 - output.add_(bias) # In-place add
319 -
320 - return output
vllm/model_executor/layers/quantization/utils/marlin_utils_test_24.py deleted
-467
@@ -1,467 +0,0 @@
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 random
6 -
7 -import numpy
8 -import torch
9 -
10 -from vllm.scalar_type import ScalarType
11 -
12 -from .marlin_utils_test import marlin_weights
13 -from .quant_utils import gptq_quantize_weights
14 -
15 -
16 -# This is PyTorch implementation of main part of reorder_meta()
17 -# function, from tools/util/include/cutlass/util/host_reorder.h file
18 -# of CUTLASS source tree. Furthermore, CUTLASS template for sparse
19 -# GEMM decides upon layout of this matrix, and at the moment for the
20 -# sparse GEMM executed on tensor cores, this is layout described by
21 -# ColumnMajorInterleaved<2> data structure, in
22 -# include/cutlass/layout/matrix.h of CUTLASS source tree. The
23 -# reordering of meta matrix into meta_reordered matrix calculated
24 -# according to these segments of CUTLASS code is re-implemented here.
25 -# Note that this calculation produces offsets for scattering metadata
26 -# matrix elements into reordered metadata matrix elements (or,
27 -# equivalently, for gathering reordered metadata matrix element back
28 -# into metadata matrix elements).
29 -def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, device):
30 - dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols)
31 - dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1)
32 -
33 - # Reorder the rows, then swizzle the 2x2 blocks.
34 - group_x = 64
35 - group_y = 32 if meta_dtype.itemsize == 2 else 16
36 -
37 - dst_rows = (
38 - dst_rows // group_x * group_x
39 - + (dst_rows % 2) * 2
40 - + (dst_rows % 8) // 4
41 - + ((dst_rows % group_y) % 4) // 2 * 32
42 - + ((dst_rows % group_x) // 8) * 4
43 - )
44 -
45 - topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8)
46 - bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8)
47 - dst_rows += topright - bottomleft
48 - dst_cols -= topright - bottomleft
49 -
50 - # Assumed that meta tensor is to be stored in CUTLASS
51 - # InterleavedColumnMajor layout, and reverse engineered
52 - # corresponding code to store values into this tensor.
53 - interleave = 2
54 - cols_maj = dst_cols // interleave
55 - cols_min = dst_cols % interleave
56 - return (cols_maj * m * interleave + dst_rows * interleave + cols_min).view(-1)
57 -
58 -
59 -# This function converts dense matrix into sparse semi-structured
60 -# representation, producing "compressed" matrix, in the layout used by
61 -# CUTLASS backend, and corresponding metadata matrix.
62 -def sparse_semi_structured_from_dense_cutlass(dense):
63 - if dense.dim() != 2:
64 - raise RuntimeError(
65 - f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" # noqa: E501
66 - )
67 -
68 - m, k = dense.shape
69 - device = dense.device
70 -
71 - meta_dtype = torch.int8
72 - if dense.dtype == torch.int8:
73 - meta_dtype = torch.int32
74 - elif dense.dtype in [torch.half, torch.bfloat16, torch.float, torch.int32]:
75 - meta_dtype = torch.int16
76 - else:
77 - raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix")
78 - quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
79 - if quadbits_per_meta_elem not in (4, 8):
80 - raise RuntimeError("Invalid number of elements per meta element calculated")
81 -
82 - if meta_dtype == torch.int32:
83 - if m % 16 != 0:
84 - raise RuntimeError(
85 - f"Number of rows of dense matrix {m} must be divisible by 16"
86 - )
87 - else:
88 - if m % 32 != 0:
89 - raise RuntimeError(
90 - f"Number of rows of dense matrix {m} must be divisible by 32"
91 - )
92 - if k % (4 * quadbits_per_meta_elem) != 0:
93 - raise RuntimeError(
94 - f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" # noqa: E501
95 - )
96 -
97 - if dense.dtype != torch.float:
98 - ksparse = 4
99 - dense_4 = dense.view(-1, k // ksparse, ksparse)
100 - m0, m1, m2, m3 = (dense_4 != 0).unbind(-1)
101 - else:
102 - ksparse = 2
103 - dense_2 = dense.view(-1, k // ksparse, ksparse)
104 - m0, m2 = m1, m3 = (dense_2 != 0).unbind(-1)
105 - meta_ncols = k // (ksparse * quadbits_per_meta_elem)
106 -
107 - # Encoding quadruples of True/False values as follows:
108 - # [True, True, False, False] -> 0b0100
109 - # [True, False, True, False] -> 0b1000
110 - # [False, True, True, False] -> 0b1001
111 - # [True, False, False, True ] -> 0b1100
112 - # [False, True, False, True ] -> 0b1101
113 - # [False, False, True, True ] -> 0b1110
114 - # Thus, lower two bits in the encoding are index of the True value
115 - # at the lowest index in the quadruple, and the higher two bits in
116 - # the encoding are index of the other True value in the quadruple.
117 - # In case there are less than two True values, than False value or
118 - # values at some index or indices are considered True for the
119 - # encoding. In case there are more than two True values, then the
120 - # excess True value(s) at some indices are considered False for
121 - # the encoding. The exact encodings used for these cases are as
122 - # follows:
123 - # [False, False, False, False] -> 0b1110
124 - # [False, False, False, True ] -> 0b1110
125 - # [False, False, True, False] -> 0b1110
126 - # [False, True, False, False] -> 0b1001
127 - # [False, True, True, True ] -> 0b1101
128 - # [True, False, False, False] -> 0b1000
129 - # [True, False, True, True ] -> 0b1100
130 - # [True, True, False, True ] -> 0b0100
131 - # [True, True, True, False] -> 0b0100
132 - # [True, True, True, True ] -> 0b0100
133 - # These particular encodings are chosen, with the help of Espresso
134 - # logic minimizer software, for the purpose of minimization of
135 - # corresponding Boolean functions, that translate non-zero flags
136 - # into encoding bits. Note also possible choices for the first
137 - # and last of these encodings were limited only to (0b0100,
138 - # 0b1110), in order to produce valid encodings for 1:2 sparsity
139 - # case.
140 -
141 - expr0 = m0 & m1
142 - expr1 = ~m0 & m1
143 - expr2 = ~m0 & ~m1
144 - bit0 = expr1
145 - bit1 = expr2
146 - bit2 = expr0 | expr2 | m3
147 - bit3 = expr1 | ~m1
148 - idxs0 = bit0 | (bit1.to(torch.int64) << 1)
149 - idxs1 = bit2 | (bit3.to(torch.int64) << 1)
150 -
151 - if dense.dtype != torch.float:
152 - sparse0 = dense_4.gather(-1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined]
153 - sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1))
154 - sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2)
155 - else:
156 - sparse = dense_2.gather(-1, idxs0.unsqueeze(-1) // 2).view(m, k // 2) # type: ignore[possibly-undefined]
157 -
158 - meta_4 = idxs0 | (idxs1 << 2)
159 - meta_n = meta_4.view((-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype)
160 -
161 - if quadbits_per_meta_elem == 4:
162 - meta = (
163 - meta_n[:, :, 0]
164 - | (meta_n[:, :, 1] << 4)
165 - | (meta_n[:, :, 2] << 8)
166 - | (meta_n[:, :, 3] << 12)
167 - )
168 - elif quadbits_per_meta_elem == 8:
169 - meta = (
170 - meta_n[:, :, 0]
171 - | (meta_n[:, :, 1] << 4)
172 - | (meta_n[:, :, 2] << 8)
173 - | (meta_n[:, :, 3] << 12)
174 - | (meta_n[:, :, 4] << 16)
175 - | (meta_n[:, :, 5] << 20)
176 - | (meta_n[:, :, 6] << 24)
177 - | (meta_n[:, :, 7] << 28)
178 - )
179 -
180 - # Reorder meta tensor elements.
181 - meta_reordered = meta.new_empty((m * meta_ncols,)) # type: ignore[possibly-undefined]
182 - meta_offsets = _calculate_meta_reordering_scatter_offsets(
183 - m, meta_ncols, meta_dtype, device
184 - )
185 - meta_reordered.scatter_(0, meta_offsets, meta.view(-1))
186 -
187 - return (sparse, meta_reordered.view(m, meta_ncols))
188 -
189 -
190 -# This function performs reverse of the function above - it
191 -# reconstructs dense matrix from a pair of "compressed" matrix, given
192 -# in the layout used by CUTLASS backend, and accompanying metadata
193 -# matrix.
194 -def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered):
195 - if sparse.dim() != 2:
196 - raise RuntimeError(
197 - f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" # noqa: E501
198 - )
199 -
200 - m, k = sparse.shape
201 - device = sparse.device
202 -
203 - if meta_reordered.dim() != 2:
204 - raise RuntimeError(
205 - f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" # noqa: E501
206 - )
207 - if meta_reordered.device != device:
208 - raise RuntimeError(
209 - f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" # noqa: E501
210 - )
211 -
212 - meta_dtype = meta_reordered.dtype
213 - if meta_dtype not in (torch.int16, torch.int32):
214 - raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix")
215 - quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4
216 -
217 - ksparse = 4 if sparse.dtype != torch.float else 2
218 -
219 - meta_nrows, meta_ncols = meta_reordered.shape
220 - if meta_nrows != m:
221 - raise RuntimeError(
222 - f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" # noqa: E501
223 - )
224 - if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k:
225 - raise RuntimeError(
226 - f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " # noqa: E501
227 - "expected according to the number of columns of meta matrix"
228 - )
229 -
230 - # Undo meta tensor elements reordering.
231 - meta_offsets = _calculate_meta_reordering_scatter_offsets(
232 - m, meta_ncols, meta_dtype, device
233 - )
234 - meta = torch.gather(meta_reordered.view(-1), 0, meta_offsets).view(m, meta_ncols)
235 -
236 - # Unpack sparse tensor back to original dense tensor, using
237 - # information provided by meta tensor. Note that torch.float
238 - # datatype is handled pretty much the same as
239 - # torch.half/torch.bfloat16, as metadata for a pair of torch.float
240 - # value is encoded as if underlying 8 bytes contain four
241 - # torch.half/torch.bfloat16 values, where either first two or last
242 - # two are zeros.
243 - meta_2 = torch.empty(
244 - (m, meta_ncols, 2 * quadbits_per_meta_elem),
245 - dtype=meta_dtype,
246 - device=device,
247 - )
248 - if quadbits_per_meta_elem == 4:
249 - meta_2[:, :, 0] = meta & 0b11
250 - meta_2[:, :, 1] = (meta >> 2) & 0b11
251 - meta_2[:, :, 2] = (meta >> 4) & 0b11
252 - meta_2[:, :, 3] = (meta >> 6) & 0b11
253 - meta_2[:, :, 4] = (meta >> 8) & 0b11
254 - meta_2[:, :, 5] = (meta >> 10) & 0b11
255 - meta_2[:, :, 6] = (meta >> 12) & 0b11
256 - meta_2[:, :, 7] = (meta >> 14) & 0b11
257 - elif quadbits_per_meta_elem == 8:
258 - meta_2[:, :, 0] = meta & 0b11
259 - meta_2[:, :, 1] = (meta >> 2) & 0b11
260 - meta_2[:, :, 2] = (meta >> 4) & 0b11
261 - meta_2[:, :, 3] = (meta >> 6) & 0b11
262 - meta_2[:, :, 4] = (meta >> 8) & 0b11
263 - meta_2[:, :, 5] = (meta >> 10) & 0b11
264 - meta_2[:, :, 6] = (meta >> 12) & 0b11
265 - meta_2[:, :, 7] = (meta >> 14) & 0b11
266 - meta_2[:, :, 8] = (meta >> 16) & 0b11
267 - meta_2[:, :, 9] = (meta >> 18) & 0b11
268 - meta_2[:, :, 10] = (meta >> 20) & 0b11
269 - meta_2[:, :, 11] = (meta >> 22) & 0b11
270 - meta_2[:, :, 12] = (meta >> 24) & 0b11
271 - meta_2[:, :, 13] = (meta >> 26) & 0b11
272 - meta_2[:, :, 14] = (meta >> 28) & 0b11
273 - meta_2[:, :, 15] = (meta >> 30) & 0b11
274 -
275 - dense_offsets = meta_2.view(-1) + (
276 - torch.arange(0, 2 * m * k // ksparse, device=device) * 4
277 - ).view(-1, 1).repeat(1, 2).view(-1)
278 -
279 - dense = torch.zeros((m * 2 * k,), dtype=sparse.dtype, device=device)
280 - if sparse.dtype != torch.float:
281 - # dense.scatter_(0, dense_offsets, sparse.view(-1))
282 - dense.scatter_(0, dense_offsets, sparse.reshape(-1))
283 - else:
284 - dense.view(torch.half).scatter_(
285 - 0, dense_offsets, sparse.view(torch.half).view(-1)
286 - )
287 -
288 - return dense.view(m, 2 * k)
289 -
290 -
291 -def mask_creator(tensor):
292 - """
293 - Class for creating N:M sparsity masks.
294 - Masks will be created using the N:M ratio, where for every block of
295 - M weights, N will be pruned based on ranked weight value. Each mask
296 - will correspond to the given tensor.
297 -
298 - :param N: The number of weights in a group to keep
299 - :param M: The size of a weight group
300 - """
301 - N = 2
302 - M = 4
303 -
304 - mask = None
305 - # for i, tensor in enumerate(tensors):
306 - if tensor.numel() % M != 0:
307 - raise ValueError(
308 - f"Tensor of size {tensor.shape} can't be evenly divided into {M} groups"
309 - )
310 -
311 - num_groups = tensor.numel() // M
312 -
313 - # N:M sparsity for linear layers
314 - tensor_temp = tensor.detach().abs().reshape(num_groups, M)
315 - index = torch.argsort(tensor_temp, dim=1)[:, : int(M - N)]
316 -
317 - w_b = torch.ones(tensor_temp.shape, device=tensor_temp.device)
318 - mask = w_b.scatter_(dim=1, index=index, value=0).reshape(tensor.shape)
319 -
320 - return mask
321 -
322 -
323 -def inject_24(w, size_k, size_n):
324 - assert w.shape == (size_k, size_n)
325 -
326 - mask = mask_creator(w.t()).t().cuda().bool()
327 -
328 - return (mask * w).contiguous(), mask.contiguous()
329 -
330 -
331 -def check_24(w, num_rows_to_sample=50, _verbose=False):
332 - BLOCK_SIZE = 4
333 - MAX_NON_ZEROS = 2
334 -
335 - w = w.t().contiguous()
336 -
337 - print("check_24: w.shape = {}".format(w.shape))
338 -
339 - num_rows, num_cols = w.shape
340 - sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample)
341 - if _verbose:
342 - print(f"Sampled row idxs = {sampled_row_idxs}")
343 -
344 - total_segments = 0
345 - non_24_segments = 0
346 - for i in sampled_row_idxs:
347 - for j in range(0, num_cols - BLOCK_SIZE, BLOCK_SIZE):
348 - total_segments += 1
349 - block = w[i, j : j + BLOCK_SIZE]
350 - num_nonzero = torch.count_nonzero(block)
351 - if num_nonzero > MAX_NON_ZEROS:
352 - print("i = {} j = {} block = {}".format(i, j, block))
353 - non_24_segments += 1
354 -
355 - print(f"{non_24_segments} / {total_segments} do not have 2:4 structure.")
356 -
357 -
358 -def compress_quantized_24_weight(q_24, size_k, size_n, wtype: ScalarType):
359 - assert q_24.shape == (size_k, size_n)
360 -
361 - # Remove bias to normalize over 0
362 - q_24_no_zp = q_24 - wtype.bias
363 -
364 - # Compress
365 - q_24_no_zp = q_24_no_zp.t().contiguous()
366 - q_24_no_zp_comp, meta = sparse_semi_structured_from_dense_cutlass(q_24_no_zp)
367 - q_24_no_zp_comp = q_24_no_zp_comp.t().contiguous()
368 -
369 - # Restore bias
370 - q_24_comp = q_24_no_zp_comp + wtype.bias
371 -
372 - # Resize meta to its actual shape (without moving any data)
373 - meta = meta.resize_(meta.shape[1] // 2, meta.shape[0] * 2)
374 -
375 - return q_24_comp, meta
376 -
377 -
378 -def get_scale_perms_24():
379 - scale_perm: list[int] = []
380 - for i in range(8):
381 - scale_perm.extend([i * 8 + j for j in [0, 4, 1, 5, 2, 6, 3, 7]])
382 - scale_perm_single: list[int] = []
383 - for i in range(8):
384 - scale_perm_single.extend([8 * i + j for j in [0, 1, 2, 3, 4, 5, 6, 7]])
385 - return scale_perm, scale_perm_single
386 -
387 -
388 -def get_weight_perm_24(num_bits: int):
389 - perm_list: list[int] = []
390 - for i in range(32):
391 - perm1: list[int] = []
392 - col = i // 4
393 - col_o = col // 2
394 - for block in [0, 1]:
395 - for row in [
396 - 2 * (i % 4),
397 - 2 * (i % 4) + 1,
398 - 2 * (i % 4 + 4),
399 - 2 * (i % 4 + 4) + 1,
400 - ]:
401 - perm1.append(16 * row + col_o * 256 + 8 * (col % 2) + 4 * block)
402 - for j in range(4):
403 - perm_list.extend([p + 1 * j for p in perm1])
404 - perm = numpy.array(perm_list)
405 -
406 - if num_bits == 4:
407 - interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
408 - elif num_bits == 8:
409 - interleave = numpy.array([0, 2, 1, 3])
410 - else:
411 - raise ValueError("num_bits must be 4 or 8, got {}".format(num_bits))
412 -
413 - perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel()
414 - perm = torch.from_numpy(perm)
415 - return perm
416 -
417 -
418 -def marlin_permute_scales_24(
419 - s: torch.Tensor, size_k: int, size_n: int, group_size: int
420 -) -> torch.Tensor:
421 - scale_perm, scale_perm_single = get_scale_perms_24()
422 - if group_size < size_k and group_size != -1:
423 - s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
424 - else:
425 - s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
426 - s = s.reshape((-1, size_n)).contiguous()
427 -
428 - return s
429 -
430 -
431 -def marlin_24_quantize(
432 - w: torch.Tensor,
433 - quant_type: ScalarType,
434 - group_size: int,
435 -):
436 - size_k, size_n = w.shape
437 -
438 - # Normalize group_size
439 - if group_size == -1:
440 - group_size = size_k
441 - assert group_size <= size_k
442 -
443 - # Inject 2:4 sparsity
444 - w_24, mask_24 = inject_24(w, size_k, size_n)
445 -
446 - # Quantize
447 - w_24_ref, q_w_24, s, g_idx, rand_perm = gptq_quantize_weights(
448 - w_24, quant_type, group_size, act_order=False
449 - )
450 -
451 - # Compress quantized weight
452 - q_w_24_comp, meta = compress_quantized_24_weight(q_w_24, size_k, size_n, quant_type)
453 - size_k_comp = size_k // 2
454 -
455 - # Reformat to marlin
456 - weight_perm = get_weight_perm_24(quant_type.size_bits)
457 - marlin_24_q_w_comp = marlin_weights(
458 - q_w_24_comp, size_k_comp, size_n, quant_type.size_bits, weight_perm
459 - )
460 - marlin_24_s = marlin_permute_scales_24(s, size_k, size_n, group_size)
461 -
462 - # Create result
463 - res_list = [w_24_ref, marlin_24_q_w_comp, meta, marlin_24_s]
464 - for i in range(len(res_list)):
465 - res_list[i] = res_list[i].to(w.device)
466 -
467 - return res_list