A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Kernel/Quant] Remove the original marlin format and qqq (#23204)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 20, 2025 at 15:13 UTC
0cdbf5e61ce3fd97d33b31b775d2faaadc99fbc5
26 files changed
+95
-3701
.buildkite/lm-eval-harness/configs/Meta-Llama-3-8B-QQQ.yaml
deleted
-12
@@ -1,12 +0,0 @@
1
-# For vllm script, with -t option (tensor parallel size).
2
-# bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m HandH1998/QQQ-Llama-3-8b-g128 -b 32 -l 1000 -f 5 -t 1
3
-model_name: "HandH1998/QQQ-Llama-3-8b-g128"
4
-tasks:
5
-- name: "gsm8k"
6
- metrics:
7
- - name: "exact_match,strict-match"
8
- value: 0.419
9
- - name: "exact_match,flexible-extract"
10
- value: 0.416
11
-limit: 1000
12
-num_fewshot: 5
.buildkite/lm-eval-harness/configs/models-large.txt
-1
@@ -3,4 +3,3 @@ Meta-Llama-3-70B-Instruct.yaml
3
Mixtral-8x7B-Instruct-v0.1.yaml
4
Qwen2-57B-A14-Instruct.yaml
5
DeepSeek-V2-Lite-Chat.yaml
6
-Meta-Llama-3-8B-QQQ.yaml
CMakeLists.txt
-2
@@ -357,9 +357,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
357
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
358
359
set(MARLIN_SRCS
360
- "csrc/quantization/marlin/dense/marlin_cuda_kernel.cu"
360
"csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
362
- "csrc/quantization/marlin/qqq/marlin_qqq_gemm_kernel.cu"
361
"csrc/quantization/gptq_marlin/gptq_marlin.cu"
362
"csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
363
"csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
benchmarks/kernels/benchmark_machete.py
+1
-22
@@ -253,28 +253,7 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
253
else:
254
assert bt.a.dtype == torch.int8
255
assert bt.wtype == scalar_types.uint4b8
256
-
257
- if bt.w_ch_s is not None:
258
- s_ch = bt.w_ch_s.to(torch.float32)
259
- else:
260
- s_ch = torch.ones(bt.w_ref.shape[1], dtype=torch.float32, device=device)
261
-
262
- if bt.w_tok_s is not None:
263
- s_tok = bt.w_tok_s.to(torch.float32)
264
- else:
265
- s_tok = torch.ones(bt.a.shape[0], dtype=torch.float32, device=device)
266
-
267
- fn = lambda: ops.marlin_qqq_gemm(
268
- a=bt.a,
269
- b_q_weight=w_q,
270
- s_group=w_s,
271
- s_tok=s_tok,
272
- s_ch=s_ch,
273
- workspace=workspace.scratch,
274
- size_m=bt.a.shape[0],
275
- size_n=bt.w_ref.shape[1],
276
- size_k=bt.w_ref.shape[0],
277
- )
256
+ raise NotImplementedError("QQQ is not supported anymore")
257
258
return fn
259
csrc/quantization/machete/generate.py
+73
-72
@@ -571,78 +571,79 @@ def generate():
571
itertools.repeat(default_heuristic))
572
]
573
574
- # Stored as "condition": ((tile_shape_mn), (cluster_shape_mnk))
575
- # TODO (LucasWilkinson): Further tuning required
576
- qqq_tile_heuristic_config = {
577
- #### M = 257+
578
- # ((128, 256), (2, 1, 1)) Broken for QQQ types
579
- # TODO (LucasWilkinson): Investigate further
580
- # "M > 256 && K <= 16384 && N <= 4096": ((128, 128), (2, 1, 1)),
581
- # "M > 256": ((128, 256), (2, 1, 1)),
582
- "M > 256": ((128, 128), (2, 1, 1)),
583
- #### M = 129-256
584
- "M > 128 && K <= 4096 && N <= 4096": ((128, 64), (2, 1, 1)),
585
- "M > 128 && K <= 8192 && N <= 8192": ((128, 128), (2, 1, 1)),
586
- # ((128, 256), (2, 1, 1)) Broken for QQQ types
587
- # TODO (LucasWilkinson): Investigate further
588
- # "M > 128": ((128, 256), (2, 1, 1)),
589
- "M > 128": ((128, 128), (2, 1, 1)),
590
- #### M = 65-128
591
- "M > 64 && K <= 4069 && N <= 4069": ((128, 32), (2, 1, 1)),
592
- "M > 64 && K <= 4069 && N <= 8192": ((128, 64), (2, 1, 1)),
593
- "M > 64 && K >= 8192 && N >= 12288": ((256, 128), (2, 1, 1)),
594
- "M > 64": ((128, 128), (2, 1, 1)),
595
- #### M = 33-64
596
- "M > 32 && K <= 6144 && N <= 6144": ((128, 16), (1, 1, 1)),
597
- # Broken for QQQ types
598
- # TODO (LucasWilkinson): Investigate further
599
- #"M > 32 && K >= 16384 && N >= 12288": ((256, 64), (2, 1, 1)),
600
- "M > 32": ((128, 64), (2, 1, 1)),
601
- #### M = 17-32
602
- "M > 16 && K <= 12288 && N <= 8192": ((128, 32), (2, 1, 1)),
603
- "M > 16": ((256, 32), (2, 1, 1)),
604
- #### M = 1-16
605
- "N >= 26624": ((256, 16), (1, 1, 1)),
606
- None: ((128, 16), (1, 1, 1)),
607
- }
608
-
609
- # For now we use the same heuristic for all types
610
- # Heuristic is currently tuned for H100s
611
- qqq_heuristic = [
612
- (cond, ScheduleConfig(*tile_config,
613
- **sch_common_params)) # type: ignore
614
- for cond, tile_config in qqq_tile_heuristic_config.items()
615
- ]
616
-
617
- QQQ_kernel_types = [
618
- *(TypeConfig(
619
- a=DataType.s8,
620
- b=VLLMDataType.u4b8,
621
- b_group_scale=b_group_scale,
622
- b_group_zeropoint=DataType.void,
623
- b_channel_scale=DataType.f32,
624
- a_token_scale=DataType.f32,
625
- out=DataType.f16,
626
- accumulator=DataType.s32,
627
- ) for b_group_scale in (DataType.f16, DataType.void)),
628
- *(TypeConfig(
629
- a=DataType.e4m3,
630
- b=VLLMDataType.u4b8,
631
- b_group_scale=b_group_scale,
632
- b_group_zeropoint=DataType.void,
633
- b_channel_scale=DataType.f32,
634
- a_token_scale=DataType.f32,
635
- out=DataType.f16,
636
- accumulator=DataType.f32,
637
- ) for b_group_scale in (DataType.f16, DataType.void)),
638
- ]
639
-
640
- impl_configs += [
641
- ImplConfig(x[0], x[1], x[2])
642
- for x in zip(QQQ_kernel_types,
643
- itertools.repeat(get_unique_schedules(qqq_heuristic)),
644
- itertools.repeat(qqq_heuristic))
645
- ]
574
+ # TODO: Support W4A8 when ready
575
+ # # Stored as "condition": ((tile_shape_mn), (cluster_shape_mnk))
576
+ # # TODO (LucasWilkinson): Further tuning required
577
+ # qqq_tile_heuristic_config = {
578
+ # #### M = 257+
579
+ # # ((128, 256), (2, 1, 1)) Broken for QQQ types
580
+ # # TODO (LucasWilkinson): Investigate further
581
+ # # "M > 256 && K <= 16384 && N <= 4096": ((128, 128), (2, 1, 1)),
582
+ # # "M > 256": ((128, 256), (2, 1, 1)),
583
+ # "M > 256": ((128, 128), (2, 1, 1)),
584
+ # #### M = 129-256
585
+ # "M > 128 && K <= 4096 && N <= 4096": ((128, 64), (2, 1, 1)),
586
+ # "M > 128 && K <= 8192 && N <= 8192": ((128, 128), (2, 1, 1)),
587
+ # # ((128, 256), (2, 1, 1)) Broken for QQQ types
588
+ # # TODO (LucasWilkinson): Investigate further
589
+ # # "M > 128": ((128, 256), (2, 1, 1)),
590
+ # "M > 128": ((128, 128), (2, 1, 1)),
591
+ # #### M = 65-128
592
+ # "M > 64 && K <= 4069 && N <= 4069": ((128, 32), (2, 1, 1)),
593
+ # "M > 64 && K <= 4069 && N <= 8192": ((128, 64), (2, 1, 1)),
594
+ # "M > 64 && K >= 8192 && N >= 12288": ((256, 128), (2, 1, 1)),
595
+ # "M > 64": ((128, 128), (2, 1, 1)),
596
+ # #### M = 33-64
597
+ # "M > 32 && K <= 6144 && N <= 6144": ((128, 16), (1, 1, 1)),
598
+ # # Broken for QQQ types
599
+ # # TODO (LucasWilkinson): Investigate further
600
+ # #"M > 32 && K >= 16384 && N >= 12288": ((256, 64), (2, 1, 1)),
601
+ # "M > 32": ((128, 64), (2, 1, 1)),
602
+ # #### M = 17-32
603
+ # "M > 16 && K <= 12288 && N <= 8192": ((128, 32), (2, 1, 1)),
604
+ # "M > 16": ((256, 32), (2, 1, 1)),
605
+ # #### M = 1-16
606
+ # "N >= 26624": ((256, 16), (1, 1, 1)),
607
+ # None: ((128, 16), (1, 1, 1)),
608
+ # }
609
+
610
+ # # For now we use the same heuristic for all types
611
+ # # Heuristic is currently tuned for H100s
612
+ # qqq_heuristic = [
613
+ # (cond, ScheduleConfig(*tile_config,
614
+ # **sch_common_params)) # type: ignore
615
+ # for cond, tile_config in qqq_tile_heuristic_config.items()
616
+ # ]
617
+
618
+ # QQQ_kernel_types = [
619
+ # *(TypeConfig(
620
+ # a=DataType.s8,
621
+ # b=VLLMDataType.u4b8,
622
+ # b_group_scale=b_group_scale,
623
+ # b_group_zeropoint=DataType.void,
624
+ # b_channel_scale=DataType.f32,
625
+ # a_token_scale=DataType.f32,
626
+ # out=DataType.f16,
627
+ # accumulator=DataType.s32,
628
+ # ) for b_group_scale in (DataType.f16, DataType.void)),
629
+ # *(TypeConfig(
630
+ # a=DataType.e4m3,
631
+ # b=VLLMDataType.u4b8,
632
+ # b_group_scale=b_group_scale,
633
+ # b_group_zeropoint=DataType.void,
634
+ # b_channel_scale=DataType.f32,
635
+ # a_token_scale=DataType.f32,
636
+ # out=DataType.f16,
637
+ # accumulator=DataType.f32,
638
+ # ) for b_group_scale in (DataType.f16, DataType.void)),
639
+ # ]
640
+
641
+ # impl_configs += [
642
+ # ImplConfig(x[0], x[1], x[2])
643
+ # for x in zip(QQQ_kernel_types,
644
+ # itertools.repeat(get_unique_schedules(qqq_heuristic)),
645
+ # itertools.repeat(qqq_heuristic))
646
+ # ]
647
648
output_dir = os.path.join(SCRIPT_DIR, "generated")
649
csrc/quantization/marlin/dense/LICENSE
deleted
-209
@@ -1,209 +0,0 @@
1
-Contains code from https://github.com/IST-DASLab/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.
204
-
205
-------------------------------------------------------------------------------------
206
-
207
-This product bundles various third-party components under other open source licenses.
208
-This section summarizes those components and their licenses. See licenses/
209
-for text of these licenses.
csrc/quantization/marlin/dense/common/base.h
deleted
-32
@@ -1,32 +0,0 @@
1
-/*
2
- * Modified by HandH1998
3
- * Modified by Neural Magic
4
- * Copyright (C) Marlin.2024 Elias Frantar
5
- *
6
- * Licensed under the Apache License, Version 2.0 (the "License");
7
- * you may not use this file except in compliance with the License.
8
- * You may obtain a copy of the License at
9
- *
10
- * http://www.apache.org/licenses/LICENSE-2.0
11
- *
12
- * Unless required by applicable law or agreed to in writing, software
13
- * distributed under the License is distributed on an "AS IS" BASIS,
14
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
- * See the License for the specific language governing permissions and
16
- * limitations under the License.
17
- */
18
-
19
-#pragma once
20
-
21
-constexpr int ceildiv(int a, int b) { return (a + b - 1) / b; }
22
-
23
-// Instances of `Vec` are used to organize groups of >>registers<<, as needed
24
-// for instance as inputs to tensor core operations. Consequently, all
25
-// corresponding index accesses must be compile-time constants, which is why we
26
-// extensively use `#pragma unroll` throughout the kernel code to guarantee
27
-// this.
28
-template <typename T, int n>
29
-struct Vec {
30
- T elems[n];
31
- __device__ T& operator[](int i) { return elems[i]; }
32
-};
csrc/quantization/marlin/dense/common/mem.h
deleted
-89
@@ -1,89 +0,0 @@
1
-/*
2
- * Modified by HandH1998
3
- * Modified by Neural Magic
4
- * Copyright (C) Marlin.2024 Elias Frantar
5
- *
6
- * Licensed under the Apache License, Version 2.0 (the "License");
7
- * you may not use this file except in compliance with the License.
8
- * You may obtain a copy of the License at
9
- *
10
- * http://www.apache.org/licenses/LICENSE-2.0
11
- *
12
- * Unless required by applicable law or agreed to in writing, software
13
- * distributed under the License is distributed on an "AS IS" BASIS,
14
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
- * See the License for the specific language governing permissions and
16
- * limitations under the License.
17
- */
18
-
19
-#pragma once
20
-
21
-// Predicated asynchronous global->shared copy; used for inputs A where we apply
22
-// predication to handle batchsizes that are not multiples of 16.
23
-__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
24
- bool pred = true) {
25
- const int BYTES = 16;
26
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
27
- asm volatile(
28
- "{\n"
29
- " .reg .pred p;\n"
30
- " setp.ne.b32 p, %0, 0;\n"
31
- " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
32
- "}\n" ::"r"((int)pred),
33
- "r"(smem), "l"(glob_ptr), "n"(BYTES));
34
-}
35
-
36
-// Asynchronous global->shared copy
37
-__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
38
- const int BYTES = 16;
39
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
40
- asm volatile(
41
- "{\n"
42
- " cp.async.cg.shared.global [%0], [%1], %2;\n"
43
- "}\n" ::"r"(smem),
44
- "l"(glob_ptr), "n"(BYTES));
45
-}
46
-
47
-// Async copy fence.
48
-__device__ inline void cp_async_fence() {
49
- asm volatile("cp.async.commit_group;\n" ::);
50
-}
51
-
52
-// Wait until at most `n` async copy stages are still pending.
53
-template <int n>
54
-__device__ inline void cp_async_wait() {
55
- asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
56
-}
57
-
58
-// Wait until barrier reaches `count`, then lock for current threadblock.
59
-__device__ inline void barrier_acquire(int* lock, int count) {
60
- if (threadIdx.x == 0) {
61
- int state = -1;
62
- do
63
- // Guarantee that subsequent writes by this threadblock will be visible
64
- // globally.
65
- asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
66
- : "=r"(state)
67
- : "l"(lock));
68
- while (state != count);
69
- }
70
- __syncthreads();
71
-}
72
-
73
-// Release barrier and increment visitation count.
74
-__device__ inline void barrier_release(int* lock, bool reset = false) {
75
- __syncthreads();
76
- if (threadIdx.x == 0) {
77
- if (reset) {
78
- lock[0] = 0;
79
- return;
80
- }
81
- int val = 1;
82
- // Make sure that all writes since acquiring this barrier are visible
83
- // globally, while releasing the barrier.
84
- asm volatile("fence.acq_rel.gpu;\n");
85
- asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
86
- :
87
- : "l"(lock), "r"(val));
88
- }
89
-}
csrc/quantization/marlin/dense/marlin_cuda_kernel.cu
deleted
-1073
@@ -1,1073 +0,0 @@
1
-/*
2
- * Modified by Neural Magic
3
- * Copyright (C) Marlin.2024 Elias Frantar
4
- *
5
- * Licensed under the Apache License, Version 2.0 (the "License");
6
- * you may not use this file except in compliance with the License.
7
- * You may obtain a copy of the License at
8
- *
9
- * http://www.apache.org/licenses/LICENSE-2.0
10
- *
11
- * Unless required by applicable law or agreed to in writing, software
12
- * distributed under the License is distributed on an "AS IS" BASIS,
13
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
- * See the License for the specific language governing permissions and
15
- * limitations under the License.
16
- */
17
-
18
-#include <torch/all.h>
19
-
20
-#include <ATen/cuda/CUDAContext.h>
21
-#include <c10/cuda/CUDAGuard.h>
22
-#include <cuda.h>
23
-#include <cuda_fp16.h>
24
-#include <cuda_runtime.h>
25
-
26
-#include <iostream>
27
-
28
-#include "common/base.h"
29
-#include "core/registration.h"
30
-
31
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
32
- #include "common/mem.h"
33
-#endif
34
-
35
-template <typename T>
36
-inline std::string str(T x) {
37
- return std::to_string(x);
38
-}
39
-
40
-namespace marlin_dense {
41
-
42
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
43
-
44
-using I4 = Vec<int, 4>;
45
-// Matrix fragments for tensor core instructions; their precise layout is
46
-// documented here:
47
-// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type
48
-using FragA = Vec<half2, 4>;
49
-using FragB = Vec<half2, 2>;
50
-using FragC = Vec<float, 4>;
51
-using FragS = Vec<half2, 1>; // quantization scales
52
-
53
-// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
54
-// output/accumulation.
55
-__device__ inline void mma(const FragA& a_frag, const FragB& frag_b,
56
- FragC& frag_c) {
57
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
58
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
59
- float* c = reinterpret_cast<float*>(&frag_c);
60
- asm volatile(
61
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
62
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
63
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
64
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
65
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
66
-}
67
-
68
-// Instruction for loading a full 16x16 matrix fragment of operand A from shared
69
-// memory, directly in tensor core layout.
70
-__device__ inline void ldsm4(FragA& frag_a, const void* smem_ptr) {
71
- uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
72
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
73
- asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
74
- : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
75
- : "r"(smem));
76
-}
77
-
78
-// Lookup-table based 3-input logical operation; explicitly used for
79
-// dequantization as the compiler does not seem to automatically recognize it in
80
-// all cases.
81
-template <int lut>
82
-__device__ inline int lop3(int a, int b, int c) {
83
- int res;
84
- asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
85
- : "=r"(res)
86
- : "r"(a), "r"(b), "r"(c), "n"(lut));
87
- return res;
88
-}
89
-
90
-// Efficiently dequantize an int32 value into a full B-fragment of 4 fp16
91
-// values. We mostly follow the strategy in the link below, with some small
92
-// changes:
93
-// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
94
-__device__ inline FragB dequant(int q) {
95
- const int LO = 0x000f000f;
96
- const int HI = 0x00f000f0;
97
- const int EX = 0x64006400;
98
- // Guarantee that the `(a & b) | c` operations are LOP3s.
99
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
100
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
101
- // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
102
- // directly into `SUB` and `ADD`.
103
- const int SUB = 0x64086408;
104
- const int MUL = 0x2c002c00;
105
- const int ADD = 0xd480d480;
106
- FragB frag_b;
107
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
108
- *reinterpret_cast<const half2*>(&SUB));
109
- frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
110
- *reinterpret_cast<const half2*>(&MUL),
111
- *reinterpret_cast<const half2*>(&ADD));
112
- return frag_b;
113
-}
114
-
115
-// Multiply dequantized values by the corresponding quantization scale; used
116
-// only for grouped quantization.
117
-__device__ inline void scale(FragB& frag_b, FragS& frag_s, int i) {
118
- half2 s = __half2half2(reinterpret_cast<__half*>(&frag_s)[i]);
119
- frag_b[0] = __hmul2(frag_b[0], s);
120
- frag_b[1] = __hmul2(frag_b[1], s);
121
-}
122
-
123
-template <const int threads, // number of threads in a threadblock
124
- const int thread_m_blocks, // number of 16x16 blocks in the m
125
- // dimension (batchsize) of the
126
- // threadblock
127
- const int thread_n_blocks, // same for n dimension (output)
128
- const int thread_k_blocks, // same for k dimension (reduction)
129
- const int stages, // number of stages for the async global->shared
130
- // fetch pipeline
131
- const int group_blocks = -1 // number of consecutive 16x16 blocks
132
- // with a separate quantization scale
133
- >
134
-__global__ void Marlin(
135
- const int4* __restrict__ A, // fp16 input matrix of shape mxk
136
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
137
- int4* __restrict__ C, // fp16 output buffer of shape mxn
138
- const int4* __restrict__ s, // fp16 quantization scales of shape
139
- // (k/groupsize)xn
140
- int prob_m, // batch dimension m
141
- int prob_n, // output dimension n
142
- int prob_k, // reduction dimension k
143
- int* locks // extra global storage for barrier synchronization
144
-) {
145
- // Each threadblock processes one "stripe" of the B matrix with (roughly) the
146
- // same size, which might involve multiple column "slices" (of width 16 *
147
- // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
148
- // example:
149
- // 0 1 3
150
- // 0 2 3
151
- // 1 2 4
152
- // While this kind of partitioning makes things somewhat more complicated, it
153
- // ensures good utilization of all SMs for many kinds of shape and GPU
154
- // configurations, while requiring as few slow global cross-threadblock
155
- // reductions as possible.
156
-
157
- // For larger GEMMs we run multiple batchsize 64 versions in parallel for a
158
- // better partitioning with less reductions
159
- int parallel = 1;
160
- if (prob_m > 16 * thread_m_blocks) {
161
- parallel = prob_m / (16 * thread_m_blocks);
162
- prob_m = 16 * thread_m_blocks;
163
- }
164
-
165
- int k_tiles = prob_k / 16 / thread_k_blocks;
166
- int n_tiles = prob_n / 16 / thread_n_blocks;
167
- int iters = ceildiv(k_tiles * n_tiles * parallel, gridDim.x);
168
- // Ensure that the number of tiles in each stripe is a multiple of the
169
- // groupsize; this avoids an annoying special case where a stripe starts in
170
- // the middle of group.
171
- if (group_blocks != -1)
172
- iters = (group_blocks / thread_k_blocks) *
173
- ceildiv(iters, (group_blocks / thread_k_blocks));
174
-
175
- int slice_row = (iters * blockIdx.x) % k_tiles;
176
- int slice_col_par = (iters * blockIdx.x) / k_tiles;
177
- int slice_col = slice_col_par;
178
- int slice_iters; // number of threadblock tiles in the current slice
179
- int slice_count =
180
- 0; // total number of active threadblocks in the current slice
181
- int slice_idx; // index of threadblock in current slice; numbered bottom to
182
- // top
183
-
184
- // We can easily implement parallel problem execution by just remapping
185
- // indices and advancing global pointers
186
- if (slice_col_par >= n_tiles) {
187
- A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8;
188
- C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
189
- locks += (slice_col_par / n_tiles) * n_tiles;
190
- slice_col = slice_col_par % n_tiles;
191
- }
192
-
193
- // Compute all information about the current slice which is required for
194
- // synchronization.
195
- auto init_slice = [&]() {
196
- slice_iters =
197
- iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
198
- if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
199
- if (slice_iters == 0) return;
200
- if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
201
- slice_count = 1;
202
- slice_idx = 0;
203
- int col_first = iters * ceildiv(k_tiles * slice_col_par, iters);
204
- if (col_first <= k_tiles * (slice_col_par + 1)) {
205
- int col_off = col_first - k_tiles * slice_col_par;
206
- slice_count = ceildiv(k_tiles - col_off, iters);
207
- if (col_off > 0) slice_count++;
208
- int delta_first = iters * blockIdx.x - col_first;
209
- if (delta_first < 0 || (col_off == 0 && delta_first == 0))
210
- slice_idx = slice_count - 1;
211
- else {
212
- slice_idx = slice_count - 1 - delta_first / iters;
213
- if (col_off > 0) slice_idx--;
214
- }
215
- }
216
- if (slice_col == n_tiles) {
217
- A += 16 * thread_m_blocks * prob_k / 8;
218
- C += 16 * thread_m_blocks * prob_n / 8;
219
- locks += n_tiles;
220
- slice_col = 0;
221
- }
222
- };
223
- init_slice();
224
-
225
- int a_gl_stride = prob_k / 8; // stride of the A matrix in global memory
226
- // We typically use `constexpr` to indicate that this value is a compile-time
227
- // constant
228
- constexpr int a_sh_stride =
229
- 16 * thread_k_blocks / 8; // stride of an A matrix tile in shared memory
230
- constexpr int a_gl_rd_delta_o =
231
- 16 * thread_k_blocks /
232
- 8; // delta between subsequent A tiles in global memory
233
- int a_gl_rd_delta_i =
234
- a_gl_stride *
235
- (threads / a_gl_rd_delta_o); // between subsequent accesses within a tile
236
- constexpr int a_sh_wr_delta =
237
- a_sh_stride *
238
- (threads / a_gl_rd_delta_o); // between shared memory writes
239
- constexpr int a_sh_rd_delta_o =
240
- 2 * ((threads / 32) /
241
- (thread_n_blocks / 4)); // between shared memory tile reads
242
- constexpr int a_sh_rd_delta_i =
243
- a_sh_stride * 16; // within a shared memory tile
244
- constexpr int a_sh_stage =
245
- a_sh_stride * (16 * thread_m_blocks); // overall size of a tile
246
- constexpr int a_sh_wr_iters =
247
- ceildiv(a_sh_stage,
248
- a_sh_wr_delta); // number of shared write iterations for a tile
249
-
250
- int b_gl_stride = 16 * prob_n / 32;
251
- constexpr int b_sh_stride = 32 * thread_n_blocks / 4;
252
- int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
253
- int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride);
254
- constexpr int b_sh_wr_delta = threads;
255
- constexpr int b_sh_rd_delta = threads;
256
- constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
257
- constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
258
-
259
- int s_gl_stride = prob_n / 8;
260
- constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
261
- constexpr int s_sh_stage = s_sh_stride;
262
- int s_gl_rd_delta = s_gl_stride;
263
-
264
- // Global A read index of current thread.
265
- int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
266
- (threadIdx.x % a_gl_rd_delta_o);
267
- a_gl_rd += a_gl_rd_delta_o * slice_row;
268
- // Shared write index of current thread.
269
- int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
270
- (threadIdx.x % a_gl_rd_delta_o);
271
- // Shared read index.
272
- int a_sh_rd =
273
- a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16;
274
- a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
275
-
276
- int b_gl_rd =
277
- b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
278
- b_gl_rd += b_sh_stride * slice_col;
279
- b_gl_rd += b_gl_rd_delta_o * slice_row;
280
- auto b_sh_wr = threadIdx.x;
281
- auto b_sh_rd = threadIdx.x;
282
-
283
- int s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
284
- s_sh_stride * slice_col + threadIdx.x;
285
- auto s_sh_wr = threadIdx.x;
286
- int s_sh_rd;
287
- // We use a different scale layout for grouped and column-wise quantization as
288
- // we scale a `half2` tile in column-major layout in the former and in
289
- // row-major in the latter case.
290
- if (group_blocks != -1)
291
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
292
- (threadIdx.x % 32) / 4;
293
- else
294
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
295
- (threadIdx.x % 32) % 4;
296
-
297
- // Precompute which thread should not read memory in which iterations; this is
298
- // needed if there are more threads than required for a certain tilesize or
299
- // when the batchsize is not a multiple of 16.
300
- bool a_sh_wr_pred[a_sh_wr_iters];
301
- #pragma unroll
302
- for (int i = 0; i < a_sh_wr_iters; i++)
303
- a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;
304
- bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
305
-
306
- // To ensure that writing and reading A tiles to/from shared memory, the
307
- // latter in fragment format, is fully bank conflict free, we need to use a
308
- // rather fancy XOR-based layout. The key here is that neither reads nor
309
- // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
310
- // same shared memory banks. Further, it seems (based on NSight-Compute) that
311
- // each warp must also write a consecutive memory segment?
312
- auto transform_a = [&](int i) {
313
- int row = i / a_gl_rd_delta_o;
314
- return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
315
- };
316
- // Since the computation of this remapping is non-trivial and, due to our main
317
- // loop unrolls, all shared memory accesses are static, we simply precompute
318
- // both transformed reads and writes.
319
- int a_sh_wr_trans[a_sh_wr_iters];
320
- #pragma unroll
321
- for (int i = 0; i < a_sh_wr_iters; i++)
322
- a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
323
- int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
324
- #pragma unroll
325
- for (int i = 0; i < b_sh_wr_iters; i++) {
326
- #pragma unroll
327
- for (int j = 0; j < thread_m_blocks; j++)
328
- a_sh_rd_trans[i][j] =
329
- transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
330
- }
331
-
332
- // Since B-accesses have non-constant stride they have to be computed at
333
- // runtime; we break dependencies between subsequent accesses with a tile by
334
- // maintining multiple pointers (we have enough registers), a tiny
335
- // optimization.
336
- const int4* B_ptr[b_sh_wr_iters];
337
- #pragma unroll
338
- for (int i = 0; i < b_sh_wr_iters; i++)
339
- B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
340
-
341
- extern __shared__ int4 sh[];
342
- // Shared memory storage for global fetch pipelines.
343
- int4* sh_a = sh;
344
- int4* sh_b = sh_a + (stages * a_sh_stage);
345
- int4* sh_s = sh_b + (stages * b_sh_stage);
346
- // Register storage for double buffer of shared memory reads.
347
- FragA frag_a[2][thread_m_blocks];
348
- I4 frag_b_quant[2];
349
- FragC frag_c[thread_m_blocks][4][2];
350
- FragS frag_s[2][4];
351
-
352
- // Zero accumulators.
353
- auto zero_accums = [&]() {
354
- #pragma unroll
355
- for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
356
- reinterpret_cast<float*>(frag_c)[i] = 0;
357
- };
358
-
359
- // Asynchronously fetch the next A, B and s tile from global to the next
360
- // shared memory pipeline location.
361
- auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
362
- if (pred) {
363
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
364
- #pragma unroll
365
- for (int i = 0; i < a_sh_wr_iters; i++) {
366
- cp_async4_pred(
367
- &sh_a_stage[a_sh_wr_trans[i]],
368
- &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
369
- a_sh_wr_pred[i]);
370
- }
371
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
372
- #pragma unroll
373
- for (int i = 0; i < b_sh_wr_iters; i++) {
374
- cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr], B_ptr[i]);
375
- B_ptr[i] += b_gl_rd_delta_o;
376
- }
377
- // Only fetch scales if this tile starts a new group
378
- if constexpr (group_blocks != -1) {
379
- // This assumes group_blocks >= thread_k_blocks
380
- // and would need to be modified to support smaller groups.
381
- static_assert(group_blocks >= thread_k_blocks);
382
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
383
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
384
- if (s_sh_wr_pred) cp_async4(&sh_s_stage[s_sh_wr], &s[s_gl_rd]);
385
- s_gl_rd += s_gl_rd_delta;
386
- }
387
- }
388
- }
389
- // Insert a fence even when we are winding down the pipeline to ensure that
390
- // waiting is also correct at this point.
391
- cp_async_fence();
392
- };
393
-
394
- // Wait until the next thread tile has been loaded to shared memory.
395
- auto wait_for_stage = [&]() {
396
- // We only have `stages - 2` active fetches since we are double buffering
397
- // and can only issue the next fetch when it is guaranteed that the previous
398
- // shared memory load is fully complete (as it may otherwise be
399
- // overwritten).
400
- cp_async_wait<stages - 2>();
401
- __syncthreads();
402
- };
403
-
404
- // Load the next sub-tile from the current location in the shared memory pipe
405
- // into the current register buffer.
406
- auto fetch_to_registers = [&](int k, int pipe) {
407
- // It may seem inefficient that we reload the groups for every sub-tile;
408
- // however, this does not seem to be a significant bottleneck, while some
409
- // theoretically better attempts have lead to bad instruction ordering by
410
- // the compiler and correspondingly a noticeable drop in performance.
411
- if constexpr (group_blocks != -1) {
412
- // This assumes group_blocks >= thread_k_blocks
413
- // and would need to be modified to support smaller groups.
414
- static_assert(group_blocks >= thread_k_blocks);
415
- int4* sh_s_stage =
416
- sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
417
- (pipe / (group_blocks / thread_k_blocks)));
418
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
419
- }
420
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
421
- #pragma unroll
422
- for (int i = 0; i < thread_m_blocks; i++)
423
- ldsm4(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
424
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
425
- frag_b_quant[k % 2] = *reinterpret_cast<I4*>(
426
- &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd]);
427
- };
428
-
429
- // Execute the actual tensor core matmul of a sub-tile.
430
- auto matmul = [&](int k) {
431
- // We have the m dimension as the inner loop in order to encourage overlapping
432
- // dequantization and matmul operations.
433
- #pragma unroll
434
- for (int j = 0; j < 4; j++) {
435
- int b_quant = frag_b_quant[k % 2][j];
436
- int b_quant_shift = b_quant >> 8;
437
- FragB frag_b0 = dequant(b_quant);
438
- // If there are no groups, we can just scale the final output once and can
439
- // avoid doing so for each weight.
440
- if (group_blocks != -1) scale(frag_b0, frag_s[k % 2][j], 0);
441
- FragB frag_b1 = dequant(b_quant_shift);
442
- if (group_blocks != -1) scale(frag_b1, frag_s[k % 2][j], 1);
443
- #pragma unroll
444
- for (int i = 0; i < thread_m_blocks; i++) {
445
- mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]);
446
- mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]);
447
- }
448
- }
449
- };
450
-
451
- // Since we slice across the k dimension of a tile in order to increase the
452
- // number of warps while keeping the n dimension of a tile reasonable, we have
453
- // multiple warps that accumulate their partial sums of the same output
454
- // location; which we have to reduce over in the end. We do in shared memory.
455
- auto thread_block_reduce = [&]() {
456
- constexpr int red_off = threads / b_sh_stride / 2;
457
- if (red_off >= 1) {
458
- auto red_idx = threadIdx.x / b_sh_stride;
459
- constexpr int red_sh_stride = b_sh_stride * 4 * 2;
460
- constexpr int red_sh_delta = b_sh_stride;
461
- int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride) +
462
- (threadIdx.x % b_sh_stride);
463
-
464
- // Parallel logarithmic shared memory reduction. We make sure to avoid any
465
- // unnecessary read or write iterations, e.g., for two warps we write only
466
- // once by warp 1 and read only once by warp 0.
467
-
468
- #pragma unroll
469
- for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
470
- #pragma unroll
471
- for (int i = red_off; i > 0; i /= 2) {
472
- if (i <= red_idx && red_idx < 2 * i) {
473
- #pragma unroll
474
- for (int j = 0; j < 4 * 2; j++) {
475
- int red_sh_wr =
476
- red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
477
- if (i < red_off) {
478
- float* c_rd =
479
- reinterpret_cast<float*>(&sh[red_sh_delta * j + red_sh_rd]);
480
- float* c_wr = reinterpret_cast<float*>(&sh[red_sh_wr]);
481
- #pragma unroll
482
- for (int k = 0; k < 4; k++)
483
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
484
- c_rd[k] + c_wr[k];
485
- }
486
- sh[red_sh_wr] =
487
- reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
488
- }
489
- }
490
- __syncthreads();
491
- }
492
- if (red_idx == 0) {
493
- #pragma unroll
494
- for (int i = 0; i < 4 * 2; i++) {
495
- float* c_rd =
496
- reinterpret_cast<float*>(&sh[red_sh_delta * i + red_sh_rd]);
497
- #pragma unroll
498
- for (int j = 0; j < 4; j++)
499
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
500
- c_rd[j];
501
- }
502
- }
503
- __syncthreads();
504
- }
505
- }
506
- };
507
-
508
- // Since multiple threadblocks may process parts of the same column slice, we
509
- // finally have to globally reduce over the results. As the striped
510
- // partitioning minimizes the number of such reductions and our outputs are
511
- // usually rather small, we perform this reduction serially in L2 cache.
512
- auto global_reduce = [&](bool first = false, bool last = false) {
513
- // We are very careful here to reduce directly in the output buffer to
514
- // maximize L2 cache utilization in this step. To do this, we write out
515
- // results in FP16 (but still reduce with FP32 compute).
516
- constexpr int active_threads = 32 * thread_n_blocks / 4;
517
- if (threadIdx.x < active_threads) {
518
- int c_gl_stride = prob_n / 8;
519
- int c_gl_wr_delta_o = 8 * c_gl_stride;
520
- int c_gl_wr_delta_i = 4 * (active_threads / 32);
521
- int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
522
- 4 * (threadIdx.x / 32) + threadIdx.x % 4;
523
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
524
- constexpr int c_sh_wr_delta = active_threads;
525
- auto c_sh_wr = threadIdx.x;
526
-
527
- int row = (threadIdx.x % 32) / 4;
528
-
529
- if (!first) {
530
- // Interestingly, doing direct global accesses here really seems to mess up
531
- // the compiler and lead to slowdowns, hence we also use async-copies even
532
- // though these fetches are not actually asynchronous.
533
- #pragma unroll
534
- for (int i = 0; i < thread_m_blocks * 4; i++) {
535
- cp_async4_pred(
536
- &sh[c_sh_wr + c_sh_wr_delta * i],
537
- &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
538
- c_gl_wr_delta_i * (i % 2)],
539
- i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
540
- }
541
- cp_async_fence();
542
- cp_async_wait<0>();
543
- }
544
-
545
- #pragma unroll
546
- for (int i = 0; i < thread_m_blocks * 4; i++) {
547
- if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) {
548
- if (!first) {
549
- int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta];
550
- #pragma unroll
551
- for (int j = 0; j < 2 * 4; j++) {
552
- reinterpret_cast<float*>(
553
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] +=
554
- __half2float(reinterpret_cast<__half*>(&c_red)[j]);
555
- }
556
- }
557
- if (!last) {
558
- int4 c;
559
- #pragma unroll
560
- for (int j = 0; j < 2 * 4; j++) {
561
- reinterpret_cast<__half*>(&c)[j] =
562
- __float2half(reinterpret_cast<float*>(
563
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]);
564
- }
565
- C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] =
566
- c;
567
- }
568
- }
569
- }
570
- }
571
- };
572
-
573
- // Write out the reduce final result in the correct layout. We only actually
574
- // reshuffle matrix fragments in this step, the reduction above is performed
575
- // in fragment layout.
576
- auto write_result = [&]() {
577
- int c_gl_stride = prob_n / 8;
578
- constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
579
- int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
580
- constexpr int c_sh_rd_delta =
581
- c_sh_stride * (threads / (2 * thread_n_blocks));
582
-
583
- int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
584
- (threadIdx.x % (2 * thread_n_blocks));
585
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
586
- int c_sh_wr =
587
- (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
588
- c_sh_wr += 32 * (threadIdx.x / 32);
589
- int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
590
- (threadIdx.x % (2 * thread_n_blocks));
591
-
592
- int c_gl_wr_end = c_gl_stride * prob_m;
593
-
594
- // We first reorder in shared memory to guarantee the most efficient final
595
- // global write patterns
596
- auto write = [&](int idx, float c0, float c1, FragS& s) {
597
- half2 res = __halves2half2(__float2half(c0), __float2half(c1));
598
- if (group_blocks ==
599
- -1) // for per-column quantization we finally apply the scale here
600
- res = __hmul2(res, s[0]);
601
- ((half2*)sh)[idx] = res;
602
- };
603
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
604
- #pragma unroll
605
- for (int i = 0; i < thread_m_blocks; i++) {
606
- #pragma unroll
607
- for (int j = 0; j < 4; j++) {
608
- int wr = c_sh_wr + 8 * j;
609
- write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
610
- frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
611
- write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
612
- frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
613
- write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
614
- frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
615
- write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
616
- frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
617
- }
618
- c_sh_wr += 16 * (4 * c_sh_stride);
619
- }
620
- }
621
- __syncthreads();
622
-
623
- #pragma unroll
624
- for (int i = 0;
625
- i < ceildiv(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
626
- i++) {
627
- if (c_gl_wr < c_gl_wr_end) {
628
- C[c_gl_wr] = sh[c_sh_rd];
629
- c_gl_wr += c_gl_wr_delta;
630
- c_sh_rd += c_sh_rd_delta;
631
- }
632
- }
633
- };
634
-
635
- // Start global fetch and register load pipelines.
636
- auto start_pipes = [&]() {
637
- #pragma unroll
638
- for (int i = 0; i < stages - 1; i++) fetch_to_shared(i, i, i < slice_iters);
639
- zero_accums();
640
- wait_for_stage();
641
- fetch_to_registers(0, 0);
642
- a_gl_rd += a_gl_rd_delta_o * (stages - 1);
643
- };
644
- start_pipes();
645
-
646
- // Main loop.
647
- while (slice_iters) {
648
- // We unroll over both the global fetch and the register load pipeline to
649
- // ensure all shared memory accesses are static. Note that both pipelines have
650
- // even length meaning that the next iteration will always start at index 0.
651
- #pragma unroll
652
- for (int pipe = 0; pipe < stages;) {
653
- #pragma unroll
654
- for (int k = 0; k < b_sh_wr_iters; k++) {
655
- fetch_to_registers(k + 1, pipe % stages);
656
- if (k == b_sh_wr_iters - 2) {
657
- fetch_to_shared((pipe + stages - 1) % stages, pipe,
658
- slice_iters >= stages);
659
- pipe++;
660
- wait_for_stage();
661
- }
662
- matmul(k);
663
- }
664
- slice_iters--;
665
- if (slice_iters == 0) break;
666
- }
667
- a_gl_rd += a_gl_rd_delta_o * stages;
668
-
669
- // Process results and, if necessary, proceed to the next column slice.
670
- // While this pattern may not be the most readable, other ways of writing
671
- // the loop seemed to noticeably worse performance after compilation.
672
- if (slice_iters == 0) {
673
- cp_async_wait<0>();
674
- bool last = slice_idx == slice_count - 1;
675
- // For per-column scales, we only fetch them here in the final step before
676
- // write-out
677
- if (group_blocks == -1 && last) {
678
- if (s_sh_wr_pred) cp_async4(&sh_s[s_sh_wr], &s[s_gl_rd]);
679
- cp_async_fence();
680
- }
681
- thread_block_reduce();
682
- if (group_blocks == -1 && last) {
683
- cp_async_wait<0>();
684
- __syncthreads();
685
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
686
- reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
687
- reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
688
- }
689
- }
690
- if (slice_count > 1) { // only globally reduce if there is more than one
691
- // block in a slice
692
- barrier_acquire(&locks[slice_col], slice_idx);
693
- global_reduce(slice_idx == 0, last);
694
- barrier_release(&locks[slice_col], last);
695
- }
696
- if (last) // only the last block in a slice actually writes the result
697
- write_result();
698
- slice_row = 0;
699
- slice_col_par++;
700
- slice_col++;
701
- init_slice();
702
- if (slice_iters) {
703
- a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
704
- (threadIdx.x % a_gl_rd_delta_o);
705
- #pragma unroll
706
- for (int i = 0; i < b_sh_wr_iters; i++)
707
- B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
708
- if (slice_col == 0) {
709
- #pragma unroll
710
- for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
711
- }
712
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
713
- start_pipes();
714
- }
715
- }
716
- }
717
-}
718
-
719
-#else
720
-
721
-template <const int threads, // number of threads in a threadblock
722
- const int thread_m_blocks, // number of 16x16 blocks in the m
723
- // dimension (batchsize) of the
724
- // threadblock
725
- const int thread_n_blocks, // same for n dimension (output)
726
- const int thread_k_blocks, // same for k dimension (reduction)
727
- const int stages, // number of stages for the async global->shared
728
- // fetch pipeline
729
- const int group_blocks = -1 // number of consecutive 16x16 blocks
730
- // with a separate quantization scale
731
- >
732
-__global__ void Marlin(
733
- const int4* __restrict__ A, // fp16 input matrix of shape mxk
734
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
735
- int4* __restrict__ C, // fp16 output buffer of shape mxn
736
- const int4* __restrict__ s, // fp16 quantization scales of shape
737
- // (k/groupsize)xn
738
- int prob_m, // batch dimension m
739
- int prob_n, // output dimension n
740
- int prob_k, // reduction dimension k
741
- int* locks // extra global storage for barrier synchronization
742
-) {
743
- // Marlin is not implemented yet for SM < 8.0
744
- assert(false);
745
- return;
746
-}
747
-
748
-#endif
749
-
750
-// 8 warps are a good choice since every SM has 4 schedulers and having more
751
-// than 1 warp per schedule allows some more latency hiding. At the same time,
752
-// we want relatively few warps to have many registers per warp and small tiles.
753
-const int USER_THREADS =
754
- 256; // Note: This is only used with user-provided thread_k/n
755
-const int STAGES = 4; // 4 pipeline stages fit into shared memory
756
-const int SHARED_MEM =
757
- 96 * 1024; // max shared memory on compute capability 8.6 (< 8.0)
758
-
759
-static constexpr int min_thread_n = 64;
760
-static constexpr int min_thread_k = 64;
761
-
762
-static constexpr int tile_size = 16;
763
-static constexpr int max_par = 16;
764
-
765
-static constexpr int pack_factor_4bit =
766
- 8; // We have 8 4-bit vals inside a 32 bit
767
-
768
-#define __CALL_IF(THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
769
- GROUP_BLOCKS, NUM_THREADS) \
770
- else if (thread_m_blocks == THREAD_M_BLOCKS && \
771
- thread_n_blocks == THREAD_N_BLOCKS && \
772
- thread_k_blocks == THREAD_K_BLOCKS && \
773
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS) { \
774
- cudaFuncSetAttribute(Marlin<NUM_THREADS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, \
775
- THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>, \
776
- cudaFuncAttributeMaxDynamicSharedMemorySize, \
777
- SHARED_MEM); \
778
- Marlin<NUM_THREADS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
779
- STAGES, GROUP_BLOCKS><<<blocks, NUM_THREADS, SHARED_MEM, stream>>>( \
780
- A_ptr, B_ptr, C_ptr, s_ptr, prob_m, prob_n, prob_k, locks); \
781
- }
782
-
783
-typedef struct {
784
- int thread_k;
785
- int thread_n;
786
- int num_threads;
787
-} thread_config_t;
788
-
789
-thread_config_t small_batch_thread_configs[] = {
790
- // Ordered by priority
791
-
792
- // thread_k, thread_n, num_threads
793
- {128, 128, 256}, // Default
794
- {128, 64, 128}, // Reduce N 2X, same K
795
- {64, 256, 256}, // Reduce K 2X, increase N 2X
796
- {64, 128, 128}, // Reduce K 2X, same N
797
-};
798
-
799
-thread_config_t large_batch_thread_configs[] = {
800
- // Ordered by priority
801
-
802
- // thread_k, thread_n, num_threads
803
- {64, 256, 256}, // Default
804
- {128, 128, 256}, // Reduce N 2X, increase K 2X
805
- {64, 128, 128}, // Reduce N 2X, same K
806
- {128, 64, 128}, // Reduce N 4X, increase K 2X
807
-};
808
-
809
-bool is_valid_config(thread_config_t const& th_config, int prob_m, int prob_n,
810
- int prob_k) {
811
- // Sanity
812
- if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
813
- th_config.num_threads == -1) {
814
- return false;
815
- }
816
-
817
- // Verify K/N are divisible by thread K/N
818
- if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) {
819
- return false;
820
- }
821
-
822
- // thread_k can be only 128 or 64 (because it must be less than groupsize
823
- // which is 128)
824
- if (th_config.thread_k != 128 && th_config.thread_k != 64) {
825
- return false;
826
- }
827
-
828
- // Verify min for thread K/N
829
- if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) {
830
- return false;
831
- }
832
-
833
- // num_threads must be at least 128 (= 4 warps)
834
- if (th_config.num_threads < 128) {
835
- return false;
836
- }
837
-
838
- return true;
839
-}
840
-
841
-thread_config_t determine_thread_config(int prob_m, int prob_n, int prob_k) {
842
- if (prob_m <= 16) {
843
- for (auto th_config : small_batch_thread_configs) {
844
- if (is_valid_config(th_config, prob_m, prob_n, prob_k)) {
845
- return th_config;
846
- }
847
- }
848
-
849
- } else {
850
- for (auto th_config : large_batch_thread_configs) {
851
- if (is_valid_config(th_config, prob_m, prob_n, prob_k)) {
852
- return th_config;
853
- }
854
- }
855
- }
856
-
857
- return thread_config_t{-1, -1, -1};
858
-}
859
-
860
-#define CALL_IF(N_BLOCKS, K_BLOCKS, NUM_THREADS) \
861
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
862
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
863
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
864
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
865
- __CALL_IF(2, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
866
- __CALL_IF(2, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
867
- __CALL_IF(3, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
868
- __CALL_IF(3, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
869
- __CALL_IF(4, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
870
- __CALL_IF(4, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS)
871
-
872
-void marlin_cuda(const void* A, const void* B, void* C, void* s, int prob_m,
873
- int prob_n, int prob_k, void* workspace, int groupsize = -1,
874
- int dev = 0, cudaStream_t stream = 0, int thread_k = -1,
875
- int thread_n = -1, int sms = -1, int max_par = 16) {
876
- int tot_m = prob_m;
877
- int tot_m_blocks = ceildiv(tot_m, 16);
878
- int pad = 16 * tot_m_blocks - tot_m;
879
-
880
- if (sms == -1)
881
- cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
882
-
883
- // Set thread config
884
- thread_config_t th_config;
885
- if (thread_k != -1 && thread_n != -1) {
886
- // User-defined config
887
- th_config = thread_config_t{thread_k, thread_n, USER_THREADS};
888
- } else {
889
- // Auto config
890
- th_config = determine_thread_config(prob_m, prob_n, prob_k);
891
- }
892
-
893
- if (!is_valid_config(th_config, prob_m, prob_n, prob_k)) {
894
- throw std::runtime_error(
895
- "Invalid thread config: thread_k = " + str(th_config.thread_k) +
896
- ", thread_n = " + str(th_config.thread_n) +
897
- ", num_threads = " + str(th_config.num_threads) + " for MKN = [" +
898
- str(prob_m) + ", " + str(prob_k) + ", " + str(prob_n) + "]");
899
- }
900
-
901
- // Uncomment for debug
902
- // std::cout << "Using thread_config: thread_k = " + str(th_config.thread_k) +
903
- // ", thread_n = " + str(th_config.thread_n) +
904
- // ", num_threads = " + str(th_config.num_threads) + " for
905
- // MKN = [" + str(prob_m) +
906
- // ", " + str(prob_k) + ", " + str(prob_n) + "]\n";
907
-
908
- int num_threads = th_config.num_threads;
909
- thread_k = th_config.thread_k;
910
- thread_n = th_config.thread_n;
911
-
912
- int thread_k_blocks = thread_k / 16;
913
- int thread_n_blocks = thread_n / 16;
914
- int group_blocks = (groupsize == -1) ? -1 : groupsize / 16;
915
- int blocks = sms;
916
-
917
- if (prob_m == 0 || prob_n == 0 || prob_k == 0) {
918
- return;
919
- }
920
-
921
- TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
922
- " is not divisible by thread_n = ", thread_n);
923
- TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
924
- " is not divisible by thread_k = ", thread_k);
925
- if (group_blocks != -1) {
926
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
927
- " is not divisible by group_blocks = ", group_blocks);
928
- }
929
-
930
- const int4* A_ptr = (const int4*)A;
931
- const int4* B_ptr = (const int4*)B;
932
- int4* C_ptr = (int4*)C;
933
- const int4* s_ptr = (const int4*)s;
934
-
935
- int* locks = (int*)workspace;
936
-
937
- for (int i = 0; i < tot_m_blocks; i += 4) {
938
- int thread_m_blocks = tot_m_blocks - i;
939
- prob_m = tot_m - 16 * i;
940
- int par = 1;
941
- if (thread_m_blocks > 4) {
942
- // Note that parallel > 1 currently only works for inputs without any
943
- // padding
944
- par = (16 * thread_m_blocks - pad) / 64;
945
- if (par > max_par) par = max_par;
946
- prob_m = 64 * par;
947
- i += 4 * (par - 1);
948
- thread_m_blocks = 4;
949
- }
950
-
951
- // For compilation speed, we only define the kernel configurations that have
952
- // seemed useful (in terms of performance) in our testing, however many more
953
- // are, in principle, possible.
954
- if (false) {
955
- }
956
- CALL_IF(8, 8, 256)
957
- CALL_IF(16, 4, 256)
958
- CALL_IF(8, 4, 128)
959
- CALL_IF(4, 8, 128)
960
- else {
961
- throw std::runtime_error("Unsupported shapes: MKN = [" + str(prob_m) +
962
- ", " + str(prob_k) + ", " + str(prob_n) + "]" +
963
- ", groupsize = " + str(groupsize) +
964
- ", thread_m_blocks = " + str(thread_m_blocks) +
965
- ", thread_n_blocks = " + str(thread_n_blocks) +
966
- ", thread_k_blocks = " + str(thread_k_blocks));
967
- }
968
-
969
- A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par;
970
- C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par;
971
- }
972
-}
973
-
974
-} // namespace marlin_dense
975
-
976
-torch::Tensor marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
977
- torch::Tensor& b_scales, torch::Tensor& workspace,
978
- int64_t size_m, int64_t size_n, int64_t size_k) {
979
- // Verify M
980
- TORCH_CHECK(size_m == a.size(0),
981
- "Shape mismatch: a.size(0) = " + str(a.size(0)) +
982
- ", size_m = " + str(size_m));
983
-
984
- // Verify K
985
- TORCH_CHECK(size_k == a.size(1),
986
- "Shape mismatch: a.size(1) = " + str(a.size(1)) +
987
- ", size_k = " + str(size_k));
988
- TORCH_CHECK(size_k % marlin_dense::tile_size == 0,
989
- "size_k = " + str(size_k) + " is not divisible by tile_size = " +
990
- str(marlin_dense::tile_size));
991
- TORCH_CHECK((size_k / marlin_dense::tile_size) == b_q_weight.size(0),
992
- "Shape mismatch: b_q_weight.size(0) = " +
993
- str(b_q_weight.size(0)) + ", size_k = " + str(size_k) +
994
- ", tile_size = " + str(marlin_dense::tile_size));
995
-
996
- // Verify N
997
- TORCH_CHECK(b_scales.size(1) == size_n,
998
- "b_scales.size(1) = " + str(b_scales.size(1)) +
999
- ", size_n = " + str(size_n));
1000
- TORCH_CHECK(
1001
- b_q_weight.size(1) % marlin_dense::tile_size == 0,
1002
- "b_q_weight.size(1) = " + str(b_q_weight.size(1)) +
1003
- " is not divisible by tile_size = " + str(marlin_dense::tile_size));
1004
-
1005
- int actual_size_n = (b_q_weight.size(1) / marlin_dense::tile_size) *
1006
- marlin_dense::pack_factor_4bit;
1007
- TORCH_CHECK(
1008
- size_n == actual_size_n,
1009
- "size_n = " + str(size_n) + ", actual_size_n = " + str(actual_size_n));
1010
-
1011
- // Verify A device and strides
1012
- TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
1013
- TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
1014
-
1015
- // Verify B device and strides
1016
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
1017
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
1018
-
1019
- // Verify scales device and strides
1020
- TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
1021
- TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
1022
-
1023
- // Alloc C matrix
1024
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
1025
- auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
1026
- torch::Tensor c = torch::empty({size_m, size_n}, options);
1027
-
1028
- // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
1029
- // auto -1)
1030
- int thread_k = -1;
1031
- // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
1032
- // auto -1)
1033
- int thread_n = -1;
1034
- // sms: number of SMs to use for the kernel (can usually be left as auto -1)
1035
- int sms = -1;
1036
-
1037
- // Detect groupsize
1038
- if (b_scales.size(0) != 1) {
1039
- TORCH_CHECK(size_k % b_scales.size(0) == 0,
1040
- "size_k = " + str(size_k) +
1041
- ", is not divisible by b_scales.size(0) = " +
1042
- str(b_scales.size(0)));
1043
- }
1044
- int groupsize = b_scales.size(0) == 1 ? -1 : size_k / b_scales.size(0);
1045
-
1046
- // Verify groupsize
1047
- TORCH_CHECK(groupsize == -1 || groupsize == 128,
1048
- "Unexpected groupsize = " + str(groupsize));
1049
-
1050
- // Verify workspace size
1051
- TORCH_CHECK(size_n % marlin_dense::min_thread_n == 0,
1052
- "size_n = " + str(size_n) +
1053
- ", is not divisible by min_thread_n = " +
1054
- str(marlin_dense::min_thread_n));
1055
- int min_workspace_size =
1056
- (size_n / marlin_dense::min_thread_n) * marlin_dense::max_par;
1057
- TORCH_CHECK(workspace.numel() >= min_workspace_size,
1058
- "workspace.numel = " + str(workspace.numel()) +
1059
- " is below min_workspace_size = " + str(min_workspace_size));
1060
-
1061
- int dev = a.get_device();
1062
- marlin_dense::marlin_cuda(a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(),
1063
- b_scales.data_ptr(), size_m, size_n, size_k,
1064
- workspace.data_ptr(), groupsize, dev,
1065
- at::cuda::getCurrentCUDAStream(dev), thread_k,
1066
- thread_n, sms, marlin_dense::max_par);
1067
-
1068
- return c;
1069
-}
1070
-
1071
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1072
- m.impl("marlin_gemm", &marlin_gemm);
1073
-}
csrc/quantization/marlin/qqq/marlin_qqq_gemm_kernel.cu
deleted
-1248
@@ -1,1248 +0,0 @@
1
-/*
2
- * Adapted from
3
- * https://github.com/IST-DASLab/marlin/blob/master/marlin/marlin_cuda_kernel.cu
4
- * https://github.com/IST-DASLab/marlin/blob/master/marlin/marlin_cuda.cpp
5
- * Modified by HandH1998
6
- * Copyright (C) 2024 HandH1998
7
- * Copyright (C) Marlin.2024 Elias Frantar
8
- *
9
- * Licensed under the Apache License, Version 2.0 (the "License");
10
- * you may not use this file except in compliance with the License.
11
- * You may obtain a copy of the License at
12
- *
13
- * http://www.apache.org/licenses/LICENSE-2.0
14
- *
15
- * Unless required by applicable law or agreed to in writing, software
16
- * distributed under the License is distributed on an "AS IS" BASIS,
17
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
- * See the License for the specific language governing permissions and
19
- * limitations under the License.
20
- */
21
-
22
-#include <torch/all.h>
23
-
24
-#include <ATen/cuda/CUDAContext.h>
25
-#include <c10/cuda/CUDAGuard.h>
26
-#include <cuda.h>
27
-#include <cuda_fp16.h>
28
-#include <cuda_runtime.h>
29
-
30
-#include <iostream>
31
-
32
-#include "../dense/common/base.h"
33
-#include "core/registration.h"
34
-
35
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
36
- #include "../dense/common/mem.h"
37
-#endif
38
-
39
-template <typename T>
40
-inline std::string str(T x) {
41
- return std::to_string(x);
42
-}
43
-
44
-namespace {
45
-
46
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
47
-
48
-using I4 = Vec<int, 4>;
49
-// Matrix fragments for tensor core instructions; their precise layout is
50
-// documented here:
51
-// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-integer-type
52
-using FragA = Vec<uint32_t, 2>;
53
-using FragB = Vec<uint32_t, 1>;
54
-using FragC = Vec<int, 4>;
55
-using FragS_GROUP = Vec<half2, 1>; // weight per-group quantization scales
56
-using FragS_CHANNEL =
57
- Vec<float, 2>; // weight per-channel quantization scales or activaton
58
- // per-token quantization scales
59
-
60
-// NOTE(HandH1998): cp.async.cg only support BYTES = 16, however,
61
-// cp.async.ca can support BYTES = 4, 8, 16;
62
-// as s_tok's shape is equal to prob_m, we need set s_tok to float type,
63
-// and cp_size = 1 float, i.e., 4 BYTES
64
-// Asynchronous global->shared copy for activation quantizaton scales s_tok
65
-__device__ inline void cp_async1(void* smem_ptr, const void* glob_ptr) {
66
- const int BYTES = 4;
67
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
68
- asm volatile(
69
- "{\n"
70
- " cp.async.ca.shared.global [%0], [%1], %2;\n"
71
- "}\n" ::"r"(smem),
72
- "l"(glob_ptr), "n"(BYTES));
73
-}
74
-
75
-// m16n8k16 tensor core mma instruction with int8 inputs and int32
76
-// output/accumulation.
77
-__device__ inline void mma(const FragA& a_frag, const FragB& frag_b,
78
- FragC& frag_c) {
79
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
80
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
81
- int* c = reinterpret_cast<int*>(&frag_c);
82
- asm volatile(
83
- "mma.sync.aligned.m16n8k16.row.col.satfinite.s32.s8.s8.s32 "
84
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
85
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
86
- : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
87
- "r"(c[3]));
88
-}
89
-
90
-// Instruction for loading a full 16x16 matrix fragment of operand A from shared
91
-// memory, directly in int8 tensor core layout.
92
-__device__ inline void ldsm4(FragA& frag_a, const void* smem_ptr) {
93
- uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
94
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
95
- asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
96
- : "=r"(a[0]), "=r"(a[1])
97
- : "r"(smem));
98
-}
99
-
100
-inline __device__ half2 float2_to_half2(float2 f) {
101
- uint32_t res;
102
- // NOTE(HandH1998): h0,h1 should be uint16_t, not half
103
- uint16_t h0, h1;
104
- asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(h0) : "f"(f.x));
105
- asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(h1) : "f"(f.y));
106
- asm volatile("mov.b32 %0, {%1, %2};\n" : "=r"(res) : "h"(h0), "h"(h1));
107
- return reinterpret_cast<half2&>(res);
108
-}
109
-
110
-inline __device__ float int32_to_float(int h) {
111
- float res;
112
- asm volatile("cvt.rn.f32.s32 %0, %1;\n" : "=f"(res) : "r"(h));
113
- return res;
114
-}
115
-
116
-// Lookup-table based 3-input logical operation; explicitly used for
117
-// dequantization as the compiler does not seem to automatically recognize it in
118
-// all cases.
119
-template <int lut>
120
-__device__ inline int lop3(int a, int b, int c) {
121
- int res;
122
- asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
123
- : "=r"(res)
124
- : "r"(a), "r"(b), "r"(c), "n"(lut));
125
- return res;
126
-}
127
-
128
-// Efficiently dequantize an int32 value into a full B-fragment of 4 int8 values
129
-// for weight per channel dequant.
130
-__device__ inline FragB dequant_per_channel(int q) {
131
- static constexpr int MASK = 0xf0f0f0f0;
132
- FragB frag_b;
133
- frag_b[0] = (q & MASK);
134
- return frag_b;
135
-}
136
-
137
-// Efficiently dequantize an int32 value into a full B-fragment of 4 int8 values
138
-// for weight per group dequant.
139
-__device__ inline FragB dequant_per_group(int q, FragS_GROUP& frag_s, int i) {
140
- static constexpr uint32_t LO = 0x000f000f;
141
- static constexpr uint32_t HI = 0x00f000f0;
142
- static constexpr uint32_t EX = 0x64006400;
143
- // Guarantee that the `(a & b) | c` operations are LOP3s.
144
- uint32_t t0 = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
145
- uint32_t t1 = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
146
- // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
147
- // directly into `SUB` and `ADD`.
148
- static constexpr uint32_t SUB = 0x64086408;
149
- static constexpr uint32_t MUL = 0x2c002c00;
150
- static constexpr uint32_t ADD = 0xd480d480;
151
- *reinterpret_cast<half2*>(&t0) = __hsub2(
152
- *reinterpret_cast<half2*>(&t0), *reinterpret_cast<const half2*>(&SUB));
153
- *reinterpret_cast<half2*>(&t1) = __hfma2(
154
- *reinterpret_cast<half2*>(&t1), *reinterpret_cast<const half2*>(&MUL),
155
- *reinterpret_cast<const half2*>(&ADD));
156
-
157
- uint16_t s = reinterpret_cast<uint16_t*>(&frag_s)[i];
158
- uint32_t double_s;
159
- // pack 2xfp16 to half2
160
- asm volatile("mov.b32 %0, {%1, %2};\n" : "=r"(double_s) : "h"(s), "h"(s));
161
- // dequant and convert 4 half to 4 uint8 (be placed at the low 8 bits of 4
162
- // half, respectively)
163
- static constexpr uint32_t MAGIC_NUM = 0x64806480;
164
- *reinterpret_cast<half2*>(&t0) = __hfma2(
165
- *reinterpret_cast<half2*>(&t0), *reinterpret_cast<half2*>(&double_s),
166
- *reinterpret_cast<const half2*>(&MAGIC_NUM));
167
- *reinterpret_cast<half2*>(&t1) = __hfma2(
168
- *reinterpret_cast<half2*>(&t1), *reinterpret_cast<half2*>(&double_s),
169
- *reinterpret_cast<const half2*>(&MAGIC_NUM));
170
- // take out the 4 uint8 from 4 half, then convert them to 4 int8 and pack 4
171
- // int8 into 1 uint32
172
- FragB frag_b;
173
- uint32_t uint8s;
174
- static constexpr uint32_t MASK_0246 = 0x6420;
175
- static constexpr uint32_t UINT8s_TO_INT8s_MASK = 0x80808080;
176
- asm volatile("prmt.b32 %0,%1,%2,%3;\n"
177
- : "=r"(uint8s)
178
- : "r"(t0), "r"(t1), "n"(MASK_0246));
179
- frag_b[0] = (uint8s ^ UINT8s_TO_INT8s_MASK);
180
- return frag_b;
181
-}
182
-
183
-template <const int threads, // number of threads in a threadblock
184
- const int thread_m_blocks, // number of 16x16 blocks in the m
185
- // dimension (batchsize) of the
186
- // threadblock
187
- const int thread_n_blocks, // same for n dimension (output)
188
- const int thread_k_blocks, // same for k dimension (reduction)
189
- const int stages, // number of stages for the async global->shared
190
- // fetch pipeline
191
- const int group_blocks = -1 // number of consecutive 16x16 blocks
192
- // with a separate quantization scale
193
- >
194
-__global__ void Marlin(
195
- const int4* __restrict__ A, // int8 input matrix of shape mxk
196
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
197
- int4* __restrict__ C, // int32 global_reduce buffer of shape
198
- // (max_par*16*4)xn, as int8 tensor core's output is
199
- // int32 dtype
200
- int4* __restrict__ D, // fp16 output buffer of shape mxn
201
- const float* __restrict__ s_tok, // fp32 activation per-token quantization
202
- // scales of shape mx1
203
- const int4* __restrict__ s_ch, // fp32 weight per-channel quantization
204
- // scales of shape 1xn
205
- const int4* __restrict__ s_group, // fp16 weight per-group quantization
206
- // scales of shape (k/groupsize)xn, when
207
- // group_blocks=-1, it should be nullptr
208
- int prob_m, // batch dimension m
209
- int prob_n, // output dimension n
210
- int prob_k, // reduction dimension k
211
- int* locks // extra global storage for barrier synchronization
212
-) {
213
- // Each threadblock processes one "stripe" of the B matrix with (roughly) the
214
- // same size, which might involve multiple column "slices" (of width 16 *
215
- // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
216
- // example:
217
- // 0 1 3
218
- // 0 2 3
219
- // 1 2 4
220
- // While this kind of partitioning makes things somewhat more complicated, it
221
- // ensures good utilization of all SMs for many kinds of shape and GPU
222
- // configurations, while requiring as few slow global cross-threadblock
223
- // reductions as possible.
224
-
225
- // For larger GEMMs we run multiple batchsize 64 versions in parallel for a
226
- // better partitioning with less reductions
227
- int parallel = 1;
228
- if (prob_m > 16 * thread_m_blocks) {
229
- parallel = prob_m / (16 * thread_m_blocks);
230
- prob_m = 16 * thread_m_blocks;
231
- }
232
-
233
- int k_tiles = prob_k / 16 / thread_k_blocks;
234
- int n_tiles = prob_n / 16 / thread_n_blocks;
235
- int iters = ceildiv(k_tiles * n_tiles * parallel, gridDim.x);
236
- // Ensure that the number of tiles in each stripe is a multiple of the
237
- // groupsize; this avoids an annoying special case where a stripe starts in
238
- // the middle of group.
239
- if constexpr (group_blocks != -1)
240
- iters = (group_blocks / thread_k_blocks) *
241
- ceildiv(iters, (group_blocks / thread_k_blocks));
242
-
243
- int slice_row = (iters * blockIdx.x) % k_tiles;
244
- int slice_col_par = (iters * blockIdx.x) / k_tiles;
245
- int slice_col = slice_col_par;
246
- int slice_iters; // number of threadblock tiles in the current slice
247
- int slice_count =
248
- 0; // total number of active threadblocks in the current slice
249
- int slice_idx; // index of threadblock in current slice; numbered bottom to
250
- // top
251
-
252
- // We can easily implement parallel problem execution by just remapping
253
- // indices and advancing global pointers
254
- if (slice_col_par >= n_tiles) {
255
- A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 16;
256
- C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 4;
257
- D += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
258
- s_tok += (slice_col_par / n_tiles) * 16 * thread_m_blocks;
259
- locks += (slice_col_par / n_tiles) * n_tiles;
260
- slice_col = slice_col_par % n_tiles;
261
- }
262
-
263
- // Compute all information about the current slice which is required for
264
- // synchronization.
265
- auto init_slice = [&]() {
266
- slice_iters =
267
- iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
268
- if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
269
- if (slice_iters == 0) return;
270
- if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
271
- slice_count = 1;
272
- slice_idx = 0;
273
- int col_first = iters * ceildiv(k_tiles * slice_col_par, iters);
274
- if (col_first <= k_tiles * (slice_col_par + 1)) {
275
- int col_off = col_first - k_tiles * slice_col_par;
276
- slice_count = ceildiv(k_tiles - col_off, iters);
277
- if (col_off > 0) slice_count++;
278
- int delta_first = iters * blockIdx.x - col_first;
279
- if (delta_first < 0 || (col_off == 0 && delta_first == 0))
280
- slice_idx = slice_count - 1;
281
- else {
282
- slice_idx = slice_count - 1 - delta_first / iters;
283
- if (col_off > 0) slice_idx--;
284
- }
285
- }
286
- if (slice_col == n_tiles) {
287
- A += 16 * thread_m_blocks * prob_k / 16;
288
- C += 16 * thread_m_blocks * prob_n / 4;
289
- D += 16 * thread_m_blocks * prob_n / 8;
290
- s_tok += 16 * thread_m_blocks;
291
- locks += n_tiles;
292
- slice_col = 0;
293
- }
294
- };
295
- init_slice();
296
-
297
- int a_gl_stride = prob_k / 16; // stride of the A matrix in global memory
298
- // We typically use `constexpr` to indicate that this value is a compile-time
299
- // constant
300
- constexpr int a_sh_stride =
301
- 16 * thread_k_blocks / 16; // stride of an A matrix tile in shared memory
302
- constexpr int a_gl_rd_delta_o =
303
- 16 * thread_k_blocks /
304
- 16; // delta between subsequent A tiles in global memory
305
- int a_gl_rd_delta_i =
306
- a_gl_stride *
307
- (threads / a_gl_rd_delta_o); // between subsequent accesses within a tile
308
- constexpr int a_sh_wr_delta =
309
- a_sh_stride *
310
- (threads / a_gl_rd_delta_o); // between shared memory writes
311
- constexpr int a_sh_rd_delta_o =
312
- 1 * ((threads / 32) /
313
- (thread_n_blocks / 4)); // between shared memory tile reads
314
- constexpr int a_sh_rd_delta_i =
315
- a_sh_stride * 16; // within a shared memory tile
316
- constexpr int a_sh_stage =
317
- a_sh_stride * (16 * thread_m_blocks); // overall size of a tile
318
- constexpr int a_sh_wr_iters =
319
- ceildiv(a_sh_stage,
320
- a_sh_wr_delta); // number of shared write iterations for a tile
321
-
322
- int b_gl_stride = 16 * prob_n / 32;
323
- constexpr int b_sh_stride = 32 * thread_n_blocks / 4;
324
- int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
325
- int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride);
326
- constexpr int b_sh_wr_delta = threads;
327
- constexpr int b_sh_rd_delta = threads;
328
- constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
329
- constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
330
-
331
- constexpr int s_tok_sh_stride = 16 * thread_m_blocks;
332
-
333
- constexpr int s_ch_sh_stride = 16 * thread_n_blocks / 4;
334
-
335
- int s_group_gl_stride = prob_n / 8;
336
- constexpr int s_group_sh_stride = 16 * thread_n_blocks / 8;
337
- constexpr int s_group_sh_stage = s_group_sh_stride;
338
- int s_group_gl_rd_delta = s_group_gl_stride;
339
-
340
- // Global A read index of current thread.
341
- int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
342
- (threadIdx.x % a_gl_rd_delta_o);
343
- a_gl_rd += a_gl_rd_delta_o * slice_row;
344
- // Shared write index of current thread.
345
- int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
346
- (threadIdx.x % a_gl_rd_delta_o);
347
- // Shared read index.
348
- // NOTE(HandH1998): int8 input a only need 16 threads to load 16x16 matrix
349
- int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % 16);
350
- a_sh_rd += 1 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
351
-
352
- int b_gl_rd =
353
- b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
354
- b_gl_rd += b_sh_stride * slice_col;
355
- b_gl_rd += b_gl_rd_delta_o * slice_row;
356
- auto b_sh_wr = threadIdx.x;
357
- auto b_sh_rd = threadIdx.x;
358
-
359
- auto s_tok_gl_rd = threadIdx.x;
360
- // NOTE(HandH1998): activation scale s_tok need shuffle to [0, 8, 1, 9, 2, 10,
361
- // 3, 11, 4, 12, 5, 13, 6, 14, 7, 15] for example, 0, 8 row scales serve for
362
- // thread 0, 1, 2, 3. For more details, refer to mma operand A layout as
363
- // s_tok's size is not fixed, we can not shuffle before inference we shuffle
364
- // it when fetching s_tok from global memory to shared memory, that's why
365
- // s_tok_sh_wr is like this
366
- int s_tok_sh_wr =
367
- (threadIdx.x / 16) * 16 + (threadIdx.x % 8) * 2 + (threadIdx.x % 16) / 8;
368
- int s_tok_sh_rd = (threadIdx.x % 32) / 4;
369
- bool s_tok_sh_wr_pred = threadIdx.x < prob_m;
370
-
371
- auto s_ch_gl_rd = s_ch_sh_stride * slice_col + threadIdx.x;
372
- auto s_ch_sh_wr = threadIdx.x;
373
- int s_ch_sh_rd = 16 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
374
- 2 * ((threadIdx.x % 32) % 4);
375
- bool s_ch_sh_wr_pred = threadIdx.x < s_ch_sh_stride;
376
-
377
- int s_group_gl_rd, s_group_sh_wr, s_group_sh_rd;
378
- bool s_group_sh_wr_pred;
379
- if constexpr (group_blocks != -1) {
380
- s_group_gl_rd =
381
- s_group_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
382
- s_group_sh_stride * slice_col + threadIdx.x;
383
- s_group_sh_wr = threadIdx.x;
384
- // NOTE(HandH1998): s_group_sh_rd is related to mma output C
385
- s_group_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
386
- (threadIdx.x % 32) / 4;
387
- s_group_sh_wr_pred = threadIdx.x < s_group_sh_stride;
388
- }
389
-
390
- // Precompute which thread should not read memory in which iterations; this is
391
- // needed if there are more threads than required for a certain tilesize or
392
- // when the batchsize is not a multiple of 16.
393
- bool a_sh_wr_pred[a_sh_wr_iters];
394
- #pragma unroll
395
- for (int i = 0; i < a_sh_wr_iters; i++)
396
- a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;
397
-
398
- // To ensure that writing and reading A tiles to/from shared memory, the
399
- // latter in fragment format, is fully bank conflict free, we need to use a
400
- // rather fancy XOR-based layout. The key here is that neither reads nor
401
- // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
402
- // same shared memory banks. Further, it seems (based on NSight-Compute) that
403
- // each warp must also write a consecutive memory segment?
404
- auto transform_a = [&](int i) {
405
- int row = i / a_gl_rd_delta_o;
406
- return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
407
- };
408
- // Since the computation of this remapping is non-trivial and, due to our main
409
- // loop unrolls, all shared memory accesses are static, we simply precompute
410
- // both transformed reads and writes.
411
- int a_sh_wr_trans[a_sh_wr_iters];
412
- #pragma unroll
413
- for (int i = 0; i < a_sh_wr_iters; i++)
414
- a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
415
- int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
416
- #pragma unroll
417
- for (int i = 0; i < b_sh_wr_iters; i++) {
418
- #pragma unroll
419
- for (int j = 0; j < thread_m_blocks; j++)
420
- a_sh_rd_trans[i][j] =
421
- transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
422
- }
423
-
424
- // Since B-accesses have non-constant stride they have to be computed at
425
- // runtime; we break dependencies between subsequent accesses with a tile by
426
- // maintining multiple pointers (we have enough registers), a tiny
427
- // optimization.
428
- const int4* B_ptr[b_sh_wr_iters];
429
- #pragma unroll
430
- for (int i = 0; i < b_sh_wr_iters; i++)
431
- B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
432
-
433
- extern __shared__ int4 sh[];
434
- // Shared memory storage for global fetch pipelines.
435
- // NOTE(HandH1998): stages need >= 4, otherwise, sh_s_tok = sh + max(stages *
436
- // a_sh_stage + stages * b_sh_stage, 4 * stages * a_sh_stage)
437
- int4* sh_a = sh;
438
- int4* sh_b = sh_a + (stages * a_sh_stage);
439
- int4* sh_s_tok = sh_b + (stages * b_sh_stage);
440
- int4* sh_s_ch = sh_s_tok + s_tok_sh_stride;
441
- int4* sh_s_group = sh_s_ch + s_ch_sh_stride;
442
-
443
- // Register storage for double buffer of shared memory reads.
444
- FragA frag_a[2][thread_m_blocks];
445
- I4 frag_b_quant[2];
446
- FragC frag_c[thread_m_blocks][4][2];
447
- FragS_GROUP frag_s_group[2][4];
448
- FragS_CHANNEL frag_s_tok[thread_m_blocks];
449
- FragS_CHANNEL frag_s_ch[2][4];
450
-
451
- // Zero accumulators.
452
- auto zero_accums = [&]() {
453
- #pragma unroll
454
- for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
455
- reinterpret_cast<int*>(frag_c)[i] = 0;
456
- };
457
-
458
- // Asynchronously fetch the next A, B and s tile from global to the next
459
- // shared memory pipeline location.
460
- auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
461
- if (pred) {
462
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
463
- #pragma unroll
464
- for (int i = 0; i < a_sh_wr_iters; i++) {
465
- cp_async4_pred(
466
- &sh_a_stage[a_sh_wr_trans[i]],
467
- &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
468
- a_sh_wr_pred[i]);
469
- }
470
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
471
- #pragma unroll
472
- for (int i = 0; i < b_sh_wr_iters; i++) {
473
- cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr], B_ptr[i]);
474
- B_ptr[i] += b_gl_rd_delta_o;
475
- }
476
- // Only fetch scales if this tile starts a new group
477
- if constexpr (group_blocks != -1) {
478
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
479
- int4* sh_s_group_stage = sh_s_group + s_group_sh_stage * pipe;
480
- if (s_group_sh_wr_pred)
481
- cp_async4(&sh_s_group_stage[s_group_sh_wr],
482
- &s_group[s_group_gl_rd]);
483
- s_group_gl_rd += s_group_gl_rd_delta;
484
- }
485
- }
486
- }
487
- // Insert a fence even when we are winding down the pipeline to ensure that
488
- // waiting is also correct at this point.
489
- cp_async_fence();
490
- };
491
-
492
- // Wait until the next thread tile has been loaded to shared memory.
493
- auto wait_for_stage = [&]() {
494
- // We only have `stages - 2` active fetches since we are double buffering
495
- // and can only issue the next fetch when it is guaranteed that the previous
496
- // shared memory load is fully complete (as it may otherwise be
497
- // overwritten).
498
- cp_async_wait<stages - 2>();
499
- __syncthreads();
500
- };
501
-
502
- // Load the next sub-tile from the current location in the shared memory pipe
503
- // into the current register buffer.
504
- auto fetch_to_registers = [&](int k, int pipe) {
505
- // It may seem inefficient that we reload the groups for every sub-tile;
506
- // however, this does not seem to be a significant bottleneck, while some
507
- // theoretically better attempts have lead to bad instruction ordering by
508
- // the compiler and correspondingly a noticeable drop in performance.
509
- if constexpr (group_blocks != -1) {
510
- int4* sh_s_group_stage =
511
- sh_s_group +
512
- s_group_sh_stage * ((group_blocks / thread_k_blocks) *
513
- (pipe / (group_blocks / thread_k_blocks)));
514
- reinterpret_cast<int4*>(&frag_s_group[k % 2])[0] =
515
- sh_s_group_stage[s_group_sh_rd];
516
- }
517
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
518
- #pragma unroll
519
- for (int i = 0; i < thread_m_blocks; i++)
520
- ldsm4(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
521
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
522
- frag_b_quant[k % 2] = *reinterpret_cast<I4*>(
523
- &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd]);
524
- };
525
-
526
- // Execute the actual tensor core matmul of a sub-tile.
527
- auto matmul = [&](int k) {
528
- // We have the m dimension as the inner loop in order to encourage overlapping
529
- // dequantization and matmul operations.
530
- #pragma unroll
531
- for (int j = 0; j < 4; j++) {
532
- int b_quant = frag_b_quant[k % 2][j];
533
- // int b_quant_shift = b_quant << 4;
534
- FragB frag_b0, frag_b1;
535
- // If there are no groups, we can just scale the final output once and can
536
- // avoid doing so for each weight.
537
- if constexpr (group_blocks != -1) {
538
- int b_quant_shift = b_quant >> 8;
539
- frag_b0 = dequant_per_group(b_quant, frag_s_group[k % 2][j], 0);
540
- frag_b1 = dequant_per_group(b_quant_shift, frag_s_group[k % 2][j], 1);
541
- } else {
542
- int b_quant_shift = b_quant << 4;
543
- frag_b0 = dequant_per_channel(b_quant);
544
- frag_b1 = dequant_per_channel(b_quant_shift);
545
- }
546
- #pragma unroll
547
- for (int i = 0; i < thread_m_blocks; i++) {
548
- mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]);
549
- mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]);
550
- }
551
- }
552
- };
553
-
554
- // Since we slice across the k dimension of a tile in order to increase the
555
- // number of warps while keeping the n dimension of a tile reasonable, we have
556
- // multiple warps that accumulate their partial sums of the same output
557
- // location; which we have to reduce over in the end. We do in shared memory.
558
- auto thread_block_reduce = [&]() {
559
- constexpr int red_off = threads / b_sh_stride / 2;
560
- if (red_off >= 1) {
561
- auto red_idx = threadIdx.x / b_sh_stride;
562
- constexpr int red_sh_stride = b_sh_stride * 4 * 2;
563
- constexpr int red_sh_delta = b_sh_stride;
564
- int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride) +
565
- (threadIdx.x % b_sh_stride);
566
-
567
- // Parallel logarithmic shared memory reduction. We make sure to avoid any
568
- // unnecessary read or write iterations, e.g., for two warps we write only
569
- // once by warp 1 and read only once by warp 0.
570
-
571
- #pragma unroll
572
- for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
573
- #pragma unroll
574
- for (int i = red_off; i > 0; i /= 2) {
575
- if (i <= red_idx && red_idx < 2 * i) {
576
- #pragma unroll
577
- for (int j = 0; j < 4 * 2; j++) {
578
- int red_sh_wr =
579
- red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
580
- if (i < red_off) {
581
- int* c_rd =
582
- reinterpret_cast<int*>(&sh[red_sh_delta * j + red_sh_rd]);
583
- int* c_wr = reinterpret_cast<int*>(&sh[red_sh_wr]);
584
- #pragma unroll
585
- for (int k = 0; k < 4; k++)
586
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
587
- c_rd[k] + c_wr[k];
588
- }
589
- sh[red_sh_wr] =
590
- reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
591
- }
592
- }
593
- __syncthreads();
594
- }
595
- if (red_idx == 0) {
596
- #pragma unroll
597
- for (int i = 0; i < 4 * 2; i++) {
598
- int* c_rd =
599
- reinterpret_cast<int*>(&sh[red_sh_delta * i + red_sh_rd]);
600
- #pragma unroll
601
- for (int j = 0; j < 4; j++)
602
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
603
- c_rd[j];
604
- }
605
- }
606
- __syncthreads();
607
- }
608
- }
609
- };
610
-
611
- // Since multiple threadblocks may process parts of the same column slice, we
612
- // finally have to globally reduce over the results. As the striped
613
- // partitioning minimizes the number of such reductions and our outputs are
614
- // usually rather small, we perform this reduction serially in L2 cache.
615
- // global_reduce works on INT32 elements, which are the results of INT8 GEMM.
616
- // This is why we need another INT32 maxtrix `C` to reduce instead of the
617
- // original half matrix `D`.
618
- auto global_reduce = [&](bool first = false, bool last = false) {
619
- // We are very careful here to reduce directly in the output buffer to
620
- // maximize L2 cache utilization in this step. To do this, we write out
621
- // results in FP16 (but still reduce with FP32 compute).
622
- constexpr int active_threads = 32 * thread_n_blocks / 4;
623
- if (threadIdx.x < active_threads) {
624
- int c_gl_stride = prob_n / 4;
625
- int c_gl_wr_delta_o = 8 * c_gl_stride;
626
- int c_gl_wr_delta_i = 8 * (active_threads / 32);
627
- int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
628
- 8 * (threadIdx.x / 32) + (threadIdx.x % 4) * 2;
629
- c_gl_wr += (4 * thread_n_blocks) * slice_col;
630
- constexpr int c_sh_wr_delta = active_threads * 2;
631
- auto c_sh_wr = 2 * threadIdx.x;
632
-
633
- int row = (threadIdx.x % 32) / 4;
634
-
635
- if (!first) {
636
- // Interestingly, doing direct global accesses here really seems to mess up
637
- // the compiler and lead to slowdowns, hence we also use async-copies even
638
- // though these fetches are not actually asynchronous.
639
- #pragma unroll
640
- for (int i = 0; i < thread_m_blocks * 4; i++) {
641
- cp_async4_pred(
642
- &sh[c_sh_wr + c_sh_wr_delta * i],
643
- &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
644
- c_gl_wr_delta_i * (i % 2)],
645
- i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
646
- cp_async4_pred(
647
- &sh[c_sh_wr + c_sh_wr_delta * i + 1],
648
- &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
649
- c_gl_wr_delta_i * (i % 2) + 1],
650
- i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
651
- }
652
- cp_async_fence();
653
- cp_async_wait<0>();
654
- }
655
-
656
- #pragma unroll
657
- for (int i = 0; i < thread_m_blocks * 4; i++) {
658
- if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) {
659
- if (!first) {
660
- int4 d_red1 = sh[c_sh_wr + i * c_sh_wr_delta];
661
- int4 d_red2 = sh[c_sh_wr + i * c_sh_wr_delta + 1];
662
- #pragma unroll
663
- for (int j = 0; j < 4; j++) {
664
- reinterpret_cast<int*>(
665
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] +=
666
- reinterpret_cast<int*>(&d_red1)[j];
667
- }
668
- #pragma unroll
669
- for (int j = 0; j < 4; j++) {
670
- reinterpret_cast<int*>(
671
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * (j + 4) + (i % 4)] +=
672
- reinterpret_cast<int*>(&d_red2)[j];
673
- }
674
- }
675
- if (!last) {
676
- int4 d1, d2;
677
- #pragma unroll
678
- for (int j = 0; j < 4; j++) {
679
- reinterpret_cast<int*>(&d1)[j] = reinterpret_cast<int*>(
680
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)];
681
- }
682
- #pragma unroll
683
- for (int j = 0; j < 4; j++) {
684
- reinterpret_cast<int*>(&d2)[j] = reinterpret_cast<int*>(
685
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * (j + 4) + (i % 4)];
686
- }
687
- C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] =
688
- d1;
689
- C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2) +
690
- 1] = d2;
691
- }
692
- }
693
- }
694
- }
695
- };
696
-
697
- // Write out the reduce final result in the correct layout. We only actually
698
- // reshuffle matrix fragments in this step, the reduction above is performed
699
- // in fragment layout.
700
- auto write_result = [&]() {
701
- int d_gl_stride = prob_n / 8;
702
- constexpr int d_sh_stride = 2 * thread_n_blocks + 1;
703
- int d_gl_wr_delta = d_gl_stride * (threads / (2 * thread_n_blocks));
704
- constexpr int d_sh_rd_delta =
705
- d_sh_stride * (threads / (2 * thread_n_blocks));
706
-
707
- int d_gl_wr = d_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
708
- (threadIdx.x % (2 * thread_n_blocks));
709
- d_gl_wr += (2 * thread_n_blocks) * slice_col;
710
- int d_sh_wr =
711
- (4 * d_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
712
- d_sh_wr += 32 * (threadIdx.x / 32);
713
- int d_sh_rd = d_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
714
- (threadIdx.x % (2 * thread_n_blocks));
715
-
716
- int d_gl_wr_end = d_gl_stride * prob_m;
717
-
718
- // We first reorder in shared memory to guarantee the most efficient final
719
- // global write patterns
720
- auto write = [&](int idx, int c0, int c1, float a_s, FragS_CHANNEL& w_s) {
721
- float2 deq_res;
722
- deq_res.x = int32_to_float(c0) * w_s[0] * a_s;
723
- deq_res.y = int32_to_float(c1) * w_s[1] * a_s;
724
- ((half2*)sh)[idx] = float2_to_half2(deq_res);
725
- };
726
-
727
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
728
- #pragma unroll
729
- for (int i = 0; i < thread_m_blocks; i++) {
730
- #pragma unroll
731
- for (int j = 0; j < 4; j++) {
732
- int wr = d_sh_wr + 8 * j;
733
- write(wr + (4 * d_sh_stride) * 0 + 0, frag_c[i][j][0][0],
734
- frag_c[i][j][0][1], frag_s_tok[i][0],
735
- frag_s_ch[j / 2][2 * (j % 2) + 0]);
736
- write(wr + (4 * d_sh_stride) * 8 + 0, frag_c[i][j][0][2],
737
- frag_c[i][j][0][3], frag_s_tok[i][1],
738
- frag_s_ch[j / 2][2 * (j % 2) + 0]);
739
- write(wr + (4 * d_sh_stride) * 0 + 4, frag_c[i][j][1][0],
740
- frag_c[i][j][1][1], frag_s_tok[i][0],
741
- frag_s_ch[j / 2][2 * (j % 2) + 1]);
742
- write(wr + (4 * d_sh_stride) * 8 + 4, frag_c[i][j][1][2],
743
- frag_c[i][j][1][3], frag_s_tok[i][1],
744
- frag_s_ch[j / 2][2 * (j % 2) + 1]);
745
- }
746
- d_sh_wr += 16 * (4 * d_sh_stride);
747
- }
748
- }
749
- __syncthreads();
750
-
751
- #pragma unroll
752
- for (int i = 0;
753
- i < ceildiv(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
754
- i++) {
755
- if (d_gl_wr < d_gl_wr_end) {
756
- D[d_gl_wr] = sh[d_sh_rd];
757
- d_gl_wr += d_gl_wr_delta;
758
- d_sh_rd += d_sh_rd_delta;
759
- }
760
- }
761
- };
762
-
763
- // Start global fetch and register load pipelines.
764
- auto start_pipes = [&]() {
765
- #pragma unroll
766
- for (int i = 0; i < stages - 1; i++) fetch_to_shared(i, i, i < slice_iters);
767
- zero_accums();
768
- wait_for_stage();
769
- fetch_to_registers(0, 0);
770
- a_gl_rd += a_gl_rd_delta_o * (stages - 1);
771
- };
772
- start_pipes();
773
-
774
- // Main loop.
775
- while (slice_iters) {
776
- // We unroll over both the global fetch and the register load pipeline to
777
- // ensure all shared memory accesses are static. Note that both pipelines have
778
- // even length meaning that the next iteration will always start at index 0.
779
- #pragma unroll
780
- for (int pipe = 0; pipe < stages;) {
781
- #pragma unroll
782
- for (int k = 0; k < b_sh_wr_iters; k++) {
783
- fetch_to_registers(k + 1, pipe % stages);
784
- if (k == b_sh_wr_iters - 2) {
785
- fetch_to_shared((pipe + stages - 1) % stages, pipe,
786
- slice_iters >= stages);
787
- pipe++;
788
- wait_for_stage();
789
- }
790
- matmul(k);
791
- }
792
- slice_iters--;
793
- if (slice_iters == 0) break;
794
- }
795
- a_gl_rd += a_gl_rd_delta_o * stages;
796
-
797
- // Process results and, if necessary, proceed to the next column slice.
798
- // While this pattern may not be the most readable, other ways of writing
799
- // the loop seemed to noticeably worse performance after compilation.
800
- if (slice_iters == 0) {
801
- cp_async_wait<0>();
802
- bool last = slice_idx == slice_count - 1;
803
- // For per-column scales, we only fetch them here in the final step before
804
- // write-out
805
- if (last) {
806
- if (s_tok_sh_wr_pred) {
807
- cp_async1(&sh_s_tok[s_tok_sh_wr], &s_tok[s_tok_gl_rd]);
808
- }
809
- if (s_ch_sh_wr_pred) {
810
- cp_async4(&sh_s_ch[s_ch_sh_wr], &s_ch[s_ch_gl_rd]);
811
- }
812
- cp_async_fence();
813
- }
814
- thread_block_reduce();
815
- if (last) {
816
- cp_async_wait<0>();
817
- __syncthreads();
818
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
819
- #pragma unroll
820
- for (int i = 0; i < thread_m_blocks; i++) {
821
- frag_s_tok[i][0] =
822
- *reinterpret_cast<float*>(&sh_s_tok[16 * i + 2 * s_tok_sh_rd]);
823
- frag_s_tok[i][1] = *reinterpret_cast<float*>(
824
- &sh_s_tok[16 * i + 2 * s_tok_sh_rd + 1]);
825
- }
826
- reinterpret_cast<int4*>(&frag_s_ch)[0] = sh_s_ch[s_ch_sh_rd + 0];
827
- reinterpret_cast<int4*>(&frag_s_ch)[1] = sh_s_ch[s_ch_sh_rd + 1];
828
- reinterpret_cast<int4*>(&frag_s_ch)[2] = sh_s_ch[s_ch_sh_rd + 8];
829
- reinterpret_cast<int4*>(&frag_s_ch)[3] = sh_s_ch[s_ch_sh_rd + 9];
830
- }
831
- }
832
- if (slice_count > 1) { // only globally reduce if there is more than one
833
- // block in a slice
834
- barrier_acquire(&locks[slice_col], slice_idx);
835
- global_reduce(slice_idx == 0, last);
836
- barrier_release(&locks[slice_col], last);
837
- }
838
- if (last) // only the last block in a slice actually writes the result
839
- write_result();
840
- slice_row = 0;
841
- slice_col_par++;
842
- slice_col++;
843
- init_slice();
844
- if (slice_iters) {
845
- a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
846
- (threadIdx.x % a_gl_rd_delta_o);
847
- #pragma unroll
848
- for (int i = 0; i < b_sh_wr_iters; i++)
849
- B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
850
- if (slice_col == 0) {
851
- #pragma unroll
852
- for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
853
- }
854
- s_group_gl_rd = s_group_sh_stride * slice_col + threadIdx.x;
855
- s_ch_gl_rd = s_ch_sh_stride * slice_col + threadIdx.x;
856
- start_pipes();
857
- }
858
- }
859
- }
860
-}
861
-
862
-#else
863
-
864
-template <const int threads, // number of threads in a threadblock
865
- const int thread_m_blocks, // number of 16x16 blocks in the m
866
- // dimension (batchsize) of the
867
- // threadblock
868
- const int thread_n_blocks, // same for n dimension (output)
869
- const int thread_k_blocks, // same for k dimension (reduction)
870
- const int stages, // number of stages for the async global->shared
871
- // fetch pipeline
872
- const int group_blocks = -1 // number of consecutive 16x16 blocks
873
- // with a separate quantization scale
874
- >
875
-__global__ void Marlin(
876
- const int4* __restrict__ A, // int8 input matrix of shape mxk
877
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
878
- int4* __restrict__ C, // int32 global_reduce buffer of shape
879
- // (max_par*16*4)xn, as int8 tensor core's output is
880
- // int32 dtype
881
- int4* __restrict__ D, // fp16 output buffer of shape mxn
882
- const float* __restrict__ s_tok, // fp32 activation per-token quantization
883
- // scales of shape mx1
884
- const int4* __restrict__ s_ch, // fp32 weight per-channel quantization
885
- // scales of shape 1xn
886
- const int4* __restrict__ s_group, // fp16 weight per-group quantization
887
- // scales of shape (k/groupsize)xn, when
888
- // group_blocks=-1, it should be nullptr
889
- int prob_m, // batch dimension m
890
- int prob_n, // output dimension n
891
- int prob_k, // reduction dimension k
892
- int* locks // extra global storage for barrier synchronization
893
-) {
894
- // Marlin is not implemented yet for SM < 8.0
895
- assert(false);
896
- return;
897
-}
898
-
899
-#endif
900
-
901
-// 8 warps are a good choice since every SM has 4 schedulers and having more
902
-// than 1 warp per schedule allows some more latency hiding. At the same time,
903
-// we want relatively few warps to have many registers per warp and small tiles.
904
-const int USER_THREADS =
905
- 256; // Note: This is only used with user-provided thread_k/n
906
-const int STAGES = 4; // 4 pipeline stages fit into shared memory
907
-
908
-static constexpr int min_thread_n = 64;
909
-static constexpr int min_thread_k = 64;
910
-
911
-static constexpr int tile_size = 16;
912
-static constexpr int max_par = 16;
913
-
914
-static constexpr int pack_factor_4bit =
915
- 8; // We have 8 4-bit vals inside a 32 bit
916
-
917
-#define __CALL_IF(THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
918
- GROUP_BLOCKS, NUM_THREADS) \
919
- else if (thread_m_blocks == THREAD_M_BLOCKS && \
920
- thread_n_blocks == THREAD_N_BLOCKS && \
921
- thread_k_blocks == THREAD_K_BLOCKS && \
922
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS) { \
923
- cudaFuncSetAttribute(Marlin<NUM_THREADS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, \
924
- THREAD_K_BLOCKS, STAGES, GROUP_BLOCKS>, \
925
- cudaFuncAttributeMaxDynamicSharedMemorySize, \
926
- max_shared_mem); \
927
- Marlin<NUM_THREADS, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
928
- STAGES, GROUP_BLOCKS> \
929
- <<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \
930
- A_ptr, B_ptr, C_ptr, D_ptr, s_tok_ptr, s_ch_ptr, s_group_ptr, \
931
- prob_m, prob_n, prob_k, locks); \
932
- }
933
-
934
-typedef struct {
935
- int thread_k;
936
- int thread_n;
937
- int num_threads;
938
-} thread_config_t;
939
-
940
-thread_config_t small_batch_thread_configs[] = {
941
- // Ordered by priority
942
-
943
- // thread_k, thread_n, num_threads
944
- {128, 128, 256}, // Default
945
- {128, 64, 128}, // Reduce N 2X, same K
946
- {64, 256, 256}, // Reduce K 2X, increase N 2X
947
- {64, 128, 128}, // Reduce K 2X, same N
948
-};
949
-
950
-thread_config_t large_batch_thread_configs[] = {
951
- // Ordered by priority
952
-
953
- // thread_k, thread_n, num_threads
954
- {64, 256, 256}, // Default
955
- {128, 128, 256}, // Reduce N 2X, increase K 2X
956
- {64, 128, 128}, // Reduce N 2X, same K
957
- {128, 64, 128}, // Reduce N 4X, increase K 2X
958
-};
959
-
960
-bool is_valid_config(thread_config_t const& th_config, int prob_m, int prob_n,
961
- int prob_k) {
962
- // Sanity
963
- if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
964
- th_config.num_threads == -1) {
965
- return false;
966
- }
967
-
968
- // Verify K/N are divisible by thread K/N
969
- if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) {
970
- return false;
971
- }
972
-
973
- // thread_k can be only 128 or 64 (because it must be less than groupsize
974
- // which is 128)
975
- if (th_config.thread_k != 128 && th_config.thread_k != 64) {
976
- return false;
977
- }
978
-
979
- // Verify min for thread K/N
980
- if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) {
981
- return false;
982
- }
983
-
984
- // num_threads must be at least 128 (= 4 warps)
985
- if (th_config.num_threads < 128) {
986
- return false;
987
- }
988
-
989
- return true;
990
-}
991
-
992
-thread_config_t determine_thread_config(int prob_m, int prob_n, int prob_k) {
993
- if (prob_m <= 16) {
994
- for (auto th_config : small_batch_thread_configs) {
995
- if (is_valid_config(th_config, prob_m, prob_n, prob_k)) {
996
- return th_config;
997
- }
998
- }
999
-
1000
- } else {
1001
- for (auto th_config : large_batch_thread_configs) {
1002
- if (is_valid_config(th_config, prob_m, prob_n, prob_k)) {
1003
- return th_config;
1004
- }
1005
- }
1006
- }
1007
-
1008
- return thread_config_t{-1, -1, -1};
1009
-}
1010
-
1011
-#define CALL_IF(N_BLOCKS, K_BLOCKS, NUM_THREADS) \
1012
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
1013
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
1014
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
1015
- __CALL_IF(1, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
1016
- __CALL_IF(2, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
1017
- __CALL_IF(2, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
1018
- __CALL_IF(3, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
1019
- __CALL_IF(3, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS) \
1020
- __CALL_IF(4, N_BLOCKS, K_BLOCKS, -1, NUM_THREADS) \
1021
- __CALL_IF(4, N_BLOCKS, K_BLOCKS, 8, NUM_THREADS)
1022
-
1023
-void marlin_qqq_cuda(const void* A, const void* B, void* C, void* D,
1024
- void* s_tok, void* s_ch, void* s_group, int prob_m,
1025
- int prob_n, int prob_k, void* workspace,
1026
- int groupsize = -1, int dev = 0, cudaStream_t stream = 0,
1027
- int thread_k = -1, int thread_n = -1, int sms = -1,
1028
- int max_par = 16) {
1029
- int tot_m = prob_m;
1030
- int tot_m_blocks = ceildiv(tot_m, 16);
1031
- int pad = 16 * tot_m_blocks - tot_m;
1032
-
1033
- if (sms == -1)
1034
- cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
1035
-
1036
- int max_shared_mem = 0;
1037
- cudaDeviceGetAttribute(&max_shared_mem,
1038
- cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
1039
- TORCH_CHECK(max_shared_mem > 0);
1040
-
1041
- // Set thread config
1042
- thread_config_t th_config;
1043
- if (thread_k != -1 && thread_n != -1) {
1044
- // User-defined config
1045
- th_config = thread_config_t{thread_k, thread_n, USER_THREADS};
1046
- } else {
1047
- // Auto config
1048
- th_config = determine_thread_config(prob_m, prob_n, prob_k);
1049
- }
1050
-
1051
- if (!is_valid_config(th_config, prob_m, prob_n, prob_k)) {
1052
- throw std::runtime_error(
1053
- "Invalid thread config: thread_k = " + str(th_config.thread_k) +
1054
- ", thread_n = " + str(th_config.thread_n) +
1055
- ", num_threads = " + str(th_config.num_threads) + " for MKN = [" +
1056
- str(prob_m) + ", " + str(prob_k) + ", " + str(prob_n) + "]");
1057
- }
1058
-
1059
- int num_threads = th_config.num_threads;
1060
- thread_k = th_config.thread_k;
1061
- thread_n = th_config.thread_n;
1062
-
1063
- int thread_k_blocks = thread_k / 16;
1064
- int thread_n_blocks = thread_n / 16;
1065
- int group_blocks = (groupsize == -1) ? -1 : groupsize / 16;
1066
- int blocks = sms;
1067
-
1068
- if (prob_m == 0 || prob_n == 0 || prob_k == 0) {
1069
- return;
1070
- }
1071
-
1072
- TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
1073
- " is not divisible by thread_n = ", thread_n);
1074
- TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
1075
- " is not divisible by thread_k = ", thread_k);
1076
- if (group_blocks != -1) {
1077
- TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
1078
- " is not divisible by group_blocks = ", group_blocks);
1079
- }
1080
-
1081
- const int4* A_ptr = (const int4*)A;
1082
- const int4* B_ptr = (const int4*)B;
1083
- int4* C_ptr = (int4*)C;
1084
- int4* D_ptr = (int4*)D;
1085
- const float* s_tok_ptr = (const float*)s_tok;
1086
- const int4* s_ch_ptr = (const int4*)s_ch;
1087
- const int4* s_group_ptr = (const int4*)s_group;
1088
-
1089
- int* locks = (int*)workspace;
1090
-
1091
- for (int i = 0; i < tot_m_blocks; i += 4) {
1092
- int thread_m_blocks = tot_m_blocks - i;
1093
- prob_m = tot_m - 16 * i;
1094
- int par = 1;
1095
- if (thread_m_blocks > 4) {
1096
- // Note that parallel > 1 currently only works for inputs without any
1097
- // padding
1098
- par = (16 * thread_m_blocks - pad) / 64;
1099
- if (par > max_par) par = max_par;
1100
- prob_m = 64 * par;
1101
- i += 4 * (par - 1);
1102
- thread_m_blocks = 4;
1103
- }
1104
-
1105
- // For compilation speed, we only define the kernel configurations that have
1106
- // seemed useful (in terms of performance) in our testing, however many more
1107
- // are, in principle, possible.
1108
- if (false) {
1109
- }
1110
- CALL_IF(8, 8, 256)
1111
- CALL_IF(16, 4, 256)
1112
- CALL_IF(8, 4, 128)
1113
- CALL_IF(4, 8, 128)
1114
- else {
1115
- throw std::runtime_error("Unsupported shapes: MKN = [" + str(prob_m) +
1116
- ", " + str(prob_k) + ", " + str(prob_n) + "]" +
1117
- ", groupsize = " + str(groupsize) +
1118
- ", thread_m_blocks = " + str(thread_m_blocks) +
1119
- ", thread_n_blocks = " + str(thread_n_blocks) +
1120
- ", thread_k_blocks = " + str(thread_k_blocks));
1121
- }
1122
-
1123
- A_ptr += 16 * thread_m_blocks * (prob_k / 16) * par;
1124
- D_ptr += 16 * thread_m_blocks * (prob_n / 8) * par;
1125
- s_tok_ptr += 16 * thread_m_blocks * par;
1126
- }
1127
-}
1128
-} // anonymous namespace
1129
-
1130
-torch::Tensor marlin_qqq_gemm(torch::Tensor const& a,
1131
- torch::Tensor const& b_q_weight,
1132
- torch::Tensor const& s_tok,
1133
- torch::Tensor const& s_ch,
1134
- torch::Tensor const& s_group,
1135
- torch::Tensor& workspace, int64_t size_m,
1136
- int64_t size_n, int64_t size_k) {
1137
- // Verify M
1138
- TORCH_CHECK(size_m == a.size(0),
1139
- "Shape mismatch: a.size(0) = " + str(a.size(0)) +
1140
- ", size_m = " + str(size_m));
1141
- TORCH_CHECK(size_m == s_tok.numel(),
1142
- "Shape mismatch: s_tok.numel() = " + str(s_tok.numel()) +
1143
- ", size_m = " + str(size_m));
1144
-
1145
- // Verify K
1146
- TORCH_CHECK(size_k == a.size(1),
1147
- "Shape mismatch: a.size(1) = " + str(a.size(1)) +
1148
- ", size_k = " + str(size_k));
1149
- TORCH_CHECK(size_k % tile_size == 0,
1150
- "size_k = " + str(size_k) +
1151
- " is not divisible by tile_size = " + str(tile_size));
1152
- TORCH_CHECK(
1153
- (size_k / tile_size) == b_q_weight.size(0),
1154
- "Shape mismatch: b_q_weight.size(0) = " + str(b_q_weight.size(0)) +
1155
- ", size_k = " + str(size_k) + ", tile_size = " + str(tile_size));
1156
-
1157
- int groupsize = (s_group.numel() == 0) ? -1 : size_k / s_group.size(0);
1158
- // Verify groupsize
1159
- TORCH_CHECK(groupsize == -1 || groupsize == 128,
1160
- "Unexpected groupsize = " + str(groupsize));
1161
-
1162
- // Verify N
1163
- TORCH_CHECK(s_ch.numel() == size_n,
1164
- "Shape mismatch: s_ch.numel() = " + str(s_ch.numel()) +
1165
- ", size_n = " + str(size_n));
1166
- TORCH_CHECK(b_q_weight.size(1) % tile_size == 0,
1167
- "b_q_weight.size(1) = " + str(b_q_weight.size(1)) +
1168
- " is not divisible by tile_size = " + str(tile_size));
1169
- if (groupsize != -1) {
1170
- TORCH_CHECK(s_group.size(1) == size_n,
1171
- "Shape mismatch: s_group.size(1) = " + str(s_group.size(1)) +
1172
- ", size_n = " + str(size_n));
1173
- TORCH_CHECK(
1174
- size_k % s_group.size(0) == 0,
1175
- "size_k = " + str(size_k) +
1176
- ", is not divisible by s_group.size(0) = " + str(s_group.size(0)));
1177
- }
1178
-
1179
- int actual_size_n = (b_q_weight.size(1) / tile_size) * pack_factor_4bit;
1180
- TORCH_CHECK(size_n == actual_size_n,
1181
- "Shape mismatch: size_n = " + str(size_n) +
1182
- ", actual_size_n = " + str(actual_size_n));
1183
-
1184
- // Verify A device and strides
1185
- TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
1186
- TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
1187
-
1188
- // Verify B device and strides
1189
- TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
1190
- TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
1191
-
1192
- // Verify s_tok device, strides and dtype
1193
- TORCH_CHECK(s_tok.device().is_cuda(), "s_tok is not on GPU");
1194
- TORCH_CHECK(s_tok.is_contiguous(), "s_tok is not contiguous");
1195
- TORCH_CHECK(s_tok.dtype() == torch::kFloat32, "s_tok's dtype is not float32");
1196
-
1197
- // Verify s_ch device, strides and dtype
1198
- TORCH_CHECK(s_ch.device().is_cuda(), "s_ch is not on GPU");
1199
- TORCH_CHECK(s_ch.is_contiguous(), "s_ch is not contiguous");
1200
- TORCH_CHECK(s_ch.dtype() == torch::kFloat32, "s_ch's dtype is not float32");
1201
-
1202
- // Verify s_group device, strides and dtype
1203
- TORCH_CHECK(s_group.device().is_cuda(), "s_group is not on GPU");
1204
- TORCH_CHECK(s_group.is_contiguous(), "s_group is not contiguous");
1205
- TORCH_CHECK(s_group.dtype() == torch::kFloat16,
1206
- "s_group's dtype is not float16");
1207
-
1208
- // Verify workspace size
1209
- TORCH_CHECK(size_n % min_thread_n == 0,
1210
- "size_n = " + str(size_n) +
1211
- ", is not divisible by min_thread_n = " + str(min_thread_n));
1212
- int min_workspace_size = (size_n / min_thread_n) * max_par;
1213
- TORCH_CHECK(workspace.numel() >= min_workspace_size,
1214
- "workspace.numel = " + str(workspace.numel()) +
1215
- " is below min_workspace_size = " + str(min_workspace_size));
1216
-
1217
- // Alloc C matrix
1218
- const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
1219
- auto options_c = torch::TensorOptions().dtype(torch::kInt).device(a.device());
1220
- torch::Tensor c = torch::empty({max_par * 64, size_n}, options_c);
1221
-
1222
- // Alloc D matrix
1223
- auto options_d =
1224
- torch::TensorOptions().dtype(torch::kFloat16).device(a.device());
1225
- torch::Tensor d = torch::empty({size_m, size_n}, options_d);
1226
-
1227
- // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
1228
- // auto -1)
1229
- int thread_k = -1;
1230
- // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
1231
- // auto -1)
1232
- int thread_n = -1;
1233
- // sms: number of SMs to use for the kernel (can usually be left as auto -1)
1234
- int sms = -1;
1235
-
1236
- int dev = a.get_device();
1237
- marlin_qqq_cuda(
1238
- a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), d.data_ptr(),
1239
- s_tok.data_ptr(), s_ch.data_ptr(), s_group.data_ptr(), size_m, size_n,
1240
- size_k, workspace.data_ptr(), groupsize, dev,
1241
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms, max_par);
1242
-
1243
- return d;
1244
-}
1245
-
1246
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1247
- m.impl("marlin_qqq_gemm", &marlin_qqq_gemm);
1248
-}
csrc/torch_bindings.cpp
-17
@@ -241,14 +241,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
241
// custom types:
242
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
243
244
- // Marlin (Dense) Optimized Quantized GEMM for GPTQ.
245
- ops.def(
246
- "marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
247
- "Tensor! workspace, SymInt size_m, SymInt size_n, SymInt size_k) -> "
248
- "Tensor",
249
- {stride_tag});
250
- // conditionally compiled so impl in source file
251
-
244
// Marlin_24 (Sparse) Optimized Quantized GEMM for GPTQ.
245
ops.def(
246
"gptq_marlin_24_gemm(Tensor a, Tensor b_q_weight, Tensor b_meta, "
@@ -353,15 +345,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
345
ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
346
347
#ifndef USE_ROCM
356
- // marlin_qqq_gemm for QQQ.
357
- ops.def(
358
- "marlin_qqq_gemm(Tensor a, Tensor b_q_weight, "
359
- "Tensor s_tok, Tensor s_ch, Tensor s_group, "
360
- "Tensor! workspace, SymInt size_m, SymInt size_n, "
361
- "SymInt size_k) -> Tensor",
362
- {stride_tag});
363
- // conditionally compiled so impl registration is in source file
364
-
348
// CUTLASS nvfp4 block scaled GEMM
349
ops.def(
350
"cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
tests/compile/test_full_graph.py
-6
@@ -53,12 +53,6 @@ def models_list(*, all: bool = True, keywords: Optional[list[str]] = None):
53
"quantization": "gptq_marlin_24"
54
}))
55
56
- if is_quant_method_supported("marlin"):
57
- TEST_MODELS.append(
58
- ("robertgshaw2/TinyLlama-1.1B-Chat-v1.0-g128-marlin", {
59
- "quantization": "marlin"
60
- }))
61
-
56
if not current_platform.is_rocm() and is_quant_method_supported("awq"):
57
TEST_MODELS.append(("TheBloke/TinyLlama-1.1B-Chat-v0.3-AWQ", {
58
"quantization": "AWQ"
tests/kernels/quantization/test_machete_mm.py
+17
-17
@@ -95,23 +95,23 @@ TEST_TYPES = [
95
token_scale_type=None)
96
for w_type in [scalar_types.uint4, scalar_types.uint8]
97
for a_type in [torch.float16, torch.bfloat16]),
98
- # QQQ style
99
- *(TypeConfig(act_type=torch.int8,
100
- weight_type=scalar_types.uint4b8,
101
- output_type=torch.float16,
102
- group_scale_type=group_scale_type,
103
- group_zero_type=None,
104
- channel_scale_type=torch.float,
105
- token_scale_type=torch.float)
106
- for group_scale_type in [None, torch.float16]),
107
- *(TypeConfig(act_type=torch.float8_e4m3fn,
108
- weight_type=scalar_types.uint4b8,
109
- output_type=torch.float16,
110
- group_scale_type=group_scale_type,
111
- group_zero_type=None,
112
- channel_scale_type=torch.float,
113
- token_scale_type=torch.float)
114
- for group_scale_type in [None, torch.float16]),
98
+ # # QQQ style
99
+ # *(TypeConfig(act_type=torch.int8,
100
+ # weight_type=scalar_types.uint4b8,
101
+ # output_type=torch.float16,
102
+ # group_scale_type=group_scale_type,
103
+ # group_zero_type=None,
104
+ # channel_scale_type=torch.float,
105
+ # token_scale_type=torch.float)
106
+ # for group_scale_type in [None, torch.float16]),
107
+ # *(TypeConfig(act_type=torch.float8_e4m3fn,
108
+ # weight_type=scalar_types.uint4b8,
109
+ # output_type=torch.float16,
110
+ # group_scale_type=group_scale_type,
111
+ # group_zero_type=None,
112
+ # channel_scale_type=torch.float,
113
+ # token_scale_type=torch.float)
114
+ # for group_scale_type in [None, torch.float16]),
115
]
116
117
# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
tests/kernels/quantization/test_marlin_gemm.py
-83
@@ -13,11 +13,7 @@ from vllm import _custom_ops as ops
13
from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
14
GPTQ_MARLIN_24_MAX_PARALLEL, GPTQ_MARLIN_24_MIN_THREAD_N,
15
GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES, GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)
16
-from vllm.model_executor.layers.quantization.qqq import (
17
- MARLIN_QQQ_MAX_PARALLEL, MARLIN_QQQ_MIN_THREAD_N,
18
- MARLIN_QQQ_SUPPORTED_GROUP_SIZES, MARLIN_QQQ_SUPPORTED_NUM_BITS)
16
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
20
- GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
17
MARLIN_SUPPORTED_GROUP_SIZES, marlin_make_empty_g_idx,
18
marlin_make_workspace_new, marlin_permute_bias, marlin_permute_scales,
19
query_marlin_supported_quant_types)
@@ -31,8 +27,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
27
marlin_weights)
28
from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
29
marlin_24_quantize)
34
-from vllm.model_executor.layers.quantization.utils.marlin_utils_test_qqq import ( # noqa: E501
35
- marlin_qqq_quantize)
30
from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
awq_pack, gptq_pack, gptq_quantize_weights, quantize_weights, sort_weights)
32
from vllm.scalar_type import scalar_types
@@ -449,68 +443,6 @@ def test_hqq_marlin_gemm(
443
assert max_diff < 0.04
444
445
452
-@pytest.mark.skipif(not is_quant_method_supported("qqq"),
453
- reason="Marlin is not supported on this GPU type.")
454
-@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
455
-@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
456
-@pytest.mark.parametrize("num_bits", MARLIN_QQQ_SUPPORTED_NUM_BITS)
457
-@pytest.mark.parametrize("group_size", MARLIN_QQQ_SUPPORTED_GROUP_SIZES)
458
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
459
-def test_marlin_qqq_gemm(
460
- k_chunk,
461
- n_chunk,
462
- num_bits,
463
- group_size,
464
- mnk_factors,
465
-):
466
- int8_traits = torch.iinfo(torch.int8)
467
- m_factor, n_factor, k_factor = mnk_factors
468
-
469
- size_m = m_factor
470
- size_k = k_chunk * k_factor
471
- size_n = n_chunk * n_factor
472
-
473
- a_input = rand_data((size_m, size_k))
474
- b_weight = rand_data((size_k, size_n))
475
-
476
- # Quantize activations
477
- s_a = a_input.abs().max(dim=-1, keepdim=True)[0].div(int8_traits.max).to(
478
- torch.float)
479
- q_a = (a_input / s_a).round().clamp(int8_traits.min,
480
- int8_traits.max).to(torch.int8)
481
-
482
- # Quantize weights
483
- w_ref, marlin_qqq_q_w, marlin_qqq_s_group, marlin_qqq_s_channel = \
484
- marlin_qqq_quantize(b_weight, num_bits, group_size)
485
-
486
- workspace = MarlinWorkspace(size_n, MARLIN_QQQ_MIN_THREAD_N,
487
- MARLIN_QQQ_MAX_PARALLEL)
488
-
489
- opcheck(torch.ops._C.marlin_qqq_gemm,
490
- (q_a, marlin_qqq_q_w, s_a, marlin_qqq_s_channel,
491
- marlin_qqq_s_group, workspace.scratch, a_input.shape[0],
492
- b_weight.shape[1], a_input.shape[1]))
493
-
494
- output = ops.marlin_qqq_gemm(
495
- q_a,
496
- marlin_qqq_q_w,
497
- s_a,
498
- marlin_qqq_s_channel,
499
- marlin_qqq_s_group,
500
- workspace.scratch,
501
- a_input.shape[0],
502
- b_weight.shape[1],
503
- a_input.shape[1],
504
- )
505
- output_ref = torch.matmul(q_a.half() * s_a.half(), w_ref)
506
-
507
- torch.cuda.synchronize()
508
-
509
- max_diff = compute_max_diff(output, output_ref)
510
-
511
- assert max_diff < 0.04
512
-
513
-
446
def test_marlin_gemm_subset_input():
447
quant_type = scalar_types.uint4b8
448
group_size = 128
@@ -602,18 +534,3 @@ def test_marlin_gemm_with_bias(size_m):
534
max_diff = compute_max_diff(output, output_ref)
535
536
assert max_diff < 0.04
605
-
606
-
607
-def test_marlin_gemm_opcheck():
608
- size_m = 2048
609
- size_n = 4096
610
- size_k = 4096
611
- a = torch.rand((size_m, size_n), device='cuda', dtype=torch.float16)
612
- w = torch.randint(-5, 5, (256, 8192), device='cuda', dtype=torch.int32)
613
- s = torch.full((32, size_k), 0.125, device='cuda', dtype=torch.float16)
614
- wk = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
615
- GPTQ_MARLIN_MAX_PARALLEL).scratch
616
- x = torch.ops._C.marlin_gemm(a, w, s, wk, size_m, size_n, size_k)
617
- y = torch.ops._C.marlin_gemm(a, w, s, wk, size_m, size_n, size_k)
618
- torch.testing.assert_close(x, y)
619
- opcheck(torch.ops._C.marlin_gemm, (a, w, s, wk, size_m, size_n, size_k))
tests/quantization/test_configs.py
-10
@@ -22,22 +22,12 @@ class ModelPair:
22
MODEL_ARG_EXPTYPES = [
23
# AUTOGPTQ
24
# compat: autogptq <=0.7.1 is_marlin_format: bool
25
- # Model Serialized in Marlin Format should always use Marlin kernel.
26
- ("neuralmagic/TinyLlama-1.1B-Chat-v1.0-marlin", None, "marlin"),
27
- ("neuralmagic/TinyLlama-1.1B-Chat-v1.0-marlin", "marlin", "marlin"),
28
- ("neuralmagic/TinyLlama-1.1B-Chat-v1.0-marlin", "gptq", "marlin"),
29
- ("neuralmagic/TinyLlama-1.1B-Chat-v1.0-marlin", "awq", "ERROR"),
25
# Model Serialized in Exllama Format.
26
("TheBloke/Llama-2-7B-Chat-GPTQ", None, "gptq_marlin"),
27
("TheBloke/Llama-2-7B-Chat-GPTQ", "marlin", "gptq_marlin"),
28
("TheBloke/Llama-2-7B-Chat-GPTQ", "gptq", "gptq"),
29
("TheBloke/Llama-2-7B-Chat-GPTQ", "awq", "ERROR"),
30
# compat: autogptq >=0.8.0 use checkpoint_format: str
36
- # Model Serialized in Marlin Format should always use Marlin kernel.
37
- ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-Marlin-4bit", None, "marlin"),
38
- ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-Marlin-4bit", "marlin", "marlin"),
39
- ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-Marlin-4bit", "gptq", "marlin"),
40
- ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-Marlin-4bit", "awq", "ERROR"),
31
# Model Serialized in Exllama Format.
32
("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", None, "gptq_marlin"),
33
("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", "marlin", "gptq_marlin"),
tests/quantization/test_lm_head.py
+1
-5
@@ -11,7 +11,6 @@ import torch
11
from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod
12
from vllm.model_executor.layers.quantization.gptq_marlin import (
13
GPTQMarlinLinearMethod)
14
-from vllm.model_executor.layers.quantization.marlin import MarlinLinearMethod
14
from vllm.model_executor.layers.vocab_parallel_embedding import (
15
UnquantizedEmbeddingMethod)
16
@@ -19,9 +18,7 @@ PROMPT = "On the surface of Mars, we found"
18
19
MODELS_QUANT = [
20
("ModelCloud/Qwen1.5-1.8B-Chat-GPTQ-4bits-dynamic-cfg-with-lm_head", True),
22
- ("ModelCloud/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit-10-25-2024", False),
21
("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", False),
24
- ("neuralmagic/Meta-Llama-3-8B-Instruct-FP8", False)
22
]
23
24
@@ -41,8 +38,7 @@ def test_lm_head(
38
lm_head_layer = model.lm_head
39
if lm_head_quantized:
40
assert isinstance(lm_head_layer.quant_method,
44
- (GPTQLinearMethod, GPTQMarlinLinearMethod,
45
- MarlinLinearMethod))
41
+ (GPTQLinearMethod, GPTQMarlinLinearMethod))
42
else:
43
assert isinstance(lm_head_layer.quant_method,
44
UnquantizedEmbeddingMethod)
tests/weight_loading/models.txt
-4
@@ -26,9 +26,5 @@ compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-W8A8-testing
26
awq, casperhansen/mixtral-instruct-awq, main
27
awq_marlin, casperhansen/mixtral-instruct-awq, main
28
fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
29
-marlin, nm-testing/zephyr-beta-7b-marlin-g128, main
30
-marlin, robertgshaw2/zephyr-7b-beta-channelwise-marlin, main
31
-qqq, HandH1998/QQQ-Llama-3-8b-g128, main
32
-qqq, HandH1998/QQQ-Llama-3-8b, main
29
hqq, nm-testing/Llama-3.2-1B-Instruct-HQQ, main
30
None, mgleize/fairseq2-dummy-Llama-3.2-1B, main
\ No newline at end of file
vllm/_custom_ops.py
-36
@@ -387,14 +387,6 @@ def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor,
387
torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
388
389
390
-# marlin
391
-def marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
392
- b_scales: torch.Tensor, workspace: torch.Tensor, size_m: int,
393
- size_n: int, size_k: int) -> torch.Tensor:
394
- return torch.ops._C.marlin_gemm(a, b_q_weight, b_scales, workspace, size_m,
395
- size_n, size_k)
396
-
397
-
390
# marlin_24
391
def gptq_marlin_24_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
392
b_meta: torch.Tensor, b_scales: torch.Tensor,
@@ -437,25 +429,6 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
429
is_zp_float: bool = False) -> torch.Tensor:
430
return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
431
440
- @register_fake("_C::marlin_qqq_gemm")
441
- def _marlin_qqq_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
442
- s_tok: torch.Tensor, s_ch: torch.Tensor,
443
- s_group: torch.Tensor, workspace: torch.Tensor,
444
- size_m: torch.SymInt, size_n: torch.SymInt,
445
- size_k: torch.SymInt) -> torch.Tensor:
446
- return torch.empty((size_m, size_n),
447
- dtype=torch.float16,
448
- device=a.device)
449
-
450
- @register_fake("_C::marlin_gemm")
451
- def _marlin_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
452
- b_scales: torch.Tensor, workspace: torch.Tensor,
453
- size_m: torch.SymInt, size_n: torch.SymInt,
454
- size_k: torch.SymInt) -> torch.Tensor:
455
- return torch.empty((size_m, size_n),
456
- dtype=torch.float16,
457
- device=a.device)
458
-
432
@register_fake("_C::awq_dequantize")
433
def _awq_dequantize_fake(qweight: torch.Tensor, scales: torch.Tensor,
434
zeros: torch.Tensor, split_k_iters: torch.SymInt,
@@ -1348,15 +1321,6 @@ def scaled_int8_quant(
1321
return output, input_scales, input_azp
1322
1323
1351
-# qqq ops
1352
-def marlin_qqq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
1353
- s_tok: torch.Tensor, s_ch: torch.Tensor,
1354
- s_group: torch.Tensor, workspace: torch.Tensor,
1355
- size_m: int, size_n: int, size_k: int) -> torch.Tensor:
1356
- return torch.ops._C.marlin_qqq_gemm(a, b_q_weight, s_tok, s_ch, s_group,
1357
- workspace, size_m, size_n, size_k)
1358
-
1359
-
1324
# gguf
1325
def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int, n: int,
1326
dtype: Optional[torch.dtype]) -> torch.Tensor:
vllm/config/__init__.py
+3
-4
@@ -1112,9 +1112,9 @@ class ModelConfig:
1112
def _verify_quantization(self) -> None:
1113
supported_quantization = me_quant.QUANTIZATION_METHODS
1114
optimized_quantization_methods = [
1115
- "fp8", "marlin", "modelopt", "gptq_marlin_24", "gptq_marlin",
1116
- "awq_marlin", "fbgemm_fp8", "compressed-tensors", "experts_int8",
1117
- "quark", "modelopt_fp4", "bitblas", "gptq_bitblas", "inc"
1115
+ "fp8", "modelopt", "gptq_marlin_24", "gptq_marlin", "awq_marlin",
1116
+ "fbgemm_fp8", "compressed-tensors", "experts_int8", "quark",
1117
+ "modelopt_fp4", "bitblas", "gptq_bitblas", "inc"
1118
]
1119
if self.quantization is not None:
1120
self.quantization = cast(me_quant.QuantizationMethods,
@@ -1137,7 +1137,6 @@ class ModelConfig:
1137
# `override_quantization_method` method) must be checked in order
1138
# of preference (this is particularly important for GPTQ).
1139
overrides = [
1140
- "marlin",
1140
"bitblas",
1141
"gptq_marlin_24",
1142
"gptq_marlin",
vllm/lora/layers.py
-3
@@ -48,9 +48,6 @@ def _get_lora_device(base_layer: nn.Module) -> torch.device:
48
# GPTQ/AWQ
49
elif hasattr(base_layer, "qweight"):
50
return base_layer.qweight.device
51
- # marlin
52
- elif hasattr(base_layer, "B"):
53
- return base_layer.B.device
51
# HQQ marlin
52
elif hasattr(base_layer, "W_q"):
53
return base_layer.W_q.device
vllm/model_executor/layers/linear.py
-1
@@ -42,7 +42,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
42
"GPTQMarlinLinearMethod",
43
"Fp8LinearMethod",
44
"MarlinLinearMethod",
45
- "QQQLinearMethod",
45
"GPTQMarlin24LinearMethod",
46
"TPUInt8LinearMethod",
47
"GPTQLinearMethod",
vllm/model_executor/layers/quantization/__init__.py
-6
@@ -15,7 +15,6 @@ QuantizationMethods = Literal[
15
"fbgemm_fp8",
16
"modelopt",
17
"modelopt_fp4",
18
- "marlin",
18
"bitblas",
19
"gguf",
20
"gptq_marlin_24",
@@ -25,7 +24,6 @@ QuantizationMethods = Literal[
24
"gptq",
25
"compressed-tensors",
26
"bitsandbytes",
28
- "qqq",
27
"hqq",
28
"experts_int8",
29
"neuron_quant",
@@ -106,13 +104,11 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
104
from .hqq_marlin import HQQMarlinConfig
105
from .inc import INCConfig
106
from .ipex_quant import IPEXConfig
109
- from .marlin import MarlinConfig
107
from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config
108
from .moe_wna16 import MoeWNA16Config
109
from .mxfp4 import Mxfp4Config
110
from .neuron_quant import NeuronQuantConfig
111
from .ptpc_fp8 import PTPCFp8Config
115
- from .qqq import QQQConfig
112
from .rtn import RTNConfig
113
from .torchao import TorchAOConfig
114
from .tpu_int8 import Int8TpuConfig
@@ -125,7 +121,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
121
"fbgemm_fp8": FBGEMMFp8Config,
122
"modelopt": ModelOptFp8Config,
123
"modelopt_fp4": ModelOptNvFp4Config,
128
- "marlin": MarlinConfig,
124
"bitblas": BitBLASConfig,
125
"gguf": GGUFConfig,
126
"gptq_marlin_24": GPTQMarlin24Config,
@@ -136,7 +131,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
131
"compressed-tensors": CompressedTensorsConfig,
132
"bitsandbytes": BitsAndBytesConfig,
133
"ptpc_fp8": PTPCFp8Config,
139
- "qqq": QQQConfig,
134
"hqq": HQQMarlinConfig,
135
"experts_int8": ExpertsInt8Config,
136
"neuron_quant": NeuronQuantConfig,
vllm/model_executor/layers/quantization/marlin.py
deleted
-263
@@ -1,263 +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 QuantizationMethods
13
-from vllm.model_executor.layers.quantization.base_config import (
14
- QuantizationConfig)
15
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
16
-from vllm.model_executor.parameter import (BasevLLMParameter,
17
- ChannelQuantScaleParameter,
18
- GroupQuantScaleParameter,
19
- PackedvLLMParameter)
20
-
21
-logger = init_logger(__name__)
22
-
23
-
24
-class MarlinConfig(QuantizationConfig):
25
- """Config class for Marlin.
26
-
27
- Reference: https://github.com/IST-DASLab/marlin/tree/master
28
- """
29
-
30
- def __init__(
31
- self,
32
- group_size: int,
33
- lm_head_quantized: bool,
34
- ) -> None:
35
- super().__init__()
36
-
37
- # Group size for the quantization.
38
- self.group_size = group_size
39
- self.lm_head_quantized = lm_head_quantized
40
- if self.group_size != 128 and self.group_size != -1:
41
- raise ValueError(
42
- "Currently, only group size 128 and -1 (channelwise) "
43
- "is supported for Marlin, but got group_size of "
44
- f"{self.group_size}")
45
-
46
- # 4 Bits packed into 32 bit datatype.
47
- self.pack_factor = 32 // 4
48
-
49
- # Tile size used by marlin kernels.
50
- self.tile_size = 16
51
-
52
- # Min out_features dim
53
- self.min_n_threads = 64
54
-
55
- # Min in_features dim
56
- self.min_k_threads = 128
57
-
58
- # Max parallel problems to solve at once (improves large
59
- # batch performance)
60
- self.max_parallel = 16
61
-
62
- # Permutation length used by the marlin kernels.
63
- self.perm_len = 1024
64
-
65
- def __repr__(self) -> str:
66
- return (f"MarlinConfig(group_size={self.group_size}, "
67
- f"lm_head_quantized={self.lm_head_quantized})")
68
-
69
- @classmethod
70
- def get_name(cls) -> QuantizationMethods:
71
- return "marlin"
72
-
73
- @classmethod
74
- def get_supported_act_dtypes(cls) -> list[torch.dtype]:
75
- return [torch.half]
76
-
77
- @classmethod
78
- # Need to figure it out
79
- def get_min_capability(cls) -> int:
80
- return 80
81
-
82
- @classmethod
83
- def get_config_filenames(cls) -> list[str]:
84
- return ["quantize_config.json"]
85
-
86
- @classmethod
87
- def from_config(cls, config: dict[str, Any]) -> "MarlinConfig":
88
- group_size = cls.get_from_keys(config, ["group_size"])
89
- lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
90
- default=False)
91
- return cls(group_size, lm_head_quantized)
92
-
93
- @classmethod
94
- def override_quantization_method(
95
- cls, hf_quant_cfg, user_quant) -> Optional[QuantizationMethods]:
96
- # compat: autogptq >=0.8.0 use checkpoint_format: str
97
- # compat: autogptq <=0.7.1 is_marlin_format: bool
98
- is_marlin_format = (hf_quant_cfg.get("checkpoint_format") == "marlin"
99
- or hf_quant_cfg.get("is_marlin_format", False))
100
-
101
- is_valid_user_quant = (user_quant is None or user_quant == "gptq"
102
- or user_quant == "marlin")
103
-
104
- if is_marlin_format and is_valid_user_quant:
105
- msg = ("The model is serialized in {} format. Using {} kernel.".
106
- format(cls.get_name(), cls.get_name()))
107
- logger.info(msg)
108
- return cls.get_name()
109
-
110
- return None
111
-
112
- def get_quant_method(self, layer: torch.nn.Module,
113
- prefix: str) -> Optional["MarlinLinearMethod"]:
114
- if (isinstance(layer, LinearBase) or
115
- (isinstance(layer, ParallelLMHead) and self.lm_head_quantized)):
116
- return MarlinLinearMethod(self)
117
- return None
118
-
119
-
120
-class MarlinLinearMethod(LinearMethodBase):
121
- """Linear method for Marlin.
122
-
123
- Args:
124
- quant_config: The Marlin quantization config.
125
- """
126
-
127
- def __init__(self, quant_config: MarlinConfig):
128
- self.quant_config = quant_config
129
-
130
- def create_weights(
131
- self,
132
- layer: torch.nn.Module,
133
- input_size_per_partition: int,
134
- output_partition_sizes: list[int],
135
- input_size: int,
136
- output_size: int,
137
- params_dtype: torch.dtype,
138
- **extra_weight_attrs,
139
- ):
140
- del output_size # Unused.
141
- weight_loader = extra_weight_attrs["weight_loader"]
142
-
143
- if params_dtype != torch.float16:
144
- raise ValueError(
145
- f"The params dtype must be float16, but got {params_dtype}")
146
-
147
- # Validate output_size_per_partition
148
- output_size_per_partition = sum(output_partition_sizes)
149
- if output_size_per_partition % self.quant_config.min_n_threads != 0:
150
- raise ValueError(
151
- f"Weight output_size_per_partition = "
152
- f"{output_size_per_partition} is not divisible by "
153
- f"min_n_threads = {self.quant_config.min_n_threads}.")
154
- if output_size_per_partition % self.quant_config.pack_factor != 0:
155
- raise ValueError(
156
- f"Weight output_size_per_partition = "
157
- f"{output_size_per_partition} is not divisible by "
158
- f"pack_factor = {self.quant_config.pack_factor}.")
159
-
160
- # Validate input_size_per_partition
161
- if input_size_per_partition % self.quant_config.min_k_threads != 0:
162
- raise ValueError(
163
- f"Weight input_size_per_partition = "
164
- f"{input_size_per_partition} is not divisible by "
165
- f"min_k_threads = {self.quant_config.min_k_threads}.")
166
- if (self.quant_config.group_size != -1 and
167
- input_size_per_partition % self.quant_config.group_size != 0):
168
- raise ValueError(f"Weight input_size_per_partition = "
169
- f"{input_size_per_partition} is not divisible by "
170
- f"group_size = {self.quant_config.group_size}.")
171
-
172
- # Check that we have at least 4 tiles horizontally in the shard
173
- num_tiles_per_perm = self.quant_config.perm_len // (
174
- self.quant_config.tile_size**2)
175
- if output_size_per_partition % num_tiles_per_perm != 0:
176
- raise ValueError(
177
- "Each permutation group must reside on the same gpu")
178
-
179
- # Quantized 4Bit weights packed into Int32.
180
- qweight = PackedvLLMParameter(
181
- data=torch.empty(
182
- input_size_per_partition // self.quant_config.tile_size,
183
- output_size_per_partition * self.quant_config.tile_size //
184
- self.quant_config.pack_factor,
185
- device="cuda",
186
- dtype=torch.int32,
187
- ),
188
- input_dim=0,
189
- output_dim=1,
190
- packed_dim=1,
191
- packed_factor=self.quant_config.pack_factor,
192
- marlin_tile_size=self.quant_config.tile_size,
193
- weight_loader=weight_loader)
194
-
195
- # Determine if channelwise or not
196
- input_groups = (1 if self.quant_config.group_size == -1 else
197
- input_size_per_partition //
198
- self.quant_config.group_size)
199
-
200
- weight_scale_args = {
201
- "data":
202
- torch.empty(
203
- input_groups,
204
- output_size_per_partition,
205
- device="cuda",
206
- dtype=params_dtype,
207
- ),
208
- "weight_loader":
209
- weight_loader
210
- }
211
- if input_groups == 1:
212
- scales = ChannelQuantScaleParameter(output_dim=1,
213
- **weight_scale_args)
214
- else:
215
- scales = GroupQuantScaleParameter(output_dim=1,
216
- input_dim=0,
217
- **weight_scale_args)
218
-
219
- # Allocate workspace (Used for internal locking mechanism)
220
- max_workspace_size = (
221
- output_size_per_partition //
222
- self.quant_config.min_n_threads) * self.quant_config.max_parallel
223
-
224
- workspace = BasevLLMParameter(data=torch.zeros(max_workspace_size,
225
- device="cuda",
226
- dtype=torch.int),
227
- weight_loader=weight_loader)
228
-
229
- layer.register_parameter("B", qweight)
230
- layer.register_parameter("s", scales)
231
- layer.register_parameter("workspace", workspace)
232
-
233
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
234
- # required by torch.compile
235
- layer.B = Parameter(layer.B.data, requires_grad=False)
236
- layer.s = Parameter(layer.s.data, requires_grad=False)
237
- layer.workspace = Parameter(layer.workspace.data, requires_grad=False)
238
-
239
- def apply(
240
- self,
241
- layer: torch.nn.Module,
242
- x: torch.Tensor,
243
- bias: Optional[torch.Tensor] = None,
244
- ) -> torch.Tensor:
245
- qweight = layer.B
246
- scales = layer.s
247
- workspace = layer.workspace
248
-
249
- x_2d = x.view(-1, x.shape[-1])
250
-
251
- size_m = x_2d.shape[0]
252
- size_k = x_2d.shape[1]
253
- size_n = scales.shape[1]
254
-
255
- output_2d = ops.marlin_gemm(x_2d, qweight, scales, workspace, size_m,
256
- size_n, size_k)
257
-
258
- output = output_2d.view(x.shape[:-1] + (output_2d.shape[1], ))
259
-
260
- if bias is not None:
261
- output.add_(bias) # In-place add
262
-
263
- return output
vllm/model_executor/layers/quantization/qqq.py
deleted
-275
@@ -1,275 +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 QuantizationMethods
13
-from vllm.model_executor.layers.quantization.base_config import (
14
- QuantizationConfig)
15
-from vllm.model_executor.parameter import (BasevLLMParameter,
16
- ChannelQuantScaleParameter,
17
- GroupQuantScaleParameter,
18
- PackedvLLMParameter)
19
-
20
-logger = init_logger(__name__)
21
-
22
-MARLIN_QQQ_TILE = 16
23
-MARLIN_QQQ_MIN_THREAD_N = 64
24
-MARLIN_QQQ_MIN_THREAD_K = 128
25
-MARLIN_QQQ_MAX_PARALLEL = 16
26
-
27
-MARLIN_QQQ_SUPPORTED_NUM_BITS = [4]
28
-MARLIN_QQQ_SUPPORTED_GROUP_SIZES = [-1, 128]
29
-MARLIN_QQQ_SUPPORTED_SYM = [True]
30
-
31
-
32
-class QQQConfig(QuantizationConfig):
33
- """Config class for QQQ
34
-
35
- Reference: https://arxiv.org/pdf/2406.09904
36
- """
37
-
38
- def __init__(
39
- self,
40
- weight_bits: int,
41
- group_size: int,
42
- is_sym: bool = True,
43
- ) -> None:
44
- super().__init__()
45
- self.weight_bits = weight_bits
46
- self.group_size = group_size
47
- self.is_sym = is_sym
48
-
49
- # Verify
50
- if self.weight_bits not in MARLIN_QQQ_SUPPORTED_NUM_BITS:
51
- raise ValueError(
52
- f"QQQ does not support weight_bits = {self.weight_bits}. "
53
- f"Only weight_bits = {MARLIN_QQQ_SUPPORTED_NUM_BITS} "
54
- "are supported.")
55
- if self.group_size not in MARLIN_QQQ_SUPPORTED_GROUP_SIZES:
56
- raise ValueError(
57
- f"QQQ does not support group_size = {self.group_size}. "
58
- f"Only group_sizes = {MARLIN_QQQ_SUPPORTED_GROUP_SIZES} "
59
- "are supported.")
60
- if self.is_sym not in MARLIN_QQQ_SUPPORTED_SYM:
61
- raise ValueError(
62
- f"QQQ does not support is_sym = {self.is_sym}. "
63
- f"Only sym = {MARLIN_QQQ_SUPPORTED_SYM} are supported.")
64
-
65
- # 4 Bits packed into 32 bit datatype.
66
- self.pack_factor = 32 // self.weight_bits
67
-
68
- # Tile size used by QQQ kernels.
69
- self.tile_size = MARLIN_QQQ_TILE
70
-
71
- # Min out_features dim
72
- self.min_n_threads = MARLIN_QQQ_MIN_THREAD_N
73
-
74
- # Min in_features dim
75
- self.min_k_threads = MARLIN_QQQ_MIN_THREAD_K
76
-
77
- # Max parallel problems to solve at once (improves large
78
- # batch performance)
79
- self.max_parallel = MARLIN_QQQ_MAX_PARALLEL
80
-
81
- # Permutation length used by the QQQ kernels.
82
- self.perm_len = 1024
83
-
84
- def __repr__(self) -> str:
85
- return "QQQConfig(weight_bits={}, group_size={})".format(
86
- self.weight_bits, self.group_size)
87
-
88
- @classmethod
89
- def get_name(cls) -> QuantizationMethods:
90
- return "qqq"
91
-
92
- @classmethod
93
- def get_supported_act_dtypes(cls) -> list[torch.dtype]:
94
- return [torch.half]
95
-
96
- @classmethod
97
- def get_min_capability(cls) -> int:
98
- return 80
99
-
100
- @classmethod
101
- def get_config_filenames(cls) -> list[str]:
102
- """List of filenames to search for in the model directory."""
103
- return [
104
- "quant_config.json",
105
- "quantize_config.json",
106
- ]
107
-
108
- @classmethod
109
- def from_config(cls, config: dict[str, Any]) -> "QQQConfig":
110
- weight_bits = cls.get_from_keys(config, ["wbits"])
111
- group_size = cls.get_from_keys(config, ["group_size"])
112
- return cls(weight_bits, group_size)
113
-
114
- def get_quant_method(self, layer: torch.nn.Module,
115
- prefix: str) -> Optional["QQQLinearMethod"]:
116
- if isinstance(layer, LinearBase):
117
- return QQQLinearMethod(self)
118
- return None
119
-
120
-
121
-class QQQLinearMethod(LinearMethodBase):
122
- """Linear method for QQQ.
123
-
124
- Args:
125
- quant_config: The QQQ quantization config.
126
- """
127
-
128
- def __init__(self, quant_config: QQQConfig):
129
- self.quant_config = quant_config
130
-
131
- def create_weights(
132
- self,
133
- layer: torch.nn.Module,
134
- input_size_per_partition: int,
135
- output_partition_sizes: list[int],
136
- input_size: int,
137
- output_size: int,
138
- params_dtype: torch.dtype,
139
- **extra_weight_attrs,
140
- ):
141
- weight_loader = extra_weight_attrs["weight_loader"]
142
- if params_dtype != torch.float16:
143
- raise ValueError(
144
- f"The params dtype must be float16, but got {params_dtype}")
145
-
146
- # Validate output_size_per_partition
147
- output_size_per_partition = sum(output_partition_sizes)
148
- if output_size_per_partition % self.quant_config.min_n_threads != 0:
149
- raise ValueError(
150
- f"Weight output_size_per_partition = "
151
- f"{output_size_per_partition} is not divisible by "
152
- f"min_n_threads = {self.quant_config.min_n_threads}.")
153
- if output_size_per_partition % self.quant_config.pack_factor != 0:
154
- raise ValueError(
155
- f"Weight output_size_per_partition = "
156
- f"{output_size_per_partition} is not divisible by "
157
- f"pack_factor = {self.quant_config.pack_factor}.")
158
-
159
- # Validate input_size_per_partition
160
- if input_size_per_partition % self.quant_config.min_k_threads != 0:
161
- raise ValueError(
162
- f"Weight input_size_per_partition = "
163
- f"{input_size_per_partition} is not divisible by "
164
- f"min_k_threads = {self.quant_config.min_k_threads}.")
165
- if (self.quant_config.group_size != -1 and
166
- input_size_per_partition % self.quant_config.group_size != 0):
167
- raise ValueError(f"Weight input_size_per_partition = "
168
- f"{input_size_per_partition} is not divisible by "
169
- f"group_size = {self.quant_config.group_size}.")
170
-
171
- # Check that we have at least 4 tiles horizontally in the shard
172
- num_tiles_per_perm = self.quant_config.perm_len // (
173
- self.quant_config.tile_size**2)
174
- if output_size_per_partition % num_tiles_per_perm != 0:
175
- raise ValueError(
176
- "Each permutation group must reside on the same gpu")
177
-
178
- # Quantized 4Bit weights packed into Int32.
179
- qweight = PackedvLLMParameter(
180
- data=torch.empty(
181
- input_size_per_partition // self.quant_config.tile_size,
182
- output_size_per_partition * self.quant_config.tile_size //
183
- self.quant_config.pack_factor,
184
- device="cuda",
185
- dtype=torch.int32,
186
- ),
187
- input_dim=0,
188
- output_dim=1,
189
- packed_dim=1,
190
- packed_factor=self.quant_config.pack_factor,
191
- marlin_tile_size=self.quant_config.tile_size,
192
- weight_loader=weight_loader)
193
-
194
- s_channel = ChannelQuantScaleParameter(data=torch.empty(
195
- 1,
196
- output_size_per_partition,
197
- device="cuda",
198
- dtype=torch.float,
199
- ),
200
- weight_loader=weight_loader,
201
- output_dim=1)
202
-
203
- if self.quant_config.group_size == -1:
204
- s_group_data = torch.tensor(
205
- [],
206
- device="cuda",
207
- dtype=torch.half,
208
- )
209
- else:
210
- s_group_data = torch.empty(
211
- input_size_per_partition // self.quant_config.group_size,
212
- output_size_per_partition,
213
- device="cuda",
214
- dtype=torch.half,
215
- )
216
-
217
- s_group_attr = {"data": s_group_data, "weight_loader": weight_loader}
218
-
219
- if self.quant_config.group_size == -1:
220
- s_group = BasevLLMParameter(**s_group_attr)
221
- else:
222
- s_group = GroupQuantScaleParameter(output_dim=1,
223
- input_dim=0,
224
- **s_group_attr)
225
-
226
- # Allocate workspace (Used for internal locking mechanism)
227
- max_workspace_size = (
228
- output_size_per_partition //
229
- self.quant_config.min_n_threads) * self.quant_config.max_parallel
230
-
231
- workspace = BasevLLMParameter(data=torch.zeros(max_workspace_size,
232
- device="cuda",
233
- dtype=torch.int),
234
- weight_loader=weight_loader)
235
-
236
- layer.register_parameter("B", qweight)
237
- layer.register_parameter("s_channel", s_channel)
238
- layer.register_parameter("s_group", s_group)
239
- layer.register_parameter("workspace", workspace)
240
-
241
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
242
- # required by torch.compile
243
- layer.B = Parameter(layer.B.data, requires_grad=False)
244
- layer.s_channel = Parameter(layer.s_channel.data, requires_grad=False)
245
- layer.s_group = Parameter(layer.s_group.data, requires_grad=False)
246
- layer.workspace = Parameter(layer.workspace.data, requires_grad=False)
247
-
248
- def apply(
249
- self,
250
- layer: torch.nn.Module,
251
- x: torch.Tensor,
252
- bias: Optional[torch.Tensor] = None,
253
- ) -> torch.Tensor:
254
- qweight = layer.B
255
- s_ch = layer.s_channel
256
- s_group = layer.s_group
257
- workspace = layer.workspace
258
-
259
- x_2d = x.view(-1, x.shape[-1])
260
-
261
- size_m = x_2d.shape[0]
262
- size_k = x_2d.shape[1]
263
- size_n = s_ch.shape[1]
264
-
265
- x_int8, s_tok, _ = ops.scaled_int8_quant(x_2d)
266
-
267
- output_2d = ops.marlin_qqq_gemm(x_int8, qweight, s_tok, s_ch, s_group,
268
- workspace, size_m, size_n, size_k)
269
-
270
- output = output_2d.view(x.shape[:-1] + (output_2d.shape[1], ))
271
-
272
- if bias is not None:
273
- output.add_(bias) # In-place add
274
-
275
- return output
vllm/model_executor/layers/quantization/utils/marlin_utils_test_qqq.py
deleted
-126
@@ -1,126 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import numpy
5
-import torch
6
-
7
-from .marlin_utils_test import marlin_permute_weights
8
-from .quant_utils import get_pack_factor, qqq_quantize_weights
9
-
10
-
11
-def marlin_qqq_weights(q_w, size_k, size_n, num_bits, perm, group_size):
12
- # Permute
13
- q_w = marlin_permute_weights(q_w, size_k, size_n, perm)
14
-
15
- # Pack
16
- pack_factor = get_pack_factor(num_bits)
17
- orig_device = q_w.device
18
-
19
- q_w = q_w.cpu().numpy().astype(numpy.uint32)
20
-
21
- q_packed = numpy.zeros((q_w.shape[0], q_w.shape[1] // pack_factor),
22
- dtype=numpy.uint32)
23
- if group_size == size_k:
24
- for i in range(pack_factor):
25
- q_packed |= (q_w[:, i::pack_factor] & 0xF) << num_bits * i
26
- else:
27
- for i in range(pack_factor):
28
- q_packed |= q_w[:, i::pack_factor] << num_bits * i
29
-
30
- q_packed = torch.from_numpy(q_packed.astype(numpy.int32)).to(orig_device)
31
-
32
- return q_packed
33
-
34
-
35
-def get_qqq_scale_perms():
36
- scale_perm: list[int] = []
37
- for i in range(8):
38
- scale_perm.extend([i + 8 * j for j in range(8)])
39
- scale_perm_single: list[int] = []
40
- for i in range(4):
41
- scale_perm_single.extend(
42
- [2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]])
43
- return scale_perm, scale_perm_single
44
-
45
-
46
-# NOTE(HandH1998): QQQ employs different perms for per-group and per-channel weight quantization. # noqa: E501
47
-def get_qqq_weight_perm(num_bits: int, quant_type: str):
48
- perm_list: list[int] = []
49
- for i in range(32):
50
- perm1: list[int] = []
51
- col = i // 4
52
- for block in [0, 1]:
53
- for row in [
54
- 4 * (i % 4),
55
- 4 * (i % 4) + 1,
56
- 4 * (i % 4) + 2,
57
- 4 * (i % 4) + 3,
58
- ]:
59
- perm1.append(16 * row + col + 8 * block)
60
- for j in range(4):
61
- perm_list.extend([p + 256 * j for p in perm1])
62
-
63
- perm = numpy.array(perm_list)
64
-
65
- assert quant_type in ["per-channel",
66
- "per-group"], "not supported quantization type"
67
- if num_bits == 4:
68
- if quant_type == "per-channel":
69
- interleave = numpy.array([4, 0, 5, 1, 6, 2, 7, 3])
70
- else:
71
- interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7])
72
- else:
73
- raise Exception("num_bits must be 4, got {}".format(num_bits))
74
-
75
- perm = perm.reshape((-1, len(interleave)))[:, interleave].ravel()
76
- perm = torch.from_numpy(perm)
77
- return perm
78
-
79
-
80
-def marlin_qqq_permute_scales(s_group, s_channel, size_k, size_n, group_size):
81
- scale_perm, scale_perm_single = get_qqq_scale_perms()
82
- if group_size < size_k and group_size != -1:
83
- s_group = s_group.reshape((-1, len(scale_perm)))[:, scale_perm]
84
- s_channel = s_channel.reshape(
85
- (-1, len(scale_perm_single)))[:, scale_perm_single]
86
- s_group = s_group.reshape((-1, size_n)).contiguous()
87
- else:
88
- s_channel = s_channel.reshape(
89
- (-1, len(scale_perm_single)))[:, scale_perm_single]
90
- s_channel = s_channel.reshape((-1, size_n)).contiguous()
91
-
92
- return s_group, s_channel
93
-
94
-
95
-def marlin_qqq_quantize(
96
- w: torch.Tensor,
97
- num_bits: int,
98
- group_size: int,
99
-):
100
- size_k, size_n = w.shape
101
-
102
- # Normalize group_size
103
- if group_size == -1:
104
- group_size = size_k
105
- assert group_size <= size_k
106
- quant_type = "per-channel" if group_size == size_k else "per-group"
107
-
108
- # Quantize
109
- w_ref, q_w, s_group, s_channel = qqq_quantize_weights(
110
- w, num_bits, group_size)
111
-
112
- # Reformat to marlin_qqq
113
- weight_perm = get_qqq_weight_perm(num_bits, quant_type)
114
- marlin_qqq_q_w = marlin_qqq_weights(q_w, size_k, size_n, num_bits,
115
- weight_perm, group_size)
116
- marlin_qqq_s_group, marlin_qqq_s_channel = marlin_qqq_permute_scales(
117
- s_group, s_channel, size_k, size_n, group_size)
118
-
119
- # Create result
120
- res_list = [
121
- w_ref, marlin_qqq_q_w, marlin_qqq_s_group, marlin_qqq_s_channel
122
- ]
123
- for i in range(len(res_list)):
124
- res_list[i] = res_list[i].to(w.device)
125
-
126
- return res_list
vllm/model_executor/layers/quantization/utils/quant_utils.py
-85
@@ -9,8 +9,6 @@ import numpy
9
import torch
10
11
from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
12
-from vllm.model_executor.layers.quantization.qqq import (
13
- MARLIN_QQQ_SUPPORTED_NUM_BITS)
12
from vllm.platforms import current_platform
13
from vllm.scalar_type import ScalarType, scalar_types
14
@@ -386,89 +384,6 @@ def gptq_quantize_weights(w: torch.Tensor,
384
return w_ref, w_q, w_s, g_idx, rand_perm
385
386
389
-# QQQ employs different quant schemes for per-group and
390
-# per-channel quantization.
391
-def qqq_quantize_weights(w: torch.Tensor, num_bits: int, group_size: int):
392
- orig_device = w.device
393
- size_k, size_n = w.shape
394
-
395
- assert w.is_floating_point(), "w must be float"
396
- assert num_bits in MARLIN_QQQ_SUPPORTED_NUM_BITS, \
397
- f"Unsupported num_bits = {num_bits}"
398
- assert group_size in SUPPORTED_GROUP_SIZES + [
399
- size_k
400
- ], f"Unsupported groupsize = {group_size}"
401
-
402
- if group_size == -1:
403
- group_size = size_k
404
- assert group_size <= size_k
405
-
406
- if group_size < size_k:
407
- # Reshape to [groupsize, -1]
408
- w = w.reshape((-1, group_size, size_n))
409
- w = w.permute(1, 0, 2)
410
- w = w.reshape((group_size, -1))
411
-
412
- max_q_val = 2**num_bits - 1
413
- half_q_val = (max_q_val + 1) // 2
414
-
415
- # Compute scale for each group
416
- s_group = torch.max(torch.abs(w), 0, keepdim=True)[0]
417
- s_group *= 2 / max_q_val # 2 => symmetric
418
-
419
- # Quantize
420
- q_w = torch.round(w / s_group).int()
421
- q_w += half_q_val
422
- q_w = torch.clamp(q_w, 0, max_q_val)
423
- # Compute ref (dequantized)
424
- w_ref = (q_w - half_q_val).half() * s_group
425
-
426
- # Restore original shapes
427
- def reshape_w(w):
428
- w = w.reshape((group_size, -1, size_n))
429
- w = w.permute(1, 0, 2)
430
- w = w.reshape((size_k, size_n)).contiguous()
431
- return w
432
-
433
- q_w = reshape_w(q_w)
434
- w_ref = reshape_w(w_ref)
435
-
436
- # Compute int8 quantization scale for each channel
437
- s_channel = torch.max(torch.abs(w_ref), 0, keepdim=True)[0]
438
- s_channel /= 127.0
439
- t_int8 = (w_ref / s_channel).round().clamp(-128, 127).to(torch.int8)
440
- w_ref = t_int8.half() * s_channel
441
- s_channel = s_channel.reshape(1, -1).to(dtype=torch.float)
442
-
443
- # Fuse scales
444
- s_group = (s_group.reshape(-1, size_n).contiguous() /
445
- s_channel).to(dtype=torch.half)
446
- else:
447
- max_q_val = 2**(num_bits - 1) - 1
448
-
449
- # Compute scale for each channel
450
- s_channel = torch.max(torch.abs(w), 0, keepdim=True)[0]
451
- s_channel /= max_q_val
452
-
453
- # Quantize
454
- q_w = torch.round(w / s_channel).int()
455
- q_w = torch.clamp(q_w, -max_q_val, max_q_val)
456
- # Compute ref (dequantized)
457
- w_ref = q_w.half() * s_channel
458
-
459
- s_group = torch.tensor([], dtype=torch.half)
460
- # div 2 ** (8 - self.bits)) to offset right shift in unpacking
461
- s_channel /= (2**(8 - num_bits))
462
- s_channel = s_channel.reshape(-1, size_n).contiguous().to(torch.float)
463
-
464
- return (
465
- w_ref.to(device=orig_device),
466
- q_w.to(device=orig_device),
467
- s_group.to(device=orig_device),
468
- s_channel.to(device=orig_device),
469
- )
470
-
471
-
387
def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor):
388
orig_device = q_w.device
389