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 15, 12:15 UTC
[Kernel] some optimizations for dense marlin and moe marlin (#16850)
Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Jinzhen Lin committed
May 6, 2025 at 00:39 UTC
1d0c9d6b2d74e6884c6c63f4856ac0a2cddd0508
26 files changed
+3501
-3257
CMakeLists.txt
+46
-2
@@ -301,8 +301,52 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
301
# are not supported by Machete yet.
302
cuda_archs_loose_intersection(MARLIN_ARCHS "8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
303
if (MARLIN_ARCHS)
304
+
305
+ #
306
+ # For the Marlin kernels we automatically generate sources for various
307
+ # preselected input type pairs and schedules.
308
+ # Generate sources:
309
+ set(MARLIN_GEN_SCRIPT
310
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
311
+ file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
312
+
313
+ message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH}")
314
+ message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH}")
315
+
316
+ if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH}
317
+ OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH})
318
+ execute_process(
319
+ COMMAND ${CMAKE_COMMAND} -E env
320
+ PYTHONPATH=$PYTHONPATH
321
+ ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT}
322
+ RESULT_VARIABLE marlin_generation_result
323
+ OUTPUT_VARIABLE marlin_generation_result
324
+ OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
325
+ ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
326
+ )
327
+
328
+ if (NOT marlin_generation_result EQUAL 0)
329
+ message(FATAL_ERROR "Marlin generation failed."
330
+ " Result: \"${marlin_generation_result}\""
331
+ "\nCheck the log for details: "
332
+ "${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
333
+ else()
334
+ set(MARLIN_GEN_SCRIPT_HASH ${MARLIN_GEN_SCRIPT_HASH}
335
+ CACHE STRING "Last run Marlin generate script hash" FORCE)
336
+ message(STATUS "Marlin generation completed successfully.")
337
+ endif()
338
+ else()
339
+ message(STATUS "Marlin generation script has not changed, skipping generation.")
340
+ endif()
341
+
342
+ file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/kernel_*.cu")
343
+ set_gencode_flags_for_srcs(
344
+ SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
345
+ CUDA_ARCHS "${MARLIN_ARCHS}")
346
+
347
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
348
+
349
set(MARLIN_SRCS
305
- "csrc/quantization/fp8/fp8_marlin.cu"
350
"csrc/quantization/marlin/dense/marlin_cuda_kernel.cu"
351
"csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
352
"csrc/quantization/marlin/qqq/marlin_qqq_gemm_kernel.cu"
@@ -644,7 +688,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
688
OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH})
689
execute_process(
690
COMMAND ${CMAKE_COMMAND} -E env
647
- PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$PYTHONPATH
691
+ PYTHONPATH=$PYTHONPATH
692
${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT}
693
RESULT_VARIABLE moe_marlin_generation_result
694
OUTPUT_VARIABLE moe_marlin_generation_output
csrc/moe/marlin_moe_wna16/.gitignore
new
+1
@@ -0,0 +1 @@
1
+kernel_*.cu
\ No newline at end of file
csrc/moe/marlin_moe_wna16/generate_kernels.py
+12
-8
@@ -25,15 +25,13 @@ TEMPLATE = ("template __global__ void Marlin<"
25
"{{thread_k_blocks}}, "
26
"{{'true' if m_block_size_8 else 'false'}}, "
27
"{{stages}}, "
28
- "{{'true' if has_act_order else 'false'}}, "
29
- "{{'true' if has_zp else 'false'}}, "
28
"{{group_blocks}}, "
29
"{{'true' if is_zp_float else 'false'}}>"
30
"( MARLIN_KERNEL_PARAMS );")
31
32
# int8 with zero point case (vllm::kU8) is also supported,
33
# we don't add it to reduce wheel size.
36
-SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128"]
34
+SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn"]
35
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
36
37
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
@@ -52,21 +50,29 @@ def remove_old_kernels():
50
51
def generate_new_kernels():
52
for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
55
- has_zp = "B" not in scalar_type
53
all_template_str_list = []
54
55
for group_blocks, m_blocks, thread_configs in itertools.product(
56
GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):
57
61
- has_act_order = group_blocks == 0
62
- if has_zp and has_act_order:
58
+ # act order case only support gptq-int4 and gptq-int8
59
+ if group_blocks == 0 and scalar_type not in [
60
+ "vllm::kU4B8", "vllm::kU8B128"
61
+ ]:
62
continue
63
if thread_configs[2] == 256:
64
+ # for small batch (m_blocks == 1), we only need (128, 128, 256)
65
+ # for large batch (m_blocks > 1), we only need (64, 256, 256)
66
if m_blocks <= 1 and thread_configs[0] != 128:
67
continue
68
if m_blocks > 1 and thread_configs[0] != 64:
69
continue
70
71
+ # we only support channelwise quantization and group_size == 128
72
+ # for fp8
73
+ if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
74
+ continue
75
+
76
k_blocks = thread_configs[0] // 16
77
n_blocks = thread_configs[1] // 16
78
threads = thread_configs[2]
@@ -82,8 +88,6 @@ def generate_new_kernels():
88
thread_k_blocks=k_blocks,
89
m_block_size_8=m_blocks == 0.5,
90
stages="pipe_stages",
85
- has_act_order=has_act_order,
86
- has_zp=has_zp,
91
group_blocks=group_blocks,
92
is_zp_float=False,
93
)
csrc/moe/marlin_moe_wna16/kernel.h
+4
-6
@@ -18,7 +18,7 @@
18
const float *__restrict__ topk_weights_ptr, int top_k, \
19
bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
20
int prob_n, int prob_k, int *locks, bool use_atomic_add, \
21
- bool use_fp32_reduce
21
+ bool use_fp32_reduce, int max_shared_mem
22
23
namespace MARLIN_NAMESPACE_NAME {
24
template <typename scalar_t, // compute dtype, half or nv_float16
@@ -33,11 +33,9 @@ template <typename scalar_t, // compute dtype, half or nv_float16
33
// only works when thread_m_blocks == 1
34
const int stages, // number of stages for the async global->shared
35
// fetch pipeline
36
- const bool has_act_order, // whether act_order is enabled
37
- const bool has_zp, // whether zero-points are enabled
38
- const int group_blocks, // number of consecutive 16x16 blocks
39
- // with a separate quantization scale
40
- const bool is_zp_float // is zero point of float16 type?
36
+ const int group_blocks, // number of consecutive 16x16 blocks
37
+ // with a separate quantization scale
38
+ const bool is_zp_float // is zero point of float16 type?
39
>
40
__global__ void Marlin(MARLIN_KERNEL_PARAMS);
41
csrc/moe/marlin_moe_wna16/marlin_template.h
+208
-260
@@ -25,6 +25,7 @@
25
26
#include "quantization/gptq_marlin/marlin.cuh"
27
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
28
+#include "quantization/gptq_marlin/dequant.h"
29
#include "core/scalar_type.hpp"
30
31
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
@@ -48,11 +49,9 @@ template <typename scalar_t, // compute dtype, half or nv_float16
49
// only works when thread_m_blocks == 1
50
const int stages, // number of stages for the async global->shared
51
// fetch pipeline
51
- const bool has_act_order, // whether act_order is enabled
52
- const bool has_zp, // whether zero-points are enabled
53
- const int group_blocks, // number of consecutive 16x16 blocks
54
- // with a separate quantization scale
55
- const bool is_zp_float // is zero point of float16 type?
52
+ const int group_blocks, // number of consecutive 16x16 blocks
53
+ // with a separate quantization scale
54
+ const bool is_zp_float // is zero point of float16 type?
55
>
56
__global__ void Marlin(
57
const int4* __restrict__ A, // fp16 input matrix of shape mxk
@@ -77,8 +76,8 @@ __global__ void Marlin(
76
int prob_k, // reduction dimension k
77
int* locks, // extra global storage for barrier synchronization
78
bool use_atomic_add, // whether to use atomic add to reduce
80
- bool use_fp32_reduce // whether to use fp32 global reduce
81
-) {}
79
+ bool use_fp32_reduce, // whether to use fp32 global reduce
80
+ int max_shared_mem) {}
81
82
} // namespace MARLIN_NAMESPACE_NAME
83
@@ -166,144 +165,6 @@ __device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
165
}
166
}
167
169
-// Lookup-table based 3-input logical operation; explicitly used for
170
-// dequantization as the compiler does not seem to automatically recognize it in
171
-// all cases.
172
-template <int lut>
173
-__device__ inline int lop3(int a, int b, int c) {
174
- int res;
175
- asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
176
- : "=r"(res)
177
- : "r"(a), "r"(b), "r"(c), "n"(lut));
178
- return res;
179
-}
180
-
181
-// Constructs destination register by taking bytes from 2 sources (based on
182
-// mask)
183
-template <int start_byte, int mask>
184
-__device__ inline uint32_t prmt(uint32_t a) {
185
- uint32_t res;
186
- asm volatile("prmt.b32 %0, %1, %2, %3;\n"
187
- : "=r"(res)
188
- : "r"(a), "n"(start_byte), "n"(mask));
189
- return res;
190
-}
191
-
192
-template <typename scalar_t, int bit>
193
-__device__ inline typename ScalarType<scalar_t>::FragB dequant(
194
- int q, typename ScalarType<scalar_t>::FragB& frag_b);
195
-
196
-//
197
-// Efficiently dequantize 4bit values packed in an int32 value into a full
198
-// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
199
-// with some small changes:
200
-// - FP16:
201
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
202
-// - BF16:
203
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
204
-//
205
-template <>
206
-__device__ inline typename ScalarType<half>::FragB dequant<half, 4>(
207
- int q, typename ScalarType<half>::FragB& frag_b) {
208
- const int LO = 0x000f000f;
209
- const int HI = 0x00f000f0;
210
- const int EX = 0x64006400;
211
- // Guarantee that the `(a & b) | c` operations are LOP3s.
212
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
213
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
214
- // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
215
- // directly into `SUB` and `ADD`.
216
- const int SUB = 0x64086408;
217
- const int MUL = 0x2c002c00;
218
- const int ADD = 0xd480d480;
219
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
220
- *reinterpret_cast<const half2*>(&SUB));
221
- frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
222
- *reinterpret_cast<const half2*>(&MUL),
223
- *reinterpret_cast<const half2*>(&ADD));
224
- return frag_b;
225
-}
226
-
227
-template <>
228
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
229
-dequant<nv_bfloat16, 4>(int q,
230
- typename ScalarType<nv_bfloat16>::FragB& frag_b) {
231
- static constexpr uint32_t MASK = 0x000f000f;
232
- static constexpr uint32_t EX = 0x43004300;
233
-
234
- // Guarantee that the `(a & b) | c` operations are LOP3s.
235
-
236
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
237
- q >>= 4;
238
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
239
-
240
- static constexpr uint32_t MUL = 0x3F803F80;
241
- static constexpr uint32_t ADD = 0xC308C308;
242
-
243
- frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
244
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
245
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
246
- frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
247
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
248
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
249
- return frag_b;
250
-}
251
-
252
-//
253
-// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
254
-// bf16 Reference:
255
-// - FP16:
256
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
257
-// - BF16:
258
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
259
-//
260
-template <>
261
-__device__ inline typename ScalarType<half>::FragB dequant<half, 8>(
262
- int q, typename ScalarType<half>::FragB& frag_b) {
263
- static constexpr uint32_t mask_for_elt_01 = 0x5250;
264
- static constexpr uint32_t mask_for_elt_23 = 0x5351;
265
- static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
266
-
267
- uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
268
- uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
269
-
270
- static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
271
-
272
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
273
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
274
- frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
275
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
276
- return frag_b;
277
-}
278
-
279
-template <>
280
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
281
-dequant<nv_bfloat16, 8>(int q,
282
- typename ScalarType<nv_bfloat16>::FragB& frag_b) {
283
- float fp32_intermediates[4];
284
- uint32_t* fp32_intermediates_casted =
285
- reinterpret_cast<uint32_t*>(fp32_intermediates);
286
-
287
- static constexpr uint32_t fp32_base = 0x4B000000;
288
- fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
289
- fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
290
- fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
291
- fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
292
-
293
- fp32_intermediates[0] -= 8388736.f;
294
- fp32_intermediates[1] -= 8388736.f;
295
- fp32_intermediates[2] -= 8388736.f;
296
- fp32_intermediates[3] -= 8388736.f;
297
-
298
- uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&frag_b);
299
- bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
300
- fp32_intermediates_casted[1], 0x7632);
301
- bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
302
- fp32_intermediates_casted[3], 0x7632);
303
-
304
- return frag_b;
305
-}
306
-
168
// Multiply dequantized values by the corresponding quantization scale; used
169
// only for grouped quantization.
170
template <typename scalar_t>
@@ -429,11 +290,9 @@ template <typename scalar_t, // compute dtype, half or nv_float16
290
// only works when thread_m_blocks == 1
291
const int stages, // number of stages for the async global->shared
292
// fetch pipeline
432
- const bool has_act_order, // whether act_order is enabled
433
- const bool has_zp, // whether zero-points are enabled
434
- const int group_blocks, // number of consecutive 16x16 blocks
435
- // with a separate quantization scale
436
- const bool is_zp_float // is zero point of float16 type?
293
+ const int group_blocks, // number of consecutive 16x16 blocks
294
+ // with a separate quantization scale
295
+ const bool is_zp_float // is zero point of float16 type?
296
>
297
__global__ void Marlin(
298
const int4* __restrict__ A, // fp16 input matrix of shape mxk
@@ -458,8 +317,8 @@ __global__ void Marlin(
317
int prob_k, // reduction dimension k
318
int* locks, // extra global storage for barrier synchronization
319
bool use_atomic_add, // whether to use atomic add to reduce
461
- bool use_fp32_reduce // whether to use fp32 global reduce
462
-) {
320
+ bool use_fp32_reduce, // whether to use fp32 global reduce
321
+ int max_shared_mem) {
322
// Each threadblock processes one "stripe" of the B matrix with (roughly) the
323
// same size, which might involve multiple column "slices" (of width 16 *
324
// `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
@@ -481,6 +340,8 @@ __global__ void Marlin(
340
341
extern __shared__ int4 sh[];
342
static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
343
+ constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
344
+ constexpr bool has_act_order = group_blocks == 0;
345
346
constexpr int pack_factor = 32 / w_type.size_bits();
347
static_assert(thread_m_blocks == 1 || !m_block_size_8);
@@ -534,13 +395,20 @@ __global__ void Marlin(
395
int64_t B_expert_off = 0;
396
397
int4* sh_block_sorted_ids_int4 = sh;
398
+ int4* sh_rd_block_sorted_ids_int4 =
399
+ sh_block_sorted_ids_int4 + moe_block_size / 4;
400
+ int4* sh_block_topk_weights_int4 =
401
+ sh_rd_block_sorted_ids_int4 + moe_block_size / 4;
402
+ // sh_block_topk_weights_int4 only need (moe_block_size / 4);
403
+ // but we pad to align to 256 bytes
404
+ int4* sh_new =
405
+ sh_block_topk_weights_int4 + moe_block_size / 2 + moe_block_size;
406
int32_t* sh_block_sorted_ids =
407
reinterpret_cast<int*>(sh_block_sorted_ids_int4);
539
- int4* sh_block_topk_weights_int4 =
540
- sh_block_sorted_ids_int4 + moe_block_size / 4;
408
+ int32_t* sh_rd_block_sorted_ids =
409
+ reinterpret_cast<int*>(sh_rd_block_sorted_ids_int4);
410
scalar_t2* sh_block_topk_weights =
411
reinterpret_cast<scalar_t2*>(sh_block_topk_weights_int4);
543
- int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 4;
412
413
int32_t block_num_valid_tokens = 0;
414
int32_t locks_off = 0;
@@ -584,6 +452,11 @@ __global__ void Marlin(
452
sh_block_sorted_ids_int4[tid4] = reinterpret_cast<const int4*>(
453
sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4];
454
455
+ #pragma unroll
456
+ for (int i = 0; i < 4; i++)
457
+ sh_rd_block_sorted_ids[tid4 * 4 + i] =
458
+ sh_block_sorted_ids[tid4 * 4 + i] / top_k;
459
+
460
if (mul_topk_weights) {
461
#pragma unroll
462
for (int i = 0; i < 4; i++) {
@@ -743,6 +616,7 @@ __global__ void Marlin(
616
constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
617
// constexpr int act_s_row_stride = 1;
618
// int act_s_col_stride = act_s_row_stride * num_groups;
619
+ constexpr int act_s_max_num_groups = 32;
620
int act_s_col_stride = 1;
621
int act_s_col_warp_stride = act_s_col_stride * 8;
622
int tb_n_warps = thread_n_blocks / 4;
@@ -758,9 +632,9 @@ __global__ void Marlin(
632
int zp_gl_rd_delta = zp_gl_stride;
633
634
// Global A read index of current thread.
761
- int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
762
- (threadIdx.x % a_gl_rd_delta_o);
763
- a_gl_rd += a_gl_rd_delta_o * slice_row;
635
+ int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o;
636
+ int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
637
+
638
// Shared write index of current thread.
639
int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
640
(threadIdx.x % a_gl_rd_delta_o);
@@ -774,8 +648,8 @@ __global__ void Marlin(
648
(threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
649
b_gl_rd += b_sh_stride * slice_col;
650
b_gl_rd += b_gl_rd_delta_o * slice_row;
777
- int b_sh_wr = threadIdx.x * b_thread_vecs;
778
- int b_sh_rd = threadIdx.x * b_thread_vecs;
651
+ auto b_sh_wr = threadIdx.x * b_thread_vecs;
652
+ auto b_sh_rd = threadIdx.x * b_thread_vecs;
653
654
// For act_order
655
constexpr int k_iter_size = tb_k / b_sh_wr_iters;
@@ -794,7 +668,7 @@ __global__ void Marlin(
668
s_sh_stride * slice_col + threadIdx.x;
669
}
670
}
797
- int s_sh_wr = threadIdx.x;
671
+ auto s_sh_wr = threadIdx.x;
672
bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
673
674
// Zero-points
@@ -807,7 +681,7 @@ __global__ void Marlin(
681
zp_sh_stride * slice_col + threadIdx.x;
682
}
683
}
810
- int zp_sh_wr = threadIdx.x;
684
+ auto zp_sh_wr = threadIdx.x;
685
bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
686
687
// We use a different scale layout for grouped and column-wise quantization as
@@ -851,7 +725,7 @@ __global__ void Marlin(
725
// each warp must also write a consecutive memory segment?
726
auto transform_a = [&](int i) {
727
int row = i / a_gl_rd_delta_o;
854
- return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
728
+ return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8);
729
};
730
// Since the computation of this remapping is non-trivial and, due to our main
731
// loop unrolls, all shared memory accesses are static, we simply precompute
@@ -879,12 +753,28 @@ __global__ void Marlin(
753
B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
754
755
// Shared memory storage for global fetch pipelines.
882
- int4* sh_a = sh_new;
883
- int4* sh_b = sh_a + (stages * a_sh_stage);
884
- int4* sh_g_idx = sh_b + (stages * b_sh_stage);
756
+ constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
757
+ constexpr int sh_b_size = stages * b_sh_stage;
758
+ int4* sh_b = sh_new;
759
+ int4* sh_red = sh_new;
760
+ int4* sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
761
int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
762
+ constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
763
+ : (stages * s_sh_stage);
764
int4* sh_s = sh_zp + (stages * zp_sh_stage);
887
- int4* sh_red = sh_b;
765
+ // shared memory reused by reduction should be smaller than
766
+ // shared memory used by weight.
767
+ static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
768
+ stages * b_sh_stage);
769
+ int4* sh_a = sh_s + sh_s_size;
770
+ constexpr int shm_size_used =
771
+ moe_block_size + stages * (g_idx_stage + zp_sh_stage) + sh_s_size +
772
+ (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
773
+
774
+ // all remaining shared memory is used to cache A (input)
775
+ // sh_a_max_row is at least ` stages * 16 * thread_m_blocks `
776
+ int sh_a_max_row =
777
+ ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2);
778
779
// Register storage for double buffer of shared memory reads.
780
FragA frag_a[2][thread_m_blocks];
@@ -905,15 +795,14 @@ __global__ void Marlin(
795
796
int sh_first_group_id = -1;
797
int sh_num_groups = -1;
908
- constexpr int sh_max_num_groups = 32;
798
799
auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
800
int last_group_id) {
801
sh_first_group_id = first_group_id;
802
sh_num_groups = last_group_id - first_group_id + 1;
803
915
- if (sh_num_groups < sh_max_num_groups) {
916
- sh_num_groups = sh_max_num_groups;
804
+ if (sh_num_groups < act_s_max_num_groups) {
805
+ sh_num_groups = act_s_max_num_groups;
806
}
807
808
if (sh_first_group_id + sh_num_groups > num_groups) {
@@ -940,27 +829,31 @@ __global__ void Marlin(
829
}
830
}
831
};
832
+
833
// Asynchronously fetch the next A, B and s tile from global to the next
834
// shared memory pipeline location.
945
- int a_remaining_load_count_in_slice = stages;
946
- auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
835
+ bool should_load_a = true;
836
+ int max_num_stage_groups =
837
+ ((sh_a_max_row - moe_block_size) / moe_block_size + 1) / stages;
838
+ max_num_stage_groups = max(max_num_stage_groups, 1);
839
+ auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true,
840
+ int pipe_a = 0) {
841
if (pred) {
948
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
949
- if (prob_k > thread_k_blocks * 16 * stages || slice_col == 0 ||
950
- a_remaining_load_count_in_slice > 0) {
951
- a_remaining_load_count_in_slice--;
842
+ if (should_load_a) {
843
+ int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
844
#pragma unroll
845
for (int i = 0; i < a_sh_wr_iters; i++) {
954
- int a_idx = a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off;
955
- int row = a_idx / a_gl_stride;
846
+ int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
847
int64_t sorted_row = 0;
848
if (!m_block_size_8 || row < 8)
958
- sorted_row = sh_block_sorted_ids[row] / top_k;
959
- int64_t true_idx = sorted_row * a_gl_stride + a_idx % a_gl_stride;
849
+ sorted_row = sh_rd_block_sorted_ids[row];
850
+ int64_t true_idx =
851
+ sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
852
cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
853
row < block_num_valid_tokens);
854
}
855
}
856
+
857
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
858
#pragma unroll
859
for (int i = 0; i < b_sh_wr_iters; i++) {
@@ -1063,8 +956,8 @@ __global__ void Marlin(
956
957
// Load the next sub-tile from the current location in the shared memory pipe
958
// into the current register buffer.
1066
- auto fetch_to_registers = [&](int k, int pipe) {
1067
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
959
+ auto fetch_to_registers = [&](int k, int pipe, int pipe_a = 0) {
960
+ int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
961
#pragma unroll
962
for (int i = 0; i < thread_m_blocks; i++)
963
ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
@@ -1109,12 +1002,17 @@ __global__ void Marlin(
1002
}
1003
} else if constexpr (group_blocks != -1) {
1004
if constexpr (group_blocks >= thread_k_blocks) {
1112
- int4* sh_s_stage =
1113
- sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
1114
- (pipe / (group_blocks / thread_k_blocks)));
1115
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1005
+ if (k % b_sh_wr_iters == 0) {
1006
+ int4* sh_s_stage =
1007
+ sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
1008
+ (pipe / (group_blocks / thread_k_blocks)));
1009
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1010
+ } else {
1011
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1012
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1013
+ }
1014
} else {
1117
- int warp_id = threadIdx.x / 32;
1015
+ auto warp_id = threadIdx.x / 32;
1016
int n_warps = thread_n_blocks / 4;
1017
1018
int warp_row = warp_id / n_warps;
@@ -1152,7 +1050,7 @@ __global__ void Marlin(
1050
1051
// Determine "position" inside the thread-block (based on warp and
1052
// thread-id)
1155
- int warp_id = threadIdx.x / 32;
1053
+ auto warp_id = threadIdx.x / 32;
1054
int n_warps =
1055
thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
1056
@@ -1161,7 +1059,7 @@ __global__ void Marlin(
1059
1060
cur_k += warp_row * 16;
1061
1164
- int th_id = threadIdx.x % 32;
1062
+ auto th_id = threadIdx.x % 32;
1063
cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1064
1065
int s_col_shift =
@@ -1222,15 +1120,18 @@ __global__ void Marlin(
1120
}
1121
1122
} else if constexpr (group_blocks >= thread_k_blocks) {
1225
- int4* sh_zp_stage =
1226
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1227
- (pipe / (group_blocks / thread_k_blocks)));
1228
- for (int i = 0; i < num_ints_per_thread; i++) {
1229
- frag_qzp[k % 2][i] =
1230
- (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1123
+ if (k % b_sh_wr_iters == 0) {
1124
+ int4* sh_zp_stage =
1125
+ sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1126
+ (pipe / (group_blocks / thread_k_blocks)));
1127
+ #pragma unroll
1128
+ for (int i = 0; i < num_ints_per_thread; i++) {
1129
+ frag_qzp[k % 2][i] =
1130
+ (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1131
+ }
1132
}
1133
} else {
1233
- int warp_id = threadIdx.x / 32;
1134
+ auto warp_id = threadIdx.x / 32;
1135
int n_warps = thread_n_blocks / 4;
1136
1137
int warp_row = warp_id / n_warps;
@@ -1251,6 +1152,7 @@ __global__ void Marlin(
1152
1153
sh_zp_stage += cur_group_id * zp_sh_stride;
1154
1155
+ #pragma unroll
1156
for (int i = 0; i < num_ints_per_thread; i++) {
1157
frag_qzp[k % 2][i] =
1158
(reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
@@ -1263,12 +1165,16 @@ __global__ void Marlin(
1165
1166
if constexpr (group_blocks != -1) {
1167
if constexpr (group_blocks >= thread_k_blocks) {
1266
- int4* sh_zp_stage =
1267
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1268
- (pipe / (group_blocks / thread_k_blocks)));
1269
- reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd];
1168
+ if (k % b_sh_wr_iters == 0) {
1169
+ int4* sh_zp_stage =
1170
+ sh_zp +
1171
+ zp_sh_stage * ((group_blocks / thread_k_blocks) *
1172
+ (pipe / (group_blocks / thread_k_blocks)));
1173
+ reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1174
+ sh_zp_stage[zp_sh_rd];
1175
+ }
1176
} else {
1271
- int warp_id = threadIdx.x / 32;
1177
+ auto warp_id = threadIdx.x / 32;
1178
int n_warps = thread_n_blocks / 4;
1179
1180
int warp_row = warp_id / n_warps;
@@ -1292,6 +1198,25 @@ __global__ void Marlin(
1198
}
1199
};
1200
1201
+ auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1202
+ if constexpr (has_zp && is_zp_float || !has_zp) {
1203
+ dequant<scalar_t2, w_type_id>(q, frag_b_ptr);
1204
+ } else {
1205
+ static_assert(has_zp && !is_zp_float);
1206
+ static_assert(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU8.id());
1207
+ // If (has_zp && !is_zp_float),
1208
+ // we use not-zp version `dequant` function
1209
+ // to improve numerical accuracy.
1210
+ // Since both weight and zero point are dequanted using this logic,
1211
+ // the final dequanted weight would be correct.
1212
+ if constexpr (w_type_id == vllm::kU4.id()) {
1213
+ dequant<scalar_t2, vllm::kU4B8.id()>(q, frag_b_ptr);
1214
+ } else if constexpr (w_type_id == vllm::kU8.id()) {
1215
+ dequant<scalar_t2, vllm::kU8B128.id()>(q, frag_b_ptr);
1216
+ }
1217
+ }
1218
+ };
1219
+
1220
// Execute the actual tensor core matmul of a sub-tile.
1221
bool is_first_matmul_in_slice = true;
1222
auto matmul = [&](int k) {
@@ -1315,15 +1240,17 @@ __global__ void Marlin(
1240
zp_quant_1 = frag_qzp[k2][1];
1241
}
1242
1318
- dequant<scalar_t, w_type.size_bits()>(zp_quant_0, frag_zp_0);
1319
- dequant<scalar_t, w_type.size_bits()>(zp_quant_1, frag_zp_1);
1320
-
1321
- frag_zp[0] = frag_zp_0[0];
1322
- frag_zp[1] = frag_zp_0[1];
1323
- frag_zp[2] = frag_zp_1[0];
1324
- frag_zp[3] = frag_zp_1[1];
1243
+ dequant_data(zp_quant_0, reinterpret_cast<scalar_t2*>(&frag_zp));
1244
+ dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1245
}
1246
}
1247
+ if constexpr (has_zp && is_zp_float) {
1248
+ if (is_new_zp) {
1249
+ reinterpret_cast<int4*>(&frag_zp)[0] =
1250
+ reinterpret_cast<int4*>(&frag_zpf[k2])[0];
1251
+ }
1252
+ }
1253
+
1254
// We have the m dimension as the inner loop in order to encourage overlapping
1255
// dequantization and matmul operations.
1256
#pragma unroll
@@ -1342,8 +1269,8 @@ __global__ void Marlin(
1269
b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1270
}
1271
1345
- dequant<scalar_t, w_type.size_bits()>(b_quant_0, frag_b0);
1346
- dequant<scalar_t, w_type.size_bits()>(b_quant_1, frag_b1);
1272
+ dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1273
+ dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1274
1275
// Apply scale to frag_b0
1276
if constexpr (has_act_order) {
@@ -1351,8 +1278,7 @@ __global__ void Marlin(
1278
scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1279
act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1280
scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1354
- act_frag_s[k][2][j], act_frag_s[k2][3][j], 1);
1355
-
1281
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1282
} else if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1283
int idx = (threadIdx.x / 4) % 2;
1284
scalar_t2 s2 = Dtype::nums2num2(
@@ -1361,18 +1287,12 @@ __global__ void Marlin(
1287
if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1288
scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1289
scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1364
- } else if constexpr (has_zp && !is_zp_float && group_blocks != -1) {
1290
+ } else if constexpr (has_zp && group_blocks != -1) {
1291
if (is_new_zp)
1292
frag_zp[j] = __hmul2(frag_zp[j],
1293
*reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1368
- scale_and_sub<scalar_t>(frag_b0, frag_s[k % 2][j][0].x, frag_zp[j].x);
1369
- scale_and_sub<scalar_t>(frag_b1, frag_s[k % 2][j][0].y, frag_zp[j].y);
1370
- } else if constexpr (has_zp && is_zp_float && group_blocks != -1) {
1371
- if (is_new_zp)
1372
- frag_zpf[k2][j] = __hmul2(
1373
- frag_zpf[k2][j], *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1374
- scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j].x, frag_zpf[k2][j].x);
1375
- scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j].y, frag_zpf[k2][j].y);
1294
+ scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1295
+ scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1296
} else if constexpr (group_blocks != -1) {
1297
scale<scalar_t>(frag_b0, frag_s[k2][j], 0);
1298
scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
@@ -1397,7 +1317,7 @@ __global__ void Marlin(
1317
auto thread_block_reduce = [&]() {
1318
constexpr int red_off = threads / b_sh_stride_threads / 2;
1319
if (red_off >= 1) {
1400
- int red_idx = threadIdx.x / b_sh_stride_threads;
1320
+ auto red_idx = threadIdx.x / b_sh_stride_threads;
1321
constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1322
constexpr int red_sh_delta = b_sh_stride_threads;
1323
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
@@ -1731,7 +1651,7 @@ __global__ void Marlin(
1651
fetch_col_scale_to_shared();
1652
}
1653
}
1734
- fetch_to_shared(i, i, i < slice_iters);
1654
+ fetch_to_shared(i, i, i < slice_iters, i);
1655
}
1656
1657
zero_accums();
@@ -1740,8 +1660,10 @@ __global__ void Marlin(
1660
fetch_to_registers(0, 0);
1661
fetch_scales_to_registers(0, 0);
1662
fetch_zp_to_registers(0, 0);
1743
- a_gl_rd += a_gl_rd_delta_o * (stages - 1);
1744
- slice_k_start_shared_fetch += tb_k * (stages - 1);
1663
+ a_gl_rd_col += a_gl_rd_delta_o * (stages - 1);
1664
+ if constexpr (has_act_order) {
1665
+ slice_k_start_shared_fetch += tb_k * (stages - 1);
1666
+ }
1667
};
1668
if (slice_iters) {
1669
start_pipes();
@@ -1754,43 +1676,56 @@ __global__ void Marlin(
1676
// have even length meaning that the next iteration will always start at
1677
// index 0.
1678
1679
+ for (int stage_group_id = 0; stage_group_id < max_num_stage_groups;
1680
+ stage_group_id++) {
1681
#pragma unroll
1758
- for (int pipe = 0; pipe < stages;) {
1682
+ for (int pipe = 0; pipe < stages;) {
1683
#pragma unroll
1760
- for (int k = 0; k < b_sh_wr_iters; k++) {
1761
- fetch_to_registers(k + 1, pipe % stages);
1762
- fetch_scales_to_registers(k + 1, pipe);
1763
- fetch_zp_to_registers(k + 1, pipe);
1764
- if (k == b_sh_wr_iters - 2) {
1765
- fetch_to_shared((pipe + stages - 1) % stages, pipe,
1766
- slice_iters >= stages);
1767
- pipe++;
1768
- wait_for_stage();
1769
- init_same_group(pipe % stages);
1684
+ for (int k = 0; k < b_sh_wr_iters; k++) {
1685
+ int idx =
1686
+ (pipe >= stages && stage_group_id == max_num_stage_groups - 1)
1687
+ ? (pipe - stages)
1688
+ : (pipe + stage_group_id * stages);
1689
+ fetch_to_registers(k + 1, pipe % stages, idx);
1690
+ fetch_scales_to_registers(k + 1, pipe);
1691
+ fetch_zp_to_registers(k + 1, pipe);
1692
+ if (k == b_sh_wr_iters - 2) {
1693
+ int idx = (pipe >= 1 && stage_group_id == max_num_stage_groups - 1)
1694
+ ? (pipe - 1)
1695
+ : (pipe + (stage_group_id + 1) * stages - 1);
1696
+ fetch_to_shared((pipe + stages - 1) % stages, pipe,
1697
+ slice_iters >= stages, idx);
1698
+ pipe++;
1699
+ wait_for_stage();
1700
+ init_same_group(pipe % stages);
1701
+ }
1702
+ matmul(k);
1703
+ }
1704
+ slice_iters--;
1705
+ if (slice_iters == 0) {
1706
+ break;
1707
}
1771
- matmul(k);
1772
- }
1773
- slice_iters--;
1774
- if (slice_iters == 0) {
1775
- break;
1708
}
1777
- }
1778
- a_remaining_load_count_in_slice = 0;
1709
1780
- a_gl_rd += a_gl_rd_delta_o * stages;
1781
- slice_k_start += tb_k * stages;
1782
- slice_k_start_shared_fetch += tb_k * stages;
1710
+ a_gl_rd_col += a_gl_rd_delta_o * stages;
1711
1784
- if constexpr (has_act_order) {
1785
- int first_group_id = g_idx[slice_k_start];
1786
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1787
- if (last_g_idx >= prob_k) {
1788
- last_g_idx = prob_k - 1;
1712
+ if constexpr (has_act_order) {
1713
+ slice_k_start += tb_k * stages;
1714
+ slice_k_start_shared_fetch += tb_k * stages;
1715
+ int first_group_id = g_idx[slice_k_start];
1716
+ int last_g_idx = slice_k_start + stages * tb_k * 2;
1717
+ if (last_g_idx >= prob_k) {
1718
+ last_g_idx = prob_k - 1;
1719
+ }
1720
+ int last_group_id = g_idx[last_g_idx];
1721
+ if (last_group_id >= sh_first_group_id + sh_num_groups) {
1722
+ fetch_act_order_scales_to_shared(false, first_group_id,
1723
+ last_group_id);
1724
+ __syncthreads();
1725
+ }
1726
}
1790
- int last_group_id = g_idx[last_g_idx];
1791
- if (last_group_id >= sh_first_group_id + sh_num_groups) {
1792
- fetch_act_order_scales_to_shared(false, first_group_id, last_group_id);
1793
- __syncthreads();
1727
+ if (slice_iters == 0) {
1728
+ break;
1729
}
1730
}
1731
@@ -1877,15 +1812,30 @@ __global__ void Marlin(
1812
if (last || use_atomic_add)
1813
// only the last block in a slice actually writes the result
1814
write_result();
1880
- if (slice_row) a_remaining_load_count_in_slice = stages;
1815
+ int old_slice_row = slice_row;
1816
slice_row = 0;
1817
slice_col_par++;
1818
slice_col++;
1819
is_first_matmul_in_slice = true;
1820
init_slice();
1821
+
1822
+ // Should we load A matrix in next slice?
1823
+ // `slice_col == 0`: when move to a new moe block
1824
+ // `old_slice_row > 0`:
1825
+ // when the last slice is not starting from k_index == 0
1826
+ // (only happen when it is the first slice of a threadblock)
1827
+ // `prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups`:
1828
+ // when the required shared memory size is larger than
1829
+ // the remaining shared memory
1830
+ if (slice_col == 0 || old_slice_row ||
1831
+ prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups) {
1832
+ should_load_a = true;
1833
+ } else {
1834
+ should_load_a = false;
1835
+ }
1836
+
1837
if (slice_iters) {
1887
- a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
1888
- (threadIdx.x % a_gl_rd_delta_o);
1838
+ a_gl_rd_col = (threadIdx.x % a_gl_rd_delta_o);
1839
#pragma unroll
1840
for (int i = 0; i < b_sh_wr_iters; i++)
1841
B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
@@ -1900,12 +1850,10 @@ __global__ void Marlin(
1850
slice_k_finish = slice_k_start + tb_k * slice_iters;
1851
slice_k_start_shared_fetch = slice_k_start;
1852
slice_n_offset = act_s_col_tb_stride * slice_col;
1903
-
1853
} else {
1854
s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1855
zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1856
}
1908
-
1857
start_pipes();
1858
}
1859
}
csrc/moe/marlin_moe_wna16/ops.cu
+155
-190
@@ -116,7 +116,7 @@ __global__ void permute_cols_kernel(
116
int base_k = 0;
117
118
for (int i = 0; i < iters; i++) {
119
- int cur_k = base_k + threadIdx.x;
119
+ auto cur_k = base_k + threadIdx.x;
120
int src_pos = perm_int_ptr[cur_k];
121
122
out_half[cur_k] = a_row_half[src_pos];
@@ -126,7 +126,7 @@ __global__ void permute_cols_kernel(
126
127
if (rest) {
128
if (threadIdx.x < rest) {
129
- int cur_k = base_k + threadIdx.x;
129
+ auto cur_k = base_k + threadIdx.x;
130
int src_pos = perm_int_ptr[cur_k];
131
132
out_half[cur_k] = a_row_half[src_pos];
@@ -195,7 +195,6 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
195
tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K
196
load_groups = max(load_groups, 32); // We load at least 32 scale groups
197
return load_groups * tb_n * 2;
198
-
198
} else {
199
int tb_scales = tb_groups * tb_n * 2;
200
@@ -203,22 +202,24 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
202
}
203
}
204
206
-int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
207
- int prob_m, int prob_n, int prob_k, int num_bits,
208
- int group_size, bool has_act_order, bool is_k_full,
209
- int has_zp, int is_zp_float) {
205
+int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
206
+ int thread_m_blocks, int prob_m, int prob_n,
207
+ int prob_k, int num_bits, int group_size,
208
+ bool has_act_order, bool is_k_full, int has_zp,
209
+ int is_zp_float) {
210
int pack_factor = 32 / num_bits;
211
212
// Get B size
213
int tb_k = th_config.thread_k;
214
int tb_n = th_config.thread_n;
215
- int tb_m = thread_m_blocks * 16;
215
+ int tb_m = thread_m_blocks * (m_block_size_8 ? 8 : 16);
216
217
- // shm size for block_sorted_ids/block_topk_weights
217
+ // shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
218
// both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
219
- int sh_block_meta_size = tb_m * 4 * 2;
219
+ int sh_block_meta_size = tb_m * 4;
220
int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
221
int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
222
+ int sh_red_size = tb_m * (tb_n + 8) * 2;
223
int sh_s_size =
224
get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
225
group_size, has_act_order, is_k_full);
@@ -233,16 +234,17 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
234
sh_zp_size = sh_s_size / 2;
235
}
236
236
- int total_size = sh_a_size + sh_b_size + sh_s_size + sh_zp_size +
237
- sh_g_idx_size + sh_block_meta_size;
237
+ int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size +
238
+ sh_zp_size + sh_g_idx_size + sh_block_meta_size;
239
240
return total_size;
241
}
242
242
-bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
243
- int prob_m, int prob_n, int prob_k, int num_bits,
244
- int group_size, bool has_act_order, bool is_k_full,
245
- int has_zp, int is_zp_float, int max_shared_mem) {
243
+bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
244
+ int thread_m_blocks, int prob_m, int prob_n, int prob_k,
245
+ int num_bits, int group_size, bool has_act_order,
246
+ bool is_k_full, int has_zp, int is_zp_float,
247
+ int max_shared_mem) {
248
// Sanity
249
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
250
th_config.num_threads == -1) {
@@ -266,143 +268,113 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
268
269
// Check that pipeline fits into cache
270
int cache_size = get_kernel_cache_size(
269
- th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
270
- has_act_order, is_k_full, has_zp, is_zp_float);
271
+ th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
272
+ num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
273
return cache_size <= max_shared_mem;
274
}
275
274
- #define __GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
275
- M_BLOCK_SIZE_8, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, \
276
- NUM_THREADS, IS_ZP_FLOAT) \
277
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
278
- thread_n_blocks == THREAD_N_BLOCKS && \
279
- thread_k_blocks == THREAD_K_BLOCKS && \
280
- m_block_size_8 == M_BLOCK_SIZE_8 && \
281
- has_act_order == HAS_ACT_ORDER && has_zp == HAS_ZP && \
282
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
283
- is_zp_float == IS_ZP_FLOAT) { \
284
- kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
285
- THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
286
- pipe_stages, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, \
287
- IS_ZP_FLOAT>; \
276
+ #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
277
+ M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
278
+ else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
279
+ thread_n_blocks == THREAD_N_BLOCKS && \
280
+ thread_k_blocks == THREAD_K_BLOCKS && \
281
+ m_block_size_8 == M_BLOCK_SIZE_8 && \
282
+ group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
283
+ is_zp_float == IS_ZP_FLOAT) { \
284
+ kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
285
+ THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
286
+ pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
287
}
288
290
- #define GPTQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
291
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, true, false, 0, NUM_THREADS, \
292
- false) \
293
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
294
- NUM_THREADS, false) \
295
- \
296
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, -1, \
297
- NUM_THREADS, false) \
298
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 2, \
299
- NUM_THREADS, false) \
300
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 4, \
301
- NUM_THREADS, false) \
302
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 8, \
303
- NUM_THREADS, false) \
304
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
305
- NUM_THREADS, false) \
306
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
307
- NUM_THREADS, false) \
308
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
309
- NUM_THREADS, false) \
310
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
311
- NUM_THREADS, false)
312
-
313
- #define GPTQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
314
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
315
- NUM_THREADS, false) \
316
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
317
- NUM_THREADS, false) \
318
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
319
- NUM_THREADS, false) \
320
- \
321
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
322
- NUM_THREADS, false) \
323
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
324
- NUM_THREADS, false) \
325
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
326
- NUM_THREADS, false) \
327
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
328
- NUM_THREADS, false) \
329
- \
330
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
331
- NUM_THREADS, false) \
332
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
333
- NUM_THREADS, false) \
334
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
335
- NUM_THREADS, false) \
336
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
337
- NUM_THREADS, false) \
338
- \
339
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
340
- NUM_THREADS, false) \
341
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
342
- NUM_THREADS, false) \
343
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
344
- NUM_THREADS, false) \
345
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
346
- NUM_THREADS, false)
347
-
348
- #define AWQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
349
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, -1, \
350
- NUM_THREADS, false) \
351
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 2, NUM_THREADS, \
352
- false) \
353
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, \
354
- false) \
355
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 8, NUM_THREADS, \
356
- false) \
357
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
358
- NUM_THREADS, false) \
359
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
360
- NUM_THREADS, false) \
361
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
362
- NUM_THREADS, false) \
363
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
364
- NUM_THREADS, false)
365
-
366
- #define AWQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
367
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
368
- NUM_THREADS, false) \
369
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
370
- NUM_THREADS, false) \
371
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
372
- NUM_THREADS, false) \
373
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
374
- NUM_THREADS, false) \
375
- \
376
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
377
- NUM_THREADS, false) \
378
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
379
- NUM_THREADS, false) \
380
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
381
- NUM_THREADS, false) \
382
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
383
- NUM_THREADS, false) \
384
- \
385
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
386
- NUM_THREADS, false) \
387
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
388
- NUM_THREADS, false) \
389
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
390
- NUM_THREADS, false) \
391
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
392
- NUM_THREADS, false)
289
+ // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
290
+ // this is the most common cases
291
+ // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
292
+ // FZP: cases for float-zero-point (is_zp_float = true)
293
+ // ACT: cases for act order case (group_blocks == 0)
294
+ #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
295
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
296
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
297
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \
298
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
299
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
300
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
301
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
302
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
303
+
304
+ #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
305
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
306
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
307
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
308
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
309
+ \
310
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
311
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
312
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
313
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
314
+ \
315
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
316
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
317
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
318
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
319
+
320
+ #define COMMON_GET_IF(W_TYPE) \
321
+ COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \
322
+ COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \
323
+ COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
324
+ COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)
325
+
326
+ #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
327
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
328
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
329
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
330
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
331
+
332
+ #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
333
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
334
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
335
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
336
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
337
+ \
338
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
339
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
340
+
341
+ #define BIGGROUP_GET_IF(W_TYPE) \
342
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
343
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
344
+ BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
345
+ BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
346
347
// We currently have 4-bit models only with group_blocks == 4
395
- #define HQQ_GET_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
396
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, \
397
- true) \
398
- __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
399
- NUM_THREADS, true) \
400
- __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
401
- NUM_THREADS, true) \
402
- __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
403
- NUM_THREADS, true) \
404
- __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
405
- NUM_THREADS, true)
348
+ #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
349
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
350
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
351
+
352
+ #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
353
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
354
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
355
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
356
+
357
+ #define FZP_GET_IF(W_TYPE) \
358
+ FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \
359
+ FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \
360
+ FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
361
+ FZP_GET_IF_M234(W_TYPE, 8, 4, 128)
362
+
363
+ // We currently have 4-bit models only with group_blocks == 4
364
+ #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
365
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
366
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
367
+
368
+ #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
369
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
370
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
371
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
372
+
373
+ #define ACT_GET_IF(W_TYPE) \
374
+ ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \
375
+ ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \
376
+ ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
377
+ ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
378
379
template <typename scalar_t>
380
MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
@@ -415,23 +387,15 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
387
auto kernel = MarlinDefault;
388
if (false) {
389
}
418
- GPTQ_GET_IF_M1(vllm::kU4B8, 8, 8, 256)
419
- GPTQ_GET_IF_M1(vllm::kU4B8, 8, 4, 128)
390
421
- GPTQ_GET_IF_M234(vllm::kU4B8, 16, 4, 256)
422
- GPTQ_GET_IF_M234(vllm::kU4B8, 8, 4, 128)
391
+ COMMON_GET_IF(vllm::kU4)
392
+ COMMON_GET_IF(vllm::kU4B8)
393
+ COMMON_GET_IF(vllm::kU8B128)
394
424
- GPTQ_GET_IF_M1(vllm::kU8B128, 8, 8, 256)
425
- GPTQ_GET_IF_M1(vllm::kU8B128, 8, 4, 128)
395
+ BIGGROUP_GET_IF(vllm::kFE4M3fn)
396
427
- GPTQ_GET_IF_M234(vllm::kU8B128, 16, 4, 256)
428
- GPTQ_GET_IF_M234(vllm::kU8B128, 8, 4, 128)
429
-
430
- AWQ_GET_IF_M1(vllm::kU4, 8, 8, 256)
431
- AWQ_GET_IF_M1(vllm::kU4, 8, 4, 128)
432
-
433
- AWQ_GET_IF_M234(vllm::kU4, 16, 4, 256)
434
- AWQ_GET_IF_M234(vllm::kU4, 8, 4, 128)
397
+ ACT_GET_IF(vllm::kU4B8)
398
+ ACT_GET_IF(vllm::kU8B128)
399
400
return kernel;
401
}
@@ -457,19 +421,19 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
421
for (int i = 0; i < thread_configs_size; i++) {
422
thread_config_t th_config = thread_configs[i];
423
460
- if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
461
- num_bits, group_size, has_act_order, is_k_full, has_zp,
462
- is_zp_float, max_shared_mem)) {
424
+ if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
425
+ prob_n, prob_k, num_bits, group_size, has_act_order,
426
+ is_k_full, has_zp, is_zp_float, max_shared_mem)) {
427
continue;
428
}
429
430
int cache_size = get_kernel_cache_size(
467
- th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits,
468
- group_size, has_act_order, is_k_full, has_zp, is_zp_float);
431
+ th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
432
+ num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
433
434
int group_blocks = 0;
435
if (!has_act_order) {
472
- group_blocks = group_size == -1 ? -1 : group_size / 16;
436
+ group_blocks = group_size == -1 ? -1 : (group_size / 16);
437
}
438
439
auto kernel = get_marlin_kernel<scalar_t>(
@@ -515,14 +479,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
479
bool m_block_size_8 = moe_block_size == 8;
480
481
if (has_zp) {
518
- TORCH_CHECK(
519
- q_type == vllm::kU4 || q_type == vllm::kU8,
520
- "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
482
+ TORCH_CHECK(q_type == vllm::kU4,
483
+ "q_type must be u4 when has_zp = True. Got = ", q_type.str());
484
} else {
522
- TORCH_CHECK(
523
- q_type == vllm::kU4B8 || q_type == vllm::kU8B128,
524
- "q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
525
- q_type.str());
485
+ TORCH_CHECK(q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
486
+ q_type == vllm::kFE4M3fn,
487
+ "q_type must be uint4b8, uint8b128 or fp8e4m3 when has_zp = "
488
+ "False. Got = ",
489
+ q_type.str());
490
}
491
492
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
@@ -631,18 +595,18 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
595
int thread_k_blocks = thread_k / 16;
596
int thread_n_blocks = thread_n / 16;
597
634
- TORCH_CHECK(is_valid_config(thread_tfg, thread_m_blocks, prob_m, prob_n,
635
- prob_k, num_bits, group_size, has_act_order,
636
- is_k_full, has_zp, is_zp_float, max_shared_mem),
637
- "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
638
- ", thread_k = ", thread_tfg.thread_k,
639
- ", thread_n = ", thread_tfg.thread_n,
640
- ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
641
- prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
642
- ", group_size = ", group_size,
643
- ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
644
- ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
645
- ", max_shared_mem = ", max_shared_mem);
598
+ TORCH_CHECK(
599
+ is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
600
+ prob_n, prob_k, num_bits, group_size, has_act_order,
601
+ is_k_full, has_zp, is_zp_float, max_shared_mem),
602
+ "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
603
+ ", thread_k = ", thread_tfg.thread_k,
604
+ ", thread_n = ", thread_tfg.thread_n,
605
+ ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
606
+ prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
607
+ ", group_size = ", group_size, ", has_act_order = ", has_act_order,
608
+ ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
609
+ ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
610
611
auto kernel = get_marlin_kernel<scalar_t>(
612
q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8,
@@ -666,7 +630,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
630
A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr,
631
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
632
topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
669
- prob_n, prob_k, locks, use_atomic_add, use_fp32_reduce);
633
+ prob_n, prob_k, locks, use_atomic_add, use_fp32_reduce, max_shared_mem);
634
// clang-format on
635
}
636
@@ -841,10 +805,11 @@ torch::Tensor moe_wna16_marlin_gemm(
805
b_q_type == vllm::kU4,
806
"b_q_type must be u4 when has_zp = True. Got = ", b_q_type.str());
807
} else {
844
- TORCH_CHECK(
845
- b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128,
846
- "b_q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
847
- b_q_type.str());
808
+ TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
809
+ b_q_type == vllm::kFE4M3fn,
810
+ "b_q_type must be uint4b8, uint8b128 or fp8e4m3 when has_zp = "
811
+ "False. Got = ",
812
+ b_q_type.str());
813
}
814
815
if (has_zp && is_zp_float) {
csrc/quantization/gptq_marlin/.gitignore
new
+1
@@ -0,0 +1 @@
1
+kernel_*.cu
\ No newline at end of file
csrc/quantization/gptq_marlin/dequant.h
new
+291
@@ -0,0 +1,291 @@
1
+
2
+#include "marlin_dtypes.cuh"
3
+
4
+namespace MARLIN_NAMESPACE_NAME {
5
+
6
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
7
+// Lookup-table based 3-input logical operation; explicitly used for
8
+// dequantization as the compiler does not seem to automatically recognize it in
9
+// all cases.
10
+template <int lut>
11
+__device__ inline int lop3(int a, int b, int c) {
12
+ int res;
13
+ asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
14
+ : "=r"(res)
15
+ : "r"(a), "r"(b), "r"(c), "n"(lut));
16
+ return res;
17
+}
18
+
19
+// Constructs destination register by taking bytes from 2 sources (based on
20
+// mask)
21
+template <int start_byte, int mask>
22
+__device__ inline uint32_t prmt(uint32_t a) {
23
+ uint32_t res;
24
+ asm volatile("prmt.b32 %0, %1, %2, %3;\n"
25
+ : "=r"(res)
26
+ : "r"(a), "n"(start_byte), "n"(mask));
27
+ return res;
28
+}
29
+
30
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id>
31
+__device__ inline void dequant(int q, scalar_t2* frag_b);
32
+
33
+//
34
+// Efficiently dequantize 4bit values packed in an int32 value into a full
35
+// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
36
+// with some small changes:
37
+// - FP16:
38
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
39
+// - BF16:
40
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
41
+//
42
+template <>
43
+__device__ inline void dequant<half2, vllm::kU4B8.id()>(int q, half2* frag_b) {
44
+ const int LO = 0x000f000f;
45
+ const int HI = 0x00f000f0;
46
+ const int EX = 0x64006400;
47
+ // Guarantee that the `(a & b) | c` operations are LOP3s.
48
+ // clang-format off
49
+ int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
50
+ int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
51
+ // clang-format on
52
+ // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
53
+ // directly into `SUB` and `ADD`.
54
+ const int SUB = 0x64086408;
55
+ const int MUL = 0x2c002c00;
56
+ const int ADD = 0xd480d480;
57
+ frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
58
+ *reinterpret_cast<const half2*>(&SUB));
59
+ frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
60
+ *reinterpret_cast<const half2*>(&MUL),
61
+ *reinterpret_cast<const half2*>(&ADD));
62
+}
63
+
64
+template <>
65
+__device__ inline void dequant<half2, vllm::kU4.id()>(int q, half2* frag_b) {
66
+ const int LO = 0x000f000f;
67
+ const int HI = 0x00f000f0;
68
+ const int EX = 0x64006400;
69
+ // Guarantee that the `(a & b) | c` operations are LOP3s.
70
+ // clang-format off
71
+ int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
72
+ int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
73
+ // clang-format on
74
+ // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
75
+ // directly into `SUB` and `ADD`.
76
+ const int SUB = 0x64006400;
77
+ const int MUL = 0x2c002c00;
78
+ const int ADD = 0xd400d400;
79
+ frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
80
+ *reinterpret_cast<const half2*>(&SUB));
81
+ frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
82
+ *reinterpret_cast<const half2*>(&MUL),
83
+ *reinterpret_cast<const half2*>(&ADD));
84
+}
85
+
86
+template <>
87
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id()>(
88
+ int q, nv_bfloat162* frag_b) {
89
+ static constexpr uint32_t MASK = 0x000f000f;
90
+ static constexpr uint32_t EX = 0x43004300;
91
+
92
+ // Guarantee that the `(a & b) | c` operations are LOP3s.
93
+ // clang-format off
94
+ int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
95
+ q >>= 4;
96
+ int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
97
+ // clang-format on
98
+
99
+ static constexpr uint32_t MUL = 0x3F803F80;
100
+ static constexpr uint32_t ADD = 0xC308C308;
101
+
102
+ frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
103
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
104
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
105
+ frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
106
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
107
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
108
+}
109
+
110
+template <>
111
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id()>(
112
+ int q, nv_bfloat162* frag_b) {
113
+ static constexpr uint32_t MASK = 0x000f000f;
114
+ static constexpr uint32_t EX = 0x43004300;
115
+
116
+ // Guarantee that the `(a & b) | c` operations are LOP3s.
117
+ // clang-format off
118
+ int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
119
+ q >>= 4;
120
+ int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
121
+ // clang-format on
122
+
123
+ static constexpr uint32_t MUL = 0x3F803F80;
124
+ static constexpr uint32_t ADD = 0xC300C300;
125
+
126
+ frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
127
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
128
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
129
+ frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
130
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
131
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
132
+}
133
+
134
+//
135
+// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
136
+// bf16 Reference:
137
+// - FP16:
138
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
139
+// - BF16:
140
+// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
141
+//
142
+template <>
143
+__device__ inline void dequant<half2, vllm::kU8B128.id()>(int q,
144
+ half2* frag_b) {
145
+ static constexpr uint32_t mask_for_elt_01 = 0x5250;
146
+ static constexpr uint32_t mask_for_elt_23 = 0x5351;
147
+ static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
148
+
149
+ uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
150
+ uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
151
+
152
+ static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
153
+
154
+ frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
155
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
156
+ frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
157
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
158
+}
159
+
160
+template <>
161
+__device__ inline void dequant<half2, vllm::kU8.id()>(int q, half2* frag_b) {
162
+ static constexpr uint32_t mask_for_elt_01 = 0x5250;
163
+ static constexpr uint32_t mask_for_elt_23 = 0x5351;
164
+ static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
165
+
166
+ uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
167
+ uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
168
+
169
+ static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
170
+
171
+ frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
172
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
173
+ frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
174
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
175
+}
176
+
177
+template <>
178
+__device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id()>(
179
+ int q, nv_bfloat162* frag_b) {
180
+ float fp32_intermediates[4];
181
+ uint32_t* fp32_intermediates_casted =
182
+ reinterpret_cast<uint32_t*>(fp32_intermediates);
183
+
184
+ static constexpr uint32_t fp32_base = 0x4B000000;
185
+ fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
186
+ fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
187
+ fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
188
+ fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
189
+
190
+ fp32_intermediates[0] -= 8388736.f;
191
+ fp32_intermediates[1] -= 8388736.f;
192
+ fp32_intermediates[2] -= 8388736.f;
193
+ fp32_intermediates[3] -= 8388736.f;
194
+
195
+ uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
196
+ bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
197
+ fp32_intermediates_casted[1], 0x7632);
198
+ bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
199
+ fp32_intermediates_casted[3], 0x7632);
200
+}
201
+
202
+template <>
203
+__device__ inline void dequant<nv_bfloat162, vllm::kU8.id()>(
204
+ int q, nv_bfloat162* frag_b) {
205
+ float fp32_intermediates[4];
206
+ uint32_t* fp32_intermediates_casted =
207
+ reinterpret_cast<uint32_t*>(fp32_intermediates);
208
+
209
+ static constexpr uint32_t fp32_base = 0x4B000000;
210
+ fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
211
+ fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
212
+ fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
213
+ fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
214
+
215
+ fp32_intermediates[0] -= 8388608.f;
216
+ fp32_intermediates[1] -= 8388608.f;
217
+ fp32_intermediates[2] -= 8388608.f;
218
+ fp32_intermediates[3] -= 8388608.f;
219
+
220
+ uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(frag_b);
221
+ bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
222
+ fp32_intermediates_casted[1], 0x7632);
223
+ bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
224
+ fp32_intermediates_casted[3], 0x7632);
225
+}
226
+
227
+template <>
228
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id()>(int q,
229
+ half2* frag_b) {
230
+ // Constants for FP8 (E4M3) and FP16 formats
231
+ constexpr int FP8_EXPONENT = 4, FP8_MANTISSA = 3, FP16_EXPONENT = 5;
232
+ constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
233
+
234
+ // Calculate MASK for extracting mantissa and exponent
235
+ constexpr int MASK1 = 0x80000000;
236
+ constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
237
+ constexpr int MASK3 = MASK2 & 0x7fffffff;
238
+ constexpr int MASK = MASK3 | (MASK3 >> 16);
239
+ // Final MASK value: 0x7F007F00
240
+
241
+ // Extract and shift FP8 values to FP16 format
242
+ int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
243
+ int Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
244
+
245
+ // Construct and apply exponent bias
246
+ constexpr int BIAS_OFFSET =
247
+ (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
248
+ const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
249
+
250
+ // Convert to half2 and apply bias
251
+ // Note: reverse indexing is intentional because weights are permuted
252
+ frag_b[1] = __hmul2(*reinterpret_cast<const half2*>(&Out1), bias_reg);
253
+ frag_b[0] = __hmul2(*reinterpret_cast<const half2*>(&Out2), bias_reg);
254
+}
255
+
256
+template <>
257
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id()>(
258
+ int q, nv_bfloat162* frag_b) {
259
+ // Constants for FP8 (E4M3) and BF16 formats
260
+ constexpr int FP8_EXPONENT = 4, FP8_MANTISSA = 3, BF16_EXPONENT = 8;
261
+ constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
262
+
263
+ // Calculate MASK for extracting mantissa and exponent
264
+ constexpr int MASK1 = 0x80000000;
265
+ constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
266
+ constexpr int MASK3 = MASK2 & 0x7fffffff;
267
+ constexpr int MASK = MASK3 | (MASK3 >> 16);
268
+ // Final MASK value: 0x7F007F00
269
+
270
+ // Extract and shift FP8 values to BF16 format
271
+ int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
272
+ int Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
273
+
274
+ // Construct and apply exponent bias
275
+ constexpr int BIAS_OFFSET =
276
+ (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
277
+ // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
278
+ // position
279
+ constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
280
+ const nv_bfloat162 bias_reg =
281
+ __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
282
+
283
+ // Convert to bfloat162 and apply bias
284
+ // Note: reverse indexing is intentional because weights are permuted
285
+ frag_b[1] = __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out1), bias_reg);
286
+ frag_b[0] = __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out2), bias_reg);
287
+}
288
+
289
+#endif
290
+
291
+} // namespace MARLIN_NAMESPACE_NAME
csrc/quantization/gptq_marlin/generate_kernels.py
new
+116
@@ -0,0 +1,116 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import glob
3
+import itertools
4
+import os
5
+import subprocess
6
+
7
+import jinja2
8
+
9
+FILE_HEAD = """
10
+// auto generated by generate.py
11
+// clang-format off
12
+
13
+#include "kernel.h"
14
+#include "marlin_template.h"
15
+
16
+namespace MARLIN_NAMESPACE_NAME {
17
+""".strip()
18
+
19
+TEMPLATE = ("template __global__ void Marlin<"
20
+ "{{scalar_t}}, "
21
+ "{{w_type_id}}, "
22
+ "{{threads}}, "
23
+ "{{thread_m_blocks}}, "
24
+ "{{thread_n_blocks}}, "
25
+ "{{thread_k_blocks}}, "
26
+ "{{'true' if m_block_size_8 else 'false'}}, "
27
+ "{{stages}}, "
28
+ "{{group_blocks}}, "
29
+ "{{'true' if is_zp_float else 'false'}}>"
30
+ "( MARLIN_KERNEL_PARAMS );")
31
+
32
+# int8 with zero point case (vllm::kU8) is also supported,
33
+# we don't add it to reduce wheel size.
34
+SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn"]
35
+THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128),
36
+ (128, 64, 128)]
37
+
38
+THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
39
+# group_blocks:
40
+# = 0 : act order case
41
+# = -1 : channelwise quantization
42
+# > 0 : group_size=16*group_blocks
43
+GROUP_BLOCKS = [0, -1, 2, 4, 8]
44
+DTYPES = ["fp16", "bf16"]
45
+
46
+
47
+def remove_old_kernels():
48
+ for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
49
+ subprocess.call(["rm", "-f", filename])
50
+
51
+
52
+def generate_new_kernels():
53
+ for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
54
+ all_template_str_list = []
55
+
56
+ for group_blocks, m_blocks, thread_configs in itertools.product(
57
+ GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):
58
+
59
+ # act order case only support gptq-int4 and gptq-int8
60
+ if group_blocks == 0 and scalar_type not in [
61
+ "vllm::kU4B8", "vllm::kU8B128"
62
+ ]:
63
+ continue
64
+ if thread_configs[2] == 256:
65
+ # for small batch (m_blocks == 1), we only need (128, 128, 256)
66
+ # for large batch (m_blocks > 1), we only need (64, 256, 256)
67
+ if m_blocks <= 1 and thread_configs[0] != 128:
68
+ continue
69
+ if m_blocks > 1 and thread_configs[0] != 64:
70
+ continue
71
+
72
+ # we only support channelwise quantization and group_size == 128
73
+ # for fp8
74
+ if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
75
+ continue
76
+
77
+ k_blocks = thread_configs[0] // 16
78
+ n_blocks = thread_configs[1] // 16
79
+ threads = thread_configs[2]
80
+
81
+ c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
82
+
83
+ is_zp_float_list = [False]
84
+ if dtype == "fp16" and scalar_type == "vllm::kU4" and \
85
+ group_blocks == 4:
86
+ # HQQ (is_zp_float = true) only supports
87
+ # 4bit quantization and fp16
88
+ is_zp_float_list.append(True)
89
+
90
+ for is_zp_float in is_zp_float_list:
91
+ template_str = jinja2.Template(TEMPLATE).render(
92
+ scalar_t=c_dtype,
93
+ w_type_id=scalar_type + ".id()",
94
+ threads=threads,
95
+ thread_m_blocks=max(m_blocks, 1),
96
+ thread_n_blocks=n_blocks,
97
+ thread_k_blocks=k_blocks,
98
+ m_block_size_8=m_blocks == 0.5,
99
+ stages="pipe_stages",
100
+ group_blocks=group_blocks,
101
+ is_zp_float=is_zp_float,
102
+ )
103
+
104
+ all_template_str_list.append(template_str)
105
+
106
+ file_content = FILE_HEAD + "\n\n"
107
+ file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
108
+ filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"
109
+
110
+ with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
111
+ f.write(file_content)
112
+
113
+
114
+if __name__ == "__main__":
115
+ remove_old_kernels()
116
+ generate_new_kernels()
csrc/quantization/gptq_marlin/gptq_marlin.cu
+442
-2063
@@ -19,10 +19,11 @@
19
* Adapted from https://github.com/IST-DASLab/marlin
20
*/
21
22
-#include "marlin.cuh"
23
-#include "marlin_dtypes.cuh"
24
-#include "core/scalar_type.hpp"
22
+#ifndef MARLIN_NAMESPACE_NAME
23
+ #define MARLIN_NAMESPACE_NAME marlin
24
+#endif
25
26
+#include "kernel.h"
27
#include "core/registration.h"
28
29
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
@@ -30,13 +31,12 @@
31
std::is_same<scalar_t, nv_bfloat16>::value, \
32
"only float16 and bfloat16 is supported");
33
33
-template <typename T>
34
-inline std::string str(T x) {
35
- return std::to_string(x);
36
-}
37
-
34
namespace marlin {
35
36
+__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
37
+
38
+using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
39
+
40
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
41
42
__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
@@ -44,46 +44,17 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
44
int4* __restrict__ out_int4_ptr, int size_m,
45
int size_k, int lda, int block_rows) {}
46
47
-template <typename scalar_t, // compute dtype, half or nv_float16
48
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
49
- const int threads, // number of threads in a threadblock
50
- const int thread_m_blocks, // number of 16x16 blocks in the m
51
- // dimension (batchsize) of the
52
- // threadblock
53
- const int thread_n_blocks, // same for n dimension (output)
54
- const int thread_k_blocks, // same for k dimension (reduction)
55
- const int stages, // number of stages for the async global->shared
56
- // fetch pipeline
57
- const bool has_act_order, // whether act_order is enabled
58
- const int group_blocks = -1, // number of consecutive 16x16 blocks
59
- // with a separate quantization scale
60
- const bool is_zp_float // is zero point of float16 type?
61
- >
62
-__global__ void Marlin(
63
- const int4* __restrict__ A, // fp16 input matrix of shape mxk
64
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
65
- int4* __restrict__ C, // fp16 output buffer of shape mxn
66
- int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
67
- const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
68
- // (k/groupsize)xn
69
- const int* __restrict__ g_idx, // int32 group indices of shape k
70
- int num_groups, // number of scale groups per output channel
71
- int prob_m, // batch dimension m
72
- int prob_n, // output dimension n
73
- int prob_k, // reduction dimension k
74
- int* locks, // extra global storage for barrier synchronization
75
- bool use_fp32_reduce // whether to use fp32 global reduce
76
-) {}
77
-
47
} // namespace marlin
48
80
-torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
81
- torch::Tensor& b_scales, torch::Tensor& b_zeros,
82
- torch::Tensor& g_idx, torch::Tensor& perm,
83
- torch::Tensor& workspace,
84
- vllm::ScalarTypeId const b_q_type_id,
85
- int64_t size_m, int64_t size_n, int64_t size_k,
86
- bool is_k_full, bool has_zp, bool is_zp_float) {
49
+torch::Tensor gptq_marlin_gemm(
50
+ torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
51
+ torch::Tensor& b_q_weight, torch::Tensor& b_scales,
52
+ std::optional<torch::Tensor> const& b_zeros_or_none,
53
+ std::optional<torch::Tensor> const& g_idx_or_none,
54
+ std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
55
+ vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
56
+ int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
57
+ bool is_zp_float) {
58
TORCH_CHECK_NOT_IMPLEMENTED(false,
59
"marlin_gemm(..) requires CUDA_ARCH >= 8.0");
60
return torch::empty({1, 1});
@@ -91,369 +62,6 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
62
63
#else
64
94
-// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
95
-// output/accumulation.
96
-template <typename scalar_t>
97
-__device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
98
- const typename ScalarType<scalar_t>::FragB& frag_b,
99
- typename ScalarType<scalar_t>::FragC& frag_c) {
100
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
101
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
102
- float* c = reinterpret_cast<float*>(&frag_c);
103
- if constexpr (std::is_same<scalar_t, half>::value) {
104
- asm volatile(
105
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
106
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
107
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
108
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
109
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
110
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
111
- asm volatile(
112
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
113
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
114
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
115
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
116
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
117
- } else {
118
- STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
119
- }
120
-}
121
-
122
-// Instruction for loading a full 16x16 matrix fragment of operand A from shared
123
-// memory, directly in tensor core layout.
124
-template <typename scalar_t>
125
-__device__ inline void ldsm4(typename ScalarType<scalar_t>::FragA& frag_a,
126
- const void* smem_ptr) {
127
- uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
128
- uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
129
- asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
130
- : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
131
- : "r"(smem));
132
-}
133
-
134
-// Lookup-table based 3-input logical operation; explicitly used for
135
-// dequantization as the compiler does not seem to automatically recognize it in
136
-// all cases.
137
-template <int lut>
138
-__device__ inline int lop3(int a, int b, int c) {
139
- int res;
140
- asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
141
- : "=r"(res)
142
- : "r"(a), "r"(b), "r"(c), "n"(lut));
143
- return res;
144
-}
145
-
146
-// Constructs destination register by taking bytes from 2 sources (based on
147
-// mask)
148
-template <int start_byte, int mask>
149
-__device__ inline uint32_t prmt(uint32_t a) {
150
- uint32_t res;
151
- asm volatile("prmt.b32 %0, %1, %2, %3;\n"
152
- : "=r"(res)
153
- : "r"(a), "n"(start_byte), "n"(mask));
154
- return res;
155
-}
156
-
157
-template <typename scalar_t, vllm::ScalarTypeId w_type_id>
158
-__device__ inline typename ScalarType<scalar_t>::FragB dequant(int q);
159
-
160
-//
161
-// Efficiently dequantize 4bit values packed in an int32 value into a full
162
-// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
163
-// with some small changes:
164
-// - FP16:
165
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
166
-// - BF16:
167
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
168
-//
169
-template <>
170
-__device__ inline typename ScalarType<half>::FragB
171
-dequant<half, vllm::kU4B8.id()>(int q) {
172
- const int LO = 0x000f000f;
173
- const int HI = 0x00f000f0;
174
- const int EX = 0x64006400;
175
- // Guarantee that the `(a & b) | c` operations are LOP3s.
176
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
177
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
178
- // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
179
- // directly into `SUB` and `ADD`.
180
- const int SUB = 0x64086408;
181
- const int MUL = 0x2c002c00;
182
- const int ADD = 0xd480d480;
183
- typename ScalarType<half>::FragB frag_b;
184
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
185
- *reinterpret_cast<const half2*>(&SUB));
186
- frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
187
- *reinterpret_cast<const half2*>(&MUL),
188
- *reinterpret_cast<const half2*>(&ADD));
189
- return frag_b;
190
-}
191
-
192
-template <>
193
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
194
-dequant<nv_bfloat16, vllm::kU4B8.id()>(int q) {
195
- static constexpr uint32_t MASK = 0x000f000f;
196
- static constexpr uint32_t EX = 0x43004300;
197
-
198
- // Guarantee that the `(a & b) | c` operations are LOP3s.
199
-
200
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
201
- q >>= 4;
202
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
203
-
204
- typename ScalarType<nv_bfloat16>::FragB frag_b;
205
- static constexpr uint32_t MUL = 0x3F803F80;
206
- static constexpr uint32_t ADD = 0xC308C308;
207
-
208
- frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
209
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
210
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
211
- frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
212
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
213
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
214
- return frag_b;
215
-}
216
-
217
-template <>
218
-__device__ inline typename ScalarType<half>::FragB
219
-dequant<half, vllm::kU4.id()>(int q) {
220
- const int LO = 0x000f000f;
221
- const int HI = 0x00f000f0;
222
- const int EX = 0x64006400;
223
- // Guarantee that the `(a & b) | c` operations are LOP3s.
224
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);
225
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);
226
-
227
- const int SUB = 0x64006400;
228
- const int MUL = 0x2c002c00;
229
- const int ADD = 0xd400d400;
230
- typename ScalarType<half>::FragB frag_b;
231
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
232
- *reinterpret_cast<const half2*>(&SUB));
233
- frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
234
- *reinterpret_cast<const half2*>(&MUL),
235
- *reinterpret_cast<const half2*>(&ADD));
236
- return frag_b;
237
-}
238
-
239
-template <>
240
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
241
-dequant<nv_bfloat16, vllm::kU4.id()>(int q) {
242
- static constexpr uint32_t MASK = 0x000f000f;
243
- static constexpr uint32_t EX = 0x43004300;
244
-
245
- // Guarantee that the `(a & b) | c` operations are LOP3s.
246
-
247
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
248
- q >>= 4;
249
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
250
-
251
- typename ScalarType<nv_bfloat16>::FragB frag_b;
252
- static constexpr uint32_t MUL = 0x3F803F80;
253
- static constexpr uint32_t ADD = 0xC300C300;
254
-
255
- frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
256
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
257
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
258
- frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
259
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
260
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
261
- return frag_b;
262
-}
263
-
264
-//
265
-// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
266
-// bf16 Reference:
267
-// - FP16:
268
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
269
-// - BF16:
270
-// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
271
-//
272
-template <>
273
-__device__ inline typename ScalarType<half>::FragB
274
-dequant<half, vllm::kU8B128.id()>(int q) {
275
- static constexpr uint32_t mask_for_elt_01 = 0x5250;
276
- static constexpr uint32_t mask_for_elt_23 = 0x5351;
277
- static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
278
-
279
- uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
280
- uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
281
-
282
- static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
283
-
284
- typename ScalarType<half>::FragB frag_b;
285
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
286
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
287
- frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
288
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
289
- return frag_b;
290
-}
291
-
292
-template <>
293
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
294
-dequant<nv_bfloat16, vllm::kU8B128.id()>(int q) {
295
- typename ScalarType<nv_bfloat16>::FragB frag_b;
296
-
297
- float fp32_intermediates[4];
298
- uint32_t* fp32_intermediates_casted =
299
- reinterpret_cast<uint32_t*>(fp32_intermediates);
300
-
301
- static constexpr uint32_t fp32_base = 0x4B000000;
302
- fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
303
- fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
304
- fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
305
- fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
306
-
307
- fp32_intermediates[0] -= 8388736.f;
308
- fp32_intermediates[1] -= 8388736.f;
309
- fp32_intermediates[2] -= 8388736.f;
310
- fp32_intermediates[3] -= 8388736.f;
311
-
312
- uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&frag_b);
313
- bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
314
- fp32_intermediates_casted[1], 0x7632);
315
- bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
316
- fp32_intermediates_casted[3], 0x7632);
317
-
318
- return frag_b;
319
-}
320
-
321
-template <>
322
-__device__ inline typename ScalarType<half>::FragB
323
-dequant<half, vllm::kU8.id()>(int q) {
324
- static constexpr uint32_t mask_for_elt_01 = 0x5250;
325
- static constexpr uint32_t mask_for_elt_23 = 0x5351;
326
- static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
327
-
328
- uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
329
- uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
330
-
331
- static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
332
-
333
- typename ScalarType<half>::FragB frag_b;
334
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
335
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
336
- frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
337
- *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
338
- return frag_b;
339
-}
340
-
341
-template <>
342
-__device__ inline typename ScalarType<nv_bfloat16>::FragB
343
-dequant<nv_bfloat16, vllm::kU8.id()>(int q) {
344
- typename ScalarType<nv_bfloat16>::FragB frag_b;
345
-
346
- float fp32_intermediates[4];
347
- uint32_t* fp32_intermediates_casted =
348
- reinterpret_cast<uint32_t*>(fp32_intermediates);
349
-
350
- static constexpr uint32_t fp32_base = 0x4B000000;
351
- fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
352
- fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
353
- fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
354
- fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
355
-
356
- fp32_intermediates[0] -= 8388608.f;
357
- fp32_intermediates[1] -= 8388608.f;
358
- fp32_intermediates[2] -= 8388608.f;
359
- fp32_intermediates[3] -= 8388608.f;
360
-
361
- uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&frag_b);
362
- bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
363
- fp32_intermediates_casted[1], 0x7632);
364
- bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
365
- fp32_intermediates_casted[3], 0x7632);
366
-
367
- return frag_b;
368
-}
369
-
370
-// Multiply dequantized values by the corresponding quantization scale; used
371
-// only for grouped quantization.
372
-template <typename scalar_t>
373
-__device__ inline void scale(typename ScalarType<scalar_t>::FragB& frag_b,
374
- typename ScalarType<scalar_t>::FragS& frag_s,
375
- int i) {
376
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
377
- scalar_t2 s =
378
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
379
- frag_b[0] = __hmul2(frag_b[0], s);
380
- frag_b[1] = __hmul2(frag_b[1], s);
381
-}
382
-
383
-template <typename scalar_t>
384
-__device__ inline void sub_zp(typename ScalarType<scalar_t>::FragB& frag_b,
385
- typename ScalarType<scalar_t>::scalar_t2& frag_zp,
386
- int i) {
387
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
388
- scalar_t2 zp =
389
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
390
- frag_b[0] = __hsub2(frag_b[0], zp);
391
- frag_b[1] = __hsub2(frag_b[1], zp);
392
-}
393
-
394
-// Same as above, but for act_order (each K is multiplied individually)
395
-template <typename scalar_t>
396
-__device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
397
- typename ScalarType<scalar_t>::FragS& frag_s_1,
398
- typename ScalarType<scalar_t>::FragS& frag_s_2,
399
- typename ScalarType<scalar_t>::FragS& frag_s_3,
400
- typename ScalarType<scalar_t>::FragS& frag_s_4,
401
- int i) {
402
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
403
- scalar_t2 s_val_1_2;
404
- s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
405
- s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
406
-
407
- scalar_t2 s_val_3_4;
408
- s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
409
- s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
410
-
411
- frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
412
- frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
413
-}
414
-
415
-// Given 2 floats multiply by 2 scales (halves)
416
-template <typename scalar_t>
417
-__device__ inline void scale_float(float* c,
418
- typename ScalarType<scalar_t>::FragS& s) {
419
- scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
420
- c[0] = __fmul_rn(c[0], ScalarType<scalar_t>::num2float(s_ptr[0]));
421
- c[1] = __fmul_rn(c[1], ScalarType<scalar_t>::num2float(s_ptr[1]));
422
-}
423
-
424
-// Wait until barrier reaches `count`, then lock for current threadblock.
425
-__device__ inline void barrier_acquire(int* lock, int count) {
426
- if (threadIdx.x == 0) {
427
- int state = -1;
428
- do
429
- // Guarantee that subsequent writes by this threadblock will be visible
430
- // globally.
431
- asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
432
- : "=r"(state)
433
- : "l"(lock));
434
- while (state != count);
435
- }
436
- __syncthreads();
437
-}
438
-
439
-// Release barrier and increment visitation count.
440
-__device__ inline void barrier_release(int* lock, bool reset = false) {
441
- __syncthreads();
442
- if (threadIdx.x == 0) {
443
- if (reset) {
444
- lock[0] = 0;
445
- return;
446
- }
447
- int val = 1;
448
- // Make sure that all writes since acquiring this barrier are visible
449
- // globally, while releasing the barrier.
450
- asm volatile("fence.acq_rel.gpu;\n");
451
- asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
452
- :
453
- : "l"(lock), "r"(val));
454
- }
455
-}
456
-
65
// For a given "a" of size [M,K] performs a permutation of the K columns based
66
// on the given "perm" indices.
67
__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
@@ -510,1304 +118,19 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
118
}
119
}
120
513
-template <typename scalar_t, // compute dtype, half or nv_float16
514
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
515
- const int threads, // number of threads in a threadblock
516
- const int thread_m_blocks, // number of 16x16 blocks in the m
517
- // dimension (batchsize) of the
518
- // threadblock
519
- const int thread_n_blocks, // same for n dimension (output)
520
- const int thread_k_blocks, // same for k dimension (reduction)
521
- const int stages, // number of stages for the async global->shared
522
- // fetch pipeline
523
- const bool has_act_order, // whether act_order is enabled
524
- const bool has_zp, // whether zero-points are enabled
525
- const int group_blocks = -1, // number of consecutive 16x16 blocks
526
- // with a separate quantization scale
527
- const bool is_zp_float // is zero point of float16 type?
528
- >
529
-__global__ void Marlin(
530
- const int4* __restrict__ A, // fp16 input matrix of shape mxk
531
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
532
- int4* __restrict__ C, // fp16 output buffer of shape mxn
533
- int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
534
- const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
535
- // (k/groupsize)xn
536
- const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
537
- // (k/groupsize)x(n/pack_factor)
538
- const int* __restrict__ g_idx, // int32 group indices of shape k
539
- int num_groups, // number of scale groups per output channel
540
- int prob_m, // batch dimension m
541
- int prob_n, // output dimension n
542
- int prob_k, // reduction dimension k
543
- int lda, // A.stride(0), equal to prob_k is A is contiguous
544
- int* locks, // extra global storage for barrier synchronization
545
- bool use_atomic_add, // whether to use atomic add to reduce
546
- bool use_fp32_reduce // whether to use fp32 global reduce
547
-) {
548
- // Each threadblock processes one "stripe" of the B matrix with (roughly) the
549
- // same size, which might involve multiple column "slices" (of width 16 *
550
- // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
551
- // example:
552
- // 0 1 3
553
- // 0 2 3
554
- // 1 2 4
555
- // While this kind of partitioning makes things somewhat more complicated, it
556
- // ensures good utilization of all SMs for many kinds of shape and GPU
557
- // configurations, while requiring as few slow global cross-threadblock
558
- // reductions as possible.
559
- using Dtype = ScalarType<scalar_t>;
560
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
561
- using FragA = typename ScalarType<scalar_t>::FragA;
562
- using FragB = typename ScalarType<scalar_t>::FragB;
563
- using FragC = typename ScalarType<scalar_t>::FragC;
564
- using FragS = typename ScalarType<scalar_t>::FragS;
565
- using FragZP = typename ScalarType<scalar_t>::FragZP;
566
-
567
- static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
568
-
569
- constexpr int pack_factor = 32 / w_type.size_bits();
570
-
571
- // For larger GEMMs we run multiple batchsize 64 versions in parallel for a
572
- // better partitioning with less reductions
573
- int parallel = 1;
574
- if (prob_m > 16 * thread_m_blocks) {
575
- parallel = prob_m / (16 * thread_m_blocks);
576
- prob_m = 16 * thread_m_blocks;
577
- }
578
-
579
- int k_tiles = prob_k / 16 / thread_k_blocks;
580
- int n_tiles = prob_n / 16 / thread_n_blocks;
581
- int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x);
582
-
583
- if constexpr (!has_act_order && group_blocks != -1) {
584
- if (group_blocks >= thread_k_blocks) {
585
- // Ensure that the number of tiles in each stripe is a multiple of the
586
- // groupsize; this avoids an annoying special case where a stripe starts
587
- // in the middle of group.
588
- iters = (group_blocks / thread_k_blocks) *
589
- div_ceil(iters, (group_blocks / thread_k_blocks));
590
- }
591
- }
592
-
593
- int slice_row = (iters * blockIdx.x) % k_tiles;
594
- int slice_col_par = (iters * blockIdx.x) / k_tiles;
595
- int slice_col = slice_col_par;
596
- int slice_iters; // number of threadblock tiles in the current slice
597
- int slice_count =
598
- 0; // total number of active threadblocks in the current slice
599
- int slice_idx; // index of threadblock in current slice; numbered bottom to
600
- // top
601
-
602
- int par_id = 0;
603
-
604
- // We can easily implement parallel problem execution by just remapping
605
- // indices and advancing global pointers
606
- if (slice_col_par >= n_tiles) {
607
- A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8;
608
- C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
609
- locks += (slice_col_par / n_tiles) * n_tiles;
610
- slice_col = slice_col_par % n_tiles;
611
- par_id = slice_col_par / n_tiles;
612
- }
613
-
614
- // Compute all information about the current slice which is required for
615
- // synchronization.
616
- auto init_slice = [&]() {
617
- slice_iters =
618
- iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
619
- if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
620
- if (slice_iters == 0) return;
621
- if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
622
- slice_count = 1;
623
- slice_idx = 0;
624
- int col_first = iters * div_ceil(k_tiles * slice_col_par, iters);
625
- if (col_first <= k_tiles * (slice_col_par + 1)) {
626
- int col_off = col_first - k_tiles * slice_col_par;
627
- slice_count = div_ceil(k_tiles - col_off, iters);
628
- if (col_off > 0) slice_count++;
629
- int delta_first = iters * blockIdx.x - col_first;
630
- if (delta_first < 0 || (col_off == 0 && delta_first == 0))
631
- slice_idx = slice_count - 1;
632
- else {
633
- slice_idx = slice_count - 1 - delta_first / iters;
634
- if (col_off > 0) slice_idx--;
635
- }
636
- }
637
- if (slice_col == n_tiles) {
638
- A += 16 * thread_m_blocks * lda / 8;
639
- C += 16 * thread_m_blocks * prob_n / 8;
640
- locks += n_tiles;
641
- slice_col = 0;
642
- par_id++;
643
- }
644
- };
645
- init_slice();
646
-
647
- // A sizes/strides
648
-
649
- // stride of the A matrix in global memory
650
- int a_gl_stride = lda / 8;
651
- // stride of an A matrix tile in shared memory
652
- constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
653
- // delta between subsequent A tiles in global memory
654
- constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;
655
- // between subsequent accesses within a tile
656
- int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
657
- // between shared memory writes
658
- constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
659
- // between shared memory tile reads
660
- constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));
661
- // within a shared memory tile
662
- constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
663
- // overall size of a tile
664
- constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);
665
- // number of shared write iterations for a tile
666
- constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
667
-
668
- // B sizes/strides
669
- int b_gl_stride = 16 * prob_n / (pack_factor * 4);
670
- constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
671
- constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;
672
- constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
673
-
674
- int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
675
- int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
676
- constexpr int b_sh_wr_delta = threads * b_thread_vecs;
677
- constexpr int b_sh_rd_delta = threads * b_thread_vecs;
678
- constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
679
- constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
680
-
681
- // Scale sizes/strides without act_order
682
- int s_gl_stride = prob_n / 8;
683
- constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
684
- constexpr int s_tb_groups =
685
- !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
686
- ? thread_k_blocks / group_blocks
687
- : 1;
688
- constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
689
- int s_gl_rd_delta = s_gl_stride;
690
-
691
- // Scale size/strides with act_order
692
- constexpr int tb_k = 16 * thread_k_blocks;
693
- constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
694
- // constexpr int act_s_row_stride = 1;
695
- // int act_s_col_stride = act_s_row_stride * num_groups;
696
- int act_s_col_stride = 1;
697
- int act_s_col_warp_stride = act_s_col_stride * 8;
698
- int tb_n_warps = thread_n_blocks / 4;
699
- int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
700
-
701
- // Zero-points sizes/strides
702
- int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
703
- constexpr int zp_sh_stride = is_zp_float
704
- ? 16 * thread_n_blocks / 8
705
- : ((16 * thread_n_blocks) / pack_factor) / 4;
706
- constexpr int zp_tb_groups = s_tb_groups;
707
- constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;
708
- int zp_gl_rd_delta = zp_gl_stride;
709
-
710
- // Global A read index of current thread.
711
- int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
712
- (threadIdx.x % a_gl_rd_delta_o);
713
- a_gl_rd += a_gl_rd_delta_o * slice_row;
714
- // Shared write index of current thread.
715
- int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
716
- (threadIdx.x % a_gl_rd_delta_o);
717
- // Shared read index.
718
- int a_sh_rd =
719
- a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16;
720
- a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
721
-
722
- int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
723
- (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
724
- b_gl_rd += b_sh_stride * slice_col;
725
- b_gl_rd += b_gl_rd_delta_o * slice_row;
726
- auto b_sh_wr = threadIdx.x * b_thread_vecs;
727
- auto b_sh_rd = threadIdx.x * b_thread_vecs;
728
-
729
- // For act_order
730
- constexpr int k_iter_size = tb_k / b_sh_wr_iters;
731
- int slice_k_start = tb_k * slice_row;
732
- int slice_k_finish = slice_k_start + tb_k * slice_iters;
733
- int slice_k_start_shared_fetch = slice_k_start;
734
- int slice_n_offset = act_s_col_tb_stride * slice_col;
735
-
736
- // No act_order
737
- int s_gl_rd;
738
- if constexpr (!has_act_order) {
739
- if constexpr (group_blocks == -1) {
740
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
741
- } else {
742
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
743
- s_sh_stride * slice_col + threadIdx.x;
744
- }
745
- }
746
- auto s_sh_wr = threadIdx.x;
747
- bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
748
-
749
- // Zero-points
750
- int zp_gl_rd;
751
- if constexpr (has_zp) {
752
- if constexpr (group_blocks == -1) {
753
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
754
- } else {
755
- zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
756
- zp_sh_stride * slice_col + threadIdx.x;
757
- }
758
- }
759
- auto zp_sh_wr = threadIdx.x;
760
- bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
761
-
762
- // We use a different scale layout for grouped and column-wise quantization as
763
- // we scale a `half2` tile in column-major layout in the former and in
764
- // row-major in the latter case.
765
- int s_sh_rd;
766
- if constexpr (group_blocks != -1)
767
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
768
- (threadIdx.x % 32) / 4;
769
- else
770
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
771
- (threadIdx.x % 32) % 4;
772
-
773
- // Zero-points have the same read layout as the scales
774
- // (without column-wise case)
775
- constexpr int num_col_threads = 8;
776
- constexpr int num_row_threads = 4;
777
- constexpr int num_ints_per_thread = 8 / pack_factor;
778
- int zp_sh_rd;
779
- if constexpr (has_zp) {
780
- if constexpr (is_zp_float) {
781
- if constexpr (group_blocks != -1) {
782
- zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
783
- (threadIdx.x % 32) / 4;
784
- }
785
- } else {
786
- zp_sh_rd = num_ints_per_thread * num_col_threads *
787
- ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
788
- num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
789
- }
790
- }
791
-
792
- // Precompute which thread should not read memory in which iterations; this is
793
- // needed if there are more threads than required for a certain tilesize or
794
- // when the batchsize is not a multiple of 16.
795
- bool a_sh_wr_pred[a_sh_wr_iters];
796
- #pragma unroll
797
- for (int i = 0; i < a_sh_wr_iters; i++)
798
- a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;
799
-
800
- // To ensure that writing and reading A tiles to/from shared memory, the
801
- // latter in fragment format, is fully bank conflict free, we need to use a
802
- // rather fancy XOR-based layout. The key here is that neither reads nor
803
- // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
804
- // same shared memory banks. Further, it seems (based on NSight-Compute) that
805
- // each warp must also write a consecutive memory segment?
806
- auto transform_a = [&](int i) {
807
- int row = i / a_gl_rd_delta_o;
808
- return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
809
- };
810
- // Since the computation of this remapping is non-trivial and, due to our main
811
- // loop unrolls, all shared memory accesses are static, we simply precompute
812
- // both transformed reads and writes.
813
- int a_sh_wr_trans[a_sh_wr_iters];
814
- #pragma unroll
815
- for (int i = 0; i < a_sh_wr_iters; i++)
816
- a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
817
- int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
818
- #pragma unroll
819
- for (int i = 0; i < b_sh_wr_iters; i++) {
820
- #pragma unroll
821
- for (int j = 0; j < thread_m_blocks; j++)
822
- a_sh_rd_trans[i][j] =
823
- transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
824
- }
825
-
826
- // Since B-accesses have non-constant stride they have to be computed at
827
- // runtime; we break dependencies between subsequent accesses with a tile by
828
- // maintining multiple pointers (we have enough registers), a tiny
829
- // optimization.
830
- const int4* B_ptr[b_sh_wr_iters];
831
- #pragma unroll
832
- for (int i = 0; i < b_sh_wr_iters; i++)
833
- B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
834
-
835
- extern __shared__ int4 sh[];
836
- // Shared memory storage for global fetch pipelines.
837
- int4* sh_a = sh;
838
- int4* sh_b = sh_a + (stages * a_sh_stage);
839
- int4* sh_g_idx = sh_b + (stages * b_sh_stage);
840
- int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
841
- int4* sh_s = sh_zp + (stages * zp_sh_stage);
842
- int4* sh_red = sh_s + (stages * s_sh_stage);
843
-
844
- // Register storage for double buffer of shared memory reads.
845
- FragA frag_a[2][thread_m_blocks];
846
- I4 frag_b_quant[2][b_thread_vecs];
847
- FragC frag_c[thread_m_blocks][4][2];
848
- FragS frag_s[2][4]; // No act-order
849
- FragS act_frag_s[2][4][4]; // For act-order
850
- int frag_qzp[2][num_ints_per_thread]; // Zero-points
851
- FragZP frag_zp; // Zero-points in fp16
852
- FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
853
-
854
- // Zero accumulators.
855
- auto zero_accums = [&]() {
856
- #pragma unroll
857
- for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
858
- reinterpret_cast<float*>(frag_c)[i] = 0;
859
- };
860
-
861
- int sh_first_group_id = -1;
862
- int sh_num_groups = -1;
863
- constexpr int sh_max_num_groups = 32;
864
-
865
- auto fetch_scales_to_shared = [&](bool is_async, int first_group_id,
866
- int last_group_id) {
867
- sh_first_group_id = first_group_id;
868
- sh_num_groups = last_group_id - first_group_id + 1;
869
-
870
- if (sh_num_groups < sh_max_num_groups) {
871
- sh_num_groups = sh_max_num_groups;
872
- }
873
-
874
- if (sh_first_group_id + sh_num_groups > num_groups) {
875
- sh_num_groups = num_groups - sh_first_group_id;
876
- }
877
-
878
- int row_offset = first_group_id * s_gl_stride;
879
-
880
- if (is_async) {
881
- for (int i = 0; i < sh_num_groups; i++) {
882
- if (threadIdx.x < s_sh_stride) {
883
- cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
884
- &scales_ptr[row_offset + (i * s_gl_stride) +
885
- slice_n_offset + threadIdx.x]);
886
- }
887
- }
888
- } else {
889
- for (int i = 0; i < sh_num_groups; i++) {
890
- if (threadIdx.x < s_sh_stride) {
891
- sh_s[(i * s_sh_stride) + threadIdx.x] =
892
- scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
893
- threadIdx.x];
894
- }
895
- }
896
- }
897
- };
898
- // Asynchronously fetch the next A, B and s tile from global to the next
899
- // shared memory pipeline location.
900
- auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
901
- if (pred) {
902
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
903
- #pragma unroll
904
- for (int i = 0; i < a_sh_wr_iters; i++) {
905
- cp_async4_pred(
906
- &sh_a_stage[a_sh_wr_trans[i]],
907
- &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
908
- a_sh_wr_pred[i]);
909
- }
910
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
911
- #pragma unroll
912
- for (int i = 0; i < b_sh_wr_iters; i++) {
913
- #pragma unroll
914
- for (int j = 0; j < b_thread_vecs; j++) {
915
- cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j);
916
- }
917
-
918
- B_ptr[i] += b_gl_rd_delta_o;
919
- }
920
-
921
- if constexpr (has_act_order) {
922
- // Fetch g_idx thread-block portion
923
- int full_pipe = a_off;
924
- int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
925
- if (cur_k < prob_k && cur_k < slice_k_finish) {
926
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
927
-
928
- int4 const* cur_g_idx_stage_ptr =
929
- reinterpret_cast<int4 const*>(&g_idx[cur_k]);
930
-
931
- if (threadIdx.x < g_idx_stage) {
932
- cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
933
- &cur_g_idx_stage_ptr[threadIdx.x]);
934
- }
935
- }
936
- } else {
937
- if constexpr (group_blocks != -1) {
938
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
939
-
940
- if constexpr (group_blocks >= thread_k_blocks) {
941
- if (s_sh_wr_pred) {
942
- cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
943
- }
944
- // Only fetch scales if this tile starts a new group
945
- if ((pipe + 1) % (group_blocks / thread_k_blocks) == 0) {
946
- s_gl_rd += s_gl_rd_delta;
947
- }
948
- } else {
949
- for (int i = 0; i < s_tb_groups; i++) {
950
- if (s_sh_wr_pred) {
951
- cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
952
- &scales_ptr[s_gl_rd]);
953
- }
954
- s_gl_rd += s_gl_rd_delta;
955
- }
956
- }
957
- }
958
-
959
- if constexpr (has_zp && group_blocks != -1) {
960
- int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
961
-
962
- if constexpr (group_blocks >= thread_k_blocks) {
963
- // Only fetch zero-points if this tile starts a new group
964
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
965
- if (zp_sh_wr_pred) {
966
- cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
967
- }
968
- zp_gl_rd += zp_gl_rd_delta;
969
- }
970
- } else {
971
- for (int i = 0; i < zp_tb_groups; i++) {
972
- if (zp_sh_wr_pred) {
973
- cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
974
- &zp_ptr[zp_gl_rd]);
975
- }
976
- zp_gl_rd += zp_gl_rd_delta;
977
- }
978
- }
979
- }
980
- }
981
- }
982
- // Insert a fence even when we are winding down the pipeline to ensure that
983
- // waiting is also correct at this point.
984
- cp_async_fence();
985
- };
986
-
987
- auto fetch_zp_to_shared = [&]() {
988
- if (zp_sh_wr_pred) {
989
- cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);
990
- }
991
- };
992
-
993
- // Wait until the next thread tile has been loaded to shared memory.
994
- auto wait_for_stage = [&]() {
995
- // We only have `stages - 2` active fetches since we are double buffering
996
- // and can only issue the next fetch when it is guaranteed that the previous
997
- // shared memory load is fully complete (as it may otherwise be
998
- // overwritten).
999
- cp_async_wait<stages - 2>();
1000
- __syncthreads();
1001
- };
1002
-
1003
- // Load the next sub-tile from the current location in the shared memory pipe
1004
- // into the current register buffer.
1005
- auto fetch_to_registers = [&](int k, int pipe) {
1006
- int4* sh_a_stage = sh_a + a_sh_stage * pipe;
1007
- #pragma unroll
1008
- for (int i = 0; i < thread_m_blocks; i++)
1009
- ldsm4<scalar_t>(frag_a[k % 2][i],
1010
- &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
1011
- int4* sh_b_stage = sh_b + b_sh_stage * pipe;
1012
-
1013
- #pragma unroll
1014
- for (int i = 0; i < b_thread_vecs; i++) {
1015
- frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
1016
- &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
1017
- }
1018
- };
1019
-
1020
- bool is_same_group[stages];
1021
- int same_group_id[stages];
1022
-
1023
- auto init_same_group = [&](int pipe) {
1024
- if constexpr (!has_act_order) {
1025
- is_same_group[pipe] = false;
1026
- same_group_id[pipe] = 0;
1027
- return;
1028
- }
1029
-
1030
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1031
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1032
-
1033
- int group_id_1 = sh_g_idx_int_ptr[0];
1034
- int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
1035
-
1036
- is_same_group[pipe] = group_id_1 == group_id_2;
1037
- same_group_id[pipe] = group_id_1;
1038
- };
1039
-
1040
- auto fetch_scales_to_registers = [&](int k, int full_pipe) {
1041
- int pipe = full_pipe % stages;
1042
-
1043
- if constexpr (!has_act_order) {
1044
- // No act-order case
1045
- if constexpr (group_blocks != -1) {
1046
- if constexpr (group_blocks >= thread_k_blocks) {
1047
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1048
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1049
- } else {
1050
- auto warp_id = threadIdx.x / 32;
1051
- int n_warps = thread_n_blocks / 4;
1052
-
1053
- int warp_row = warp_id / n_warps;
1054
-
1055
- int cur_k = warp_row * 16;
1056
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1057
-
1058
- int k_blocks = cur_k / 16;
1059
- int cur_group_id = k_blocks / group_blocks;
1060
-
1061
- int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1062
-
1063
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1064
- sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1065
- }
1066
- }
1067
-
1068
- return;
1069
- }
1070
-
1071
- // Act-order case
1072
-
1073
- // Determine K of the "current" thread-block
1074
- int cur_k = slice_k_start + tb_k * full_pipe;
1075
- if (cur_k >= prob_k || cur_k >= slice_k_finish) {
1076
- return;
1077
- }
1078
-
1079
- // Reset (to current thread-block) since we read g_idx portion from the
1080
- // shared memory
1081
- cur_k = 0;
1082
-
1083
- // Progress to current iteration
1084
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1085
-
1086
- // Determine "position" inside the thread-block (based on warp and
1087
- // thread-id)
1088
- auto warp_id = threadIdx.x / 32;
1089
- int n_warps =
1090
- thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
1091
-
1092
- int warp_row = warp_id / n_warps;
1093
- int warp_col = warp_id % n_warps;
1094
-
1095
- cur_k += warp_row * 16;
1096
-
1097
- auto th_id = threadIdx.x % 32;
1098
- cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1099
-
1100
- int s_col_shift =
1101
- /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
1102
- (th_id / 4) * act_s_col_stride;
1103
-
1104
- if (is_same_group[pipe]) {
1105
- if (k % 2 == 0) {
1106
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1107
- sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
1108
- s_col_shift];
1109
- } else {
1110
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1111
- *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
1112
- }
1113
-
1114
- for (int i = 1; i < 4; i++) {
1115
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1116
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
1117
- }
1118
- return;
1119
- }
1120
-
1121
- int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1122
- int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1123
-
1124
- constexpr int k_frag_offsets[4] = {0, 1, 8,
1125
- 9}; // Tensor core offsets per thread
1126
-
1127
- #pragma unroll
1128
- for (int i = 0; i < 4; i++) {
1129
- int actual_k = cur_k + k_frag_offsets[i];
1130
-
1131
- int group_id = sh_g_idx_int_ptr[actual_k];
1132
- int rel_group_id = group_id - sh_first_group_id;
1133
-
1134
- *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1135
- sh_s[rel_group_id * s_sh_stride + s_col_shift];
1136
- }
1137
- };
1138
-
1139
- auto fetch_zp_to_registers = [&](int k, int full_pipe) {
1140
- // This code does not handle group_blocks == 0,
1141
- // which signifies act_order.
1142
- // has_zp implies AWQ, which doesn't have act_order,
1143
- static_assert(!has_zp || group_blocks != 0);
1144
-
1145
- if constexpr (has_zp && !is_zp_float) {
1146
- int pipe = full_pipe % stages;
1147
-
1148
- if constexpr (group_blocks == -1) {
1149
- for (int i = 0; i < num_ints_per_thread; i++) {
1150
- frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
1151
- }
1152
-
1153
- } else if constexpr (group_blocks >= thread_k_blocks) {
1154
- int4* sh_zp_stage =
1155
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1156
- (pipe / (group_blocks / thread_k_blocks)));
1157
- for (int i = 0; i < num_ints_per_thread; i++) {
1158
- frag_qzp[k % 2][i] =
1159
- (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1160
- }
1161
- } else {
1162
- auto warp_id = threadIdx.x / 32;
1163
- int n_warps = thread_n_blocks / 4;
1164
-
1165
- int warp_row = warp_id / n_warps;
1166
-
1167
- int cur_k = warp_row * 16;
1168
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1169
-
1170
- int k_blocks = cur_k / 16;
1171
- int cur_group_id = 0;
1172
-
1173
- // Suppress bogus and persistent divide-by-zero warning
1174
- #pragma nv_diagnostic push
1175
- #pragma nv_diag_suppress divide_by_zero
1176
- cur_group_id = k_blocks / group_blocks;
1177
- #pragma nv_diagnostic pop
1178
-
1179
- int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1180
-
1181
- sh_zp_stage += cur_group_id * zp_sh_stride;
1182
-
1183
- for (int i = 0; i < num_ints_per_thread; i++) {
1184
- frag_qzp[k % 2][i] =
1185
- (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1186
- }
1187
- }
1188
- }
1189
-
1190
- else if constexpr (has_zp && is_zp_float) {
1191
- int pipe = full_pipe % stages;
1192
-
1193
- if constexpr (group_blocks != -1) {
1194
- if constexpr (group_blocks >= thread_k_blocks) {
1195
- int4* sh_zp_stage =
1196
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1197
- (pipe / (group_blocks / thread_k_blocks)));
1198
- reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd];
1199
- } else {
1200
- auto warp_id = threadIdx.x / 32;
1201
- int n_warps = thread_n_blocks / 4;
1202
-
1203
- int warp_row = warp_id / n_warps;
1204
-
1205
- int cur_k = warp_row * 16;
1206
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1207
-
1208
- int k_blocks = cur_k / 16;
1209
- // Suppress bogus and persistent divide-by-zero warning
1210
- #pragma nv_diagnostic push
1211
- #pragma nv_diag_suppress divide_by_zero
1212
- int cur_group_id = k_blocks / group_blocks;
1213
- #pragma nv_diagnostic pop
1214
-
1215
- int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1216
-
1217
- reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1218
- sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride];
1219
- }
1220
- }
1221
- }
1222
- };
1223
-
1224
- // Execute the actual tensor core matmul of a sub-tile.
1225
- auto matmul = [&](int k) {
1226
- if constexpr (has_zp && !is_zp_float) {
1227
- FragB frag_zp_0;
1228
- FragB frag_zp_1;
1229
- int zp_quant_0, zp_quant_1;
1230
-
1231
- if constexpr (w_type.size_bits() == 4) {
1232
- zp_quant_0 = frag_qzp[k % 2][0];
1233
- zp_quant_1 = zp_quant_0 >> 8;
1234
- } else {
1235
- static_assert(w_type.size_bits() == 8);
1236
- zp_quant_0 = frag_qzp[k % 2][0];
1237
- zp_quant_1 = frag_qzp[k % 2][1];
1238
- }
1239
-
1240
- frag_zp_0 = dequant<scalar_t, w_type_id>(zp_quant_0);
1241
- frag_zp_1 = dequant<scalar_t, w_type_id>(zp_quant_1);
1242
-
1243
- frag_zp[0] = frag_zp_0[0];
1244
- frag_zp[1] = frag_zp_0[1];
1245
- frag_zp[2] = frag_zp_1[0];
1246
- frag_zp[3] = frag_zp_1[1];
1247
- }
1248
-
1249
- // We have the m dimension as the inner loop in order to encourage overlapping
1250
- // dequantization and matmul operations.
1251
- #pragma unroll
1252
- for (int j = 0; j < 4; j++) {
1253
- FragB frag_b0;
1254
- FragB frag_b1;
1255
- int b_quant_0, b_quant_1;
1256
-
1257
- if constexpr (w_type.size_bits() == 4) {
1258
- b_quant_0 = frag_b_quant[k % 2][0][j];
1259
- b_quant_1 = b_quant_0 >> 8;
1260
- } else {
1261
- static_assert(w_type.size_bits() == 8);
1262
- int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k % 2]);
1263
- b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
1264
- b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1265
- }
1266
-
1267
- frag_b0 = dequant<scalar_t, w_type_id>(b_quant_0);
1268
- frag_b1 = dequant<scalar_t, w_type_id>(b_quant_1);
1269
-
1270
- // Apply zero-point to frag_b0
1271
- if constexpr (has_zp && !is_zp_float) {
1272
- sub_zp<scalar_t>(frag_b0, frag_zp[j], 0);
1273
- }
1274
-
1275
- else if constexpr (has_zp && is_zp_float && group_blocks != -1) {
1276
- sub_zp<scalar_t>(frag_b0, frag_zpf[k % 2][j], 0);
1277
- }
1278
-
1279
- // Apply scale to frag_b0
1280
- if constexpr (has_act_order) {
1281
- scale4<scalar_t>(frag_b0, act_frag_s[k % 2][0][j],
1282
- act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j],
1283
- act_frag_s[k % 2][3][j], 0);
1284
- } else {
1285
- if constexpr (group_blocks != -1) {
1286
- scale<scalar_t>(frag_b0, frag_s[k % 2][j], 0);
1287
- }
1288
- }
1289
-
1290
- // Apply zero-point to frag_b1
1291
- if constexpr (has_zp && !is_zp_float) {
1292
- sub_zp<scalar_t>(frag_b1, frag_zp[j], 1);
1293
- }
1294
-
1295
- else if constexpr (has_zp && is_zp_float && group_blocks != -1) {
1296
- sub_zp<scalar_t>(frag_b1, frag_zpf[k % 2][j], 1);
1297
- }
1298
-
1299
- // Apply scale to frag_b1
1300
- if constexpr (has_act_order) {
1301
- scale4<scalar_t>(frag_b1, act_frag_s[k % 2][0][j],
1302
- act_frag_s[k % 2][1][j], act_frag_s[k % 2][2][j],
1303
- act_frag_s[k % 2][3][j], 1);
1304
-
1305
- } else {
1306
- if constexpr (group_blocks != -1) {
1307
- scale<scalar_t>(frag_b1, frag_s[k % 2][j], 1);
1308
- }
1309
- }
1310
-
1311
- #pragma unroll
1312
- for (int i = 0; i < thread_m_blocks; i++) {
1313
- mma<scalar_t>(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]);
1314
- mma<scalar_t>(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]);
1315
- }
1316
- }
1317
- };
1318
-
1319
- // Since we slice across the k dimension of a tile in order to increase the
1320
- // number of warps while keeping the n dimension of a tile reasonable, we have
1321
- // multiple warps that accumulate their partial sums of the same output
1322
- // location; which we have to reduce over in the end. We do in shared memory.
1323
- auto thread_block_reduce = [&]() {
1324
- constexpr int red_off = threads / b_sh_stride_threads / 2;
1325
- if (red_off >= 1) {
1326
- auto red_idx = threadIdx.x / b_sh_stride_threads;
1327
- constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1328
- constexpr int red_sh_delta = b_sh_stride_threads;
1329
- int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
1330
- (threadIdx.x % b_sh_stride_threads);
1331
-
1332
- // Parallel logarithmic shared memory reduction. We make sure to avoid any
1333
- // unnecessary read or write iterations, e.g., for two warps we write only
1334
- // once by warp 1 and read only once by warp 0.
1335
-
1336
- #pragma unroll
1337
- for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
1338
- #pragma unroll
1339
- for (int i = red_off; i > 0; i /= 2) {
1340
- if (i <= red_idx && red_idx < 2 * i) {
1341
- #pragma unroll
1342
- for (int j = 0; j < 4 * 2; j++) {
1343
- int red_sh_wr =
1344
- red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
1345
- if (i < red_off) {
1346
- float* c_rd = reinterpret_cast<float*>(
1347
- &sh_red[red_sh_delta * j + red_sh_rd]);
1348
- float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
1349
- #pragma unroll
1350
- for (int k = 0; k < 4; k++)
1351
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
1352
- c_rd[k] + c_wr[k];
1353
- }
1354
- sh_red[red_sh_wr] =
1355
- reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
1356
- }
1357
- }
1358
- __syncthreads();
1359
- }
1360
- if (red_idx == 0) {
1361
- #pragma unroll
1362
- for (int i = 0; i < 4 * 2; i++) {
1363
- float* c_rd =
1364
- reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
1365
- #pragma unroll
1366
- for (int j = 0; j < 4; j++)
1367
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
1368
- c_rd[j];
1369
- }
1370
- }
1371
- __syncthreads();
1372
- }
1373
- }
1374
- };
1375
-
1376
- // Since multiple threadblocks may process parts of the same column slice, we
1377
- // finally have to globally reduce over the results. As the striped
1378
- // partitioning minimizes the number of such reductions and our outputs are
1379
- // usually rather small, we perform this reduction serially in L2 cache.
1380
- auto global_reduce_fp16 = [&](bool first = false, bool last = false) {
1381
- // We are very careful here to reduce directly in the output buffer to
1382
- // maximize L2 cache utilization in this step. To do this, we write out
1383
- // results in FP16 (but still reduce with FP32 compute).
1384
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1385
- if (threadIdx.x < active_threads) {
1386
- int c_gl_stride = prob_n / 8;
1387
- int c_gl_wr_delta_o = 8 * c_gl_stride;
1388
- int c_gl_wr_delta_i = 4 * (active_threads / 32);
1389
- int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
1390
- 4 * (threadIdx.x / 32) + threadIdx.x % 4;
1391
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
1392
- constexpr int c_sh_wr_delta = active_threads;
1393
- auto c_sh_wr = threadIdx.x;
1394
-
1395
- int row = (threadIdx.x % 32) / 4;
1396
-
1397
- if (!first) {
1398
- // Interestingly, doing direct global accesses here really seems to mess up
1399
- // the compiler and lead to slowdowns, hence we also use async-copies even
1400
- // though these fetches are not actually asynchronous.
1401
- #pragma unroll
1402
- for (int i = 0; i < thread_m_blocks * 4; i++) {
1403
- cp_async4_pred(
1404
- &sh_red[c_sh_wr + c_sh_wr_delta * i],
1405
- &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1406
- c_gl_wr_delta_i * (i % 2)],
1407
- i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
1408
- }
1409
- cp_async_fence();
1410
- cp_async_wait<0>();
1411
- }
1412
-
1413
- #pragma unroll
1414
- for (int i = 0; i < thread_m_blocks * 4; i++) {
1415
- if (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) {
1416
- if (!first) {
1417
- int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
1418
- #pragma unroll
1419
- for (int j = 0; j < 2 * 4; j++) {
1420
- reinterpret_cast<float*>(
1421
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] +=
1422
- Dtype::num2float(reinterpret_cast<scalar_t*>(&c_red)[j]);
1423
- }
1424
- }
1425
- if (!last) {
1426
- int4 c;
1427
- #pragma unroll
1428
- for (int j = 0; j < 2 * 4; j++) {
1429
- reinterpret_cast<scalar_t*>(&c)[j] =
1430
- Dtype::float2num(reinterpret_cast<float*>(
1431
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]);
1432
- }
1433
- C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] =
1434
- c;
1435
- }
1436
- }
1437
- }
1438
- }
1439
- };
1440
-
1441
- // Globally reduce over threadblocks that compute the same column block.
1442
- // We use a tmp C buffer to reduce in full fp32 precision.
1443
- auto global_reduce_fp32 = [&](bool first = false, bool last = false) {
1444
- constexpr int tb_m = thread_m_blocks * 16;
1445
- constexpr int tb_n = thread_n_blocks * 16;
1446
-
1447
- constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
1448
-
1449
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1450
- bool is_th_active = threadIdx.x < active_threads;
1451
-
1452
- int par_offset = c_size * n_tiles * par_id;
1453
- int slice_offset = c_size * slice_col;
1454
-
1455
- constexpr int num_floats = thread_m_blocks * 4 * 2 * 4;
1456
- constexpr int th_size = num_floats * sizeof(float) / 16;
1457
-
1458
- int c_cur_offset = par_offset + slice_offset;
1459
-
1460
- if (!is_th_active) {
1461
- return;
1462
- }
1463
-
1464
- if (!first) {
1465
- float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
1466
- #pragma unroll
1467
- for (int k = 0; k < th_size; k++) {
1468
- sh_red[threadIdx.x] =
1469
- C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
1470
-
1471
- float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
1472
- #pragma unroll
1473
- for (int f = 0; f < 4; f++) {
1474
- frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
1475
- }
1476
- }
1477
- }
1478
-
1479
- if (!last) {
1480
- int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
1481
- #pragma unroll
1482
- for (int k = 0; k < th_size; k++) {
1483
- C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k];
1484
- }
1485
- }
1486
- };
1487
-
1488
- // Write out the reduce final result in the correct layout. We only actually
1489
- // reshuffle matrix fragments in this step, the reduction above is performed
1490
- // in fragment layout.
1491
- auto write_result = [&]() {
1492
- int c_gl_stride = prob_n / 8;
1493
- constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
1494
- int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
1495
- constexpr int c_sh_rd_delta =
1496
- c_sh_stride * (threads / (2 * thread_n_blocks));
1497
-
1498
- int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1499
- (threadIdx.x % (2 * thread_n_blocks));
1500
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
1501
- int c_sh_wr =
1502
- (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
1503
- c_sh_wr += 32 * (threadIdx.x / 32);
1504
- int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1505
- (threadIdx.x % (2 * thread_n_blocks));
1506
-
1507
- int c_gl_wr_end = c_gl_stride * prob_m;
1508
-
1509
- // We first reorder in shared memory to guarantee the most efficient final
1510
- // global write patterns
1511
- auto write = [&](int idx, float c0, float c1, FragS& s) {
1512
- scalar_t2 res =
1513
- Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1514
-
1515
- // For per-column quantization we finally apply the scale here (only for
1516
- // 4-bit)
1517
- if constexpr (!has_act_order && group_blocks == -1 &&
1518
- w_type.size_bits() == 4) {
1519
- res = __hmul2(res, s[0]);
1520
- }
1521
-
1522
- ((scalar_t2*)sh_red)[idx] = res;
1523
- };
1524
-
1525
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1526
- #pragma unroll
1527
- for (int i = 0; i < thread_m_blocks; i++) {
1528
- #pragma unroll
1529
- for (int j = 0; j < 4; j++) {
1530
- int wr = c_sh_wr + 8 * j;
1531
- write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
1532
- frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
1533
- write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
1534
- frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
1535
- write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
1536
- frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
1537
- write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
1538
- frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
1539
- }
1540
- c_sh_wr += 16 * (4 * c_sh_stride);
1541
- }
1542
- }
1543
- __syncthreads();
1544
-
1545
- #pragma unroll
1546
- for (int i = 0;
1547
- i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
1548
- i++) {
1549
- if (c_gl_wr < c_gl_wr_end) {
1550
- if (use_atomic_add && slice_count > 1) {
1551
- scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[c_gl_wr]);
1552
- scalar_t2* sh_red_half2 =
1553
- reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
1554
- #pragma unroll
1555
- for (int a = 0; a < 4; a++) {
1556
- atomicAdd(&C_half2[a], sh_red_half2[a]);
1557
- }
1558
- } else {
1559
- C[c_gl_wr] = sh_red[c_sh_rd];
1560
- }
1561
- c_gl_wr += c_gl_wr_delta;
1562
- c_sh_rd += c_sh_rd_delta;
1563
- }
1564
- }
1565
- };
1566
-
1567
- // Start global fetch and register load pipelines.
1568
- auto start_pipes = [&]() {
1569
-
1570
- #pragma unroll
1571
- for (int i = 0; i < stages - 1; i++) {
1572
- if (has_act_order && i == 0) {
1573
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1574
- if (last_g_idx >= prob_k) {
1575
- last_g_idx = prob_k - 1;
1576
- }
1577
- fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]);
1578
- }
1579
-
1580
- if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1581
- if (i == 0) {
1582
- fetch_zp_to_shared();
1583
- }
1584
- }
1585
- fetch_to_shared(i, i, i < slice_iters);
1586
- }
1587
-
1588
- zero_accums();
1589
- wait_for_stage();
1590
- init_same_group(0);
1591
- fetch_to_registers(0, 0);
1592
- fetch_scales_to_registers(0, 0);
1593
- fetch_zp_to_registers(0, 0);
1594
- a_gl_rd += a_gl_rd_delta_o * (stages - 1);
1595
- slice_k_start_shared_fetch += tb_k * (stages - 1);
1596
- };
1597
- if (slice_iters) {
1598
- start_pipes();
1599
- }
1600
-
1601
- // Main loop.
1602
- while (slice_iters) {
1603
- // We unroll over both the global fetch and the register load pipeline to
1604
- // ensure all shared memory accesses are static. Note that both pipelines
1605
- // have even length meaning that the next iteration will always start at
1606
- // index 0.
1607
-
1608
- #pragma unroll
1609
- for (int pipe = 0; pipe < stages;) {
1610
- #pragma unroll
1611
- for (int k = 0; k < b_sh_wr_iters; k++) {
1612
- fetch_to_registers(k + 1, pipe % stages);
1613
- fetch_scales_to_registers(k + 1, pipe);
1614
- fetch_zp_to_registers(k + 1, pipe);
1615
- if (k == b_sh_wr_iters - 2) {
1616
- fetch_to_shared((pipe + stages - 1) % stages, pipe,
1617
- slice_iters >= stages);
1618
- pipe++;
1619
- wait_for_stage();
1620
- init_same_group(pipe % stages);
1621
- }
1622
- matmul(k);
1623
- }
1624
- slice_iters--;
1625
- if (slice_iters == 0) {
1626
- break;
1627
- }
1628
- }
1629
-
1630
- a_gl_rd += a_gl_rd_delta_o * stages;
1631
- slice_k_start += tb_k * stages;
1632
- slice_k_start_shared_fetch += tb_k * stages;
1633
-
1634
- if constexpr (has_act_order) {
1635
- int first_group_id = g_idx[slice_k_start];
1636
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1637
- if (last_g_idx >= prob_k) {
1638
- last_g_idx = prob_k - 1;
1639
- }
1640
- int last_group_id = g_idx[last_g_idx];
1641
- if (last_group_id >= sh_first_group_id + sh_num_groups) {
1642
- fetch_scales_to_shared(false, first_group_id, last_group_id);
1643
- __syncthreads();
1644
- }
1645
- }
1646
-
1647
- // Process results and, if necessary, proceed to the next column slice.
1648
- // While this pattern may not be the most readable, other ways of writing
1649
- // the loop seemed to noticeably worse performance after compilation.
1650
- if (slice_iters == 0) {
1651
- cp_async_wait<0>();
1652
- bool last = slice_idx == slice_count - 1;
1653
- // For per-column scales, we only fetch them here in the final step before
1654
- // write-out
1655
- if constexpr (!has_act_order && group_blocks == -1) {
1656
- if constexpr (w_type.size_bits() == 8) {
1657
- if (s_sh_wr_pred) {
1658
- cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
1659
- }
1660
- cp_async_fence();
1661
- } else {
1662
- if (last || use_atomic_add) {
1663
- if (s_sh_wr_pred) {
1664
- cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
1665
- }
1666
- cp_async_fence();
1667
- }
1668
- }
1669
- }
1670
-
1671
- thread_block_reduce();
1672
- if constexpr (!has_act_order && group_blocks == -1) {
1673
- if constexpr (w_type.size_bits() == 8) {
1674
- cp_async_wait<0>();
1675
- __syncthreads();
1676
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1677
- reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
1678
- reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1679
- }
1680
-
1681
- } else {
1682
- if (last || use_atomic_add) {
1683
- cp_async_wait<0>();
1684
- __syncthreads();
1685
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1686
- reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
1687
- reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1688
- }
1689
- }
1690
- }
1691
- }
1692
-
1693
- // For 8-bit channelwise, we apply the scale before the global reduction
1694
- // that converts the fp32 results to fp16 (so that we avoid possible
1695
- // overflow in fp16)
1696
- if constexpr (!has_act_order && group_blocks == -1 &&
1697
- w_type.size_bits() == 8) {
1698
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1699
- #pragma unroll
1700
- for (int i = 0; i < thread_m_blocks; i++) {
1701
- #pragma unroll
1702
- for (int j = 0; j < 4; j++) {
1703
- scale_float<scalar_t>(
1704
- reinterpret_cast<float*>(&frag_c[i][j][0][0]),
1705
- frag_s[j / 2][2 * (j % 2) + 0]);
1706
- scale_float<scalar_t>(
1707
- reinterpret_cast<float*>(&frag_c[i][j][0][2]),
1708
- frag_s[j / 2][2 * (j % 2) + 0]);
1709
-
1710
- scale_float<scalar_t>(
1711
- reinterpret_cast<float*>(&frag_c[i][j][1][0]),
1712
- frag_s[j / 2][2 * (j % 2) + 1]);
1713
- scale_float<scalar_t>(
1714
- reinterpret_cast<float*>(&frag_c[i][j][1][2]),
1715
- frag_s[j / 2][2 * (j % 2) + 1]);
1716
- }
1717
- }
1718
- }
1719
- }
1720
-
1721
- if (slice_count > 1 && !use_atomic_add) {
1722
- // only globally reduce if there is more than one block in a slice
1723
- barrier_acquire(&locks[slice_col], slice_idx);
1724
- if (use_fp32_reduce) {
1725
- global_reduce_fp32(slice_idx == 0, last);
1726
- } else {
1727
- global_reduce_fp16(slice_idx == 0, last);
1728
- }
1729
- barrier_release(&locks[slice_col], last);
1730
- }
1731
- if (last || use_atomic_add)
1732
- // only the last block in a slice actuallywrites the result
1733
- write_result();
1734
- slice_row = 0;
1735
- slice_col_par++;
1736
- slice_col++;
1737
- init_slice();
1738
- if (slice_iters) {
1739
- a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
1740
- (threadIdx.x % a_gl_rd_delta_o);
1741
- #pragma unroll
1742
- for (int i = 0; i < b_sh_wr_iters; i++)
1743
- B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
1744
- if (slice_col == 0) {
1745
- #pragma unroll
1746
- for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1747
- }
1748
-
1749
- // Update slice k/n for scales loading
1750
- if constexpr (has_act_order) {
1751
- slice_k_start = tb_k * slice_row;
1752
- slice_k_finish = slice_k_start + tb_k * slice_iters;
1753
- slice_k_start_shared_fetch = slice_k_start;
1754
- slice_n_offset = act_s_col_tb_stride * slice_col;
1755
-
1756
- } else {
1757
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1758
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1759
- }
1760
-
1761
- start_pipes();
1762
- }
1763
- }
1764
- }
1765
-}
1766
-
1767
- #define __CALL_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
1768
- HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, NUM_THREADS, \
1769
- IS_ZP_FLOAT) \
1770
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
1771
- thread_n_blocks == THREAD_N_BLOCKS && \
1772
- thread_k_blocks == THREAD_K_BLOCKS && \
1773
- has_act_order == HAS_ACT_ORDER && has_zp == HAS_ZP && \
1774
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
1775
- is_zp_float == IS_ZP_FLOAT) { \
1776
- if constexpr (!IS_ZP_FLOAT || std::is_same<scalar_t, half>::value) { \
1777
- cudaFuncSetAttribute( \
1778
- Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
1779
- THREAD_N_BLOCKS, THREAD_K_BLOCKS, pipe_stages, \
1780
- HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, IS_ZP_FLOAT>, \
1781
- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
1782
- Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
1783
- THREAD_N_BLOCKS, THREAD_K_BLOCKS, pipe_stages, HAS_ACT_ORDER, \
1784
- HAS_ZP, GROUP_BLOCKS, IS_ZP_FLOAT> \
1785
- <<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \
1786
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, \
1787
- num_groups, prob_m, prob_n, prob_k, lda, locks, \
1788
- part_use_atomic_add, use_fp32_reduce); \
1789
- } \
1790
- }
1791
-
121
typedef struct {
122
int thread_k;
123
int thread_n;
124
int num_threads;
125
} thread_config_t;
126
1798
-typedef struct {
1799
- int max_m_blocks;
1800
- thread_config_t tb_cfg;
1801
-} exec_config_t;
1802
-
127
thread_config_t small_batch_thread_configs[] = {
128
// Ordered by priority
129
130
// thread_k, thread_n, num_threads
131
{128, 128, 256},
132
{64, 128, 128},
1809
- {128, 64, 128},
1810
-};
133
+ {128, 64, 128}};
134
135
thread_config_t large_batch_thread_configs[] = {
136
// Ordered by priority
@@ -1815,9 +138,12 @@ thread_config_t large_batch_thread_configs[] = {
138
// thread_k, thread_n, num_threads
139
{64, 256, 256},
140
{64, 128, 128},
1818
- {128, 64, 128},
141
+ {128, 64, 128}};
142
1820
-};
143
+typedef struct {
144
+ int blocks_per_sm;
145
+ thread_config_t tb_cfg;
146
+} exec_config_t;
147
148
int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
149
int prob_n, int prob_k, int num_bits, int group_size,
@@ -1842,7 +168,6 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
168
tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K
169
load_groups = max(load_groups, 32); // We load at least 32 scale groups
170
return load_groups * tb_n * 2;
1845
-
171
} else {
172
int tb_scales = tb_groups * tb_n * 2;
173
@@ -1850,49 +175,43 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
175
}
176
}
177
1853
-bool is_valid_cache_size(thread_config_t const& th_config, int max_m_blocks,
1854
- int prob_m, int prob_n, int prob_k, int num_bits,
1855
- int scales_cache_size, int max_shared_mem) {
178
+int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
179
+ int prob_m, int prob_n, int prob_k, int num_bits,
180
+ int group_size, bool has_act_order, bool is_k_full,
181
+ int has_zp, int is_zp_float) {
182
int pack_factor = 32 / num_bits;
183
184
// Get B size
185
int tb_k = th_config.thread_k;
186
int tb_n = th_config.thread_n;
1861
-
1862
- int b_size = (tb_k * tb_n / pack_factor) * 4;
1863
-
1864
- // Get A size
1865
- int m_blocks = div_ceil(prob_m, 16);
1866
- int tb_max_m = 16;
1867
-
1868
- while (true) {
1869
- if (m_blocks >= max_m_blocks) {
1870
- tb_max_m *= max_m_blocks;
1871
- break;
1872
- }
1873
-
1874
- max_m_blocks--;
1875
- if (max_m_blocks == 0) {
1876
- TORCH_CHECK(false, "Unexpected m_blocks = ", m_blocks);
1877
- }
187
+ int tb_m = thread_m_blocks * 16;
188
+ int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
189
+ int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
190
+ int sh_red_size = tb_m * (tb_n + 8);
191
+ int sh_s_size =
192
+ get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
193
+ group_size, has_act_order, is_k_full);
194
+ int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0;
195
+ int sh_zp_size = 0;
196
+ if (has_zp) {
197
+ if (is_zp_float)
198
+ sh_zp_size = sh_s_size;
199
+ else if (num_bits == 4)
200
+ sh_zp_size = sh_s_size / 4;
201
+ else if (num_bits == 8)
202
+ sh_zp_size = sh_s_size / 2;
203
}
204
1880
- int a_size = (tb_max_m * tb_k) * 2;
1881
-
1882
- float pipe_size = (a_size + b_size) * pipe_stages;
1883
-
1884
- float reduce_size = max(th_config.num_threads * 32 * 4,
1885
- (tb_n / 64) * 32 * (tb_max_m / 16) * 4 * 2 * 4 * 2);
205
+ int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size +
206
+ sh_zp_size + sh_g_idx_size;
207
1887
- TORCH_CHECK(max_shared_mem / 2 > scales_cache_size); // Sanity
1888
-
1889
- return pipe_size + reduce_size < 0.95f * (max_shared_mem - scales_cache_size);
208
+ return total_size;
209
}
210
1892
-bool is_valid_config(thread_config_t const& th_config, int max_m_blocks,
211
+bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
212
int prob_m, int prob_n, int prob_k, int num_bits,
213
int group_size, bool has_act_order, bool is_k_full,
1895
- int max_shared_mem) {
214
+ int has_zp, int is_zp_float, int max_shared_mem) {
215
// Sanity
216
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
217
th_config.num_threads == -1) {
@@ -1914,163 +233,204 @@ bool is_valid_config(thread_config_t const& th_config, int max_m_blocks,
233
return false;
234
}
235
1917
- // Determine cache for scales
1918
- int scales_cache_size =
1919
- get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
1920
- group_size, has_act_order, is_k_full);
1921
-
236
// Check that pipeline fits into cache
1923
- if (!is_valid_cache_size(th_config, max_m_blocks, prob_m, prob_n, prob_k,
1924
- num_bits, scales_cache_size, max_shared_mem)) {
1925
- return false;
1926
- }
1927
-
1928
- return true;
237
+ int cache_size = get_kernel_cache_size(
238
+ th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
239
+ has_act_order, is_k_full, has_zp, is_zp_float);
240
+ return cache_size <= max_shared_mem;
241
}
242
1931
-int determine_reduce_max_m(int prob_m, int max_par) {
1932
- constexpr int tile_m_size = 16;
243
+ #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
244
+ M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
245
+ else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
246
+ thread_n_blocks == THREAD_N_BLOCKS && \
247
+ thread_k_blocks == THREAD_K_BLOCKS && \
248
+ m_block_size_8 == M_BLOCK_SIZE_8 && \
249
+ group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
250
+ is_zp_float == IS_ZP_FLOAT) { \
251
+ kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
252
+ THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
253
+ pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
254
+ }
255
+
256
+ // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
257
+ // this is the most common cases
258
+ // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
259
+ // FZP: cases for float-zero-point (is_zp_float = true)
260
+ // ACT: cases for act order case (group_blocks == 0)
261
+ #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
262
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
263
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
264
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \
265
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
266
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
267
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
268
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
269
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
270
+
271
+ #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
272
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
273
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
274
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
275
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
276
+ \
277
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
278
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
279
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
280
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
281
+ \
282
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
283
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
284
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
285
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
286
+
287
+ #define COMMON_GET_IF(W_TYPE) \
288
+ COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \
289
+ COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \
290
+ COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \
291
+ COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
292
+ COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \
293
+ COMMON_GET_IF_M234(W_TYPE, 4, 8, 128)
294
+
295
+ #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
296
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
297
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
298
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
299
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
300
+
301
+ #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
302
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
303
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
304
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
305
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
306
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
307
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
308
+
309
+ #define BIGGROUP_GET_IF(W_TYPE) \
310
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
311
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
312
+ BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \
313
+ BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
314
+ BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \
315
+ BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
316
+
317
+ // We currently have 4-bit models only with group_blocks == 4
318
+ #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
319
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
320
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
321
+
322
+ #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
323
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
324
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
325
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
326
+
327
+ #define FZP_GET_IF(W_TYPE) \
328
+ FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \
329
+ FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \
330
+ FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \
331
+ FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
332
+ FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \
333
+ FZP_GET_IF_M234(W_TYPE, 4, 8, 128)
334
+
335
+ // We currently have 4-bit models only with group_blocks == 4
336
+ #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
337
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
338
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
339
+
340
+ #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
341
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
342
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
343
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
344
+
345
+ #define ACT_GET_IF(W_TYPE) \
346
+ ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \
347
+ ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \
348
+ ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \
349
+ ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
350
+ ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \
351
+ ACT_GET_IF_M234(W_TYPE, 4, 8, 128)
352
1934
- if (prob_m <= tile_m_size) {
1935
- return tile_m_size;
353
+template <typename scalar_t>
354
+MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
355
+ int thread_m_blocks, int thread_n_blocks,
356
+ int thread_k_blocks, bool m_block_size_8,
357
+ bool has_act_order, bool has_zp,
358
+ int group_blocks, int num_threads,
359
+ bool is_zp_float) {
360
+ int num_bits = q_type.size_bits();
361
+ auto kernel = MarlinDefault;
362
+ if (false) {
363
+ }
364
1937
- } else if (prob_m <= tile_m_size * 2) {
1938
- return tile_m_size * 2;
365
+ COMMON_GET_IF(vllm::kU4)
366
+ COMMON_GET_IF(vllm::kU4B8)
367
+ COMMON_GET_IF(vllm::kU8B128)
368
1940
- } else if (prob_m <= tile_m_size * 3) {
1941
- return tile_m_size * 3;
369
+ BIGGROUP_GET_IF(vllm::kFE4M3fn)
370
1943
- } else if (prob_m <= tile_m_size * 4) {
1944
- return tile_m_size * 4;
371
+ ACT_GET_IF(vllm::kU4B8)
372
+ ACT_GET_IF(vllm::kU8B128)
373
1946
- } else {
1947
- int cur_par = min(div_ceil(prob_m, tile_m_size * 4), max_par);
1948
- return tile_m_size * 4 * cur_par;
374
+ if (std::is_same<scalar_t, half>::value) {
375
+ if (false) {
376
+ }
377
+ FZP_GET_IF(vllm::kU4)
378
}
379
+
380
+ return kernel;
381
}
382
1952
-exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k,
1953
- int num_bits, int group_size,
1954
- bool has_act_order, bool is_k_full,
1955
- int max_shared_mem) {
1956
- int max_m_blocks = 4;
1957
- while (max_m_blocks > 0) {
1958
- if (prob_m <= 16) {
1959
- for (auto th_config : small_batch_thread_configs) {
1960
- if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k,
1961
- num_bits, group_size, has_act_order, is_k_full,
1962
- max_shared_mem)) {
1963
- return exec_config_t{max_m_blocks, th_config};
1964
- }
1965
- }
1966
- } else {
1967
- for (auto th_config : large_batch_thread_configs) {
1968
- if (is_valid_config(th_config, max_m_blocks, prob_m, prob_n, prob_k,
1969
- num_bits, group_size, has_act_order, is_k_full,
1970
- max_shared_mem)) {
1971
- return exec_config_t{max_m_blocks, th_config};
1972
- }
1973
- }
383
+template <typename scalar_t>
384
+exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
385
+ int prob_n, int prob_k, int thread_m_blocks,
386
+ bool m_block_size_8, int num_bits,
387
+ int group_size, bool has_act_order,
388
+ bool is_k_full, bool has_zp,
389
+ bool is_zp_float, int max_shared_mem,
390
+ int sms) {
391
+ exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
392
+ thread_config_t* thread_configs = thread_m_blocks > 1
393
+ ? large_batch_thread_configs
394
+ : small_batch_thread_configs;
395
+ int thread_configs_size =
396
+ thread_m_blocks > 1
397
+ ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t)
398
+ : sizeof(small_batch_thread_configs) / sizeof(thread_config_t);
399
+
400
+ for (int i = 0; i < thread_configs_size; i++) {
401
+ thread_config_t th_config = thread_configs[i];
402
+
403
+ if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
404
+ num_bits, group_size, has_act_order, is_k_full, has_zp,
405
+ is_zp_float, max_shared_mem)) {
406
+ continue;
407
}
408
1976
- max_m_blocks--; // Process less M blocks per invocation to reduce cache
1977
- // usage
1978
- }
409
+ int cache_size = get_kernel_cache_size(
410
+ th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits,
411
+ group_size, has_act_order, is_k_full, has_zp, is_zp_float);
412
1980
- return exec_config_t{0, {-1, -1, -1}};
1981
-}
413
+ int group_blocks = 0;
414
+ if (!has_act_order) {
415
+ group_blocks = group_size == -1 ? -1 : group_size / 16;
416
+ }
417
1983
- #define GPTQ_CALL_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
1984
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, 0, NUM_THREADS, \
1985
- false) \
1986
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, true, false, 0, NUM_THREADS, \
1987
- false) \
1988
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, true, false, 0, NUM_THREADS, \
1989
- false) \
1990
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, true, false, 0, NUM_THREADS, \
1991
- false) \
1992
- \
1993
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, -1, NUM_THREADS, \
1994
- false) \
1995
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, 2, NUM_THREADS, \
1996
- false) \
1997
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, 4, NUM_THREADS, \
1998
- false) \
1999
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, 8, NUM_THREADS, \
2000
- false) \
2001
- \
2002
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, -1, NUM_THREADS, \
2003
- false) \
2004
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, 2, NUM_THREADS, \
2005
- false) \
2006
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, 4, NUM_THREADS, \
2007
- false) \
2008
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, 8, NUM_THREADS, \
2009
- false) \
2010
- \
2011
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, -1, NUM_THREADS, \
2012
- false) \
2013
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, 2, NUM_THREADS, \
2014
- false) \
2015
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, 4, NUM_THREADS, \
2016
- false) \
2017
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, 8, NUM_THREADS, \
2018
- false) \
2019
- \
2020
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, -1, NUM_THREADS, \
2021
- false) \
2022
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, 2, NUM_THREADS, \
2023
- false) \
2024
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, 4, NUM_THREADS, \
2025
- false) \
2026
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, 8, NUM_THREADS, \
2027
- false)
2028
-
2029
- #define AWQ_CALL_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
2030
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, -1, NUM_THREADS, \
2031
- false) \
2032
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, 2, NUM_THREADS, \
2033
- false) \
2034
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2035
- false) \
2036
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, 8, NUM_THREADS, \
2037
- false) \
2038
- \
2039
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, -1, NUM_THREADS, \
2040
- false) \
2041
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, 2, NUM_THREADS, \
2042
- false) \
2043
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2044
- false) \
2045
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, 8, NUM_THREADS, \
2046
- false) \
2047
- \
2048
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, -1, NUM_THREADS, \
2049
- false) \
2050
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, 2, NUM_THREADS, \
2051
- false) \
2052
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2053
- false) \
2054
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, 8, NUM_THREADS, \
2055
- false) \
2056
- \
2057
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, -1, NUM_THREADS, \
2058
- false) \
2059
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, 2, NUM_THREADS, \
2060
- false) \
2061
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2062
- false) \
2063
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, 8, NUM_THREADS, false)
418
+ auto kernel = get_marlin_kernel<scalar_t>(
419
+ q_type, thread_m_blocks, th_config.thread_n / 16,
420
+ th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
421
+ group_blocks, th_config.num_threads, is_zp_float);
422
2065
- // We currently have 4-bit models only with group_blocks == 4
2066
- #define HQQ_CALL_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
2067
- __CALL_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2068
- true) \
2069
- __CALL_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2070
- true) \
2071
- __CALL_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, \
2072
- true) \
2073
- __CALL_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS, true)
423
+ if (kernel == MarlinDefault) continue;
424
+
425
+ // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16);
426
+ // int n_tiles = prob_n / th_config.thread_n;
427
+ // int k_tiles = prob_k / th_config.thread_k;
428
+
429
+ return {1, th_config};
430
+ }
431
+
432
+ return exec_cfg;
433
+}
434
435
template <typename scalar_t>
436
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
@@ -2078,78 +438,24 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
438
int prob_n, int prob_k, int lda, void* workspace,
439
vllm::ScalarType const& q_type, bool has_act_order,
440
bool is_k_full, bool has_zp, int num_groups, int group_size,
2081
- int dev, cudaStream_t stream, int thread_k, int thread_n,
2082
- int sms, int max_par, bool use_atomic_add, bool use_fp32_reduce,
2083
- bool is_zp_float) {
441
+ int dev, cudaStream_t stream, int thread_k_init,
442
+ int thread_n_init, int sms, bool use_atomic_add,
443
+ bool use_fp32_reduce, bool is_zp_float) {
444
if (has_zp) {
445
TORCH_CHECK(
446
q_type == vllm::kU4 || q_type == vllm::kU8,
447
"q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
448
} else {
2089
- TORCH_CHECK(
2090
- q_type == vllm::kU4B8 || q_type == vllm::kU8B128,
2091
- "q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
2092
- q_type.str());
449
+ TORCH_CHECK(q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
450
+ q_type == vllm::kFE4M3fn,
451
+ "q_type must be uint4b8, uint8b128 or float8_e4m3fn when "
452
+ "has_zp = False. Got = ",
453
+ q_type.str());
454
}
455
456
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
457
", ", prob_n, ", ", prob_k, "]");
458
2098
- // TODO: remove alias when we start supporting other 8bit types
2099
- int num_bits = q_type.size_bits();
2100
- int tot_m = prob_m;
2101
- int tot_m_blocks = div_ceil(tot_m, 16);
2102
- int pad = 16 * tot_m_blocks - tot_m;
2103
-
2104
- if (sms == -1) {
2105
- cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
2106
- }
2107
-
2108
- int max_shared_mem = 0;
2109
- cudaDeviceGetAttribute(&max_shared_mem,
2110
- cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
2111
- TORCH_CHECK(max_shared_mem > 0);
2112
-
2113
- // Set thread config
2114
- exec_config_t exec_cfg;
2115
- if (thread_k != -1 && thread_n != -1) {
2116
- // User-defined config
2117
- exec_cfg =
2118
- exec_config_t{4, thread_config_t{thread_k, thread_n, default_threads}};
2119
- } else {
2120
- // Auto config
2121
- exec_cfg =
2122
- determine_thread_config(prob_m, prob_n, prob_k, num_bits, group_size,
2123
- has_act_order, is_k_full, max_shared_mem);
2124
- }
2125
-
2126
- TORCH_CHECK(exec_cfg.max_m_blocks > 0 &&
2127
- is_valid_config(exec_cfg.tb_cfg, exec_cfg.max_m_blocks,
2128
- prob_m, prob_n, prob_k, num_bits, group_size,
2129
- has_act_order, is_k_full, max_shared_mem),
2130
- "Invalid thread config: max_m_blocks = ", exec_cfg.max_m_blocks,
2131
- ", thread_k = ", exec_cfg.tb_cfg.thread_k,
2132
- ", thread_n = ", exec_cfg.tb_cfg.thread_n,
2133
- ", num_threads = ", exec_cfg.tb_cfg.num_threads, " for MKN = [",
2134
- prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
2135
- ", group_size = ", group_size,
2136
- ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
2137
- ", max_shared_mem = ", max_shared_mem);
2138
-
2139
- int num_threads = exec_cfg.tb_cfg.num_threads;
2140
- thread_k = exec_cfg.tb_cfg.thread_k;
2141
- thread_n = exec_cfg.tb_cfg.thread_n;
2142
-
2143
- int thread_k_blocks = thread_k / 16;
2144
- int thread_n_blocks = thread_n / 16;
2145
-
2146
- int blocks = sms;
2147
-
2148
- TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
2149
- " is not divisible by thread_n = ", thread_n);
2150
- TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
2151
- " is not divisible by thread_k = ", thread_k);
2152
-
459
int group_blocks = 0;
460
if (has_act_order) {
461
if (is_k_full) {
@@ -2161,7 +467,6 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
467
TORCH_CHECK(group_size == 0);
468
group_blocks = 0;
469
}
2164
-
470
} else {
471
if (group_size == -1) {
472
group_blocks = -1;
@@ -2172,6 +477,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
477
}
478
}
479
480
+ int num_bits = q_type.size_bits();
481
const int4* A_ptr = (const int4*)A;
482
const int4* B_ptr = (const int4*)B;
483
int4* C_ptr = (int4*)C;
@@ -2186,106 +492,138 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
492
493
if (has_act_order) {
494
// Permute A columns
2189
- int block_rows = div_ceil(prob_m, blocks);
2190
- permute_cols_kernel<<<blocks, default_threads, 0, stream>>>(
495
+ int block_rows = div_ceil(prob_m, sms);
496
+ // avoid ">>>" being formatted to "> > >"
497
+ // clang-format off
498
+ permute_cols_kernel<<<sms, default_threads, 0, stream>>>(
499
A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows);
500
+ // clang-format on
501
A_ptr = a_tmp_ptr;
502
lda = prob_k;
2194
- }
503
2196
- // If we have a full K, then we can run the non-act-order version of Marlin
2197
- // (since the weight rows are reordered by increasing group ids, and by having
2198
- // a full K, we have full original groups)
2199
- if (is_k_full) {
2200
- has_act_order = false;
504
+ // If we have a full K, then we can run the non-act-order version of Marlin
505
+ // (since the weight rows are reordered by increasing group ids, and by
506
+ // having a full K, we have full original groups)
507
+ if (is_k_full) has_act_order = false;
508
}
509
2203
- // Main loop
2204
- for (int i = 0; i < tot_m_blocks; i += exec_cfg.max_m_blocks) {
2205
- int thread_m_blocks = tot_m_blocks - i;
2206
- prob_m = tot_m - 16 * i;
2207
- int par = 1;
2208
- if (thread_m_blocks > exec_cfg.max_m_blocks) {
2209
- // Note that parallel > 1 currently only works for inputs without any
2210
- // padding
2211
- par = (16 * thread_m_blocks - pad) / (16 * exec_cfg.max_m_blocks);
2212
- if (par > max_par) par = max_par;
2213
- prob_m = (16 * exec_cfg.max_m_blocks) * par;
2214
- i += exec_cfg.max_m_blocks * (par - 1);
2215
- thread_m_blocks = exec_cfg.max_m_blocks;
510
+ int max_shared_mem = 0;
511
+ cudaDeviceGetAttribute(&max_shared_mem,
512
+ cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
513
+ TORCH_CHECK(max_shared_mem > 0);
514
+
515
+ int max_par = 16;
516
+ if (prob_n <= 4096) max_par = 16 * 8;
517
+ int max_shared_mem_new = max_shared_mem;
518
+ int rest_m = prob_m;
519
+ int max_thread_m_blocks = 4;
520
+ while (rest_m) {
521
+ int par_count = rest_m / (max_thread_m_blocks * 16);
522
+ if (par_count > max_par) par_count = max_par;
523
+ int prob_m_split =
524
+ par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m;
525
+
526
+ int thread_k = thread_k_init;
527
+ int thread_n = thread_n_init;
528
+
529
+ int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks);
530
+ int m_block_size_8 = prob_m_split <= 8;
531
+
532
+ // Set thread config
533
+ exec_config_t exec_cfg;
534
+ thread_config_t thread_tfg;
535
+ if (thread_k != -1 && thread_n != -1) {
536
+ thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
537
+ exec_cfg = exec_config_t{1, thread_tfg};
538
+ TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
539
+ " is not divisible by thread_n = ", thread_n);
540
+ TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
541
+ " is not divisible by thread_k = ", thread_k);
542
+ } else {
543
+ // Auto config
544
+ exec_cfg = determine_exec_config<scalar_t>(
545
+ q_type, prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8,
546
+ num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
547
+ max_shared_mem, sms);
548
+ thread_tfg = exec_cfg.tb_cfg;
549
+ if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) {
550
+ max_thread_m_blocks--;
551
+ continue;
552
+ }
553
}
554
2218
- // atomic add reduce have better performance only when m * n is small
2219
- bool part_use_atomic_add =
2220
- use_atomic_add && div_ceil(prob_m, 64) * prob_n <= 2048;
555
+ int num_threads = thread_tfg.num_threads;
556
+ thread_k = thread_tfg.thread_k;
557
+ thread_n = thread_tfg.thread_n;
558
+ int blocks = sms * exec_cfg.blocks_per_sm;
559
+ if (exec_cfg.blocks_per_sm > 1)
560
+ max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024;
561
2222
- if (false) {
2223
- }
2224
- GPTQ_CALL_IF(vllm::kU4B8, 16, 4, 256)
2225
- GPTQ_CALL_IF(vllm::kU4B8, 8, 8, 256)
2226
- GPTQ_CALL_IF(vllm::kU4B8, 8, 4, 128)
2227
- GPTQ_CALL_IF(vllm::kU4B8, 4, 8, 128)
2228
- GPTQ_CALL_IF(vllm::kU8B128, 16, 4, 256)
2229
- GPTQ_CALL_IF(vllm::kU8B128, 8, 8, 256)
2230
- GPTQ_CALL_IF(vllm::kU8B128, 8, 4, 128)
2231
- GPTQ_CALL_IF(vllm::kU8B128, 4, 8, 128)
2232
-
2233
- AWQ_CALL_IF(vllm::kU4, 16, 4, 256)
2234
- AWQ_CALL_IF(vllm::kU4, 8, 8, 256)
2235
- AWQ_CALL_IF(vllm::kU4, 8, 4, 128)
2236
- AWQ_CALL_IF(vllm::kU4, 4, 8, 128)
2237
- AWQ_CALL_IF(vllm::kU8, 16, 4, 256)
2238
- AWQ_CALL_IF(vllm::kU8, 8, 8, 256)
2239
- AWQ_CALL_IF(vllm::kU8, 8, 4, 128)
2240
- AWQ_CALL_IF(vllm::kU8, 4, 8, 128)
2241
-
2242
- HQQ_CALL_IF(vllm::kU4, 16, 4, 256)
2243
- HQQ_CALL_IF(vllm::kU4, 8, 8, 256)
2244
- HQQ_CALL_IF(vllm::kU4, 8, 4, 128)
2245
- HQQ_CALL_IF(vllm::kU4, 4, 8, 128)
2246
- else {
562
+ int thread_k_blocks = thread_k / 16;
563
+ int thread_n_blocks = thread_n / 16;
564
+
565
+ TORCH_CHECK(
566
+ is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n,
567
+ prob_k, num_bits, group_size, has_act_order, is_k_full,
568
+ has_zp, is_zp_float, max_shared_mem_new),
569
+ "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
570
+ ", thread_k = ", thread_tfg.thread_k,
571
+ ", thread_n = ", thread_tfg.thread_n,
572
+ ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m,
573
+ ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
574
+ ", prob_m_split = ", prob_m_split, ", group_size = ", group_size,
575
+ ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
576
+ ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
577
+ ", max_shared_mem_new = ", max_shared_mem_new);
578
+
579
+ auto kernel = get_marlin_kernel<scalar_t>(
580
+ q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks,
581
+ m_block_size_8, has_act_order, has_zp, group_blocks, num_threads,
582
+ is_zp_float);
583
+
584
+ if (kernel == MarlinDefault) {
585
TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
586
", ", prob_k, "]", ", has_act_order = ", has_act_order,
587
", num_groups = ", num_groups, ", group_size = ", group_size,
588
+ ", prob_m_split = ", prob_m_split,
589
", thread_m_blocks = ", thread_m_blocks,
590
", thread_n_blocks = ", thread_n_blocks,
591
", thread_k_blocks = ", thread_k_blocks,
2253
- ", num_bits = ", num_bits);
592
+ ", num_threads = ", num_threads, ", num_bits = ", num_bits);
593
}
594
2256
- A_ptr += 16 * thread_m_blocks * (lda / 8) * par;
2257
- C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par;
595
+ cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
596
+ max_shared_mem_new);
597
+
598
+ bool part_use_atomic_add =
599
+ use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048;
600
+
601
+ // avoid ">>>" being formatted to "> > >"
602
+ // clang-format off
603
+ kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
604
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, num_groups,
605
+ prob_m_split, prob_n, prob_k, lda, locks, part_use_atomic_add,
606
+ use_fp32_reduce, max_shared_mem_new);
607
+ // clang-format on
608
+
609
+ A_ptr += prob_m_split * (lda / 8);
610
+ C_ptr += prob_m_split * (prob_n / 8);
611
+ rest_m -= prob_m_split;
612
}
613
}
614
615
} // namespace marlin
616
2263
-torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
2264
- torch::Tensor& b_scales, torch::Tensor& b_zeros,
2265
- torch::Tensor& g_idx, torch::Tensor& perm,
2266
- torch::Tensor& workspace,
2267
- vllm::ScalarTypeId const& b_q_type_id,
2268
- int64_t size_m, int64_t size_n, int64_t size_k,
2269
- bool is_k_full, bool has_zp, bool use_atomic_add,
2270
- bool use_fp32_reduce, bool is_zp_float) {
617
+torch::Tensor gptq_marlin_gemm(
618
+ torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
619
+ torch::Tensor& b_q_weight, torch::Tensor& b_scales,
620
+ std::optional<torch::Tensor> const& b_zeros_or_none,
621
+ std::optional<torch::Tensor> const& g_idx_or_none,
622
+ std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
623
+ vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
624
+ int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
625
+ bool is_zp_float) {
626
vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
2272
- if (has_zp) {
2273
- TORCH_CHECK(
2274
- b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
2275
- "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
2276
- } else {
2277
- TORCH_CHECK(
2278
- b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128,
2279
- "b_q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
2280
- b_q_type.str());
2281
- }
2282
-
2283
- if (has_zp && is_zp_float) {
2284
- TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
2285
- "Computation type must be float16 (half) when using float zero "
2286
- "points.");
2287
- }
2288
-
627
int pack_factor = 32 / b_q_type.size_bits();
628
629
// Verify A
@@ -2295,15 +633,19 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
633
", size_k = ", size_k);
634
635
// Verify B
2298
- TORCH_CHECK(size_k % marlin::tile_size == 0, "size_k = ", size_k,
2299
- " is not divisible by tile_size = ", marlin::tile_size);
2300
- TORCH_CHECK((size_k / marlin::tile_size) == b_q_weight.size(0),
636
+ TORCH_CHECK(
637
+ size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
638
+ " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
639
+ TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0),
640
"Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
2302
- ", size_k = ", size_k, ", tile_size = ", marlin::tile_size);
2303
- TORCH_CHECK(b_q_weight.size(1) % marlin::tile_size == 0,
2304
- "b_q_weight.size(1) = ", b_q_weight.size(1),
2305
- " is not divisible by tile_size = ", marlin::tile_size);
2306
- int actual_size_n = (b_q_weight.size(1) / marlin::tile_size) * pack_factor;
641
+ ", size_k = ", size_k,
642
+ ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
643
+ TORCH_CHECK(
644
+ b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0,
645
+ "b_q_weight.size(1) = ", b_q_weight.size(1),
646
+ " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
647
+ int actual_size_n =
648
+ (b_q_weight.size(1) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
649
TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
650
", actual_size_n = ", actual_size_n);
651
@@ -2320,63 +662,47 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
662
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
663
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
664
2323
- TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
2324
- TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
2325
-
2326
- TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
2327
- TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
2328
-
2329
- TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
2330
- TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
665
+ // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
666
+ // auto -1)
667
+ int thread_k = -1;
668
+ // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
669
+ // auto -1)
670
+ int thread_n = -1;
671
+ // sms: number of SMs to use for the kernel
672
+ int sms = -1;
673
+ cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
674
675
// Alloc buffers
676
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
677
auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
678
torch::Tensor c;
2336
- if (use_atomic_add) {
2337
- c = torch::zeros({size_m, size_n}, options);
679
+ if (c_or_none.has_value()) {
680
+ c = c_or_none.value();
681
+ TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
682
+ TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
683
+ TORCH_CHECK(c.size(0) == size_m, "Shape mismatch: c.size(0) = ", c.size(0),
684
+ ", size_m = ", size_m);
685
+ TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
686
+ ", size_n = ", size_n);
687
} else {
688
c = torch::empty({size_m, size_n}, options);
689
}
2341
-
2342
- torch::Tensor a_tmp;
2343
- bool has_act_order = g_idx.size(0) != 0;
2344
- if (has_act_order) {
2345
- a_tmp = torch::empty({size_m, size_k}, options);
2346
- } else {
2347
- a_tmp = torch::empty({0}, options);
2348
- }
690
+ if (size_m == 0) return c;
691
692
// Alloc C tmp buffer that is going to be used for the global reduce
693
torch::Tensor c_tmp;
2352
- int reduce_max_m = marlin::determine_reduce_max_m(size_m, marlin::max_par);
2353
- int reduce_n = size_n;
694
auto options_fp32 =
695
torch::TensorOptions().dtype(at::kFloat).device(a.device());
696
if (use_fp32_reduce) {
2357
- c_tmp = torch::empty({reduce_max_m, reduce_n}, options_fp32);
697
+ int max_m_block_size = (size_m + 16 - 1) / 16 * 16;
698
+ max_m_block_size = min(max_m_block_size, 64);
699
+ int max_c_tmp_size =
700
+ sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n;
701
+ c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
702
} else {
2359
- reduce_max_m = 0;
2360
- reduce_n = 0;
703
c_tmp = torch::empty({0}, options_fp32);
704
}
705
2364
- // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
2365
- // auto -1)
2366
- int thread_k = -1;
2367
- // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
2368
- // auto -1)
2369
- int thread_n = -1;
2370
- // sms: number of SMs to use for the kernel (can usually be left as auto -1)
2371
- int sms = -1;
2372
-
2373
- // Verify g_idx and perm
2374
- TORCH_CHECK((g_idx.size(0) == 0 && perm.size(0) == 0) ||
2375
- (g_idx.size(0) == size_k && perm.size(0) == size_k),
2376
- "Unexpected g_idx.size(0) = ", g_idx.size(0),
2377
- " and perm.size(0) = ", perm.size(0),
2378
- ", where size_k = ", size_k);
2379
-
706
// Detect groupsize and act_order
707
int num_groups = -1;
708
int group_size = -1;
@@ -2387,7 +713,31 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
713
" is not size_n = ", size_n);
714
num_groups = b_scales.size(0);
715
716
+ torch::Tensor g_idx, perm, a_tmp;
717
+ if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
718
+ g_idx = g_idx_or_none.value();
719
+ perm = perm_or_none.value();
720
+
721
+ TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
722
+ TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
723
+ TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
724
+ TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
725
+
726
+ // Verify g_idx and perm
727
+ TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
728
+ (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
729
+ "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
730
+ " and perm.size(-1) = ", perm.size(-1),
731
+ ", where size_k = ", size_k);
732
+ } else {
733
+ g_idx = torch::empty({0}, options);
734
+ perm = torch::empty({0}, options);
735
+ a_tmp = torch::empty({0}, options);
736
+ }
737
+ bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
738
+
739
if (has_act_order) {
740
+ a_tmp = torch::empty({size_m, size_k}, options);
741
if (is_k_full) {
742
TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
743
TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
@@ -2398,6 +748,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
748
}
749
750
} else {
751
+ a_tmp = torch::empty({0}, options);
752
if (num_groups > 1) {
753
TORCH_CHECK(
754
size_k % num_groups == 0, "size_k = ", size_k,
@@ -2408,6 +759,33 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
759
}
760
}
761
762
+ torch::Tensor b_zeros;
763
+ if (b_zeros_or_none.has_value()) {
764
+ b_zeros = b_zeros_or_none.value();
765
+ TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
766
+ TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
767
+ } else {
768
+ b_zeros = torch::empty({0}, options);
769
+ }
770
+ bool has_zp = b_zeros.size(-1) > 0;
771
+ if (has_zp) {
772
+ TORCH_CHECK(
773
+ b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
774
+ "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
775
+ } else {
776
+ TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
777
+ b_q_type == vllm::kFE4M3fn,
778
+ "b_q_type must be uint4b8, uint8b128 or float8_e4m3fn when "
779
+ "has_zp = False. Got = ",
780
+ b_q_type.str());
781
+ }
782
+
783
+ if (has_zp && is_zp_float) {
784
+ TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
785
+ "Computation type must be float16 (half) when using float zero "
786
+ "points.");
787
+ }
788
+
789
// Verify b_zeros
790
if (has_zp) {
791
int rank = b_zeros.sizes().size();
@@ -2431,9 +809,11 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
809
}
810
811
// Verify workspace size
2434
- TORCH_CHECK(size_n % marlin::min_thread_n == 0, "size_n = ", size_n,
2435
- ", is not divisible by min_thread_n = ", marlin::min_thread_n);
2436
- int min_workspace_size = (size_n / marlin::min_thread_n) * marlin::max_par;
812
+ TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
813
+ "size_n = ", size_n, ", is not divisible by min_thread_n = ",
814
+ MARLIN_NAMESPACE_NAME::min_thread_n);
815
+
816
+ int min_workspace_size = sms;
817
TORCH_CHECK(workspace.numel() >= min_workspace_size,
818
"workspace.numel = ", workspace.numel(),
819
" is below min_workspace_size = ", min_workspace_size);
@@ -2447,8 +827,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
827
a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k, a.stride(0),
828
workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
829
num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
2450
- thread_k, thread_n, sms, marlin::max_par, use_atomic_add,
2451
- use_fp32_reduce, is_zp_float);
830
+ thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float);
831
} else if (a.scalar_type() == at::ScalarType::BFloat16) {
832
marlin::marlin_mm<nv_bfloat16>(
833
a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
@@ -2458,7 +837,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
837
a.stride(0), workspace.data_ptr(), b_q_type, has_act_order, is_k_full,
838
has_zp, num_groups, group_size, dev,
839
at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
2461
- marlin::max_par, use_atomic_add, use_fp32_reduce, is_zp_float);
840
+ use_atomic_add, use_fp32_reduce, is_zp_float);
841
} else {
842
TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16");
843
}
csrc/quantization/gptq_marlin/kernel.h
new
+37
@@ -0,0 +1,37 @@
1
+
2
+#ifndef MARLIN_NAMESPACE_NAME
3
+ #define MARLIN_NAMESPACE_NAME marlin
4
+#endif
5
+
6
+#include "marlin.cuh"
7
+#include "marlin_dtypes.cuh"
8
+#include "core/scalar_type.hpp"
9
+
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ scales_ptr, const int4 *__restrict__ zp_ptr, \
14
+ const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, \
15
+ int prob_k, int lda, int *locks, bool use_atomic_add, \
16
+ bool use_fp32_reduce, int max_shared_mem
17
+
18
+namespace MARLIN_NAMESPACE_NAME {
19
+template <typename scalar_t, // compute dtype, half or nv_float16
20
+ const vllm::ScalarTypeId w_type_id, // weight ScalarType id
21
+ const int threads, // number of threads in a threadblock
22
+ const int thread_m_blocks, // number of 16x16 blocks in the m
23
+ // dimension (batchsize) of the
24
+ // threadblock
25
+ const int thread_n_blocks, // same for n dimension (output)
26
+ const int thread_k_blocks, // same for k dimension (reduction)
27
+ const bool m_block_size_8, // whether m_block_size == 8
28
+ // only works when thread_m_blocks == 1
29
+ const int stages, // number of stages for the async global->shared
30
+ // fetch pipeline
31
+ const int group_blocks, // number of consecutive 16x16 blocks
32
+ // with a separate quantization scale
33
+ const bool is_zp_float // is zero point of float16 type?
34
+ >
35
+__global__ void Marlin(MARLIN_KERNEL_PARAMS);
36
+
37
+}
csrc/quantization/gptq_marlin/marlin_template.h
new
+1678
@@ -0,0 +1,1678 @@
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
+/*
19
+ * Adapted from https://github.com/IST-DASLab/marlin
20
+ */
21
+
22
+#ifndef MARLIN_NAMESPACE_NAME
23
+ #define MARLIN_NAMESPACE_NAME marlin
24
+#endif
25
+
26
+#include "marlin.cuh"
27
+#include "marlin_dtypes.cuh"
28
+#include "dequant.h"
29
+#include "core/scalar_type.hpp"
30
+
31
+#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
32
+ static_assert(std::is_same<scalar_t, half>::value || \
33
+ std::is_same<scalar_t, nv_bfloat16>::value, \
34
+ "only float16 and bfloat16 is supported");
35
+
36
+namespace MARLIN_NAMESPACE_NAME {
37
+
38
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
39
+
40
+template <typename scalar_t, // compute dtype, half or nv_float16
41
+ const vllm::ScalarTypeId w_type_id, // weight ScalarType id
42
+ const int threads, // number of threads in a threadblock
43
+ const int thread_m_blocks, // number of 16x16 blocks in the m
44
+ // dimension (batchsize) of the
45
+ // threadblock
46
+ const int thread_n_blocks, // same for n dimension (output)
47
+ const int thread_k_blocks, // same for k dimension (reduction)
48
+ const bool m_block_size_8, // whether m_block_size == 8
49
+ // only works when thread_m_blocks == 1
50
+ const int stages, // number of stages for the async global->shared
51
+ // fetch pipeline
52
+ const bool has_act_order, // whether act_order is enabled
53
+ const int group_blocks, // number of consecutive 16x16 blocks
54
+ // with a separate quantization scale
55
+ const bool is_zp_float // is zero point of float16 type?
56
+ >
57
+__global__ void Marlin(
58
+ const int4* __restrict__ A, // fp16 input matrix of shape mxk
59
+ const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
60
+ int4* __restrict__ C, // fp16 output buffer of shape mxn
61
+ int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
62
+ const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
63
+ // (k/groupsize)xn
64
+ const int* __restrict__ g_idx, // int32 group indices of shape k
65
+ int num_groups, // number of scale groups per output channel
66
+ int prob_m, // batch dimension m
67
+ int prob_n, // output dimension n
68
+ int prob_k, // reduction dimension k
69
+ int* locks, // extra global storage for barrier synchronization
70
+ bool use_fp32_reduce // whether to use fp32 global reduce
71
+) {}
72
+
73
+} // namespace marlin
74
+
75
+#else
76
+
77
+// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
78
+// output/accumulation.
79
+template <typename scalar_t>
80
+__device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
81
+ const typename ScalarType<scalar_t>::FragB& frag_b,
82
+ typename ScalarType<scalar_t>::FragC& frag_c) {
83
+ const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
84
+ const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
85
+ float* c = reinterpret_cast<float*>(&frag_c);
86
+ if constexpr (std::is_same<scalar_t, half>::value) {
87
+ asm volatile(
88
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
89
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
90
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
91
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
92
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
93
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
94
+ asm volatile(
95
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
96
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
97
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
98
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
99
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
100
+ } else {
101
+ STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
102
+ }
103
+}
104
+
105
+template <typename scalar_t>
106
+__device__ inline void mma_trans(
107
+ const typename ScalarType<scalar_t>::FragA& a_frag,
108
+ const typename ScalarType<scalar_t>::FragB& frag_b,
109
+ const typename ScalarType<scalar_t>::FragB& frag_b2,
110
+ typename ScalarType<scalar_t>::FragC& frag_c) {
111
+ const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
112
+ const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
113
+ const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
114
+ float* c = reinterpret_cast<float*>(&frag_c);
115
+ if constexpr (std::is_same<scalar_t, half>::value) {
116
+ asm volatile(
117
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
118
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
119
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
120
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
121
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
122
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
123
+ asm volatile(
124
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
125
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
126
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
127
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
128
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
129
+ } else {
130
+ STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
131
+ }
132
+}
133
+
134
+// Instruction for loading a full 16x16 matrix fragment of operand A from shared
135
+// memory, directly in tensor core layout.
136
+template <int count, typename scalar_t>
137
+__device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
138
+ const void* smem_ptr) {
139
+ uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
140
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
141
+ if constexpr (count == 4) {
142
+ asm volatile(
143
+ "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
144
+ : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
145
+ : "r"(smem));
146
+ } else if constexpr (count == 2) {
147
+ asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
148
+ : "=r"(a[0]), "=r"(a[1])
149
+ : "r"(smem));
150
+ } else if constexpr (count == 1) {
151
+ asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n"
152
+ : "=r"(a[0])
153
+ : "r"(smem));
154
+ } else {
155
+ static_assert(count == 1 || count == 2 || count == 4, "invalid count");
156
+ }
157
+}
158
+
159
+// Multiply dequantized values by the corresponding quantization scale; used
160
+// only for grouped quantization.
161
+template <typename scalar_t>
162
+__device__ inline void scale(typename ScalarType<scalar_t>::FragB& frag_b,
163
+ typename ScalarType<scalar_t>::FragS& frag_s,
164
+ int i) {
165
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
166
+ scalar_t2 s =
167
+ ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
168
+ frag_b[0] = __hmul2(frag_b[0], s);
169
+ frag_b[1] = __hmul2(frag_b[1], s);
170
+}
171
+
172
+template <typename scalar_t>
173
+__device__ inline void scale_and_sub(
174
+ typename ScalarType<scalar_t>::FragB& frag_b, scalar_t s, scalar_t zp) {
175
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
176
+ scalar_t2 s2 = ScalarType<scalar_t>::num2num2(s);
177
+ scalar_t2 zp2 = ScalarType<scalar_t>::num2num2(zp);
178
+ frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
179
+ frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
180
+}
181
+
182
+template <typename scalar_t>
183
+__device__ inline void sub_zp(typename ScalarType<scalar_t>::FragB& frag_b,
184
+ typename ScalarType<scalar_t>::scalar_t2& frag_zp,
185
+ int i) {
186
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
187
+ scalar_t2 zp =
188
+ ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
189
+ frag_b[0] = __hsub2(frag_b[0], zp);
190
+ frag_b[1] = __hsub2(frag_b[1], zp);
191
+}
192
+
193
+// Same as above, but for act_order (each K is multiplied individually)
194
+template <typename scalar_t>
195
+__device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
196
+ typename ScalarType<scalar_t>::FragS& frag_s_1,
197
+ typename ScalarType<scalar_t>::FragS& frag_s_2,
198
+ typename ScalarType<scalar_t>::FragS& frag_s_3,
199
+ typename ScalarType<scalar_t>::FragS& frag_s_4,
200
+ int i) {
201
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
202
+ scalar_t2 s_val_1_2;
203
+ s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
204
+ s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
205
+
206
+ scalar_t2 s_val_3_4;
207
+ s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
208
+ s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
209
+
210
+ frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
211
+ frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
212
+}
213
+
214
+// Given 2 floats multiply by 2 scales (halves)
215
+template <typename scalar_t>
216
+__device__ inline void scale_float(float* c,
217
+ typename ScalarType<scalar_t>::FragS& s) {
218
+ scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
219
+ c[0] = __fmul_rn(c[0], ScalarType<scalar_t>::num2float(s_ptr[0]));
220
+ c[1] = __fmul_rn(c[1], ScalarType<scalar_t>::num2float(s_ptr[1]));
221
+}
222
+
223
+// Wait until barrier reaches `count`, then lock for current threadblock.
224
+__device__ inline void barrier_acquire(int* lock, int count) {
225
+ if (threadIdx.x == 0) {
226
+ int state = -1;
227
+ do
228
+ // Guarantee that subsequent writes by this threadblock will be visible
229
+ // globally.
230
+ asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
231
+ : "=r"(state)
232
+ : "l"(lock));
233
+ while (state != count);
234
+ }
235
+ __syncthreads();
236
+}
237
+
238
+// Release barrier and increment visitation count.
239
+__device__ inline void barrier_release(int* lock, bool reset = false) {
240
+ __syncthreads();
241
+ if (threadIdx.x == 0) {
242
+ if (reset) {
243
+ lock[0] = 0;
244
+ return;
245
+ }
246
+ int val = 1;
247
+ // Make sure that all writes since acquiring this barrier are visible
248
+ // globally, while releasing the barrier.
249
+ asm volatile("fence.acq_rel.gpu;\n");
250
+ asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
251
+ :
252
+ : "l"(lock), "r"(val));
253
+ }
254
+}
255
+
256
+// Wait until value of lock to be negative, and then add 1
257
+__device__ inline void wait_negative_and_add(int* lock) {
258
+ if (threadIdx.x == 0) {
259
+ int state = 0;
260
+ do
261
+ // Guarantee that subsequent writes by this threadblock will be visible
262
+ // globally.
263
+ asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
264
+ : "=r"(state)
265
+ : "l"(lock));
266
+ while (state >= 0);
267
+ atomicAdd(lock, 1);
268
+ }
269
+ __syncthreads();
270
+}
271
+
272
+template <typename scalar_t, // compute dtype, half or nv_float16
273
+ const vllm::ScalarTypeId w_type_id, // weight ScalarType id
274
+ const int threads, // number of threads in a threadblock
275
+ const int thread_m_blocks, // number of 16x16 blocks in the m
276
+ // dimension (batchsize) of the
277
+ // threadblock
278
+ const int thread_n_blocks, // same for n dimension (output)
279
+ const int thread_k_blocks, // same for k dimension (reduction)
280
+ const bool m_block_size_8, // whether m_block_size == 8
281
+ // only works when thread_m_blocks == 1
282
+ const int stages, // number of stages for the async global->shared
283
+ // fetch pipeline
284
+ const int group_blocks, // number of consecutive 16x16 blocks
285
+ // with a separate quantization scale
286
+ const bool is_zp_float // is zero point of float16 type?
287
+ >
288
+__global__ void Marlin(
289
+ const int4* __restrict__ A, // fp16 input matrix of shape mxk
290
+ const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
291
+ int4* __restrict__ C, // fp16 output buffer of shape mxn
292
+ int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
293
+ const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
294
+ // (k/groupsize)xn
295
+ const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
296
+ // (k/groupsize)x(n/pack_factor)
297
+ const int* __restrict__ g_idx, // int32 group indices of shape k
298
+ int num_groups, // number of scale groups per output channel
299
+ int prob_m, // batch dimension m
300
+ int prob_n, // output dimension n
301
+ int prob_k, // reduction dimension k
302
+ int lda, // A.stride(0), equal to prob_k is A is contiguous
303
+ int* locks, // extra global storage for barrier synchronization
304
+ bool use_atomic_add, // whether to use atomic add to reduce
305
+ bool use_fp32_reduce, // whether to use fp32 global reduce
306
+ int max_shared_mem) {
307
+ // Each threadblock processes one "stripe" of the B matrix with (roughly) the
308
+ // same size, which might involve multiple column "slices" (of width 16 *
309
+ // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
310
+ // example:
311
+ // 0 1 3
312
+ // 0 2 3
313
+ // 1 2 4
314
+ // While this kind of partitioning makes things somewhat more complicated, it
315
+ // ensures good utilization of all SMs for many kinds of shape and GPU
316
+ // configurations, while requiring as few slow global cross-threadblock
317
+ // reductions as possible.
318
+ using Dtype = ScalarType<scalar_t>;
319
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
320
+ using FragA = typename ScalarType<scalar_t>::FragA;
321
+ using FragB = typename ScalarType<scalar_t>::FragB;
322
+ using FragC = typename ScalarType<scalar_t>::FragC;
323
+ using FragS = typename ScalarType<scalar_t>::FragS;
324
+ using FragZP = typename ScalarType<scalar_t>::FragZP;
325
+
326
+ static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
327
+ constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
328
+ constexpr bool has_act_order = group_blocks == 0;
329
+ constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
330
+
331
+ constexpr int pack_factor = 32 / w_type.size_bits();
332
+ static_assert(thread_m_blocks == 1 || !m_block_size_8);
333
+
334
+ // For larger GEMMs we run multiple batchsize 64 versions in parallel for a
335
+ // better partitioning with less reductions
336
+ int parallel = 1;
337
+ if (prob_m > m_block_size) {
338
+ parallel = prob_m / m_block_size;
339
+ prob_m = m_block_size;
340
+ }
341
+
342
+ int k_tiles = prob_k / 16 / thread_k_blocks;
343
+ int n_tiles = prob_n / 16 / thread_n_blocks;
344
+ int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x);
345
+
346
+ if constexpr (!has_act_order && group_blocks != -1) {
347
+ if (group_blocks >= thread_k_blocks) {
348
+ // Ensure that the number of tiles in each stripe is a multiple of the
349
+ // groupsize; this avoids an annoying special case where a stripe starts
350
+ // in the middle of group.
351
+ iters = (group_blocks / thread_k_blocks) *
352
+ div_ceil(iters, (group_blocks / thread_k_blocks));
353
+ }
354
+ }
355
+
356
+ int slice_row = (iters * blockIdx.x) % k_tiles;
357
+ int slice_col_par = (iters * blockIdx.x) / k_tiles;
358
+ int slice_col = slice_col_par;
359
+ int slice_iters; // number of threadblock tiles in the current slice
360
+ int slice_count =
361
+ 0; // total number of active threadblocks in the current slice
362
+ int slice_idx; // index of threadblock in current slice; numbered bottom to
363
+ // top
364
+
365
+ int par_id = 0;
366
+ int locks_off = 0;
367
+
368
+ // We can easily implement parallel problem execution by just remapping
369
+ // indices and advancing global pointers
370
+ if (slice_col_par >= n_tiles) {
371
+ A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8;
372
+ C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
373
+ slice_col = slice_col_par % n_tiles;
374
+ par_id = slice_col_par / n_tiles;
375
+ }
376
+ if (parallel * n_tiles >= gridDim.x) {
377
+ // when parallel * n_tiles >= sms
378
+ // then there are at most $sms$ conflict tile blocks
379
+ locks_off = blockIdx.x;
380
+ } else {
381
+ locks_off = (iters * blockIdx.x) / k_tiles - 1;
382
+ }
383
+
384
+ // Compute all information about the current slice which is required for
385
+ // synchronization.
386
+ auto init_slice = [&](bool first_init = false) {
387
+ slice_iters =
388
+ iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
389
+ if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
390
+ if (slice_iters == 0) return;
391
+ if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
392
+ slice_count = 1;
393
+ slice_idx = 0;
394
+ int col_first = iters * div_ceil(k_tiles * slice_col_par, iters);
395
+ if (col_first <= k_tiles * (slice_col_par + 1)) {
396
+ int col_off = col_first - k_tiles * slice_col_par;
397
+ slice_count = div_ceil(k_tiles - col_off, iters);
398
+ if (col_off > 0) slice_count++;
399
+ int delta_first = iters * blockIdx.x - col_first;
400
+ if (delta_first < 0 || (col_off == 0 && delta_first == 0))
401
+ slice_idx = slice_count - 1;
402
+ else {
403
+ slice_idx = slice_count - 1 - delta_first / iters;
404
+ if (col_off > 0) slice_idx--;
405
+ }
406
+ }
407
+ if (parallel * n_tiles >= gridDim.x) {
408
+ if (slice_count > 1 && slice_idx == slice_count - 1) {
409
+ locks_off++;
410
+ }
411
+ } else {
412
+ locks_off++;
413
+ }
414
+
415
+ if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) {
416
+ constexpr int threads_per_m = 16 * thread_n_blocks / 8;
417
+ int m_per_thread =
418
+ div_ceil(thread_m_blocks * 16, threads / threads_per_m);
419
+ if (m_block_size_8) m_per_thread = div_ceil(8, threads / threads_per_m);
420
+ for (int i = 0; i < m_per_thread; i++) {
421
+ int row = threads / threads_per_m * i + threadIdx.x / threads_per_m;
422
+ if (row < prob_m) {
423
+ int col = slice_col * 16 * thread_n_blocks / 8 +
424
+ threadIdx.x % threads_per_m;
425
+ C[row * prob_n / 8 + col] = {0, 0, 0, 0};
426
+ }
427
+ }
428
+ // After write zero to output, write a negative value to lock.
429
+ // Every SM that processes the same slice would wait for
430
+ // the negative value, and then atomicAdd 1 to it.
431
+ // After all SMs are processed, the lock value would back to 0 again.
432
+ __syncthreads();
433
+ if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count;
434
+ }
435
+
436
+ if (slice_col == n_tiles) {
437
+ A += 16 * thread_m_blocks * lda / 8;
438
+ C += 16 * thread_m_blocks * prob_n / 8;
439
+ slice_col = 0;
440
+ par_id++;
441
+ }
442
+ };
443
+ init_slice(true);
444
+
445
+ // A sizes/strides
446
+
447
+ // stride of the A matrix in global memory
448
+ int a_gl_stride = lda / 8;
449
+ // stride of an A matrix tile in shared memory
450
+ constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
451
+ // delta between subsequent A tiles in global memory
452
+ constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;
453
+ // between subsequent accesses within a tile
454
+ int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
455
+ // between shared memory writes
456
+ constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
457
+ // between shared memory tile reads
458
+ constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));
459
+ // within a shared memory tile
460
+ constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
461
+ // overall size of a tile
462
+ constexpr int a_sh_stage = a_sh_stride * m_block_size;
463
+ // number of shared write iterations for a tile
464
+ constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
465
+
466
+ // B sizes/strides
467
+ int b_gl_stride = 16 * prob_n / (pack_factor * 4);
468
+ constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
469
+ constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;
470
+ constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
471
+
472
+ int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
473
+ int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
474
+ constexpr int b_sh_wr_delta = threads * b_thread_vecs;
475
+ constexpr int b_sh_rd_delta = threads * b_thread_vecs;
476
+ constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
477
+ constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
478
+
479
+ // Scale sizes/strides without act_order
480
+ int s_gl_stride = prob_n / 8;
481
+ constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
482
+ constexpr int s_tb_groups =
483
+ !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
484
+ ? thread_k_blocks / group_blocks
485
+ : 1;
486
+ constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
487
+ int s_gl_rd_delta = s_gl_stride;
488
+
489
+ // Scale size/strides with act_order
490
+ constexpr int tb_k = 16 * thread_k_blocks;
491
+ constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
492
+ // constexpr int act_s_row_stride = 1;
493
+ // int act_s_col_stride = act_s_row_stride * num_groups;
494
+ constexpr int act_s_max_num_groups = 32;
495
+ int act_s_col_stride = 1;
496
+ int act_s_col_warp_stride = act_s_col_stride * 8;
497
+
498
+ int tb_n_warps = thread_n_blocks / 4;
499
+ int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
500
+
501
+ // Zero-points sizes/strides
502
+ int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
503
+ constexpr int zp_sh_stride = is_zp_float
504
+ ? 16 * thread_n_blocks / 8
505
+ : ((16 * thread_n_blocks) / pack_factor) / 4;
506
+ constexpr int zp_tb_groups = s_tb_groups;
507
+ constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;
508
+ int zp_gl_rd_delta = zp_gl_stride;
509
+
510
+ // Global A read index of current thread.
511
+ int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
512
+ (threadIdx.x % a_gl_rd_delta_o);
513
+ a_gl_rd += a_gl_rd_delta_o * slice_row;
514
+ // Shared write index of current thread.
515
+ int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
516
+ (threadIdx.x % a_gl_rd_delta_o);
517
+ // Shared read index.
518
+ int a_sh_rd =
519
+ a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
520
+ (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
521
+ a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
522
+
523
+ int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
524
+ (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
525
+ b_gl_rd += b_sh_stride * slice_col;
526
+ b_gl_rd += b_gl_rd_delta_o * slice_row;
527
+ auto b_sh_wr = threadIdx.x * b_thread_vecs;
528
+ auto b_sh_rd = threadIdx.x * b_thread_vecs;
529
+
530
+ // For act_order
531
+ constexpr int k_iter_size = tb_k / b_sh_wr_iters;
532
+ int slice_k_start = tb_k * slice_row;
533
+ int slice_k_finish = slice_k_start + tb_k * slice_iters;
534
+ int slice_k_start_shared_fetch = slice_k_start;
535
+ int slice_n_offset = act_s_col_tb_stride * slice_col;
536
+
537
+ // No act_order
538
+ int s_gl_rd;
539
+ if constexpr (!has_act_order) {
540
+ if constexpr (group_blocks == -1) {
541
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
542
+ } else {
543
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
544
+ s_sh_stride * slice_col + threadIdx.x;
545
+ }
546
+ }
547
+ auto s_sh_wr = threadIdx.x;
548
+ bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
549
+
550
+ // Zero-points
551
+ int zp_gl_rd;
552
+ if constexpr (has_zp) {
553
+ if constexpr (group_blocks == -1) {
554
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
555
+ } else {
556
+ zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
557
+ zp_sh_stride * slice_col + threadIdx.x;
558
+ }
559
+ }
560
+ auto zp_sh_wr = threadIdx.x;
561
+ bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
562
+
563
+ // We use a different scale layout for grouped and column-wise quantization as
564
+ // we scale a `half2` tile in column-major layout in the former and in
565
+ // row-major in the latter case.
566
+ int s_sh_rd;
567
+ if constexpr (group_blocks != -1)
568
+ s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
569
+ (threadIdx.x % 32) / 4;
570
+ else if constexpr (group_blocks == -1 && (m_block_size_8 || has_zp))
571
+ s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
572
+ (threadIdx.x % 32) / 8;
573
+ else
574
+ s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
575
+ (threadIdx.x % 32) % 4;
576
+
577
+ // Zero-points have the same read layout as the scales
578
+ // (without column-wise case)
579
+ constexpr int num_col_threads = 8;
580
+ constexpr int num_row_threads = 4;
581
+ constexpr int num_ints_per_thread = 8 / pack_factor;
582
+ int zp_sh_rd;
583
+ if constexpr (has_zp) {
584
+ if constexpr (is_zp_float) {
585
+ if constexpr (group_blocks != -1) {
586
+ zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
587
+ (threadIdx.x % 32) / 4;
588
+ }
589
+ } else {
590
+ zp_sh_rd = num_ints_per_thread * num_col_threads *
591
+ ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
592
+ num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
593
+ }
594
+ }
595
+
596
+ // Precompute which thread should not read memory in which iterations; this is
597
+ // needed if there are more threads than required for a certain tilesize or
598
+ // when the batchsize is not a multiple of 16.
599
+ bool a_sh_wr_pred[a_sh_wr_iters];
600
+ #pragma unroll
601
+ for (int i = 0; i < a_sh_wr_iters; i++)
602
+ a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;
603
+
604
+ // To ensure that writing and reading A tiles to/from shared memory, the
605
+ // latter in fragment format, is fully bank conflict free, we need to use a
606
+ // rather fancy XOR-based layout. The key here is that neither reads nor
607
+ // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
608
+ // same shared memory banks. Further, it seems (based on NSight-Compute) that
609
+ // each warp must also write a consecutive memory segment?
610
+ auto transform_a = [&](int i) {
611
+ int row = i / a_gl_rd_delta_o;
612
+ return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8);
613
+ };
614
+ // Since the computation of this remapping is non-trivial and, due to our main
615
+ // loop unrolls, all shared memory accesses are static, we simply precompute
616
+ // both transformed reads and writes.
617
+ int a_sh_wr_trans[a_sh_wr_iters];
618
+ #pragma unroll
619
+ for (int i = 0; i < a_sh_wr_iters; i++)
620
+ a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
621
+ int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
622
+ #pragma unroll
623
+ for (int i = 0; i < b_sh_wr_iters; i++) {
624
+ #pragma unroll
625
+ for (int j = 0; j < thread_m_blocks; j++)
626
+ a_sh_rd_trans[i][j] =
627
+ transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
628
+ }
629
+
630
+ // Since B-accesses have non-constant stride they have to be computed at
631
+ // runtime; we break dependencies between subsequent accesses with a tile by
632
+ // maintining multiple pointers (we have enough registers), a tiny
633
+ // optimization.
634
+ const int4* B_ptr[b_sh_wr_iters];
635
+ #pragma unroll
636
+ for (int i = 0; i < b_sh_wr_iters; i++)
637
+ B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
638
+
639
+ extern __shared__ int4 sh[];
640
+ // Shared memory storage for global fetch pipelines.
641
+ constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
642
+ constexpr int sh_b_size = stages * b_sh_stage;
643
+ int4* sh_b = sh;
644
+ int4* sh_red = sh;
645
+ int4* sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
646
+ int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
647
+ constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
648
+ : (stages * s_sh_stage);
649
+ int4* sh_s = sh_zp + (stages * zp_sh_stage);
650
+ // shared memory reused by reduction should be smaller than
651
+ // shared memory used by weight.
652
+ static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
653
+ stages * b_sh_stage);
654
+ int4* sh_a = sh_s + sh_s_size;
655
+ // constexpr int shm_size_used =
656
+ // stages * (g_idx_stage + zp_sh_stage) + sh_s_size +
657
+ // (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
658
+
659
+ // Register storage for double buffer of shared memory reads.
660
+ FragA frag_a[2][thread_m_blocks];
661
+ I4 frag_b_quant[2][b_thread_vecs];
662
+ FragC frag_c[thread_m_blocks][4][2];
663
+ FragS frag_s[2][4]; // No act-order
664
+ FragS act_frag_s[2][4][4]; // For act-order
665
+ int frag_qzp[2][num_ints_per_thread]; // Zero-points
666
+ FragZP frag_zp; // Zero-points in fp16
667
+ FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
668
+
669
+ // Zero accumulators.
670
+ auto zero_accums = [&]() {
671
+ #pragma unroll
672
+ for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
673
+ reinterpret_cast<float*>(frag_c)[i] = 0;
674
+ };
675
+
676
+ int sh_first_group_id = -1;
677
+ int sh_num_groups = -1;
678
+
679
+ auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
680
+ int last_group_id) {
681
+ sh_first_group_id = first_group_id;
682
+ sh_num_groups = last_group_id - first_group_id + 1;
683
+
684
+ if (sh_num_groups < act_s_max_num_groups) {
685
+ sh_num_groups = act_s_max_num_groups;
686
+ }
687
+
688
+ if (sh_first_group_id + sh_num_groups > num_groups) {
689
+ sh_num_groups = num_groups - sh_first_group_id;
690
+ }
691
+
692
+ int row_offset = first_group_id * s_gl_stride;
693
+
694
+ if (is_async) {
695
+ for (int i = 0; i < sh_num_groups; i++) {
696
+ if (threadIdx.x < s_sh_stride) {
697
+ cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
698
+ &scales_ptr[row_offset + (i * s_gl_stride) +
699
+ slice_n_offset + threadIdx.x]);
700
+ }
701
+ }
702
+ } else {
703
+ for (int i = 0; i < sh_num_groups; i++) {
704
+ if (threadIdx.x < s_sh_stride) {
705
+ sh_s[(i * s_sh_stride) + threadIdx.x] =
706
+ scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
707
+ threadIdx.x];
708
+ }
709
+ }
710
+ }
711
+ };
712
+ // Asynchronously fetch the next A, B and s tile from global to the next
713
+ // shared memory pipeline location.
714
+ auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
715
+ if (pred) {
716
+ int4* sh_a_stage = sh_a + a_sh_stage * pipe;
717
+ #pragma unroll
718
+ for (int i = 0; i < a_sh_wr_iters; i++) {
719
+ cp_async4_pred(
720
+ &sh_a_stage[a_sh_wr_trans[i]],
721
+ &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
722
+ a_sh_wr_pred[i]);
723
+ }
724
+ int4* sh_b_stage = sh_b + b_sh_stage * pipe;
725
+ #pragma unroll
726
+ for (int i = 0; i < b_sh_wr_iters; i++) {
727
+ #pragma unroll
728
+ for (int j = 0; j < b_thread_vecs; j++) {
729
+ cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j);
730
+ }
731
+
732
+ B_ptr[i] += b_gl_rd_delta_o;
733
+ }
734
+
735
+ if constexpr (has_act_order) {
736
+ // Fetch g_idx thread-block portion
737
+ int full_pipe = a_off;
738
+ int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
739
+ if (cur_k < prob_k && cur_k < slice_k_finish) {
740
+ int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
741
+
742
+ int4 const* cur_g_idx_stage_ptr =
743
+ reinterpret_cast<int4 const*>(&g_idx[cur_k]);
744
+
745
+ if (threadIdx.x < g_idx_stage) {
746
+ cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
747
+ &cur_g_idx_stage_ptr[threadIdx.x]);
748
+ }
749
+ }
750
+ } else {
751
+ if constexpr (group_blocks != -1) {
752
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
753
+
754
+ if constexpr (group_blocks >= thread_k_blocks) {
755
+ // Only fetch scales if this tile starts a new group
756
+ if (pipe % (group_blocks / thread_k_blocks) == 0) {
757
+ if (s_sh_wr_pred) {
758
+ cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
759
+ }
760
+ s_gl_rd += s_gl_rd_delta;
761
+ }
762
+ } else {
763
+ for (int i = 0; i < s_tb_groups; i++) {
764
+ if (s_sh_wr_pred) {
765
+ cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
766
+ &scales_ptr[s_gl_rd]);
767
+ }
768
+ s_gl_rd += s_gl_rd_delta;
769
+ }
770
+ }
771
+ }
772
+
773
+ if constexpr (has_zp && group_blocks != -1) {
774
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
775
+
776
+ if constexpr (group_blocks >= thread_k_blocks) {
777
+ // Only fetch zero-points if this tile starts a new group
778
+ if (pipe % (group_blocks / thread_k_blocks) == 0) {
779
+ if (zp_sh_wr_pred) {
780
+ cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
781
+ }
782
+ zp_gl_rd += zp_gl_rd_delta;
783
+ }
784
+ } else {
785
+ for (int i = 0; i < zp_tb_groups; i++) {
786
+ if (zp_sh_wr_pred) {
787
+ cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
788
+ &zp_ptr[zp_gl_rd]);
789
+ }
790
+ zp_gl_rd += zp_gl_rd_delta;
791
+ }
792
+ }
793
+ }
794
+ }
795
+ }
796
+ // Insert a fence even when we are winding down the pipeline to ensure that
797
+ // waiting is also correct at this point.
798
+ cp_async_fence();
799
+ };
800
+
801
+ auto fetch_col_zp_to_shared = [&]() {
802
+ if (zp_sh_wr_pred) {
803
+ cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);
804
+ }
805
+ };
806
+
807
+ auto fetch_col_scale_to_shared = [&]() {
808
+ if (s_sh_wr_pred) {
809
+ cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
810
+ }
811
+ };
812
+
813
+ // Wait until the next thread tile has been loaded to shared memory.
814
+ auto wait_for_stage = [&]() {
815
+ // We only have `stages - 2` active fetches since we are double buffering
816
+ // and can only issue the next fetch when it is guaranteed that the previous
817
+ // shared memory load is fully complete (as it may otherwise be
818
+ // overwritten).
819
+ cp_async_wait<stages - 2>();
820
+ __syncthreads();
821
+ };
822
+
823
+ // Load the next sub-tile from the current location in the shared memory pipe
824
+ // into the current register buffer.
825
+ auto fetch_to_registers = [&](int k, int pipe) {
826
+ int4* sh_a_stage = sh_a + a_sh_stage * pipe;
827
+ #pragma unroll
828
+ for (int i = 0; i < thread_m_blocks; i++)
829
+ ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
830
+ frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
831
+ int4* sh_b_stage = sh_b + b_sh_stage * pipe;
832
+
833
+ #pragma unroll
834
+ for (int i = 0; i < b_thread_vecs; i++) {
835
+ frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
836
+ &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
837
+ }
838
+ };
839
+
840
+ bool is_same_group[stages];
841
+ int same_group_id[stages];
842
+
843
+ auto init_same_group = [&](int pipe) {
844
+ if constexpr (!has_act_order) {
845
+ return;
846
+ }
847
+
848
+ int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
849
+ int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
850
+
851
+ int group_id_1 = sh_g_idx_int_ptr[0];
852
+ int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
853
+
854
+ is_same_group[pipe] = group_id_1 == group_id_2;
855
+ same_group_id[pipe] = group_id_1;
856
+ };
857
+
858
+ auto fetch_scales_to_registers = [&](int k, int full_pipe) {
859
+ int pipe = full_pipe % stages;
860
+
861
+ if constexpr (!has_act_order) {
862
+ // No act-order case
863
+ if constexpr (group_blocks == -1) {
864
+ // load only when starting a new slice
865
+ if (k == 0 && full_pipe == 0) {
866
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
867
+ reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
868
+ }
869
+ } else if constexpr (group_blocks != -1) {
870
+ if constexpr (group_blocks >= thread_k_blocks) {
871
+ if (k % b_sh_wr_iters == 0) {
872
+ int4* sh_s_stage =
873
+ sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
874
+ (pipe / (group_blocks / thread_k_blocks)));
875
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
876
+ } else {
877
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
878
+ reinterpret_cast<int4*>(&frag_s[0])[0];
879
+ }
880
+ } else {
881
+ auto warp_id = threadIdx.x / 32;
882
+ int n_warps = thread_n_blocks / 4;
883
+
884
+ int warp_row = warp_id / n_warps;
885
+
886
+ int cur_k = warp_row * 16;
887
+ cur_k += k_iter_size * (k % b_sh_wr_iters);
888
+
889
+ int k_blocks = cur_k / 16;
890
+ int cur_group_id = k_blocks / group_blocks;
891
+
892
+ int4* sh_s_stage = sh_s + s_sh_stage * pipe;
893
+
894
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
895
+ sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
896
+ }
897
+ }
898
+
899
+ return;
900
+ }
901
+
902
+ // Act-order case
903
+
904
+ // Determine K of the "current" thread-block
905
+ int cur_k = slice_k_start + tb_k * full_pipe;
906
+ if (cur_k >= prob_k || cur_k >= slice_k_finish) {
907
+ return;
908
+ }
909
+
910
+ // Reset (to current thread-block) since we read g_idx portion from the
911
+ // shared memory
912
+ cur_k = 0;
913
+
914
+ // Progress to current iteration
915
+ cur_k += k_iter_size * (k % b_sh_wr_iters);
916
+
917
+ // Determine "position" inside the thread-block (based on warp and
918
+ // thread-id)
919
+ auto warp_id = threadIdx.x / 32;
920
+ int n_warps =
921
+ thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
922
+
923
+ int warp_row = warp_id / n_warps;
924
+ int warp_col = warp_id % n_warps;
925
+
926
+ cur_k += warp_row * 16;
927
+
928
+ auto th_id = threadIdx.x % 32;
929
+ cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
930
+
931
+ int s_col_shift =
932
+ /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
933
+ (th_id / 4) * act_s_col_stride;
934
+
935
+ if (is_same_group[pipe]) {
936
+ if (k % 2 == 0) {
937
+ *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
938
+ sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
939
+ s_col_shift];
940
+ } else {
941
+ *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
942
+ *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
943
+ }
944
+
945
+ for (int i = 1; i < 4; i++) {
946
+ *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
947
+ *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
948
+ }
949
+ return;
950
+ }
951
+
952
+ int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
953
+ int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
954
+
955
+ constexpr int k_frag_offsets[4] = {0, 1, 8,
956
+ 9}; // Tensor core offsets per thread
957
+
958
+ #pragma unroll
959
+ for (int i = 0; i < 4; i++) {
960
+ int actual_k = cur_k + k_frag_offsets[i];
961
+
962
+ int group_id = sh_g_idx_int_ptr[actual_k];
963
+ int rel_group_id = group_id - sh_first_group_id;
964
+
965
+ *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
966
+ sh_s[rel_group_id * s_sh_stride + s_col_shift];
967
+ }
968
+ };
969
+
970
+ auto fetch_zp_to_registers = [&](int k, int full_pipe) {
971
+ // This code does not handle group_blocks == 0,
972
+ // which signifies act_order.
973
+ // has_zp implies AWQ, which doesn't have act_order,
974
+ static_assert(!has_zp || group_blocks != 0);
975
+
976
+ if constexpr (has_zp && !is_zp_float) {
977
+ int pipe = full_pipe % stages;
978
+
979
+ if constexpr (group_blocks == -1) {
980
+ // load only when starting a new slice
981
+ if (k == 0 && full_pipe == 0) {
982
+ #pragma unroll
983
+ for (int i = 0; i < num_ints_per_thread; i++) {
984
+ frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
985
+ }
986
+ }
987
+
988
+ } else if constexpr (group_blocks >= thread_k_blocks) {
989
+ if (k % b_sh_wr_iters == 0) {
990
+ int4* sh_zp_stage =
991
+ sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
992
+ (pipe / (group_blocks / thread_k_blocks)));
993
+ #pragma unroll
994
+ for (int i = 0; i < num_ints_per_thread; i++) {
995
+ frag_qzp[k % 2][i] =
996
+ (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
997
+ }
998
+ }
999
+ } else {
1000
+ auto warp_id = threadIdx.x / 32;
1001
+ int n_warps = thread_n_blocks / 4;
1002
+
1003
+ int warp_row = warp_id / n_warps;
1004
+
1005
+ int cur_k = warp_row * 16;
1006
+ cur_k += k_iter_size * (k % b_sh_wr_iters);
1007
+
1008
+ int k_blocks = cur_k / 16;
1009
+ int cur_group_id = 0;
1010
+
1011
+ // Suppress bogus and persistent divide-by-zero warning
1012
+ #pragma nv_diagnostic push
1013
+ #pragma nv_diag_suppress divide_by_zero
1014
+ cur_group_id = k_blocks / group_blocks;
1015
+ #pragma nv_diagnostic pop
1016
+
1017
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1018
+
1019
+ sh_zp_stage += cur_group_id * zp_sh_stride;
1020
+
1021
+ #pragma unroll
1022
+ for (int i = 0; i < num_ints_per_thread; i++) {
1023
+ frag_qzp[k % 2][i] =
1024
+ (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1025
+ }
1026
+ }
1027
+ }
1028
+
1029
+ else if constexpr (has_zp && is_zp_float) {
1030
+ int pipe = full_pipe % stages;
1031
+
1032
+ if constexpr (group_blocks != -1) {
1033
+ if constexpr (group_blocks >= thread_k_blocks) {
1034
+ if (k % b_sh_wr_iters == 0) {
1035
+ int4* sh_zp_stage =
1036
+ sh_zp +
1037
+ zp_sh_stage * ((group_blocks / thread_k_blocks) *
1038
+ (pipe / (group_blocks / thread_k_blocks)));
1039
+ reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1040
+ sh_zp_stage[zp_sh_rd];
1041
+ }
1042
+ } else {
1043
+ auto warp_id = threadIdx.x / 32;
1044
+ int n_warps = thread_n_blocks / 4;
1045
+
1046
+ int warp_row = warp_id / n_warps;
1047
+
1048
+ int cur_k = warp_row * 16;
1049
+ cur_k += k_iter_size * (k % b_sh_wr_iters);
1050
+
1051
+ int k_blocks = cur_k / 16;
1052
+ // Suppress bogus and persistent divide-by-zero warning
1053
+ #pragma nv_diagnostic push
1054
+ #pragma nv_diag_suppress divide_by_zero
1055
+ int cur_group_id = k_blocks / group_blocks;
1056
+ #pragma nv_diagnostic pop
1057
+
1058
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1059
+
1060
+ reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1061
+ sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride];
1062
+ }
1063
+ }
1064
+ }
1065
+ };
1066
+
1067
+ auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1068
+ if constexpr (has_zp && is_zp_float || !has_zp) {
1069
+ dequant<scalar_t2, w_type_id>(q, frag_b_ptr);
1070
+ } else {
1071
+ static_assert(has_zp && !is_zp_float);
1072
+ static_assert(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU8.id());
1073
+ // If (has_zp && !is_zp_float),
1074
+ // we use not-zp version `dequant` function
1075
+ // to improve numerical accuracy.
1076
+ // Since both weight and zero point are dequanted using this logic,
1077
+ // the final dequanted weight would be correct.
1078
+ if constexpr (w_type_id == vllm::kU4.id()) {
1079
+ dequant<scalar_t2, vllm::kU4B8.id()>(q, frag_b_ptr);
1080
+ } else if constexpr (w_type_id == vllm::kU8.id()) {
1081
+ dequant<scalar_t2, vllm::kU8B128.id()>(q, frag_b_ptr);
1082
+ }
1083
+ }
1084
+ };
1085
+
1086
+ // Execute the actual tensor core matmul of a sub-tile.
1087
+ bool is_first_matmul_in_slice = true;
1088
+ auto matmul = [&](int k) {
1089
+ int k2 = k % 2;
1090
+ const bool is_new_zp =
1091
+ ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) ||
1092
+ (group_blocks == -1 && is_first_matmul_in_slice);
1093
+ if constexpr (has_zp && !is_zp_float) {
1094
+ if (is_new_zp) {
1095
+ if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
1096
+ FragB frag_zp_0;
1097
+ FragB frag_zp_1;
1098
+ int zp_quant_0, zp_quant_1;
1099
+
1100
+ if constexpr (w_type.size_bits() == 4) {
1101
+ zp_quant_0 = frag_qzp[k2][0];
1102
+ zp_quant_1 = zp_quant_0 >> 8;
1103
+ } else {
1104
+ static_assert(w_type.size_bits() == 8);
1105
+ zp_quant_0 = frag_qzp[k2][0];
1106
+ zp_quant_1 = frag_qzp[k2][1];
1107
+ }
1108
+
1109
+ dequant_data(zp_quant_0, reinterpret_cast<scalar_t2*>(&frag_zp));
1110
+ dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1111
+ }
1112
+ }
1113
+ if constexpr (has_zp && is_zp_float) {
1114
+ if (is_new_zp) {
1115
+ reinterpret_cast<int4*>(&frag_zp)[0] =
1116
+ reinterpret_cast<int4*>(&frag_zpf[k2])[0];
1117
+ }
1118
+ }
1119
+
1120
+ // We have the m dimension as the inner loop in order to encourage overlapping
1121
+ // dequantization and matmul operations.
1122
+ #pragma unroll
1123
+ for (int j = 0; j < 4; j++) {
1124
+ FragB frag_b0;
1125
+ FragB frag_b1;
1126
+ int b_quant_0, b_quant_1;
1127
+
1128
+ if constexpr (w_type.size_bits() == 4) {
1129
+ b_quant_0 = frag_b_quant[k2][0][j];
1130
+ b_quant_1 = b_quant_0 >> 8;
1131
+ } else {
1132
+ static_assert(w_type.size_bits() == 8);
1133
+ int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
1134
+ b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
1135
+ b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1136
+ }
1137
+
1138
+ dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1139
+ dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1140
+
1141
+ // Apply scale to frag_b0
1142
+ if constexpr (has_act_order) {
1143
+ static_assert(group_blocks != -1);
1144
+ scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1145
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1146
+ scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1147
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1148
+ } else if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1149
+ int idx = (threadIdx.x / 4) % 2;
1150
+ scalar_t2 s2 = Dtype::nums2num2(
1151
+ reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
1152
+ reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
1153
+ if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1154
+ scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1155
+ scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1156
+ } else if constexpr (has_zp && group_blocks != -1) {
1157
+ if (is_new_zp)
1158
+ frag_zp[j] = __hmul2(frag_zp[j],
1159
+ *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1160
+ scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1161
+ scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1162
+ } else if constexpr (group_blocks != -1) {
1163
+ scale<scalar_t>(frag_b0, frag_s[k2][j], 0);
1164
+ scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
1165
+ }
1166
+
1167
+ #pragma unroll
1168
+ for (int i = 0; i < thread_m_blocks; i++) {
1169
+ if constexpr (m_block_size_8) {
1170
+ mma_trans<scalar_t>(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]);
1171
+ } else {
1172
+ mma<scalar_t>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1173
+ mma<scalar_t>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
1174
+ }
1175
+ }
1176
+ }
1177
+ };
1178
+
1179
+ // Since we slice across the k dimension of a tile in order to increase the
1180
+ // number of warps while keeping the n dimension of a tile reasonable, we have
1181
+ // multiple warps that accumulate their partial sums of the same output
1182
+ // location; which we have to reduce over in the end. We do in shared memory.
1183
+ auto thread_block_reduce = [&]() {
1184
+ constexpr int red_off = threads / b_sh_stride_threads / 2;
1185
+ if (red_off >= 1) {
1186
+ auto red_idx = threadIdx.x / b_sh_stride_threads;
1187
+ constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1188
+ constexpr int red_sh_delta = b_sh_stride_threads;
1189
+ int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
1190
+ (threadIdx.x % b_sh_stride_threads);
1191
+
1192
+ // Parallel logarithmic shared memory reduction. We make sure to avoid any
1193
+ // unnecessary read or write iterations, e.g., for two warps we write only
1194
+ // once by warp 1 and read only once by warp 0.
1195
+
1196
+ #pragma unroll
1197
+ for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
1198
+ #pragma unroll
1199
+ for (int i = red_off; i > 0; i /= 2) {
1200
+ if (i <= red_idx && red_idx < 2 * i) {
1201
+ #pragma unroll
1202
+ for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) {
1203
+ int red_sh_wr =
1204
+ red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
1205
+ if (i < red_off) {
1206
+ float* c_rd = reinterpret_cast<float*>(
1207
+ &sh_red[red_sh_delta * j + red_sh_rd]);
1208
+ float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
1209
+ #pragma unroll
1210
+ for (int k = 0; k < 4; k++)
1211
+ reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
1212
+ c_rd[k] + c_wr[k];
1213
+ }
1214
+ sh_red[red_sh_wr] =
1215
+ reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
1216
+ }
1217
+ }
1218
+ __syncthreads();
1219
+ }
1220
+ if (red_idx == 0) {
1221
+ #pragma unroll
1222
+ for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) {
1223
+ float* c_rd =
1224
+ reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
1225
+ #pragma unroll
1226
+ for (int j = 0; j < 4; j++)
1227
+ reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
1228
+ c_rd[j];
1229
+ }
1230
+ }
1231
+ __syncthreads();
1232
+ }
1233
+ }
1234
+ };
1235
+
1236
+ // Since multiple threadblocks may process parts of the same column slice, we
1237
+ // finally have to globally reduce over the results. As the striped
1238
+ // partitioning minimizes the number of such reductions and our outputs are
1239
+ // usually rather small, we perform this reduction serially in L2 cache.
1240
+ auto global_reduce_fp16 = [&](bool first = false, bool last = false) {
1241
+ // We are very careful here to reduce directly in the output buffer to
1242
+ // maximize L2 cache utilization in this step. To do this, we write out
1243
+ // results in FP16 (but still reduce with FP32 compute).
1244
+ constexpr int active_threads = 32 * thread_n_blocks / 4;
1245
+ if (threadIdx.x < active_threads) {
1246
+ int c_gl_stride = prob_n / 8;
1247
+ int c_gl_wr_delta_o = 8 * c_gl_stride;
1248
+ int c_gl_wr_delta_i = 4 * (active_threads / 32);
1249
+ int c_gl_wr;
1250
+ if constexpr (m_block_size_8) {
1251
+ c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) +
1252
+ 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8;
1253
+ c_gl_wr += (2 * thread_n_blocks) * slice_col;
1254
+ } else {
1255
+ c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
1256
+ 4 * (threadIdx.x / 32) + threadIdx.x % 4;
1257
+ c_gl_wr += (2 * thread_n_blocks) * slice_col;
1258
+ }
1259
+ constexpr int c_sh_wr_delta = active_threads;
1260
+ auto c_sh_wr = threadIdx.x;
1261
+
1262
+ int row = (threadIdx.x % 32) / 4;
1263
+
1264
+ if (!first) {
1265
+ // Interestingly, doing direct global accesses here really seems to mess up
1266
+ // the compiler and lead to slowdowns, hence we also use async-copies even
1267
+ // though these fetches are not actually asynchronous.
1268
+ #pragma unroll
1269
+ for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
1270
+ if constexpr (m_block_size_8) {
1271
+ cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i],
1272
+ &C[c_gl_wr + i * c_gl_stride +
1273
+ (threadIdx.x % 8) / 4 * c_gl_wr_delta_i],
1274
+ (threadIdx.x % 4) * 2 + i < prob_m);
1275
+ } else {
1276
+ cp_async4_pred(
1277
+ &sh_red[c_sh_wr + c_sh_wr_delta * i],
1278
+ &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1279
+ c_gl_wr_delta_i * (i % 2)],
1280
+ i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
1281
+ }
1282
+ }
1283
+ cp_async_fence();
1284
+ cp_async_wait<0>();
1285
+ }
1286
+
1287
+ #pragma unroll
1288
+ for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
1289
+ bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 ||
1290
+ 8 * (i / 2) + row < prob_m) ||
1291
+ (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m);
1292
+ if (mask) {
1293
+ if (!first) {
1294
+ int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
1295
+ #pragma unroll
1296
+ for (int j = 0; j < 2 * 4; j++) {
1297
+ int delta = 0;
1298
+ if constexpr (m_block_size_8) {
1299
+ delta = j % 2 == 1 ? -2 : 0;
1300
+ }
1301
+ reinterpret_cast<float*>(
1302
+ &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] +=
1303
+ Dtype::num2float(reinterpret_cast<scalar_t*>(&c_red)[j]);
1304
+ }
1305
+ }
1306
+ if (!last) {
1307
+ int4 c;
1308
+ #pragma unroll
1309
+ for (int j = 0; j < 2 * 4; j++) {
1310
+ int delta = 0;
1311
+ if constexpr (m_block_size_8) {
1312
+ delta = j % 2 == 1 ? -2 : 0;
1313
+ }
1314
+ reinterpret_cast<scalar_t*>(&c)[j] =
1315
+ Dtype::float2num(reinterpret_cast<float*>(
1316
+ &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]);
1317
+ }
1318
+ if constexpr (m_block_size_8)
1319
+ C[c_gl_wr + i * c_gl_stride +
1320
+ (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = c;
1321
+ else
1322
+ C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1323
+ c_gl_wr_delta_i * (i % 2)] = c;
1324
+ }
1325
+ }
1326
+ }
1327
+ }
1328
+ };
1329
+
1330
+ // Globally reduce over threadblocks that compute the same column block.
1331
+ // We use a tmp C buffer to reduce in full fp32 precision.
1332
+ auto global_reduce_fp32 = [&](bool first = false, bool last = false) {
1333
+ constexpr int tb_m = thread_m_blocks * 16;
1334
+ constexpr int tb_n = thread_n_blocks * 16;
1335
+
1336
+ constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
1337
+
1338
+ constexpr int active_threads = 32 * thread_n_blocks / 4;
1339
+ bool is_th_active = threadIdx.x < active_threads;
1340
+
1341
+ constexpr int num_floats = thread_m_blocks * 4 * 2 * 4;
1342
+ constexpr int th_size = num_floats * sizeof(float) / 16;
1343
+
1344
+ int c_cur_offset = locks_off * c_size;
1345
+
1346
+ if (!is_th_active) {
1347
+ return;
1348
+ }
1349
+
1350
+ if (!first) {
1351
+ float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
1352
+ #pragma unroll
1353
+ for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) {
1354
+ sh_red[threadIdx.x] =
1355
+ C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
1356
+
1357
+ float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
1358
+ #pragma unroll
1359
+ for (int f = 0; f < 4; f++) {
1360
+ frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
1361
+ }
1362
+ }
1363
+ }
1364
+
1365
+ if (!last) {
1366
+ int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
1367
+ #pragma unroll
1368
+ for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) {
1369
+ C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k];
1370
+ }
1371
+ }
1372
+ };
1373
+
1374
+ // Write out the reduce final result in the correct layout. We only actually
1375
+ // reshuffle matrix fragments in this step, the reduction above is performed
1376
+ // in fragment layout.
1377
+ auto write_result = [&]() {
1378
+ int c_gl_stride = prob_n / 8;
1379
+ constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
1380
+ int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
1381
+ constexpr int c_sh_rd_delta =
1382
+ c_sh_stride * (threads / (2 * thread_n_blocks));
1383
+
1384
+ int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1385
+ (threadIdx.x % (2 * thread_n_blocks));
1386
+ c_gl_wr += (2 * thread_n_blocks) * slice_col;
1387
+ int c_sh_wr;
1388
+ if constexpr (m_block_size_8) {
1389
+ c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) +
1390
+ (threadIdx.x % 32) / 4;
1391
+ c_sh_wr += 64 * (threadIdx.x / 32);
1392
+ } else {
1393
+ c_sh_wr =
1394
+ (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
1395
+ c_sh_wr += 32 * (threadIdx.x / 32);
1396
+ }
1397
+
1398
+ int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1399
+ (threadIdx.x % (2 * thread_n_blocks));
1400
+
1401
+ int c_gl_wr_end = c_gl_stride * prob_m;
1402
+ // We first reorder in shared memory to guarantee the most efficient final
1403
+ // global write patterns
1404
+ auto write = [&](int idx, float c0, float c1, FragS& s) {
1405
+ scalar_t2 res =
1406
+ Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1407
+
1408
+ // For per-column quantization we finally apply the scale here (only for
1409
+ // 4-bit)
1410
+ if constexpr (!has_act_order && group_blocks == -1 &&
1411
+ w_type.size_bits() == 4 && !has_zp) {
1412
+ res = __hmul2(res, s[0]);
1413
+ }
1414
+
1415
+ if constexpr (m_block_size_8) {
1416
+ ((scalar_t*)sh_red)[idx] = res.x;
1417
+ ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1418
+ } else {
1419
+ ((scalar_t2*)sh_red)[idx] = res;
1420
+ }
1421
+ };
1422
+
1423
+ if (threadIdx.x / 32 < thread_n_blocks / 4) {
1424
+ #pragma unroll
1425
+ for (int i = 0; i < thread_m_blocks; i++) {
1426
+ #pragma unroll
1427
+ for (int j = 0; j < 4; j++) {
1428
+ if constexpr (m_block_size_8) {
1429
+ int wr = c_sh_wr + 16 * j;
1430
+ write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
1431
+ frag_s[j / 2][2 * (j % 2) + 0]);
1432
+ write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
1433
+ frag_s[j / 2][2 * (j % 2) + 1]);
1434
+ } else {
1435
+ int wr = c_sh_wr + 8 * j;
1436
+ write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
1437
+ frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
1438
+ write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
1439
+ frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
1440
+ write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
1441
+ frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
1442
+ write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
1443
+ frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
1444
+ }
1445
+ }
1446
+ c_sh_wr += 16 * (4 * c_sh_stride);
1447
+ }
1448
+ }
1449
+ __syncthreads();
1450
+
1451
+ #pragma unroll
1452
+ for (int i = 0;
1453
+ i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
1454
+ i++) {
1455
+ if (c_gl_wr < c_gl_wr_end) {
1456
+ if (use_atomic_add && slice_count > 1) {
1457
+ scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[c_gl_wr]);
1458
+ scalar_t2* sh_red_half2 =
1459
+ reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
1460
+ #pragma unroll
1461
+ for (int a = 0; a < 4; a++) {
1462
+ atomicAdd(&C_half2[a], sh_red_half2[a]);
1463
+ }
1464
+ } else {
1465
+ C[c_gl_wr] = sh_red[c_sh_rd];
1466
+ }
1467
+ c_gl_wr += c_gl_wr_delta;
1468
+ c_sh_rd += c_sh_rd_delta;
1469
+ }
1470
+ }
1471
+ __syncthreads();
1472
+ };
1473
+
1474
+ // Start global fetch and register load pipelines.
1475
+ auto start_pipes = [&]() {
1476
+
1477
+ #pragma unroll
1478
+ for (int i = 0; i < stages - 1; i++) {
1479
+ if (has_act_order && i == 0) {
1480
+ int last_g_idx = slice_k_start + stages * tb_k * 2;
1481
+ if (last_g_idx >= prob_k) {
1482
+ last_g_idx = prob_k - 1;
1483
+ }
1484
+ fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
1485
+ g_idx[last_g_idx]);
1486
+ }
1487
+
1488
+ if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1489
+ if (i == 0) {
1490
+ fetch_col_zp_to_shared();
1491
+ fetch_col_scale_to_shared();
1492
+ }
1493
+ }
1494
+ fetch_to_shared(i, i, i < slice_iters);
1495
+ }
1496
+
1497
+ zero_accums();
1498
+ wait_for_stage();
1499
+ init_same_group(0);
1500
+ fetch_to_registers(0, 0);
1501
+ fetch_scales_to_registers(0, 0);
1502
+ fetch_zp_to_registers(0, 0);
1503
+ a_gl_rd += a_gl_rd_delta_o * (stages - 1);
1504
+ if constexpr (has_act_order) {
1505
+ slice_k_start_shared_fetch += tb_k * (stages - 1);
1506
+ }
1507
+ };
1508
+ if (slice_iters) {
1509
+ start_pipes();
1510
+ }
1511
+
1512
+ // Main loop.
1513
+ while (slice_iters) {
1514
+ // We unroll over both the global fetch and the register load pipeline to
1515
+ // ensure all shared memory accesses are static. Note that both pipelines
1516
+ // have even length meaning that the next iteration will always start at
1517
+ // index 0.
1518
+
1519
+ #pragma unroll
1520
+ for (int pipe = 0; pipe < stages;) {
1521
+ #pragma unroll
1522
+ for (int k = 0; k < b_sh_wr_iters; k++) {
1523
+ fetch_to_registers(k + 1, pipe % stages);
1524
+ fetch_scales_to_registers(k + 1, pipe);
1525
+ fetch_zp_to_registers(k + 1, pipe);
1526
+ if (k == b_sh_wr_iters - 2) {
1527
+ fetch_to_shared((pipe + stages - 1) % stages, pipe,
1528
+ slice_iters >= stages);
1529
+ pipe++;
1530
+ wait_for_stage();
1531
+ init_same_group(pipe % stages);
1532
+ }
1533
+ matmul(k);
1534
+ }
1535
+ slice_iters--;
1536
+ if (slice_iters == 0) {
1537
+ break;
1538
+ }
1539
+ }
1540
+
1541
+ a_gl_rd += a_gl_rd_delta_o * stages;
1542
+
1543
+ if constexpr (has_act_order) {
1544
+ slice_k_start += tb_k * stages;
1545
+ slice_k_start_shared_fetch += tb_k * stages;
1546
+ int first_group_id = g_idx[slice_k_start];
1547
+ int last_g_idx = slice_k_start + stages * tb_k * 2;
1548
+ if (last_g_idx >= prob_k) {
1549
+ last_g_idx = prob_k - 1;
1550
+ }
1551
+ int last_group_id = g_idx[last_g_idx];
1552
+ if (last_group_id >= sh_first_group_id + sh_num_groups) {
1553
+ fetch_act_order_scales_to_shared(false, first_group_id, last_group_id);
1554
+ __syncthreads();
1555
+ }
1556
+ }
1557
+
1558
+ // Process results and, if necessary, proceed to the next column slice.
1559
+ // While this pattern may not be the most readable, other ways of writing
1560
+ // the loop seemed to noticeably worse performance after compilation.
1561
+ if (slice_iters == 0) {
1562
+ cp_async_wait<0>();
1563
+ bool last = slice_idx == slice_count - 1;
1564
+ // For per-column scales, we only fetch them here in the final step before
1565
+ // write-out
1566
+ if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1567
+ if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1568
+ if (s_sh_wr_pred) {
1569
+ cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
1570
+ }
1571
+ cp_async_fence();
1572
+ }
1573
+ }
1574
+
1575
+ thread_block_reduce();
1576
+ if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1577
+ if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1578
+ cp_async_wait<0>();
1579
+ __syncthreads();
1580
+ if (threadIdx.x / 32 < thread_n_blocks / 4) {
1581
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
1582
+ reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1583
+ if constexpr (m_block_size_8) {
1584
+ int idx = (threadIdx.x / 4) % 2;
1585
+ scalar_t2* frag_s_half2 = reinterpret_cast<scalar_t2*>(frag_s);
1586
+ #pragma unroll
1587
+ for (int i = 0; i < 8; i++) {
1588
+ frag_s_half2[i] = Dtype::num2num2(
1589
+ reinterpret_cast<scalar_t*>(&frag_s_half2[i])[idx]);
1590
+ }
1591
+ }
1592
+ }
1593
+ }
1594
+ }
1595
+
1596
+ // For 8-bit channelwise, we apply the scale before the global reduction
1597
+ // that converts the fp32 results to fp16 (so that we avoid possible
1598
+ // overflow in fp16)
1599
+ if constexpr (!has_act_order && group_blocks == -1 &&
1600
+ w_type.size_bits() == 8 && !has_zp) {
1601
+ if (threadIdx.x / 32 < thread_n_blocks / 4) {
1602
+ #pragma unroll
1603
+ for (int i = 0; i < thread_m_blocks; i++) {
1604
+ #pragma unroll
1605
+ for (int j = 0; j < 4; j++) {
1606
+ scale_float<scalar_t>(
1607
+ reinterpret_cast<float*>(&frag_c[i][j][0][0]),
1608
+ frag_s[j / 2][2 * (j % 2) + 0]);
1609
+ scale_float<scalar_t>(
1610
+ reinterpret_cast<float*>(&frag_c[i][j][0][2]),
1611
+ frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
1612
+
1613
+ if constexpr (!m_block_size_8) {
1614
+ scale_float<scalar_t>(
1615
+ reinterpret_cast<float*>(&frag_c[i][j][1][0]),
1616
+ frag_s[j / 2][2 * (j % 2) + 1]);
1617
+ scale_float<scalar_t>(
1618
+ reinterpret_cast<float*>(&frag_c[i][j][1][2]),
1619
+ frag_s[j / 2][2 * (j % 2) + 1]);
1620
+ }
1621
+ }
1622
+ }
1623
+ }
1624
+ }
1625
+
1626
+ if (slice_count > 1 && !use_atomic_add) {
1627
+ // only globally reduce if there is more than one block in a slice
1628
+ barrier_acquire(&locks[locks_off], slice_idx);
1629
+ if (use_fp32_reduce) {
1630
+ global_reduce_fp32(slice_idx == 0, last);
1631
+ } else {
1632
+ global_reduce_fp16(slice_idx == 0, last);
1633
+ }
1634
+ barrier_release(&locks[locks_off], last);
1635
+ }
1636
+ if (use_atomic_add && slice_count > 1 && slice_idx != 0)
1637
+ wait_negative_and_add(&locks[locks_off]);
1638
+ if (last || use_atomic_add)
1639
+ // only the last block in a slice actually writes the result
1640
+ write_result();
1641
+ slice_row = 0;
1642
+ slice_col_par++;
1643
+ slice_col++;
1644
+ is_first_matmul_in_slice = true;
1645
+ init_slice();
1646
+
1647
+ if (slice_iters) {
1648
+ a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
1649
+ (threadIdx.x % a_gl_rd_delta_o);
1650
+ #pragma unroll
1651
+ for (int i = 0; i < b_sh_wr_iters; i++)
1652
+ B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
1653
+ if (slice_col == 0) {
1654
+ #pragma unroll
1655
+ for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1656
+ }
1657
+
1658
+ // Update slice k/n for scales loading
1659
+ if constexpr (has_act_order) {
1660
+ slice_k_start = tb_k * slice_row;
1661
+ slice_k_finish = slice_k_start + tb_k * slice_iters;
1662
+ slice_k_start_shared_fetch = slice_k_start;
1663
+ slice_n_offset = act_s_col_tb_stride * slice_col;
1664
+
1665
+ } else {
1666
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1667
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1668
+ }
1669
+
1670
+ start_pipes();
1671
+ }
1672
+ }
1673
+ }
1674
+}
1675
+
1676
+} // namespace MARLIN_NAMESPACE_NAME
1677
+
1678
+#endif
csrc/torch_bindings.cpp
+4
-13
@@ -291,12 +291,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
291
292
// gptq_marlin Optimized Quantized GEMM for GPTQ.
293
ops.def(
294
- "gptq_marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
295
- "Tensor b_zeros, Tensor g_idx, Tensor perm, Tensor workspace, "
296
- "int b_q_type, "
294
+ "gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
295
+ "Tensor b_scales, Tensor? b_zeros_or_none, Tensor? g_idx_or_none, "
296
+ "Tensor? perm_or_none, Tensor workspace, int b_q_type, "
297
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
298
- "bool has_zp, bool use_atomic_add, bool use_fp32_reduce, "
299
- "bool is_zp_float) -> Tensor",
298
+ "bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
299
{stride_tag});
300
// conditionally compiled so impl registration is in source file
301
@@ -341,14 +340,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
340
ops.def("ggml_moe_get_block_size", &ggml_moe_get_block_size);
341
342
#ifndef USE_ROCM
344
- // fp8_marlin Optimized Quantized GEMM for FP8 weight-only.
345
- ops.def(
346
- "fp8_marlin_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, "
347
- "Tensor! workspace, int num_bits, SymInt size_m, SymInt size_n, "
348
- "SymInt size_k) -> Tensor",
349
- {stride_tag});
350
- // conditionally compiled so impl registration is in source file
351
-
343
// marlin_qqq_gemm for QQQ.
344
ops.def(
345
"marlin_qqq_gemm(Tensor a, Tensor b_q_weight, "
tests/kernels/moe/test_moe.py
+44
-121
@@ -11,19 +11,20 @@ from transformers import MixtralConfig
11
from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
12
13
import vllm.model_executor.layers.fused_moe # noqa
14
-from tests.kernels.utils import (opcheck, stack_and_dev, torch_moe,
15
- torch_moe_single)
14
+from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
15
from vllm.model_executor.layers.fused_moe import fused_moe
16
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
17
from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
18
fused_moe as iterative_moe)
19
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
20
+ marlin_quant_fp8_torch)
21
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
22
awq_marlin_quantize, marlin_quantize)
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
quantize_weights)
25
from vllm.model_executor.models.mixtral import MixtralMoE
26
from vllm.platforms import current_platform
26
-from vllm.scalar_type import scalar_types
27
+from vllm.scalar_type import ScalarType, scalar_types
28
29
NUM_EXPERTS = [8, 64]
30
EP_SIZE = [1, 4]
@@ -285,7 +286,7 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
286
atol=mixtral_moe_tol[dtype])
287
288
288
-@pytest.mark.parametrize("m", [1, 33, 123])
289
+@pytest.mark.parametrize("m", [1, 123, 666])
290
@pytest.mark.parametrize("n", [128, 1024])
291
@pytest.mark.parametrize("k", [256, 2048])
292
@pytest.mark.parametrize("e", [4, 12])
@@ -294,8 +295,10 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
295
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
296
@pytest.mark.parametrize("group_size", [-1, 32, 128])
297
@pytest.mark.parametrize("act_order", [True, False])
297
-@pytest.mark.parametrize("num_bits", [4, 8])
298
-@pytest.mark.parametrize("has_zp", [True, False])
298
+@pytest.mark.parametrize("quant_type", [
299
+ scalar_types.uint4, scalar_types.uint8b128, scalar_types.uint4b8,
300
+ scalar_types.float8_e4m3fn
301
+])
302
@pytest.mark.parametrize("is_k_full", [True, False])
303
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
304
def test_fused_marlin_moe(
@@ -308,14 +311,22 @@ def test_fused_marlin_moe(
311
dtype: torch.dtype,
312
group_size: int,
313
act_order: bool,
311
- num_bits: int,
312
- has_zp: bool,
314
+ quant_type: ScalarType,
315
is_k_full: bool,
316
):
315
- current_platform.seed_everything(7)
317
+ torch.cuda.manual_seed(0)
318
+ has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
319
+
320
+ if quant_type == scalar_types.float8_e4m3fn:
321
+ if group_size not in [-1, 128]:
322
+ return
323
+ if act_order:
324
+ return
325
326
# Filter act_order
327
if act_order:
328
+ if quant_type == scalar_types.float8_e4m3fn:
329
+ return
330
if group_size == -1:
331
return
332
if group_size in (k, n):
@@ -326,17 +337,9 @@ def test_fused_marlin_moe(
337
if not is_k_full:
338
return
339
329
- if has_zp:
330
- # we don't build kernel for int8 with zero
331
- if num_bits == 8:
332
- return
333
- quant_type = scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
334
- else:
335
- quant_type = scalar_types.uint4b8 \
336
- if num_bits == 4 else scalar_types.uint8b128
340
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
338
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
339
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
341
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
342
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
343
344
if ep_size > 1:
345
local_e = e // ep_size
@@ -364,17 +367,23 @@ def test_fused_marlin_moe(
367
qweight1_l.append(qweight1)
368
scales1_l.append(scales1)
369
zeros1_l.append(zeros1)
367
- else:
370
+ elif quant_type != scalar_types.float8_e4m3fn:
371
test_perm = torch.randperm(k)
369
- quant_res = marlin_quantize(w1[i].transpose(1, 0), quant_type,
370
- group_size, act_order, test_perm)
371
- w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = quant_res
372
+ w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = \
373
+ marlin_quantize(w1[i].transpose(1, 0), quant_type,
374
+ group_size, act_order, test_perm)
375
376
w_ref1_l.append(w_ref1.T)
377
qweight1_l.append(qweight1)
378
scales1_l.append(scales1)
379
g_idx1_l.append(g_idx1)
380
sort_indices1_l.append(sort_indices1)
381
+ else:
382
+ w_ref1, qweight1, scales1 = marlin_quant_fp8_torch(
383
+ w1[i], group_size)
384
+ w_ref1_l.append(w_ref1.T)
385
+ qweight1_l.append(qweight1)
386
+ scales1_l.append(scales1)
387
388
w_ref1 = stack_and_dev(w_ref1_l)
389
qweight1 = stack_and_dev(qweight1_l).contiguous()
@@ -399,17 +408,23 @@ def test_fused_marlin_moe(
408
qweight2_l.append(qweight2)
409
scales2_l.append(scales2)
410
zeros2_l.append(zeros2)
402
- else:
411
+ elif quant_type != scalar_types.float8_e4m3fn:
412
test_perm = torch.randperm(n)
404
- quant_res = marlin_quantize(w2[i].transpose(1, 0), quant_type,
405
- group_size, act_order, test_perm)
406
- w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = quant_res
413
+ w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = \
414
+ marlin_quantize(w2[i].transpose(1, 0), quant_type,
415
+ group_size, act_order, test_perm)
416
417
w_ref2_l.append(w_ref2.T)
418
qweight2_l.append(qweight2)
419
scales2_l.append(scales2)
420
g_idx2_l.append(g_idx2)
421
sort_indices2_l.append(sort_indices2)
422
+ else:
423
+ w_ref2, qweight2, scales2 = marlin_quant_fp8_torch(
424
+ w2[i], group_size)
425
+ w_ref2_l.append(w_ref2.T)
426
+ qweight2_l.append(qweight2)
427
+ scales2_l.append(scales2)
428
429
w_ref2 = stack_and_dev(w_ref2_l)
430
qweight2 = stack_and_dev(qweight2_l).contiguous()
@@ -442,102 +457,10 @@ def test_fused_marlin_moe(
457
sort_indices2=sort_indices2,
458
w1_zeros=zeros1,
459
w2_zeros=zeros2,
445
- num_bits=num_bits,
460
+ quant_type_id=quant_type.id,
461
is_k_full=is_k_full)
462
448
- torch.testing.assert_close(marlin_output, torch_output, atol=2e-2, rtol=0)
449
-
450
-
451
-@pytest.mark.skip("This test is here for the sake of debugging, "
452
- "don't run it in automated tests.")
453
-@pytest.mark.parametrize("m", [1, 33, 123])
454
-@pytest.mark.parametrize("n", [128, 1024])
455
-@pytest.mark.parametrize("k", [256, 2048])
456
-@pytest.mark.parametrize("e", [4, 12])
457
-@pytest.mark.parametrize("topk", [2, 3])
458
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
459
-@pytest.mark.parametrize("group_size", [-1, 32, 128])
460
-@pytest.mark.parametrize("act_order", [True, False])
461
-@pytest.mark.parametrize("num_bits", [4, 8])
462
-@pytest.mark.parametrize("has_zp", [True, False])
463
-@pytest.mark.parametrize("is_k_full", [True, False])
464
-def test_single_marlin_moe_multiply(m: int, n: int, k: int, e: int, topk: int,
465
- dtype: torch.dtype, group_size: int,
466
- act_order: bool, num_bits: int,
467
- has_zp: bool, is_k_full: bool):
468
- # Filter act_order
469
- if act_order:
470
- if group_size == -1:
471
- return
472
- if group_size in (k, n):
473
- return
474
- if has_zp:
475
- return
476
- else:
477
- if not is_k_full:
478
- return
479
-
480
- if has_zp:
481
- quant_type = scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
482
- else:
483
- quant_type = scalar_types.uint4b8 \
484
- if num_bits == 4 else scalar_types.uint8b128
485
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
486
- w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10
487
-
488
- w_ref_l = []
489
- qweight_l = []
490
- scales_l = []
491
- zeros_l = []
492
- g_idx_l = []
493
- sort_indices_l = []
494
-
495
- for i in range(w.shape[0]):
496
- if has_zp:
497
- w_ref, qweight, scales, zeros = awq_marlin_quantize(
498
- w[i].transpose(1, 0), quant_type, group_size)
499
-
500
- w_ref_l.append(w_ref.T)
501
- qweight_l.append(qweight)
502
- scales_l.append(scales)
503
- zeros_l.append(zeros)
504
- else:
505
- test_perm = torch.randperm(k)
506
- w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
507
- w[i].transpose(1, 0), quant_type, group_size, act_order,
508
- test_perm)
509
-
510
- w_ref_l.append(w_ref.T)
511
- qweight_l.append(qweight)
512
- scales_l.append(scales)
513
- g_idx_l.append(g_idx)
514
- sort_indices_l.append(sort_indices)
515
-
516
- w_ref = stack_and_dev(w_ref_l)
517
- qweight = stack_and_dev(qweight_l).contiguous()
518
- scales = stack_and_dev(scales_l)
519
- g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
520
- zeros = stack_and_dev(zeros_l) if zeros_l else None
521
- sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
522
-
523
- score = torch.randn((m, e), device="cuda", dtype=dtype)
524
- marlin_output = torch.ops.vllm.single_marlin_moe(
525
- a,
526
- qweight,
527
- scales,
528
- score,
529
- topk,
530
- renormalize=False,
531
- g_idx=g_idx,
532
- sort_indices=sort_indices,
533
- w_zeros=zeros,
534
- num_bits=num_bits,
535
- is_k_full=is_k_full,
536
- )
537
-
538
- torch_output = torch_moe_single(a, w_ref, score, topk)
539
-
540
- torch.testing.assert_close(marlin_output, torch_output, atol=2e-2, rtol=0)
463
+ torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
464
465
466
def test_moe_align_block_size_opcheck():
tests/kernels/quantization/test_awq_marlin.py
deleted
-164
@@ -1,164 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-"""Test AWQ with fused MoE Marlin kernels.
3
-
4
-Run `pytest tests/kernels/test_awq_marlin.py`.
5
-"""
6
-import pytest
7
-import torch
8
-
9
-import vllm.model_executor.layers.fused_moe # noqa
10
-from tests.kernels.utils import (compute_max_diff, stack_and_dev, torch_moe,
11
- torch_moe_single)
12
-from vllm import _custom_ops as ops
13
-from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14
-from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
15
- awq_marlin_quantize)
16
-from vllm.scalar_type import scalar_types
17
-
18
-NUM_EXPERTS = [8, 64]
19
-TOP_KS = [2, 6]
20
-GROUP_SIZES = [-1, 32, 128]
21
-
22
-
23
-@pytest.mark.parametrize("m", [1, 33, 64, 222])
24
-@pytest.mark.parametrize("n", [128, 2048])
25
-@pytest.mark.parametrize("k", [128, 1024])
26
-@pytest.mark.parametrize("e", NUM_EXPERTS)
27
-@pytest.mark.parametrize("topk", TOP_KS)
28
-@pytest.mark.parametrize("group_size", GROUP_SIZES)
29
-@pytest.mark.skipif(not (ops.supports_moe_ops
30
- and hasattr(torch.ops._moe_C, "marlin_gemm_moe")),
31
- reason="Marlin is not supported on this GPU type.")
32
-def test_fused_marlin_moe_awq(
33
- m: int,
34
- n: int,
35
- k: int,
36
- e: int,
37
- topk: int,
38
- group_size: int,
39
-):
40
- torch.manual_seed(7)
41
-
42
- num_bits = 4
43
- quant_type = scalar_types.uint4
44
- dtype = torch.float16
45
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
46
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
47
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
48
-
49
- w_ref1_l = []
50
- qweights1_l = []
51
- scales1_l = []
52
- zp1_l = []
53
-
54
- for i in range(w1.shape[0]):
55
- w_ref1, qweight1, scales1, zp1 = awq_marlin_quantize(
56
- w1[i].transpose(1, 0), quant_type, group_size)
57
- w_ref1_l.append(w_ref1)
58
- qweights1_l.append(qweight1)
59
- scales1_l.append(scales1)
60
- zp1_l.append(zp1)
61
-
62
- w_ref1 = stack_and_dev(w_ref1_l)
63
- qweight1 = stack_and_dev(qweights1_l).contiguous()
64
- scales1 = stack_and_dev(scales1_l)
65
- zp1 = stack_and_dev(zp1_l)
66
-
67
- w_ref2_l = []
68
- qweights2_l = []
69
- scales2_l = []
70
- zp2_l = []
71
-
72
- for i in range(w2.shape[0]):
73
- w_ref2, qweight2, scales2, zp2 = awq_marlin_quantize(
74
- w2[i].transpose(1, 0), quant_type, group_size)
75
- w_ref2_l.append(w_ref2)
76
- qweights2_l.append(qweight2)
77
- scales2_l.append(scales2)
78
- zp2_l.append(zp2)
79
-
80
- w_ref2 = stack_and_dev(w_ref2_l)
81
- qweight2 = stack_and_dev(qweights2_l).contiguous()
82
- scales2 = stack_and_dev(scales2_l)
83
- zp2 = stack_and_dev(zp2_l)
84
-
85
- score = torch.randn((m, e), device="cuda", dtype=dtype)
86
-
87
- topk_weights, topk_ids, token_expert_indices = fused_topk(
88
- a, score, topk, False)
89
- marlin_output = torch.ops.vllm.fused_marlin_moe(
90
- a,
91
- qweight1,
92
- qweight2,
93
- scales1,
94
- scales2,
95
- score,
96
- topk_weights,
97
- topk_ids,
98
- w1_zeros=zp1,
99
- w2_zeros=zp2,
100
- num_bits=num_bits,
101
- )
102
-
103
- torch_output = torch_moe(a, w_ref1.transpose(1, 2), w_ref2.transpose(1, 2),
104
- score, topk, None)
105
-
106
- assert compute_max_diff(marlin_output, torch_output) < 4e-2
107
-
108
-
109
-@pytest.mark.skip("This test is here for the sake of debugging, "
110
- "don't run it in automated tests.")
111
-@pytest.mark.parametrize("m", [64, 512, 222, 33, 1])
112
-@pytest.mark.parametrize("n", [128, 2048, 256, 1024])
113
-@pytest.mark.parametrize("k", [128, 1024, 512])
114
-@pytest.mark.parametrize("e", [8, 64])
115
-@pytest.mark.parametrize("topk", [2, 6])
116
-@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
117
-def test_single_marlin_moe_multiply_awq(
118
- m: int,
119
- n: int,
120
- k: int,
121
- e: int,
122
- topk: int,
123
- group_size: int,
124
-):
125
- torch.manual_seed(7)
126
-
127
- num_bits = 4
128
- quant_type = scalar_types.uint4
129
- dtype = torch.float16
130
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
131
- w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10
132
-
133
- w_ref_l = []
134
- qweights_l = []
135
- scales_l = []
136
- zp_l = []
137
-
138
- for i in range(w.shape[0]):
139
- w_ref, qweight, scales, zp = awq_marlin_quantize(
140
- w[i].transpose(1, 0), quant_type, group_size)
141
- w_ref_l.append(w_ref)
142
- qweights_l.append(qweight)
143
- scales_l.append(scales)
144
- zp_l.append(zp)
145
-
146
- w_ref = stack_and_dev(w_ref_l)
147
- qweight = stack_and_dev(qweights_l).contiguous()
148
- scales = stack_and_dev(scales_l).contiguous()
149
- zp = stack_and_dev(zp_l).contiguous()
150
-
151
- score = torch.randn((m, e), device="cuda", dtype=dtype)
152
-
153
- marlin_output = torch.ops.vllm.single_marlin_moe(a,
154
- qweight,
155
- scales,
156
- score,
157
- topk,
158
- renormalize=False,
159
- w_zeros=zp,
160
- num_bits=num_bits)
161
-
162
- torch_output = torch_moe_single(a, w_ref.transpose(1, 2), score, topk)
163
-
164
- assert compute_max_diff(marlin_output, torch_output) < 1e-2
tests/kernels/quantization/test_marlin_gemm.py
+38
-103
@@ -18,9 +18,10 @@ from vllm.model_executor.layers.quantization.qqq import (
18
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
19
GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
20
MARLIN_SUPPORTED_GROUP_SIZES, marlin_make_empty_g_idx,
21
- marlin_permute_scales, query_marlin_supported_quant_types)
21
+ marlin_make_workspace_new, marlin_permute_scales,
22
+ query_marlin_supported_quant_types)
23
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
23
- pack_fp8_to_int32)
24
+ marlin_quant_fp8_torch)
25
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
26
MarlinWorkspace, awq_marlin_quantize, get_weight_perm, marlin_quantize,
27
marlin_weights)
@@ -73,7 +74,7 @@ def rand_data(shape, dtype=torch.float16):
74
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
75
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
76
@pytest.mark.parametrize("quant_type",
76
- query_marlin_supported_quant_types(False))
77
+ query_marlin_supported_quant_types(False, False))
78
@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
79
@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
80
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
@@ -138,7 +139,7 @@ def test_gptq_marlin_repack(k_chunk, n_chunk, quant_type, group_size,
139
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
140
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
141
@pytest.mark.parametrize("quant_type",
141
- query_marlin_supported_quant_types(False))
142
+ query_marlin_supported_quant_types(True))
143
@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
144
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
145
def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size,
@@ -220,38 +221,50 @@ def test_gptq_marlin_gemm(
221
if group_size == size_k:
222
return
223
224
+ if size_k % group_size != 0:
225
+ return
226
+
227
a_input = rand_data((size_m, size_k))
228
b_weight = rand_data((size_k, size_n))
229
226
- w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
227
- b_weight, quant_type, group_size, act_order)
230
+ if quant_type == scalar_types.float8_e4m3fn:
231
+ if group_size not in [-1, 128]:
232
+ return
233
+ if act_order:
234
+ return
235
+ w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(
236
+ b_weight.T, group_size)
237
+ g_idx = None
238
+ sort_indices = None
239
+ else:
240
+ w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
241
+ b_weight, quant_type, group_size, act_order)
242
243
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
244
231
- workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
232
- GPTQ_MARLIN_MAX_PARALLEL)
245
+ workspace = marlin_make_workspace_new(w_ref.device)
246
234
- opcheck(torch.ops._C.gptq_marlin_gemm,
235
- (a_input, marlin_q_w, marlin_s, marlin_zp, g_idx, sort_indices,
236
- workspace.scratch, quant_type.id, a_input.shape[0],
237
- b_weight.shape[1], a_input.shape[1], is_k_full, False,
238
- use_atomic_add, use_fp32_reduce, False),
239
- test_utils=DEFAULT_OPCHECK_TEST_UTILS)
247
+ opcheck(
248
+ torch.ops._C.gptq_marlin_gemm,
249
+ (a_input, None, marlin_q_w, marlin_s, marlin_zp, g_idx, sort_indices,
250
+ workspace, quant_type.id, a_input.shape[0], b_weight.shape[1],
251
+ a_input.shape[1], is_k_full, use_atomic_add, use_fp32_reduce, False),
252
+ test_utils=DEFAULT_OPCHECK_TEST_UTILS)
253
254
output = ops.gptq_marlin_gemm(
255
a_input,
256
+ None,
257
marlin_q_w,
258
marlin_s,
259
marlin_zp,
260
g_idx,
261
sort_indices,
248
- workspace.scratch,
262
+ workspace,
263
quant_type,
264
a_input.shape[0],
265
b_weight.shape[1],
266
a_input.shape[1],
267
is_k_full=is_k_full,
254
- has_zp=False,
268
use_atomic_add=use_atomic_add,
269
use_fp32_reduce=use_fp32_reduce,
270
is_zp_float=False,
@@ -326,80 +339,6 @@ def test_gptq_marlin_24_gemm(k_chunk, n_chunk, quant_type, group_size,
339
assert max_diff < 0.04
340
341
329
-@pytest.mark.skipif(not is_quant_method_supported("fp8"),
330
- reason="Marlin is not supported on this GPU type.")
331
-@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
332
-@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
333
-@pytest.mark.parametrize("num_bits", [8])
334
-@pytest.mark.parametrize("group_size", [-1])
335
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
336
-@pytest.mark.parametrize("dtype", DTYPES)
337
-def test_fp8_marlin_gemm(
338
- k_chunk,
339
- n_chunk,
340
- num_bits,
341
- group_size,
342
- mnk_factors,
343
- dtype,
344
-):
345
- m_factor, n_factor, k_factor = mnk_factors
346
-
347
- size_m = m_factor
348
- size_k = k_chunk * k_factor
349
- size_n = n_chunk * n_factor
350
-
351
- a_input = rand_data((size_m, size_k), dtype=dtype)
352
- b_weight = rand_data((size_k, size_n), dtype=dtype)
353
-
354
- # WEIGHTS
355
- fp8_weight, weight_scale = ops.scaled_fp8_quant(b_weight, scale=None)
356
- # Repack weights to gptq format (packed int32 elements)
357
- packed_gptq_qweight = pack_fp8_to_int32(fp8_weight)
358
- # Repack weights to marlin format
359
- marlin_qweight = ops.gptq_marlin_repack(
360
- b_q_weight=packed_gptq_qweight,
361
- perm=torch.empty(0, dtype=torch.int, device="cuda"),
362
- size_k=size_k,
363
- size_n=size_n,
364
- num_bits=8,
365
- )
366
-
367
- # WEIGHT SCALES
368
- # Currently Marlin doesn't support per-tensor scales, so we
369
- # expand it to channelwise
370
- scales = weight_scale.repeat(1, size_n).to(a_input.dtype).to("cuda")
371
- # Permute scales
372
- marlin_scales = marlin_permute_scales(s=scales,
373
- size_k=size_k,
374
- size_n=size_n,
375
- group_size=-1)
376
-
377
- workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
378
- GPTQ_MARLIN_MAX_PARALLEL)
379
-
380
- opcheck(torch.ops._C.fp8_marlin_gemm,
381
- (a_input, marlin_qweight, marlin_scales, workspace.scratch,
382
- num_bits, a_input.shape[0], b_weight.shape[1], a_input.shape[1]))
383
-
384
- output = ops.fp8_marlin_gemm(
385
- a=a_input,
386
- b_q_weight=marlin_qweight,
387
- b_scales=marlin_scales,
388
- workspace=workspace.scratch,
389
- num_bits=num_bits,
390
- size_m=a_input.shape[0],
391
- size_n=b_weight.shape[1],
392
- size_k=a_input.shape[1],
393
- )
394
- output_ref = torch.matmul(a_input, b_weight)
395
-
396
- torch.cuda.synchronize()
397
-
398
- max_diff = compute_max_diff(output, output_ref)
399
-
400
- assert max_diff < 0.04
401
-
402
-
342
@pytest.mark.skipif(not is_quant_method_supported("gptq_marlin"),
343
reason="Marlin is not supported on this GPU type.")
344
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
@@ -432,25 +371,23 @@ def test_awq_marlin_gemm(
371
g_idx = torch.empty(0, dtype=torch.int, device=marlin_q_w.device)
372
sort_indices = torch.empty(0, dtype=torch.int, device=marlin_q_w.device)
373
is_k_full = True
435
- has_zp = True
374
437
- workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
438
- GPTQ_MARLIN_MAX_PARALLEL)
375
+ workspace = marlin_make_workspace_new(a_input.device)
376
377
output = ops.gptq_marlin_gemm(
378
a_input,
379
+ None,
380
marlin_q_w,
381
marlin_s,
382
marlin_zp,
383
g_idx,
384
sort_indices,
447
- workspace.scratch,
385
+ workspace,
386
quant_type,
387
a_input.shape[0],
388
b_weight.shape[1],
389
a_input.shape[1],
390
is_k_full=is_k_full,
453
- has_zp=has_zp,
391
use_fp32_reduce=use_fp32_reduce,
392
is_zp_float=False,
393
)
@@ -508,23 +445,22 @@ def test_hqq_marlin_gemm(
445
g_idx = marlin_make_empty_g_idx(dev)
446
g_idx_sort_indices = marlin_make_empty_g_idx(dev)
447
511
- workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
512
- GPTQ_MARLIN_MAX_PARALLEL)
448
+ workspace = marlin_make_workspace_new(b_weight.device)
449
450
output = ops.gptq_marlin_gemm(
451
a_input,
452
+ None,
453
marlin_w_q,
454
marlin_s,
455
marlin_zp,
456
g_idx,
457
g_idx_sort_indices,
521
- workspace.scratch,
458
+ workspace,
459
quant_type,
460
a_input.shape[0],
461
b_weight.shape[0],
462
a_input.shape[1],
463
is_k_full=True,
527
- has_zp=True,
464
use_fp32_reduce=use_fp32_reduce,
465
is_zp_float=True,
466
)
@@ -621,23 +557,22 @@ def test_marlin_gemm_subset_input():
557
b_weight, quant_type, group_size, False)
558
559
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
624
- workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
625
- GPTQ_MARLIN_MAX_PARALLEL)
560
+ workspace = marlin_make_workspace_new(a_input.device)
561
562
output = ops.gptq_marlin_gemm(
563
a_input,
564
+ None,
565
marlin_q_w,
566
marlin_s,
567
marlin_zp,
568
g_idx,
569
sort_indices,
634
- workspace.scratch,
570
+ workspace,
571
quant_type,
572
a_input.shape[0],
573
b_weight.shape[1],
574
a_input.shape[1],
575
is_k_full=True,
640
- has_zp=False,
576
use_atomic_add=False,
577
use_fp32_reduce=True,
578
is_zp_float=False,
vllm/_custom_ops.py
+14
-31
@@ -325,18 +325,18 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
325
326
@register_fake("_C::gptq_marlin_gemm")
327
def _gptq_marlin_gemm_fake(a: torch.Tensor,
328
+ c: Optional[torch.Tensor],
329
b_q_weight: torch.Tensor,
330
b_scales: torch.Tensor,
330
- b_zeros: torch.Tensor,
331
- g_idx: torch.Tensor,
332
- perm: torch.Tensor,
331
+ b_zeros: Optional[torch.Tensor],
332
+ g_idx: Optional[torch.Tensor],
333
+ perm: Optional[torch.Tensor],
334
workspace: torch.Tensor,
334
- b_q_type: ScalarType,
335
+ b_q_type_id: int,
336
size_m: torch.SymInt,
337
size_n: torch.SymInt,
338
size_k: torch.SymInt,
338
- is_k_full: bool,
339
- has_zp: bool = False,
339
+ is_k_full: bool = True,
340
use_atomic_add: bool = False,
341
use_fp32_reduce: bool = False,
342
is_zp_float: bool = False) -> torch.Tensor:
@@ -407,14 +407,6 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
407
dtype=codebooks.dtype,
408
device=codebooks.device)
409
410
- @register_fake("_C::fp8_marlin_gemm")
411
- def _fp8_marlin_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
412
- b_scales: torch.Tensor, workspace: torch.Tensor,
413
- num_bits: int, size_m: torch.SymInt,
414
- size_n: torch.SymInt,
415
- size_k: torch.SymInt) -> torch.Tensor:
416
- return torch.empty((size_m, size_n), dtype=a.dtype, device=a.device)
417
-
410
@register_fake("_C::machete_mm")
411
def machete_mm_fake(
412
a: torch.Tensor,
@@ -815,35 +807,26 @@ def awq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
807
808
809
def gptq_marlin_gemm(a: torch.Tensor,
810
+ c: Optional[torch.Tensor],
811
b_q_weight: torch.Tensor,
812
b_scales: torch.Tensor,
820
- b_zeros: torch.Tensor,
821
- g_idx: torch.Tensor,
822
- perm: torch.Tensor,
813
+ b_zeros: Optional[torch.Tensor],
814
+ g_idx: Optional[torch.Tensor],
815
+ perm: Optional[torch.Tensor],
816
workspace: torch.Tensor,
817
b_q_type: ScalarType,
818
size_m: int,
819
size_n: int,
820
size_k: int,
828
- is_k_full: bool,
829
- has_zp: bool = False,
821
+ is_k_full: bool = True,
822
use_atomic_add: bool = False,
823
use_fp32_reduce: bool = False,
824
is_zp_float: bool = False) -> torch.Tensor:
833
- return torch.ops._C.gptq_marlin_gemm(a, b_q_weight, b_scales, b_zeros,
825
+ return torch.ops._C.gptq_marlin_gemm(a, c, b_q_weight, b_scales, b_zeros,
826
g_idx, perm, workspace, b_q_type.id,
827
size_m, size_n, size_k, is_k_full,
836
- has_zp, use_atomic_add,
837
- use_fp32_reduce, is_zp_float)
838
-
839
-
840
-# fp8 marlin
841
-def fp8_marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
842
- b_scales: torch.Tensor, workspace: torch.Tensor,
843
- num_bits: int, size_m: int, size_n: int,
844
- size_k: int) -> torch.Tensor:
845
- return torch.ops._C.fp8_marlin_gemm(a, b_q_weight, b_scales, workspace,
846
- num_bits, size_m, size_n, size_k)
828
+ use_atomic_add, use_fp32_reduce,
829
+ is_zp_float)
830
831
832
# machete
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+19
-170
@@ -7,163 +7,13 @@ import torch
7
8
import vllm._custom_ops as ops
9
from vllm.model_executor.layers.fused_moe.fused_moe import (
10
- fused_topk, moe_align_block_size, try_get_optimal_moe_config)
11
-from vllm.scalar_type import scalar_types
10
+ moe_align_block_size, try_get_optimal_moe_config)
11
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
12
+ marlin_make_workspace_new, maybe_warn_marlin_atomic_add)
13
+from vllm.scalar_type import ScalarType, scalar_types
14
from vllm.utils import direct_register_custom_op
15
16
15
-def get_scalar_type(num_bits: int, has_zp: bool):
16
- if has_zp:
17
- return scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
18
- else:
19
- return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128
20
-
21
-
22
-def single_marlin_moe(
23
- hidden_states: torch.Tensor,
24
- w: torch.Tensor,
25
- scales: torch.Tensor,
26
- gating_output: torch.Tensor,
27
- topk: int,
28
- renormalize: bool,
29
- global_num_experts: int = -1,
30
- expert_map: Optional[torch.Tensor] = None,
31
- g_idx: Optional[torch.Tensor] = None,
32
- sort_indices: Optional[torch.Tensor] = None,
33
- w_zeros: Optional[torch.Tensor] = None,
34
- workspace: Optional[torch.Tensor] = None,
35
- num_bits: int = 8,
36
- is_k_full: bool = True,
37
-) -> torch.Tensor:
38
- """
39
- This function computes the multiplication of hidden_states with expert
40
- weights used in Marlin MoE, using weights w and top-k gating mechanism.
41
- Its purpose is testing and debugging the fused MoE kernel.
42
-
43
- Parameters:
44
- - hidden_states (torch.Tensor): The input tensor to the Marlin Mul.
45
- - w (torch.Tensor): The set of expert weights.
46
- - scales (torch.Tensor): The quantization scales.
47
- - gating_output (torch.Tensor): The output of the gating operation
48
- (before softmax).
49
- - g_idx (Optional[torch.Tensor]): Optional act_order indices.
50
- - sort_indices (Optional[torch.Tensor]): Optional act_order input
51
- permutation.
52
- - topk (int): The number of top-k experts to select.
53
- - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
54
- - w_zeros (Optional[torch.Tensor]): Optional zero points to be used for w.
55
- - num_bits (bool): The number of bits in expert weights quantization.
56
-
57
- Returns:
58
- - torch.Tensor: The output tensor after applying the MoE layer.
59
- """
60
- # Check constraints.
61
- assert hidden_states.shape[0] == gating_output.shape[0], (
62
- "Number of tokens mismatch")
63
- assert hidden_states.shape[1] == w.shape[1] * 16, "Hidden size mismatch"
64
- assert gating_output.shape[1] == w.shape[0], "Number of experts mismatch"
65
- assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
66
- assert w.is_contiguous(), "Expert weights must be contiguous"
67
- assert hidden_states.dtype in [torch.float16, torch.bfloat16]
68
- assert num_bits in [4, 8]
69
-
70
- M, K = hidden_states.shape
71
- E = w.shape[0]
72
- N = w.shape[2] // (num_bits // 2)
73
-
74
- topk_weights, topk_ids, token_expert_indices = fused_topk(
75
- hidden_states, gating_output, topk, renormalize)
76
-
77
- # This might not be an optimal config for a single MMM
78
- get_config_func = functools.partial(try_get_optimal_moe_config,
79
- w.shape,
80
- w.shape,
81
- topk_ids.shape[1],
82
- None,
83
- is_marlin=True)
84
- config = get_config_func(M)
85
-
86
- block_size_m = config['BLOCK_SIZE_M']
87
-
88
- if global_num_experts == -1:
89
- global_num_experts = E
90
- sorted_token_ids, expert_ids, num_tokens_post_padded = \
91
- moe_align_block_size(topk_ids, block_size_m, E, expert_map)
92
-
93
- if workspace is None:
94
- max_workspace_size = (max(2 * N, K) // 64) * \
95
- (sorted_token_ids.size(0) // block_size_m)
96
- device = hidden_states.device
97
- sms = torch.cuda.get_device_properties(device).multi_processor_count
98
- max_workspace_size = min(max_workspace_size, sms)
99
- workspace = torch.zeros(max_workspace_size,
100
- dtype=torch.int,
101
- device=device,
102
- requires_grad=False)
103
-
104
- scalar_type = get_scalar_type(num_bits, w_zeros is not None)
105
- intermediate_cache = torch.empty(
106
- (M * topk_ids.shape[1], N),
107
- device=hidden_states.device,
108
- dtype=hidden_states.dtype,
109
- )
110
-
111
- ops.moe_wna16_marlin_gemm(hidden_states,
112
- intermediate_cache,
113
- w,
114
- scales,
115
- w_zeros,
116
- g_idx,
117
- sort_indices,
118
- workspace,
119
- sorted_token_ids,
120
- expert_ids,
121
- num_tokens_post_padded,
122
- topk_weights,
123
- moe_block_size=block_size_m,
124
- top_k=topk,
125
- mul_topk_weights=False,
126
- is_ep=expert_map is not None,
127
- b_q_type=scalar_type,
128
- size_m=M,
129
- size_n=N,
130
- size_k=K,
131
- is_k_full=is_k_full,
132
- use_atomic_add=False,
133
- use_fp32_reduce=True,
134
- is_zp_float=False)
135
- intermediate_cache = intermediate_cache.view(-1, topk, N)
136
-
137
- return torch.sum(intermediate_cache.view(*intermediate_cache.shape), dim=1)
138
-
139
-
140
-def single_marlin_moe_fake(
141
- hidden_states: torch.Tensor,
142
- w: torch.Tensor,
143
- scales: torch.Tensor,
144
- gating_output: torch.Tensor,
145
- topk: int,
146
- renormalize: bool,
147
- global_num_experts: int = -1,
148
- expert_map: Optional[torch.Tensor] = None,
149
- g_idx: Optional[torch.Tensor] = None,
150
- sort_indices: Optional[torch.Tensor] = None,
151
- w_zeros: Optional[torch.Tensor] = None,
152
- workspace: Optional[torch.Tensor] = None,
153
- num_bits: int = 8,
154
- is_k_full: bool = True,
155
-) -> torch.Tensor:
156
- return torch.empty_like(hidden_states)
157
-
158
-
159
-direct_register_custom_op(
160
- op_name="single_marlin_moe",
161
- op_func=single_marlin_moe,
162
- mutates_args=[],
163
- fake_impl=single_marlin_moe_fake,
164
-)
165
-
166
-
17
def fused_marlin_moe(hidden_states: torch.Tensor,
18
w1: torch.Tensor,
19
w2: torch.Tensor,
@@ -172,6 +22,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
22
gating_output: torch.Tensor,
23
topk_weights: torch.Tensor,
24
topk_ids: torch.Tensor,
25
+ quant_type_id: int,
26
global_num_experts: int = -1,
27
expert_map: Optional[torch.Tensor] = None,
28
g_idx1: Optional[torch.Tensor] = None,
@@ -181,7 +32,6 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
32
w1_zeros: Optional[torch.Tensor] = None,
33
w2_zeros: Optional[torch.Tensor] = None,
34
workspace: Optional[torch.Tensor] = None,
184
- num_bits: int = 8,
35
is_k_full: bool = True,
36
inplace: bool = False) -> torch.Tensor:
37
"""
@@ -211,6 +61,15 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
61
Returns:
62
- torch.Tensor: The output tensor after applying the MoE layer.
63
"""
64
+ quant_type = ScalarType.from_id(quant_type_id)
65
+ assert quant_type in [
66
+ scalar_types.uint4, scalar_types.uint8b128, scalar_types.uint4b8,
67
+ scalar_types.float8_e4m3fn
68
+ ]
69
+
70
+ int4_scalar_types = [scalar_types.uint4, scalar_types.uint4b8]
71
+ num_bits = 4 if quant_type in int4_scalar_types else 8
72
+
73
# Check constraints.
74
assert hidden_states.shape[0] == gating_output.shape[
75
0], "Number of tokens mismatch"
@@ -248,18 +107,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
107
expert_map)
108
109
if workspace is None:
251
- max_workspace_size = (max(2 * N, K) // 64) * \
252
- (sorted_token_ids.size(0) // block_size_m)
253
- device = hidden_states.device
254
- sms = torch.cuda.get_device_properties(device).multi_processor_count
255
- max_workspace_size = min(max_workspace_size, sms * 4)
256
- workspace = torch.zeros(max_workspace_size,
257
- dtype=torch.int,
258
- device=device,
259
- requires_grad=False)
260
-
261
- scalar_type1 = get_scalar_type(num_bits, w1_zeros is not None)
262
- scalar_type2 = get_scalar_type(num_bits, w2_zeros is not None)
110
+ workspace = marlin_make_workspace_new(hidden_states.device, 4)
111
112
intermediate_cache2 = torch.empty(
113
(M * topk_ids.shape[1], N),
@@ -276,6 +124,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
124
intermediate_cache3 = intermediate_cache13[:M * topk_ids.shape[1] * K]
125
intermediate_cache3 = intermediate_cache3.view(-1, K)
126
127
+ maybe_warn_marlin_atomic_add(hidden_states.device, hidden_states.dtype)
128
use_atomic_add = hidden_states.dtype == torch.half or \
129
torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
130
@@ -296,7 +145,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
145
top_k=topk,
146
mul_topk_weights=False,
147
is_ep=expert_map is not None,
299
- b_q_type=scalar_type1,
148
+ b_q_type=quant_type,
149
size_m=M,
150
size_n=2 * N,
151
size_k=K,
@@ -328,7 +177,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
177
top_k=1,
178
mul_topk_weights=True,
179
is_ep=expert_map is not None,
331
- b_q_type=scalar_type2,
180
+ b_q_type=quant_type,
181
size_m=M * topk,
182
size_n=K,
183
size_k=N,
@@ -351,6 +200,7 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
200
gating_output: torch.Tensor,
201
topk_weights: torch.Tensor,
202
topk_ids: torch.Tensor,
203
+ quant_type_id: int,
204
global_num_experts: int = -1,
205
expert_map: Optional[torch.Tensor] = None,
206
g_idx1: Optional[torch.Tensor] = None,
@@ -360,7 +210,6 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
210
w1_zeros: Optional[torch.Tensor] = None,
211
w2_zeros: Optional[torch.Tensor] = None,
212
workspace: Optional[torch.Tensor] = None,
363
- num_bits: int = 8,
213
is_k_full: bool = True,
214
inplace: bool = False) -> torch.Tensor:
215
return torch.empty_like(hidden_states)
vllm/model_executor/layers/quantization/awq_marlin.py
+11
-13
@@ -22,9 +22,10 @@ from vllm.model_executor.layers.quantization.utils import replace_parameter
22
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
23
apply_awq_marlin_linear, awq_to_marlin_zero_points, check_marlin_supported,
24
check_marlin_supports_layer, check_moe_marlin_supports_layer,
25
- marlin_make_empty_g_idx, marlin_make_workspace, marlin_moe_permute_scales,
26
- marlin_permute_scales, moe_awq_to_marlin_zero_points,
27
- verify_marlin_supported, verify_marlin_supports_shape)
25
+ marlin_make_empty_g_idx, marlin_make_workspace_new,
26
+ marlin_moe_permute_scales, marlin_permute_scales,
27
+ moe_awq_to_marlin_zero_points, verify_marlin_supported,
28
+ verify_marlin_supports_shape)
29
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
30
from vllm.model_executor.parameter import (GroupQuantScaleParameter,
31
PackedvLLMParameter)
@@ -267,8 +268,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):
268
requires_grad=False)
269
270
# Allocate marlin workspace
270
- layer.workspace = marlin_make_workspace(
271
- layer.output_size_per_partition, device)
271
+ layer.workspace = marlin_make_workspace_new(device)
272
273
# Repack weights from AWQ format to marlin format.
274
marlin_qweight = ops.awq_marlin_repack(
@@ -322,6 +322,9 @@ class AWQMoEMethod(FusedMoEMethodBase):
322
323
def __init__(self, quant_config: AWQMarlinConfig):
324
self.quant_config = quant_config
325
+ if self.quant_config.weight_bits != 4:
326
+ raise ValueError("AWQMoEMethod only supports 4bit now.")
327
+ self.quant_type = scalar_types.uint4
328
329
def create_weights(self, layer: torch.nn.Module, num_experts: int,
330
hidden_size: int, intermediate_size_per_partition: int,
@@ -396,11 +399,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
399
set_weight_attrs(w2_qzeros, extra_weight_attrs)
400
401
device = layer.w13_qweight.device
399
- sms = torch.cuda.get_device_properties(device).multi_processor_count
400
- layer.workspace = torch.zeros((sms * 4, ),
401
- dtype=torch.int,
402
- device=device,
403
- requires_grad=False)
402
+ layer.workspace = marlin_make_workspace_new(device, 4)
403
404
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
405
num_experts = layer.w13_qweight.shape[0]
@@ -511,10 +510,9 @@ class AWQMoEMethod(FusedMoEMethodBase):
510
router_logits,
511
topk_weights,
512
topk_ids,
513
+ quant_type_id=self.quant_type.id,
514
global_num_experts=global_num_experts,
515
expert_map=expert_map,
516
w1_zeros=layer.w13_qzeros,
517
w2_zeros=layer.w2_qzeros,
518
- workspace=layer.workspace,
519
- num_bits=self.quant_config.weight_bits,
520
- )
518
+ workspace=layer.workspace)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
+2
-1
@@ -55,7 +55,7 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
55
# required by torch.compile to be torch.nn.Parameter
56
layer.input_scale = torch.nn.Parameter(layer.input_scale.data,
57
requires_grad=False)
58
- prepare_fp8_layer_for_marlin(layer, strategy="channel")
58
+ prepare_fp8_layer_for_marlin(layer)
59
60
def create_weights(self, layer: torch.nn.Module, input_size: int,
61
output_partition_sizes: List[int],
@@ -68,6 +68,7 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
68
layer.input_size_per_partition = input_size_per_partition
69
layer.output_size_per_partition = output_size_per_partition
70
layer.orig_dtype = params_dtype
71
+ layer.weight_block_size = None
72
73
# WEIGHT
74
weight = ModelWeightParameter(data=torch.empty(
vllm/model_executor/layers/quantization/fp8.py
+57
-42
@@ -21,19 +21,21 @@ from vllm.model_executor.layers.quantization.base_config import (
21
QuantizationConfig, QuantizeMethodBase)
22
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
23
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
24
- apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin)
24
+ apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,
25
+ prepare_moe_fp8_layer_for_marlin)
26
from vllm.model_executor.layers.quantization.utils.quant_utils import (
27
is_layer_skipped)
28
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
28
- Fp8LinearOp, all_close_1d, convert_to_channelwise,
29
- cutlass_block_fp8_supported, cutlass_fp8_supported,
30
- maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz,
31
- per_tensor_dequantize, requantize_with_max_scale)
29
+ Fp8LinearOp, all_close_1d, cutlass_block_fp8_supported,
30
+ cutlass_fp8_supported, maybe_create_device_identity,
31
+ normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize,
32
+ requantize_with_max_scale)
33
from vllm.model_executor.parameter import (BlockQuantScaleParameter,
34
ModelWeightParameter,
35
PerTensorScaleParameter)
36
from vllm.model_executor.utils import set_weight_attrs
37
from vllm.platforms import current_platform
38
+from vllm.scalar_type import scalar_types
39
40
ACTIVATION_SCHEMES = ["static", "dynamic"]
41
@@ -181,10 +183,6 @@ class Fp8LinearMethod(LinearMethodBase):
183
self.use_marlin = False
184
185
self.block_quant = self.quant_config.weight_block_size is not None
184
- if self.block_quant:
185
- # Marlin doesn't support block-wise fp8
186
- self.use_marlin = False
187
-
186
self.fp8_linear = Fp8LinearOp(
187
# Default to using per_token quantization if cutlass is supported
188
use_per_token_if_dynamic=cutlass_fp8_supported())
@@ -203,10 +201,16 @@ class Fp8LinearMethod(LinearMethodBase):
201
202
output_size_per_partition = sum(output_partition_sizes)
203
weight_loader = extra_weight_attrs.get("weight_loader")
204
+ layer.logical_widths = output_partition_sizes
205
+ layer.input_size_per_partition = input_size_per_partition
206
+ layer.output_size_per_partition = output_size_per_partition
207
+ layer.orig_dtype = params_dtype
208
+ layer.weight_block_size = None
209
210
if self.block_quant:
211
tp_size = get_tensor_model_parallel_world_size()
212
assert self.quant_config.weight_block_size is not None
213
+ layer.weight_block_size = self.quant_config.weight_block_size
214
block_n, block_k = (
215
self.quant_config.weight_block_size[0],
216
self.quant_config.weight_block_size[1],
@@ -229,12 +233,6 @@ class Fp8LinearMethod(LinearMethodBase):
233
f"{output_partition_size} is not divisible by "
234
f"weight quantization block_n = {block_n}.")
235
232
- layer.logical_widths = output_partition_sizes
233
-
234
- layer.input_size_per_partition = input_size_per_partition
235
- layer.output_size_per_partition = output_size_per_partition
236
- layer.orig_dtype = params_dtype
237
-
236
# WEIGHT
237
weight_dtype = (torch.float8_e4m3fn
238
if self.quant_config.is_checkpoint_fp8_serialized else
@@ -303,9 +301,11 @@ class Fp8LinearMethod(LinearMethodBase):
301
return weight
302
303
def process_weights_after_loading(self, layer: Module) -> None:
304
+ size_k_first = True
305
# TODO(rob): refactor block quant into separate class.
306
if self.block_quant:
307
assert self.quant_config.activation_scheme == "dynamic"
308
+ size_k_first = False
309
if current_platform.is_fp8_fnuz():
310
weight, weight_scale_inv, _ = \
311
normalize_e4m3fn_to_e4m3fnuz(
@@ -321,21 +321,12 @@ class Fp8LinearMethod(LinearMethodBase):
321
layer.weight = Parameter(weight, requires_grad=False)
322
layer.weight_scale_inv = Parameter(weight_scale_inv,
323
requires_grad=False)
324
- return
324
325
# If checkpoint not serialized fp8, quantize the weights.
327
- if not self.quant_config.is_checkpoint_fp8_serialized:
326
+ elif not self.quant_config.is_checkpoint_fp8_serialized:
327
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight,
328
scale=None)
329
331
- # If using marlin (w8a16), kernel uses channelwise weights,
332
- # so extend the weight scales to be channelwise.
333
- if self.use_marlin:
334
- assert weight_scale.numel() == 1
335
- weight_scale = convert_to_channelwise(
336
- weight_scale.expand(len(layer.logical_widths)),
337
- layer.logical_widths)
338
-
330
# Update the layer with the new values.
331
layer.weight = Parameter(qweight.t(), requires_grad=False)
332
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
@@ -349,20 +340,14 @@ class Fp8LinearMethod(LinearMethodBase):
340
if self.quant_config.activation_scheme == "static":
341
layer.input_scale = torch.nn.Parameter(layer.input_scale.data,
342
requires_grad=False)
352
- # If using marlin (w8a16), kernel uses channelwise weights,
353
- # so extend the weight scales to be channelwise.
354
- if self.use_marlin:
355
- weight = layer.weight
356
- weight_scale = convert_to_channelwise(layer.weight_scale,
357
- layer.logical_widths)
343
+
344
+ weight = layer.weight
345
+ weight_scale = layer.weight_scale
346
347
# If using w8a8, torch._scaled_mm needs per tensor, so
348
# requantize the logical shards as a single weight.
361
- else:
349
+ if not self.use_marlin:
350
# Dequant -> Quant with max scale so we can run per tensor.
363
- weight = layer.weight
364
- weight_scale = layer.weight_scale
365
-
351
if current_platform.is_fp8_fnuz():
352
weight, weight_scale, input_scale = \
353
normalize_e4m3fn_to_e4m3fnuz(
@@ -388,7 +373,7 @@ class Fp8LinearMethod(LinearMethodBase):
373
requires_grad=False)
374
375
if self.use_marlin:
391
- prepare_fp8_layer_for_marlin(layer)
376
+ prepare_fp8_layer_for_marlin(layer, size_k_first)
377
# Activations not quantized for marlin.
378
del layer.input_scale
379
@@ -444,6 +429,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
429
self.quant_config = quant_config
430
self.block_quant = self.quant_config.weight_block_size is not None
431
432
+ # For GPUs that lack FP8 hardware support, we can leverage the Marlin
433
+ # kernel for fast weight-only FP8 quantization
434
+ self.use_marlin = (not current_platform.has_device_capability(89)
435
+ or envs.VLLM_TEST_FORCE_FP8_MARLIN)
436
+ # Disable marlin for rocm
437
+ if current_platform.is_rocm():
438
+ self.use_marlin = False
439
+
440
# Check for DeepGemm support.
441
self.allow_deep_gemm = False
442
if envs.VLLM_USE_DEEP_GEMM:
@@ -461,10 +454,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
454
intermediate_size_per_partition: int,
455
params_dtype: torch.dtype, **extra_weight_attrs):
456
457
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
458
+ layer.hidden_size = hidden_size
459
+ layer.num_experts = num_experts
460
+ layer.orig_dtype = params_dtype
461
+ layer.weight_block_size = None
462
+
463
if self.quant_config.is_checkpoint_fp8_serialized:
464
params_dtype = torch.float8_e4m3fn
465
if self.block_quant:
466
assert self.quant_config.weight_block_size is not None
467
+ layer.weight_block_size = self.quant_config.weight_block_size
468
tp_size = get_tensor_model_parallel_world_size()
469
block_n, block_k = (
470
self.quant_config.weight_block_size[0],
@@ -630,10 +630,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
630
layer.w2_weight_scale_inv = \
631
dg.get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()
632
633
- return
634
-
633
# If checkpoint is fp16, quantize in place.
636
- if not self.quant_config.is_checkpoint_fp8_serialized:
634
+ elif not self.quant_config.is_checkpoint_fp8_serialized:
635
fp8_dtype = current_platform.fp8_dtype()
636
w13_weight = torch.empty_like(layer.w13_weight.data,
637
dtype=fp8_dtype)
@@ -677,8 +675,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
675
requires_grad=False)
676
layer.w2_weight = torch.nn.Parameter(shuffled_w2,
677
requires_grad=False)
680
- return
681
-
678
# If checkpoint is fp8, we need to handle that the
679
# MoE kernels require single activation scale and single weight
680
# scale for w13 per expert.
@@ -766,7 +762,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
762
763
layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
764
requires_grad=False)
769
- return
765
+
766
+ if self.use_marlin:
767
+ prepare_moe_fp8_layer_for_marlin(layer, False)
768
+ # Activations not quantized for marlin.
769
+ del layer.w13_input_scale
770
+ del layer.w2_input_scale
771
772
def apply(
773
self,
@@ -801,6 +802,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
802
e_score_correction_bias=e_score_correction_bias,
803
)
804
805
+ if self.use_marlin:
806
+ return torch.ops.vllm.fused_marlin_moe(
807
+ x,
808
+ layer.w13_weight,
809
+ layer.w2_weight,
810
+ layer.w13_weight_scale,
811
+ layer.w2_weight_scale,
812
+ router_logits,
813
+ topk_weights,
814
+ topk_ids,
815
+ quant_type_id=scalar_types.float8_e4m3fn.id,
816
+ global_num_experts=global_num_experts,
817
+ expert_map=expert_map)
818
+
819
return fused_experts(
820
x,
821
layer.w13_weight,
vllm/model_executor/layers/quantization/gptq_marlin.py
+11
-8
@@ -21,8 +21,8 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import (
21
get_linear_quant_method)
22
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
23
check_marlin_supported, check_moe_marlin_supports_layer,
24
- marlin_moe_permute_scales, marlin_repeat_scales_on_all_ranks,
25
- verify_marlin_supported)
24
+ marlin_make_workspace_new, marlin_moe_permute_scales,
25
+ marlin_repeat_scales_on_all_ranks, verify_marlin_supported)
26
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
27
GroupQuantScaleParameter,
28
PackedColumnParameter,
@@ -350,6 +350,13 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
350
351
def __init__(self, quant_config: GPTQMarlinConfig) -> None:
352
self.quant_config = quant_config
353
+ if self.quant_config.quant_type.size_bits == 4:
354
+ self.quant_type = scalar_types.uint4b8
355
+ elif self.quant_config.quant_type.size_bits == 8:
356
+ self.quant_type = scalar_types.uint8b128
357
+ else:
358
+ raise ValueError(
359
+ "GPTQMarlinMoEMethod only supports int4 and int8 now.")
360
361
def create_weights(
362
self,
@@ -498,11 +505,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
505
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
506
507
device = layer.w13_qweight.device
501
- sms = torch.cuda.get_device_properties(device).multi_processor_count
502
- layer.workspace = torch.zeros((sms * 4, ),
503
- dtype=torch.int,
504
- device=device,
505
- requires_grad=False)
508
+ layer.workspace = marlin_make_workspace_new(device, 4)
509
510
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
511
@@ -633,12 +636,12 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
636
router_logits,
637
topk_weights,
638
topk_ids,
639
+ quant_type_id=self.quant_type.id,
640
global_num_experts=global_num_experts,
641
expert_map=expert_map,
642
g_idx1=layer.w13_g_idx,
643
g_idx2=layer.w2_g_idx,
644
sort_indices1=layer.w13_g_idx_sort_indices,
645
sort_indices2=layer.w2_g_idx_sort_indices,
642
- num_bits=self.quant_config.quant_type.size_bits,
646
workspace=layer.workspace,
647
is_k_full=self.is_k_full)
vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py
+2
-4
@@ -8,7 +8,7 @@ from vllm import _custom_ops as ops
8
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
9
MARLIN_SUPPORTED_GROUP_SIZES, apply_gptq_marlin_linear,
10
check_marlin_supports_shape, marlin_is_k_full, marlin_make_empty_g_idx,
11
- marlin_make_workspace, marlin_permute_scales, marlin_sort_g_idx,
11
+ marlin_make_workspace_new, marlin_permute_scales, marlin_sort_g_idx,
12
marlin_zero_points, query_marlin_supported_quant_types, unpack_cols)
13
from vllm.model_executor.parameter import (BasevLLMParameter,
14
permute_param_layout_)
@@ -53,8 +53,7 @@ class MarlinLinearKernel(MPLinearKernel):
53
self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
54
55
# Allocate marlin workspace.
56
- self.workspace = marlin_make_workspace(c.partition_weight_shape[1],
57
- device)
56
+ self.workspace = marlin_make_workspace_new(device)
57
58
# Default names since marlin requires empty parameters for these,
59
# TODO: remove this requirement from marlin (allow optional tensors)
@@ -127,6 +126,5 @@ class MarlinLinearKernel(MPLinearKernel):
126
wtype=c.weight_type,
127
input_size_per_partition=c.partition_weight_shape[0],
128
output_size_per_partition=c.partition_weight_shape[1],
130
- has_zp=self.config.zero_points,
129
is_k_full=self.is_k_full,
130
bias=bias)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+60
-16
@@ -7,12 +7,15 @@ import torch
7
8
import vllm.envs as envs
9
from vllm import _custom_ops as ops
10
+from vllm.logger import init_logger
11
from vllm.model_executor.layers.linear import LinearBase
12
from vllm.platforms import current_platform
13
from vllm.scalar_type import ScalarType, scalar_types
14
15
from .quant_utils import pack_cols, unpack_cols
16
17
+logger = init_logger(__name__)
18
+
19
GPTQ_MARLIN_TILE = 16
20
GPTQ_MARLIN_MIN_THREAD_N = 64
21
GPTQ_MARLIN_MIN_THREAD_K = 128
@@ -29,9 +32,11 @@ USE_FP32_REDUCE_DEFAULT = True
32
# For binary size and compile time, we don't support the same types for with and
33
# without runtime zero-point. We support common cases, i.e. AWQ and GPTQ.
34
# TODO: we may want to move this into the C++ so its closer to the actual impl
32
-def query_marlin_supported_quant_types(has_zp: bool,
33
- device_capability: Optional[int] = None
34
- ):
35
+def query_marlin_supported_quant_types(
36
+ has_zp: bool,
37
+ include_fp_type: bool = True,
38
+ device_capability: Optional[int] = None,
39
+):
40
if device_capability is None:
41
capability_tuple = current_platform.get_device_capability()
42
device_capability = (-1 if capability_tuple is None else
@@ -42,12 +47,13 @@ def query_marlin_supported_quant_types(has_zp: bool,
47
48
if has_zp:
49
# AWQ style, unsigned + runtime zero-point
45
- return [scalar_types.uint4, scalar_types.uint8]
50
+ return [scalar_types.uint4]
51
else:
52
# GPTQ style, unsigned + symmetric bias
48
- # TODO: once fp8_marlin is merged into "gptq_marlin" we should be able
49
- # to add `scalar_types.float8_e4m3fn` here
50
- return [scalar_types.uint4b8, scalar_types.uint8b128]
53
+ res = [scalar_types.uint4b8, scalar_types.uint8b128]
54
+ if include_fp_type:
55
+ res += [scalar_types.float8_e4m3fn]
56
+ return res
57
58
59
def _check_marlin_supported(
@@ -62,7 +68,7 @@ def _check_marlin_supported(
68
capability_tuple.to_int())
69
70
supported_types = query_marlin_supported_quant_types(
65
- has_zp, device_capability)
71
+ has_zp, True, device_capability)
72
73
if quant_type not in supported_types:
74
return (False, f"Marlin does not support weight_bits = {quant_type}. "
@@ -175,6 +181,17 @@ def marlin_make_workspace(output_size_per_partition: int,
181
requires_grad=False)
182
183
184
+def marlin_make_workspace_new(device: torch.device,
185
+ max_blocks_per_sm: int = 1) -> torch.Tensor:
186
+ # In the new marlin kernel, we use the num of threadblocks as workspace
187
+ # size. The num of threadblocks is is sms_count * max_blocks_per_sm.
188
+ sms = torch.cuda.get_device_properties(device).multi_processor_count
189
+ return torch.zeros(sms * max_blocks_per_sm,
190
+ dtype=torch.int,
191
+ device=device,
192
+ requires_grad=False)
193
+
194
+
195
def marlin_is_k_full(act_order: bool, is_row_parallel: bool) -> bool:
196
return (not act_order) or (act_order and not is_row_parallel)
197
@@ -304,21 +321,50 @@ def moe_awq_to_marlin_zero_points(q_zp_packed: torch.Tensor, size_k: int,
321
return output
322
323
324
+def maybe_warn_marlin_atomic_add(device, dtype):
325
+ if torch.compiler.is_dynamo_compiling():
326
+ return
327
+ device_capability = torch.cuda.get_device_capability(device)
328
+ if device_capability[0] < 9 and dtype == torch.bfloat16:
329
+ logger.info_once(
330
+ "You are running Marlin kernel with bf16 on GPUs before SM90. "
331
+ "You can consider change to fp16 to achieve better performance "
332
+ "if possible.")
333
+
334
+
335
+def maybe_warn_marlin_atomic_add_env():
336
+ if torch.compiler.is_dynamo_compiling():
337
+ return
338
+ if envs.VLLM_MARLIN_USE_ATOMIC_ADD:
339
+ return
340
+ logger.info_once(
341
+ "Marlin kernel can achieve better performance for small size_n "
342
+ "with experimental use_atomic_add feature. "
343
+ "You can consider set environment variable "
344
+ "VLLM_MARLIN_USE_ATOMIC_ADD to 1 if possible.")
345
+
346
+
347
def should_use_atomic_add_reduce(m: int, n: int, k: int, device: torch.device,
348
dtype: torch.dtype) -> bool:
349
+
350
+ # the performance of atomicAdd is better than global reduce
351
+ # only when m*n is small and k is large
352
+ if n >= 2048 or k < 2048 or device.type != "cuda":
353
+ return False
354
+
355
# disable atomicAdd reduce by default,
356
# one can enable it with VLLM_MARLIN_USE_ATOMIC_ADD=1
311
- if not envs.VLLM_MARLIN_USE_ATOMIC_ADD or device.type != "cuda":
357
+ if not envs.VLLM_MARLIN_USE_ATOMIC_ADD:
358
+ maybe_warn_marlin_atomic_add_env()
359
return False
360
361
# sm8x doesn't support atomicAdd + bfloat16 natively
362
device_capability = torch.cuda.get_device_capability(device)
363
if device_capability[0] < 9 and dtype == torch.bfloat16:
364
+ maybe_warn_marlin_atomic_add(device, dtype)
365
return False
366
319
- # the performance of atomicAdd is better than global reduce
320
- # only when m*n is small and k is large
321
- return n < 2048 and k >= 2048
367
+ return True
368
369
370
def apply_gptq_marlin_linear(
@@ -332,7 +378,6 @@ def apply_gptq_marlin_linear(
378
wtype: ScalarType,
379
output_size_per_partition: int,
380
input_size_per_partition: int,
335
- has_zp: bool,
381
is_k_full: bool,
382
bias: Optional[torch.Tensor] = None,
383
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT) -> torch.Tensor:
@@ -346,6 +391,7 @@ def apply_gptq_marlin_linear(
391
dtype=input.dtype)
392
393
output = ops.gptq_marlin_gemm(reshaped_x,
394
+ None,
395
weight,
396
weight_scale,
397
weight_zp,
@@ -358,7 +404,6 @@ def apply_gptq_marlin_linear(
404
size_k=input_size_per_partition,
405
is_k_full=is_k_full,
406
use_atomic_add=use_atomic_add,
361
- has_zp=has_zp,
407
use_fp32_reduce=use_fp32_reduce,
408
is_zp_float=False)
409
@@ -391,6 +436,7 @@ def apply_awq_marlin_linear(
436
dtype=input.dtype)
437
438
output = ops.gptq_marlin_gemm(reshaped_x,
439
+ None,
440
weight,
441
weight_scale,
442
weight_zp,
@@ -401,8 +447,6 @@ def apply_awq_marlin_linear(
447
size_m=reshaped_x.shape[0],
448
size_n=output_size_per_partition,
449
size_k=input_size_per_partition,
404
- is_k_full=True,
405
- has_zp=True,
450
use_atomic_add=use_atomic_add,
451
use_fp32_reduce=use_fp32_reduce,
452
is_zp_float=False)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+237
-42
@@ -6,9 +6,11 @@ import torch
6
7
import vllm._custom_ops as ops
8
from vllm.logger import init_logger
9
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
10
+ USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_scales,
11
+ should_use_atomic_add_reduce)
12
from vllm.platforms import current_platform
10
-
11
-from .marlin_utils import marlin_make_workspace, marlin_permute_scales
13
+from vllm.scalar_type import scalar_types
14
15
logger = init_logger(__name__)
16
@@ -18,30 +20,40 @@ def is_fp8_marlin_supported():
20
21
22
def apply_fp8_marlin_linear(
21
- input: torch.Tensor,
22
- weight: torch.Tensor,
23
- weight_scale: torch.Tensor,
24
- workspace: torch.Tensor,
25
- size_n: int,
26
- size_k: int,
27
- bias: Optional[torch.Tensor],
28
-) -> torch.Tensor:
23
+ input: torch.Tensor,
24
+ weight: torch.Tensor,
25
+ weight_scale: torch.Tensor,
26
+ workspace: torch.Tensor,
27
+ size_n: int,
28
+ size_k: int,
29
+ bias: Optional[torch.Tensor],
30
+ use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT) -> torch.Tensor:
31
# For GPUs that lack FP8 hardware support, we can leverage the
32
# Marlin kernel for fast weight-only FP8 quantization
33
34
reshaped_x = input.reshape(-1, input.shape[-1])
35
out_shape = input.shape[:-1] + (size_n, )
36
35
- output = ops.fp8_marlin_gemm(
36
- a=reshaped_x,
37
- b_q_weight=weight,
38
- b_scales=weight_scale,
39
- workspace=workspace,
40
- num_bits=8,
41
- size_m=reshaped_x.shape[0],
42
- size_n=size_n,
43
- size_k=size_k,
44
- )
37
+ use_atomic_add = should_use_atomic_add_reduce(m=reshaped_x.size(0),
38
+ n=size_n,
39
+ k=size_k,
40
+ device=input.device,
41
+ dtype=input.dtype)
42
+
43
+ output = ops.gptq_marlin_gemm(a=reshaped_x,
44
+ c=None,
45
+ b_q_weight=weight,
46
+ b_scales=weight_scale,
47
+ b_zeros=None,
48
+ g_idx=None,
49
+ perm=None,
50
+ workspace=workspace,
51
+ b_q_type=scalar_types.float8_e4m3fn,
52
+ size_m=reshaped_x.size(0),
53
+ size_n=size_n,
54
+ size_k=size_k,
55
+ use_atomic_add=use_atomic_add,
56
+ use_fp32_reduce=use_fp32_reduce)
57
58
if bias is not None:
59
output.add_(bias) # In-place add
@@ -50,7 +62,7 @@ def apply_fp8_marlin_linear(
62
63
64
def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
53
- strategy: str = "tensor") -> None:
65
+ size_k_first: bool = True) -> None:
66
logger.warning_once(
67
"Your GPU does not have native support for FP8 computation but "
68
"FP8 quantization is being used. Weight-only FP8 compression will "
@@ -60,51 +72,234 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
72
part_size_n = layer.output_size_per_partition
73
part_size_k = layer.input_size_per_partition
74
75
+ if size_k_first:
76
+ assert layer.weight.shape == (part_size_k, part_size_n)
77
+ else:
78
+ assert layer.weight.shape == (part_size_n, part_size_k)
79
+
80
device = layer.weight.device
81
82
# WORKSPACE
66
- layer.workspace = marlin_make_workspace(part_size_n, device)
83
+ layer.workspace = marlin_make_workspace_new(device)
84
85
# WEIGHT
86
# Repack weights to marlin format
70
- marlin_qweight = ops.gptq_marlin_repack(b_q_weight=pack_fp8_to_int32(
71
- layer.weight),
72
- perm=torch.empty(0,
73
- dtype=torch.int,
74
- device=device),
87
+ perm = torch.empty(0, dtype=torch.int, device=device)
88
+ qweight = pack_fp8_to_int32(layer.weight, size_k_first)
89
+ if not size_k_first:
90
+ qweight = qweight.T.contiguous()
91
+
92
+ marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
93
+ perm=perm,
94
size_k=part_size_k,
95
size_n=part_size_n,
96
num_bits=8)
97
layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
98
99
# WEIGHT SCALES
81
- scales = layer.weight_scale.to(layer.orig_dtype)
100
# Permute scales
101
+ if "weight_scale" in dir(layer):
102
+ scales = layer.weight_scale.to(layer.orig_dtype)
103
+ elif "weight_scale_inv" in dir(layer):
104
+ scales = layer.weight_scale_inv.to(layer.orig_dtype)
105
+ del layer.weight_scale_inv
106
+
107
+ if layer.weight_block_size is None:
108
+ group_size = -1
109
+ else:
110
+ group_size = layer.weight_block_size[1]
111
+
112
+ # marlin kernel only support channel-wise and group-wise quantization
113
+ # we need to convert the scales
114
+ if layer.weight_block_size is None:
115
+ if scales.nelement() == 1:
116
+ # tensor-wise quantization -> channel-wise quantization
117
+ # (1, 1) =>(repeat)=> (1, size_n)
118
+ scales = scales.view(1, 1).repeat_interleave(part_size_n, 1)
119
+ elif scales.nelement() > 1 and scales.nelement() != part_size_n:
120
+ assert part_size_n % scales.nelement() == 0
121
+ s_size = scales.nelement()
122
+ # tensor-wise quantization (for gate-up proj)
123
+ # -> channel-wise quantization
124
+ # (1, s_size) =>(repeat)=> (1, size_n)
125
+ scales = scales.view(1, s_size)
126
+ scales = scales.repeat_interleave(part_size_n // s_size, 1)
127
+ else:
128
+ # channel-wise quantization
129
+ # (1, size_n)
130
+ scales = scales.view(1, part_size_n)
131
+ else:
132
+ # block-wise quantization -> group-wise quantization
133
+ # (size_k // block_size[1], ceil(size_n / block_size[0]))
134
+ # =>(repeat)=> (size_k // block_size[1], size_n)
135
+ block_n = layer.weight_block_size[0]
136
+ scales = scales.T.repeat_interleave(block_n, 1)
137
+ # size_n may not divisible by block_size[0]
138
+ scales = scales[:, :part_size_n]
139
+
140
marlin_scales = marlin_permute_scales(s=scales,
141
size_k=part_size_k,
142
size_n=part_size_n,
86
- group_size=-1)
143
+ group_size=group_size)
144
layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)
145
146
90
-def pack_fp8_to_int32(fp8_tensor: torch.Tensor) -> torch.Tensor:
147
+def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
148
+ size_k_first: bool = True) -> None:
149
+ logger.warning_once(
150
+ "Your GPU does not have native support for FP8 computation but "
151
+ "FP8 quantization is being used. Weight-only FP8 compression will "
152
+ "be used leveraging the Marlin kernel. This may degrade "
153
+ "performance for compute-heavy workloads.")
154
+
155
+ e = layer.num_experts
156
+ k = layer.hidden_size
157
+ n = layer.intermediate_size_per_partition
158
+
159
+ # WORKSPACE
160
+ device = layer.w13_weight.device
161
+ layer.workspace = marlin_make_workspace_new(device, 4)
162
+ perm = torch.empty(0, dtype=torch.int, device=device)
163
+
164
+ # WEIGHT
165
+ # Repack weights to marlin format
166
+ for name in ["w13_weight", "w2_weight"]:
167
+ weight = getattr(layer, name)
168
+ tensor_list = []
169
+ if "w13" in name:
170
+ size_n, size_k = n * 2, k
171
+ else:
172
+ size_n, size_k = k, n
173
+
174
+ if size_k_first:
175
+ assert weight.shape == (e, size_k, size_n)
176
+ else:
177
+ assert weight.shape == (e, size_n, size_k)
178
+
179
+ for i in range(e):
180
+ qweight = pack_fp8_to_int32(weight[i], size_k_first)
181
+ if not size_k_first:
182
+ qweight = qweight.T.contiguous()
183
+
184
+ marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
185
+ perm=perm,
186
+ size_k=size_k,
187
+ size_n=size_n,
188
+ num_bits=8)
189
+ tensor_list.append(marlin_qweight)
190
+
191
+ weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
192
+ weight = torch.nn.Parameter(weight, requires_grad=False)
193
+
194
+ setattr(layer, name, weight)
195
+
196
+ # WEIGHT SCALES
197
+ # Permute scales
198
+ if layer.weight_block_size is None:
199
+ group_size = -1
200
+ else:
201
+ group_size = layer.weight_block_size[1]
202
+
203
+ for name in ["w13", "w2"]:
204
+ if name + "_weight_scale" in dir(layer):
205
+ new_name = name + "_weight_scale"
206
+ scales = getattr(layer, new_name).to(layer.orig_dtype)
207
+ delattr(layer, new_name)
208
+ elif name + "_weight_scale_inv" in dir(layer):
209
+ new_name = name + "_weight_scale_inv"
210
+ scales = getattr(layer, new_name).to(layer.orig_dtype)
211
+ delattr(layer, new_name)
212
+
213
+ tensor_list = []
214
+ if "w13" in name:
215
+ size_n, size_k = n * 2, k
216
+ else:
217
+ size_n, size_k = k, n
218
+
219
+ # marlin kernel only support channel-wise and group-wise quantization
220
+ # we need to convert the scales
221
+ if layer.weight_block_size is None:
222
+ if scales.nelement() == e:
223
+ # tensor-wise quantization -> channel-wise quantization
224
+ # (e, 1, 1) =>(repeat)=> (e, 1, size_n)
225
+ scales = scales.view(e, 1, 1).repeat_interleave(size_n, 2)
226
+ elif scales.nelement() > e and scales.nelement() != e * size_n:
227
+ assert (e * size_n) % scales.nelement() == 0
228
+ s_size = scales.nelement() // e
229
+ # tensor-wise quantization (for gate-up proj)
230
+ # -> channel-wise quantization
231
+ # (e, 1, s_size) =>(repeat)=> (e, 1, size_n)
232
+ scales = scales.view(e, 1, s_size)
233
+ scales = scales.repeat_interleave(size_n // s_size, 2)
234
+ else:
235
+ # channel-wise quantization
236
+ # (e, 1, size_n)
237
+ scales = scales.view(e, 1, size_n)
238
+ else:
239
+ # block-wise quantization -> group-wise quantization
240
+ # (e, size_k // block_size[1], ceil(size_n / block_size[0]))
241
+ # =>(repeat)=> (e, size_k // block_size[1], size_n)
242
+ block_n = layer.weight_block_size[0]
243
+ scales = scales.permute(0, 2, 1).repeat_interleave(block_n, 2)
244
+ # size_n may not divisible by block_size[0]
245
+ scales = scales[..., :size_n].contiguous()
246
+
247
+ for i in range(e):
248
+ marlin_scales = marlin_permute_scales(s=scales[i],
249
+ size_k=size_k,
250
+ size_n=size_n,
251
+ group_size=group_size)
252
+ tensor_list.append(marlin_scales)
253
+
254
+ scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
255
+ scales = torch.nn.Parameter(scales, requires_grad=False)
256
+
257
+ setattr(layer, name + "_weight_scale", scales)
258
+
259
+
260
+def pack_fp8_to_int32(fp8_tensor: torch.Tensor,
261
+ size_k_first: bool = True) -> torch.Tensor:
262
"""
263
Repack FP8 weights to gptq format (packed int32 elements)
264
"""
265
assert fp8_tensor.dtype == torch.float8_e4m3fn
95
- assert fp8_tensor.shape[0] % 4 == 0
266
+ assert fp8_tensor.ndim == 2
267
+
268
+ fp8_tensor = fp8_tensor.T if size_k_first else fp8_tensor
269
+ fp8_tensor = fp8_tensor.contiguous()
270
+ # fp8_tensor is contiguous and have shape (N, K) now
271
+ # with `.view(torch.int32)`, it become (N, K // 4)
272
+ int32_tensor = fp8_tensor.view(torch.int32)
273
+ return int32_tensor.T.contiguous() if size_k_first else int32_tensor
274
+
275
97
- # Reshape to prepare for packing
98
- reshaped = fp8_tensor.reshape(-1, 4, *fp8_tensor.shape[1:])
276
+def marlin_quant_fp8_torch(weight, group_size):
277
+ size_n, size_k = weight.shape
278
+ device = weight.device
279
100
- # Convert fp8 to uint8 (byte) representation
101
- byte_tensor = reshaped.view(torch.uint8)
280
+ if group_size != -1:
281
+ scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 448
282
+ repeated_scales = scales.repeat_interleave(group_size, 1)
283
+ fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn)
284
+ weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
285
+ else:
286
+ scales = weight.view(size_n, 1, group_size).abs().max(-1)[0] / 448
287
+ repeated_scales = scales.repeat_interleave(size_k, 1)
288
+ fp8_weight = (weight / repeated_scales).to(torch.float8_e4m3fn)
289
+ weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
290
+
291
+ packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous()
292
+ marlin_qweight = ops.gptq_marlin_repack(
293
+ b_q_weight=packed_weight,
294
+ perm=torch.empty(0, dtype=torch.int, device=device),
295
+ size_k=size_k,
296
+ size_n=size_n,
297
+ num_bits=8,
298
+ )
299
103
- # Pack 4 uint8 values into one int32
104
- packed = (byte_tensor[:, 0].to(torch.int32) |
105
- (byte_tensor[:, 1].to(torch.int32) << 8) |
106
- (byte_tensor[:, 2].to(torch.int32) << 16) |
107
- (byte_tensor[:, 3].to(torch.int32) << 24))
300
+ marlin_scales = marlin_permute_scales(s=scales.T,
301
+ size_k=size_k,
302
+ size_n=size_n,
303
+ group_size=group_size)
304
109
- return packed.view(fp8_tensor.shape[0] // 4,
110
- *fp8_tensor.shape[1:]).contiguous()
305
+ return weight_ref.T, marlin_qweight, marlin_scales
vllm/scalar_type.py
+11
@@ -6,6 +6,8 @@ from dataclasses import dataclass
6
from enum import Enum
7
from typing import Optional, Union
8
9
+_SCALAR_TYPES_ID_MAP = {}
10
+
11
12
# Mirrors enum in `core/scalar_type.hpp`
13
class NanRepr(Enum):
@@ -158,6 +160,8 @@ class ScalarType:
160
assert offset <= 64, \
161
f"ScalarType fields too big {offset} to fit into an int64"
162
163
+ _SCALAR_TYPES_ID_MAP[val] = self
164
+
165
return val
166
167
@property
@@ -295,6 +299,13 @@ class ScalarType:
299
ret.id # noqa B018: make sure the id is cached
300
return ret
301
302
+ @classmethod
303
+ def from_id(cls, scalar_type_id: int):
304
+ if scalar_type_id not in _SCALAR_TYPES_ID_MAP:
305
+ raise ValueError(
306
+ f"scalar_type_id {scalar_type_id} doesn't exists.")
307
+ return _SCALAR_TYPES_ID_MAP[scalar_type_id]
308
+
309
310
# naming generally follows: https://github.com/jax-ml/ml_dtypes
311
# for floating point types (leading f) the scheme is: