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 11, 08:45 UTC
[Kernel][Quantization][MoE] add marlin kernel support for turing (sm75) (#29901)
Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Jinzhen Lin committed
Dec 17, 2025 at 06:35 UTC
ce96857fdd2bf2390aaa2183561fd1a0f5c464c7
16 files changed
+728
-512
CMakeLists.txt
+76
-31
@@ -357,6 +357,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
357
358
# marlin arches for fp16 output
359
cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
360
+ # marlin has limited support for turing
361
+ cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
362
# marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
363
cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
364
# marlin arches for fp8 input
@@ -364,8 +366,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
366
# - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
367
# so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
368
cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
369
+ # marlin arches for other files
370
+ cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
371
368
- if (MARLIN_ARCHS)
372
+ if (MARLIN_OTHER_ARCHS)
373
374
#
375
# For the Marlin kernels we automatically generate sources for various
@@ -406,25 +410,39 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
410
message(STATUS "Marlin generation script has not changed, skipping generation.")
411
endif()
412
409
- file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
410
- set_gencode_flags_for_srcs(
411
- SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
412
- CUDA_ARCHS "${MARLIN_ARCHS}")
413
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
414
- set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
415
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
413
+ if (MARLIN_ARCHS)
414
+ file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
415
+ set_gencode_flags_for_srcs(
416
+ SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
417
+ CUDA_ARCHS "${MARLIN_ARCHS}")
418
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
419
+ set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
420
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
421
+ endif()
422
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
423
+
424
+ file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
425
+ set_gencode_flags_for_srcs(
426
+ SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
427
+ CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
428
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
429
+ set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
430
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
431
+ endif()
432
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
433
endif()
417
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
434
419
- file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
420
- set_gencode_flags_for_srcs(
421
- SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
422
- CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
423
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
424
- set_source_files_properties(${MARLIN_TEMPLATE_BF16_KERNEL_SRC}
425
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
435
+ if (MARLIN_SM75_ARCHS)
436
+ file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/gptq_marlin/sm75_kernel_*.cu")
437
+ set_gencode_flags_for_srcs(
438
+ SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
439
+ CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
440
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
441
+ set_source_files_properties(${MARLIN_TEMPLATE_SM75_KERNEL_SRC}
442
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
443
+ endif()
444
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
445
endif()
427
- list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
446
447
if (MARLIN_FP8_ARCHS)
448
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
@@ -446,14 +464,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
464
"csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
465
set_gencode_flags_for_srcs(
466
SRCS "${MARLIN_SRCS}"
449
- CUDA_ARCHS "${MARLIN_ARCHS}")
467
+ CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
468
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
451
- set_source_files_properties("csrc/quantization/gptq_marlin/gptq_marlin.cu"
469
+ set_source_files_properties(${MARLIN_SRCS}
470
PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
471
endif()
472
list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}")
473
456
- message(STATUS "Building Marlin kernels for archs: ${MARLIN_ARCHS}")
474
+ message(STATUS "Building Marlin kernels for archs: ${MARLIN_OTHER_ARCHS}")
475
else()
476
message(STATUS "Not building Marlin kernels as no compatible archs found"
477
" in CUDA target architectures")
@@ -980,12 +998,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
998
# note that we always set `use_atomic_add=False` for moe marlin now,
999
# so we don't need 9.0 for bf16 atomicAdd PTX
1000
cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
1001
+ # moe marlin has limited support for turing
1002
+ cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
1003
# moe marlin arches for fp8 input
1004
# - sm80 doesn't support fp8 computation
1005
# - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
1006
# so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
1007
cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
988
- if (MARLIN_MOE_ARCHS)
1008
+ # moe marlin arches for other files
1009
+ cuda_archs_loose_intersection(MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
1010
+ if (MARLIN_MOE_OTHER_ARCHS)
1011
1012
#
1013
# For the Marlin MOE kernels we automatically generate sources for various
@@ -1026,16 +1048,29 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1048
message(STATUS "Marlin MOE generation script has not changed, skipping generation.")
1049
endif()
1050
1029
- file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
1030
- list(APPEND MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
1031
- set_gencode_flags_for_srcs(
1032
- SRCS "${MARLIN_MOE_SRC}"
1033
- CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
1034
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
1035
- set_source_files_properties(${MARLIN_MOE_SRC}
1036
- PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1051
+ if (MARLIN_MOE_ARCHS)
1052
+ file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
1053
+ set_gencode_flags_for_srcs(
1054
+ SRCS "${MARLIN_MOE_SRC}"
1055
+ CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
1056
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
1057
+ set_source_files_properties(${MARLIN_MOE_SRC}
1058
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1059
+ endif()
1060
+ list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
1061
+ endif()
1062
+
1063
+ if (MARLIN_MOE_SM75_ARCHS)
1064
+ file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu")
1065
+ set_gencode_flags_for_srcs(
1066
+ SRCS "${MARLIN_MOE_SM75_SRC}"
1067
+ CUDA_ARCHS "${MARLIN_MOE_SM75_ARCHS}")
1068
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
1069
+ set_source_files_properties(${MARLIN_MOE_SM75_SRC}
1070
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1071
+ endif()
1072
+ list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SM75_SRC})
1073
endif()
1038
- list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
1074
1075
if (MARLIN_MOE_FP8_ARCHS)
1076
file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu")
@@ -1049,7 +1084,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1084
list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC})
1085
endif()
1086
1052
- message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
1087
+ set(MARLIN_MOE_OTHER_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
1088
+ set_gencode_flags_for_srcs(
1089
+ SRCS "${MARLIN_MOE_OTHER_SRC}"
1090
+ CUDA_ARCHS "${MARLIN_MOE_OTHER_ARCHS}")
1091
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
1092
+ set_source_files_properties(${MARLIN_MOE_OTHER_SRC}
1093
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1094
+ endif()
1095
+ list(APPEND VLLM_MOE_EXT_SRC "${MARLIN_MOE_OTHER_SRC}")
1096
+
1097
+ message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_OTHER_ARCHS}")
1098
else()
1099
message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
1100
" in CUDA target architectures")
csrc/moe/marlin_moe_wna16/.gitignore
+1
@@ -1,2 +1,3 @@
1
sm*_kernel_*.cu
2
kernel_selector.h
3
+kernel_*.cu
csrc/moe/marlin_moe_wna16/generate_kernels.py
+76
-56
@@ -10,6 +10,8 @@ import jinja2
10
11
ARCHS = []
12
SUPPORT_FP8 = False
13
+SUPPORT_SM75 = False
14
+SUPPORT_SM80 = False
15
for arch in sys.argv[1].split(","):
16
arch = arch[: arch.index(".") + 2].replace(".", "")
17
arch = int(arch)
@@ -19,6 +21,10 @@ for arch in sys.argv[1].split(","):
21
# with FP16 MMA, so it cannot achieve any acceleration.
22
if arch in [89, 120]:
23
SUPPORT_FP8 = True
24
+ if arch >= 80:
25
+ SUPPORT_SM80 = True
26
+ if arch == 75:
27
+ SUPPORT_SM75 = True
28
29
FILE_HEAD_COMMENT = """
30
// auto generated by generate_kernels.py
@@ -157,6 +163,7 @@ def remove_old_kernels():
163
164
def generate_new_kernels():
165
result_dict = {}
166
+ sm_75_result_dict = {}
167
168
for quant_config in QUANT_CONFIGS:
169
c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
@@ -174,6 +181,8 @@ def generate_new_kernels():
181
s_type = quant_config.get("s_type", c_type)
182
if (a_type, b_type, c_type) not in result_dict:
183
result_dict[(a_type, b_type, c_type)] = []
184
+ if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16":
185
+ sm_75_result_dict[(a_type, b_type, c_type)] = []
186
187
for group_blocks, m_blocks, thread_configs in itertools.product(
188
all_group_blocks, all_m_blocks, all_thread_configs
@@ -197,78 +206,89 @@ def generate_new_kernels():
206
"thread_k_blocks": thread_k // 16,
207
"thread_n_blocks": thread_n // 16,
208
"m_block_size_8": "true" if m_blocks == 0.5 else "false",
200
- "stages": "pipe_stages",
209
+ "stages": 4,
210
"group_blocks": group_blocks,
211
"is_zp_float": "false",
212
}
213
205
- result_dict[(a_type, b_type, c_type)].append(config)
214
+ if SUPPORT_SM80:
215
+ result_dict[(a_type, b_type, c_type)].append(config)
216
+ if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75:
217
+ config_sm75 = config.copy()
218
+ config_sm75["stages"] = 2
219
+ sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75)
220
221
kernel_selector_str = FILE_HEAD_COMMENT
222
209
- for (a_type, b_type, c_type), config_list in result_dict.items():
210
- all_template_str_list = []
211
- for config in config_list:
212
- s_type = config["s_type"]
213
- template_str = jinja2.Template(TEMPLATE).render(
214
- a_type_id=f"vllm::{a_type}.id()",
215
- b_type_id=f"vllm::{b_type}.id()",
216
- c_type_id=f"vllm::{c_type}.id()",
217
- s_type_id=f"vllm::{s_type}.id()",
218
- **config,
219
- )
220
- all_template_str_list.append(template_str)
221
-
222
- conditions = [
223
- f"a_type == vllm::{a_type}",
224
- f"b_type == vllm::{b_type}",
225
- f"c_type == vllm::{c_type}",
226
- f"s_type == vllm::{s_type}",
227
- f"threads == {config['threads']}",
228
- f"thread_m_blocks == {config['thread_m_blocks']}",
229
- f"thread_n_blocks == {config['thread_n_blocks']}",
230
- f"thread_k_blocks == {config['thread_k_blocks']}",
231
- f"m_block_size_8 == {config['m_block_size_8']}",
232
- f"group_blocks == {config['group_blocks']}",
233
- f"is_zp_float == {config['is_zp_float']}",
234
- ]
235
- conditions = " && ".join(conditions)
236
-
237
- if kernel_selector_str == FILE_HEAD_COMMENT:
238
- kernel_selector_str += f"if ({conditions})\n kernel = "
239
- else:
240
- kernel_selector_str += f"else if ({conditions})\n kernel = "
241
-
242
- kernel_template2 = (
243
- "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
244
- "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
245
- "{{thread_n_blocks}}, {{thread_k_blocks}}, "
246
- "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
247
- "{{is_zp_float}}>;"
248
- )
249
-
250
- kernel_selector_str += (
251
- jinja2.Template(kernel_template2).render(
223
+ for result_dict_tmp in [result_dict, sm_75_result_dict]:
224
+ for (a_type, b_type, c_type), config_list in result_dict_tmp.items():
225
+ all_template_str_list = []
226
+ if not config_list:
227
+ continue
228
+ for config in config_list:
229
+ s_type = config["s_type"]
230
+ template_str = jinja2.Template(TEMPLATE).render(
231
a_type_id=f"vllm::{a_type}.id()",
232
b_type_id=f"vllm::{b_type}.id()",
233
c_type_id=f"vllm::{c_type}.id()",
234
s_type_id=f"vllm::{s_type}.id()",
235
**config,
236
)
258
- + "\n"
259
- )
237
+ all_template_str_list.append(template_str)
238
+
239
+ conditions = [
240
+ f"a_type == vllm::{a_type}",
241
+ f"b_type == vllm::{b_type}",
242
+ f"c_type == vllm::{c_type}",
243
+ f"s_type == vllm::{s_type}",
244
+ f"threads == {config['threads']}",
245
+ f"thread_m_blocks == {config['thread_m_blocks']}",
246
+ f"thread_n_blocks == {config['thread_n_blocks']}",
247
+ f"thread_k_blocks == {config['thread_k_blocks']}",
248
+ f"m_block_size_8 == {config['m_block_size_8']}",
249
+ f"stages == {config['stages']}",
250
+ f"group_blocks == {config['group_blocks']}",
251
+ f"is_zp_float == {config['is_zp_float']}",
252
+ ]
253
+ conditions = " && ".join(conditions)
254
+
255
+ if kernel_selector_str == FILE_HEAD_COMMENT:
256
+ kernel_selector_str += f"if ({conditions})\n kernel = "
257
+ else:
258
+ kernel_selector_str += f"else if ({conditions})\n kernel = "
259
+
260
+ kernel_template2 = (
261
+ "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
262
+ "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
263
+ "{{thread_n_blocks}}, {{thread_k_blocks}}, "
264
+ "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
265
+ "{{is_zp_float}}>;"
266
+ )
267
261
- file_content = FILE_HEAD + "\n\n"
262
- file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
263
- if a_type == "kFE4M3fn":
264
- filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
265
- else:
266
- filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
268
+ kernel_selector_str += (
269
+ jinja2.Template(kernel_template2).render(
270
+ a_type_id=f"vllm::{a_type}.id()",
271
+ b_type_id=f"vllm::{b_type}.id()",
272
+ c_type_id=f"vllm::{c_type}.id()",
273
+ s_type_id=f"vllm::{s_type}.id()",
274
+ **config,
275
+ )
276
+ + "\n"
277
+ )
278
+
279
+ file_content = FILE_HEAD + "\n\n"
280
+ file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
281
+ if a_type == "kFE4M3fn":
282
+ filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
283
+ elif result_dict_tmp is sm_75_result_dict:
284
+ filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
285
+ else:
286
+ filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
287
268
- filename = filename.lower()
288
+ filename = filename.lower()
289
270
- with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
271
- f.write(file_content)
290
+ with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
291
+ f.write(file_content)
292
293
if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
294
kernel_selector_str += (
csrc/moe/marlin_moe_wna16/marlin_template.h
+55
-153
@@ -26,6 +26,7 @@
26
#include "quantization/gptq_marlin/marlin.cuh"
27
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
28
#include "quantization/gptq_marlin/dequant.h"
29
+#include "quantization/gptq_marlin/marlin_mma.h"
30
#include "core/scalar_type.hpp"
31
32
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
@@ -35,7 +36,7 @@
36
37
namespace MARLIN_NAMESPACE_NAME {
38
38
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
39
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
40
41
template <typename scalar_t, // compute dtype, half or nv_float16
42
const vllm::ScalarTypeId b_type_id, // weight MarlinScalarType id
@@ -84,146 +85,6 @@ __global__ void Marlin(
85
86
#else
87
87
-// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
88
-// output/accumulation.
89
-template <vllm::ScalarTypeId type_id, int k_size = 16>
90
-__device__ inline void mma(
91
- const typename MarlinScalarType<type_id>::FragA& a_frag,
92
- const typename MarlinScalarType<type_id>::FragB& frag_b,
93
- typename MarlinScalarType<type_id>::FragC& frag_c, int idx = 0) {
94
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
95
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
96
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
97
- if constexpr (k_size == 16) {
98
- if constexpr (std::is_same<scalar_t, half>::value) {
99
- float* c = reinterpret_cast<float*>(&frag_c);
100
- asm volatile(
101
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
102
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
103
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
104
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
105
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
106
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
107
- float* c = reinterpret_cast<float*>(&frag_c);
108
- asm volatile(
109
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
110
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
111
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
112
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
113
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
114
- } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
115
- float* c = reinterpret_cast<float*>(&frag_c);
116
- asm volatile(
117
- "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
118
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
119
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
120
- : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]),
121
- "f"(c[1]), "f"(c[2]), "f"(c[3]));
122
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
123
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
124
- asm volatile(
125
- "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
126
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
127
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
128
- : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]),
129
- "r"(c[1]), "r"(c[2]), "r"(c[3]));
130
- }
131
- } else if (k_size == 32) {
132
- if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
133
- float* c = reinterpret_cast<float*>(&frag_c);
134
- asm volatile(
135
- "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
136
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
137
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
138
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
139
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
140
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
141
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
142
- asm volatile(
143
- "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
144
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
145
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
146
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
147
- "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
148
- }
149
- }
150
-}
151
-
152
-template <vllm::ScalarTypeId type_id, int k_size = 16>
153
-__device__ inline void mma_trans(
154
- const typename MarlinScalarType<type_id>::FragA& a_frag,
155
- const typename MarlinScalarType<type_id>::FragB& frag_b,
156
- const typename MarlinScalarType<type_id>::FragB& frag_b2,
157
- typename MarlinScalarType<type_id>::FragC& frag_c) {
158
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
159
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
160
- const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
161
- float* c = reinterpret_cast<float*>(&frag_c);
162
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
163
- if constexpr (k_size == 16) {
164
- if constexpr (std::is_same<scalar_t, half>::value) {
165
- float* c = reinterpret_cast<float*>(&frag_c);
166
- asm volatile(
167
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
168
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
169
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
170
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
171
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
172
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
173
- float* c = reinterpret_cast<float*>(&frag_c);
174
- asm volatile(
175
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
176
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
177
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
178
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
179
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
180
- } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
181
- float* c = reinterpret_cast<float*>(&frag_c);
182
- asm volatile(
183
- "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
184
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
185
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
186
- : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
187
- "f"(c[3]));
188
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
189
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
190
- asm volatile(
191
- "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
192
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
193
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
194
- : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
195
- "r"(c[3]));
196
- }
197
- } else {
198
- if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
199
- float* c = reinterpret_cast<float*>(&frag_c);
200
- #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 1200
201
- asm volatile(
202
- "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
203
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
204
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
205
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
206
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
207
- #else
208
- asm volatile(
209
- "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
210
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
211
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
212
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
213
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
214
- #endif
215
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
216
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
217
- asm volatile(
218
- "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
219
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
220
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
221
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
222
- "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
223
- }
224
- }
225
-}
226
-
88
// Instruction for loading a full 16x16 matrix fragment of operand A from shared
89
// memory, directly in tensor core layout.
90
template <int count, vllm::ScalarTypeId type_id>
@@ -439,9 +300,20 @@ __global__ void Marlin(
300
if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
301
#endif
302
303
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
304
+ // Turing TensorCore only supports fp16 and int8
305
+ if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id())
306
+ return;
307
+ #endif
308
+
309
int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
310
constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
311
312
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
313
+ constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
314
+ #else
315
+ constexpr bool use_fp16_accum = false;
316
+ #endif
317
using Adtype = MarlinScalarType<a_type_id>;
318
using Cdtype = MarlinScalarType<c_type_id>;
319
@@ -618,7 +490,22 @@ __global__ void Marlin(
490
}
491
}
492
493
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
494
+
495
+ if constexpr (moe_block_size >= 16)
496
+ local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 16);
497
+ if constexpr (moe_block_size >= 8)
498
+ local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 8);
499
+ if constexpr (moe_block_size >= 4)
500
+ local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 4);
501
+ if constexpr (moe_block_size >= 2)
502
+ local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 2);
503
+
504
+ local_count += __shfl_down_sync(0xFFFFFFFF, local_count, 1);
505
+ block_num_valid_tokens = local_count;
506
+ #else
507
block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count);
508
+ #endif
509
510
if (lane_id == 0)
511
reinterpret_cast<int*>(sh_new)[0] = block_num_valid_tokens;
@@ -1018,10 +905,6 @@ __global__ void Marlin(
905
constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
906
: (stages * s_sh_stage);
907
int4* sh_s = sh_zp + (stages * zp_sh_stage);
1021
- // shared memory reused by reduction should be smaller than
1022
- // shared memory used by weight.
1023
- static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
1024
- stages * b_sh_stage);
908
int4* sh_a = sh_s + sh_s_size;
909
910
// Register storage for double buffer of shared memory reads.
@@ -1545,11 +1428,13 @@ __global__ void Marlin(
1428
#pragma unroll
1429
for (int i = 0; i < thread_m_blocks; i++) {
1430
if constexpr (m_block_size_8) {
1548
- mma_trans<a_type_id>(frag_a[k2][i], frag_b0, frag_b1,
1549
- frag_c[i][j][0]);
1431
+ mma_trans<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, frag_b1,
1432
+ frag_c[i][j][0]);
1433
} else {
1551
- mma<a_type_id>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1552
- mma<a_type_id>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
1434
+ mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0,
1435
+ frag_c[i][j][0]);
1436
+ mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b1,
1437
+ frag_c[i][j][1]);
1438
}
1439
}
1440
}
@@ -1583,10 +1468,12 @@ __global__ void Marlin(
1468
1469
#pragma unroll
1470
for (int i = 0; i < thread_m_blocks; i++) {
1586
- mma<a_type_id, 32>(frag_a[k2][i], frag_b[0],
1587
- (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
1588
- mma<a_type_id, 32>(frag_a[k2][i], frag_b[1],
1589
- (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
1471
+ mma<a_type_id, false, 32>(
1472
+ frag_a[k2][i], frag_b[0],
1473
+ (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
1474
+ mma<a_type_id, false, 32>(
1475
+ frag_a[k2][i], frag_b[1],
1476
+ (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
1477
}
1478
1479
if constexpr (group_blocks != -1) {
@@ -2132,6 +2019,21 @@ __global__ void Marlin(
2019
// While this pattern may not be the most readable, other ways of writing
2020
// the loop seemed to noticeably worse performance after compilation.
2021
if (slice_iters == 0) {
2022
+ // convert fp16 accum to fp32 for reduction
2023
+ if constexpr (use_fp16_accum) {
2024
+ #pragma unroll
2025
+ for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) {
2026
+ float* frag_c_part_float = reinterpret_cast<float*>(frag_c) + i * 4;
2027
+ scalar_t* frag_c_part_half =
2028
+ reinterpret_cast<scalar_t*>(frag_c_part_float);
2029
+
2030
+ #pragma unroll
2031
+ for (int i = 3; i >= 0; i--) {
2032
+ frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]);
2033
+ }
2034
+ }
2035
+ }
2036
+
2037
if constexpr (is_a_8bit) {
2038
float frag_a_s[2 * thread_m_blocks];
2039
csrc/moe/marlin_moe_wna16/ops.cu
+30
-24
@@ -142,7 +142,7 @@ typedef struct {
142
143
int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
144
int prob_n, int prob_k, int num_bits, int group_size,
145
- bool has_act_order, bool is_k_full) {
145
+ bool has_act_order, bool is_k_full, int stages) {
146
bool cache_scales_chunk = has_act_order && !is_k_full;
147
148
int tb_n = th_config.thread_n;
@@ -160,13 +160,13 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
160
161
if (cache_scales_chunk) {
162
int load_groups =
163
- tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K
163
+ tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K
164
load_groups = max(load_groups, 32); // We load at least 32 scale groups
165
return load_groups * tb_n * 2;
166
} else {
167
int tb_scales = tb_groups * tb_n * 2;
168
169
- return tb_scales * pipe_stages;
169
+ return tb_scales * stages;
170
}
171
}
172
@@ -174,7 +174,7 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
174
int thread_m_blocks, int prob_m, int prob_n,
175
int prob_k, int num_bits, int group_size,
176
bool has_act_order, bool is_k_full, int has_zp,
177
- int is_zp_float, bool is_a_8bit) {
177
+ int is_zp_float, bool is_a_8bit, int stages) {
178
int pack_factor = 32 / num_bits;
179
180
// Get B size
@@ -185,8 +185,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
185
// shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
186
// both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
187
int sh_block_meta_size = tb_m * 16;
188
- int sh_a_size = pipe_stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
189
- int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
188
+ int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
189
+ int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4;
190
int sh_red_size = tb_m * (tb_n + 8) * 2;
191
int sh_bias_size = tb_n * 2;
192
int tmp_size =
@@ -195,8 +195,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
195
196
int sh_s_size =
197
get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
198
- group_size, has_act_order, is_k_full);
199
- int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0;
198
+ group_size, has_act_order, is_k_full, stages);
199
+ int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
200
int sh_zp_size = 0;
201
if (has_zp) {
202
if (is_zp_float)
@@ -217,7 +217,7 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
217
int thread_m_blocks, int prob_m, int prob_n, int prob_k,
218
int num_bits, int group_size, bool has_act_order,
219
bool is_k_full, int has_zp, int is_zp_float,
220
- int max_shared_mem, bool is_a_8bit) {
220
+ bool is_a_8bit, int stages, int max_shared_mem) {
221
// Sanity
222
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
223
th_config.num_threads == -1) {
@@ -243,7 +243,7 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
243
int cache_size =
244
get_kernel_cache_size(th_config, m_block_size_8, thread_m_blocks, prob_m,
245
prob_n, prob_k, num_bits, group_size, has_act_order,
246
- is_k_full, has_zp, is_zp_float, is_a_8bit);
246
+ is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
247
return cache_size <= max_shared_mem;
248
}
249
@@ -252,7 +252,7 @@ MarlinFuncPtr get_marlin_kernel(
252
const vllm::ScalarType c_type, const vllm::ScalarType s_type,
253
int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
254
bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
255
- int threads, bool is_zp_float) {
255
+ int threads, bool is_zp_float, int stages) {
256
int num_bits = b_type.size_bits();
257
auto kernel = MarlinDefault;
258
@@ -266,8 +266,8 @@ exec_config_t determine_exec_config(
266
const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
267
int prob_n, int prob_k, int num_experts, int top_k, int thread_m_blocks,
268
bool m_block_size_8, int num_bits, int group_size, bool has_act_order,
269
- bool is_k_full, bool has_zp, bool is_zp_float, int max_shared_mem, int sms,
270
- bool is_a_8bit) {
269
+ bool is_k_full, bool has_zp, bool is_zp_float, bool is_a_8bit, int stages,
270
+ int max_shared_mem, int sms) {
271
exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
272
thread_config_t* thread_configs = thread_m_blocks > 1
273
? large_batch_thread_configs
@@ -284,15 +284,15 @@ exec_config_t determine_exec_config(
284
285
if (!is_valid_config(th_config, m_block_size_8, thread_m_blocks, prob_m,
286
prob_n, prob_k, num_bits, group_size, has_act_order,
287
- is_k_full, has_zp, is_zp_float, max_shared_mem - 512,
288
- is_a_8bit)) {
287
+ is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
288
+ max_shared_mem - 512)) {
289
continue;
290
}
291
292
int cache_size = get_kernel_cache_size(
293
th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
294
num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
295
- is_a_8bit);
295
+ is_a_8bit, stages);
296
297
int group_blocks = 0;
298
if (!has_act_order) {
@@ -303,7 +303,7 @@ exec_config_t determine_exec_config(
303
get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
304
th_config.thread_n / 16, th_config.thread_k / 16,
305
m_block_size_8, has_act_order, has_zp, group_blocks,
306
- th_config.num_threads, is_zp_float);
306
+ th_config.num_threads, is_zp_float, stages);
307
308
if (kernel == MarlinDefault) continue;
309
@@ -433,8 +433,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
433
dev);
434
cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
435
dev);
436
- TORCH_CHECK(major_capability * 10 + minor_capability >= 80,
437
- "marlin kernel only support Ampere or newer GPUs.");
436
+ TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
437
+ "marlin kernel only support Turing or newer GPUs.");
438
+ int stages = 4;
439
+ if (major_capability == 7 && minor_capability == 5) {
440
+ stages = 2;
441
+ TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
442
+ "Turing only support FP16 or INT8 activation.");
443
+ }
444
if (a_type == vllm::kFE4M3fn) {
445
TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
446
"FP8 only support Ada Lovelace or newer GPUs.");
@@ -461,8 +467,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
467
exec_cfg = determine_exec_config(
468
a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
469
top_k, thread_m_blocks, m_block_size_8, num_bits, group_size,
464
- has_act_order, is_k_full, has_zp, is_zp_float, max_shared_mem, sms,
465
- is_a_8bit);
470
+ has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages,
471
+ max_shared_mem, sms);
472
thread_tfg = exec_cfg.tb_cfg;
473
}
474
@@ -479,7 +485,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
485
TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks,
486
prob_m, prob_n, prob_k, num_bits, group_size,
487
has_act_order, is_k_full, has_zp, is_zp_float,
482
- max_shared_mem, is_a_8bit),
488
+ is_a_8bit, stages, max_shared_mem),
489
"Invalid thread config: thread_m_blocks = ", thread_m_blocks,
490
", thread_k = ", thread_tfg.thread_k,
491
", thread_n = ", thread_tfg.thread_n,
@@ -493,12 +499,12 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
499
int sh_cache_size =
500
get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
501
prob_n, prob_k, num_bits, group_size, has_act_order,
496
- is_k_full, has_zp, is_zp_float, is_a_8bit);
502
+ is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
503
504
auto kernel = get_marlin_kernel(
505
a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
506
thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
501
- num_threads, is_zp_float);
507
+ num_threads, is_zp_float, stages);
508
509
if (kernel == MarlinDefault) {
510
TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
csrc/quantization/gptq_marlin/.gitignore
+1
@@ -1,2 +1,3 @@
1
sm*_kernel_*.cu
2
kernel_selector.h
3
+kernel_*.cu
csrc/quantization/gptq_marlin/dequant.h
+1
-1
@@ -67,7 +67,7 @@ where `scale_factor * multiplier` can be computed at weight loading.
67
68
namespace MARLIN_NAMESPACE_NAME {
69
70
-#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
70
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750
71
// Lookup-table based 3-input logical operation; explicitly used for
72
// dequantization as the compiler does not seem to automatically recognize it in
73
// all cases.
csrc/quantization/gptq_marlin/generate_kernels.py
+76
-56
@@ -10,6 +10,8 @@ import jinja2
10
11
ARCHS = []
12
SUPPORT_FP8 = False
13
+SUPPORT_SM75 = False
14
+SUPPORT_SM80 = False
15
for arch in sys.argv[1].split(","):
16
arch = arch[: arch.index(".") + 2].replace(".", "")
17
arch = int(arch)
@@ -19,6 +21,10 @@ for arch in sys.argv[1].split(","):
21
# with FP16 MMA, so it cannot achieve any acceleration.
22
if arch in [89, 120]:
23
SUPPORT_FP8 = True
24
+ if arch >= 80:
25
+ SUPPORT_SM80 = True
26
+ if arch == 75:
27
+ SUPPORT_SM75 = True
28
29
FILE_HEAD_COMMENT = """
30
// auto generated by generate_kernels.py
@@ -166,6 +172,7 @@ def remove_old_kernels():
172
173
def generate_new_kernels():
174
result_dict = {}
175
+ sm_75_result_dict = {}
176
177
for quant_config in QUANT_CONFIGS:
178
c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
@@ -184,6 +191,8 @@ def generate_new_kernels():
191
s_type = quant_config.get("s_type", c_type)
192
if (a_type, b_type, c_type) not in result_dict:
193
result_dict[(a_type, b_type, c_type)] = []
194
+ if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16":
195
+ sm_75_result_dict[(a_type, b_type, c_type)] = []
196
197
for group_blocks, m_blocks, thread_configs in itertools.product(
198
all_group_blocks, all_m_blocks, all_thread_configs
@@ -207,78 +216,89 @@ def generate_new_kernels():
216
"thread_k_blocks": thread_k // 16,
217
"thread_n_blocks": thread_n // 16,
218
"m_block_size_8": "true" if m_blocks == 0.5 else "false",
210
- "stages": "pipe_stages",
219
+ "stages": 4,
220
"group_blocks": group_blocks,
221
"is_zp_float": "true" if is_zp_float else "false",
222
}
223
215
- result_dict[(a_type, b_type, c_type)].append(config)
224
+ if SUPPORT_SM80:
225
+ result_dict[(a_type, b_type, c_type)].append(config)
226
+ if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75:
227
+ config_sm75 = config.copy()
228
+ config_sm75["stages"] = 2
229
+ sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75)
230
231
kernel_selector_str = FILE_HEAD_COMMENT
232
219
- for (a_type, b_type, c_type), config_list in result_dict.items():
220
- all_template_str_list = []
221
- for config in config_list:
222
- s_type = config["s_type"]
223
- template_str = jinja2.Template(TEMPLATE).render(
224
- a_type_id=f"vllm::{a_type}.id()",
225
- b_type_id=f"vllm::{b_type}.id()",
226
- c_type_id=f"vllm::{c_type}.id()",
227
- s_type_id=f"vllm::{s_type}.id()",
228
- **config,
229
- )
230
- all_template_str_list.append(template_str)
231
-
232
- conditions = [
233
- f"a_type == vllm::{a_type}",
234
- f"b_type == vllm::{b_type}",
235
- f"c_type == vllm::{c_type}",
236
- f"s_type == vllm::{s_type}",
237
- f"threads == {config['threads']}",
238
- f"thread_m_blocks == {config['thread_m_blocks']}",
239
- f"thread_n_blocks == {config['thread_n_blocks']}",
240
- f"thread_k_blocks == {config['thread_k_blocks']}",
241
- f"m_block_size_8 == {config['m_block_size_8']}",
242
- f"group_blocks == {config['group_blocks']}",
243
- f"is_zp_float == {config['is_zp_float']}",
244
- ]
245
- conditions = " && ".join(conditions)
246
-
247
- if kernel_selector_str == FILE_HEAD_COMMENT:
248
- kernel_selector_str += f"if ({conditions})\n kernel = "
249
- else:
250
- kernel_selector_str += f"else if ({conditions})\n kernel = "
251
-
252
- kernel_template2 = (
253
- "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
254
- "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
255
- "{{thread_n_blocks}}, {{thread_k_blocks}}, "
256
- "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
257
- "{{is_zp_float}}>;"
258
- )
259
-
260
- kernel_selector_str += (
261
- jinja2.Template(kernel_template2).render(
233
+ for result_dict_tmp in [result_dict, sm_75_result_dict]:
234
+ for (a_type, b_type, c_type), config_list in result_dict_tmp.items():
235
+ all_template_str_list = []
236
+ if not config_list:
237
+ continue
238
+ for config in config_list:
239
+ s_type = config["s_type"]
240
+ template_str = jinja2.Template(TEMPLATE).render(
241
a_type_id=f"vllm::{a_type}.id()",
242
b_type_id=f"vllm::{b_type}.id()",
243
c_type_id=f"vllm::{c_type}.id()",
244
s_type_id=f"vllm::{s_type}.id()",
245
**config,
246
)
268
- + "\n"
269
- )
247
+ all_template_str_list.append(template_str)
248
+
249
+ conditions = [
250
+ f"a_type == vllm::{a_type}",
251
+ f"b_type == vllm::{b_type}",
252
+ f"c_type == vllm::{c_type}",
253
+ f"s_type == vllm::{s_type}",
254
+ f"threads == {config['threads']}",
255
+ f"thread_m_blocks == {config['thread_m_blocks']}",
256
+ f"thread_n_blocks == {config['thread_n_blocks']}",
257
+ f"thread_k_blocks == {config['thread_k_blocks']}",
258
+ f"m_block_size_8 == {config['m_block_size_8']}",
259
+ f"stages == {config['stages']}",
260
+ f"group_blocks == {config['group_blocks']}",
261
+ f"is_zp_float == {config['is_zp_float']}",
262
+ ]
263
+ conditions = " && ".join(conditions)
264
+
265
+ if kernel_selector_str == FILE_HEAD_COMMENT:
266
+ kernel_selector_str += f"if ({conditions})\n kernel = "
267
+ else:
268
+ kernel_selector_str += f"else if ({conditions})\n kernel = "
269
+
270
+ kernel_template2 = (
271
+ "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
272
+ "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
273
+ "{{thread_n_blocks}}, {{thread_k_blocks}}, "
274
+ "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
275
+ "{{is_zp_float}}>;"
276
+ )
277
271
- file_content = FILE_HEAD + "\n\n"
272
- file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
273
- if a_type == "kFE4M3fn":
274
- filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
275
- else:
276
- filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
278
+ kernel_selector_str += (
279
+ jinja2.Template(kernel_template2).render(
280
+ a_type_id=f"vllm::{a_type}.id()",
281
+ b_type_id=f"vllm::{b_type}.id()",
282
+ c_type_id=f"vllm::{c_type}.id()",
283
+ s_type_id=f"vllm::{s_type}.id()",
284
+ **config,
285
+ )
286
+ + "\n"
287
+ )
288
+
289
+ file_content = FILE_HEAD + "\n\n"
290
+ file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
291
+ if a_type == "kFE4M3fn":
292
+ filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
293
+ elif result_dict_tmp is sm_75_result_dict:
294
+ filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
295
+ else:
296
+ filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
297
278
- filename = filename.lower()
298
+ filename = filename.lower()
299
280
- with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
281
- f.write(file_content)
300
+ with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
301
+ f.write(file_content)
302
303
if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
304
kernel_selector_str += (
csrc/quantization/gptq_marlin/gptq_marlin.cu
+39
-29
@@ -37,7 +37,7 @@ __global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
37
38
using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
39
40
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
40
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
41
42
__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
43
int const* __restrict__ perm_int_ptr,
@@ -148,7 +148,7 @@ typedef struct {
148
149
int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
150
int prob_n, int prob_k, int num_bits, int group_size,
151
- bool has_act_order, bool is_k_full) {
151
+ bool has_act_order, bool is_k_full, int stages) {
152
bool cache_scales_chunk = has_act_order && !is_k_full;
153
154
int tb_n = th_config.thread_n;
@@ -166,28 +166,29 @@ int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
166
167
if (cache_scales_chunk) {
168
int load_groups =
169
- tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K
169
+ tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K
170
load_groups = max(load_groups, 32); // We load at least 32 scale groups
171
return load_groups * tb_n * 2;
172
} else {
173
int tb_scales = tb_groups * tb_n * 2;
174
175
- return tb_scales * pipe_stages;
175
+ return tb_scales * stages;
176
}
177
}
178
179
int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
180
int prob_m, int prob_n, int prob_k, int num_bits,
181
int group_size, bool has_act_order, bool is_k_full,
182
- int has_zp, int is_zp_float) {
182
+ int has_zp, bool is_zp_float, bool is_a_8bit,
183
+ int stages) {
184
int pack_factor = 32 / num_bits;
185
186
// Get B size
187
int tb_k = th_config.thread_k;
188
int tb_n = th_config.thread_n;
189
int tb_m = thread_m_blocks * 16;
189
- int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
190
- int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
190
+ int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2);
191
+ int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4;
192
int sh_red_size = tb_m * (tb_n + 8) * 2;
193
int sh_bias_size = tb_n * 2;
194
int tmp_size =
@@ -196,8 +197,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
197
198
int sh_s_size =
199
get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
199
- group_size, has_act_order, is_k_full);
200
- int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0;
200
+ group_size, has_act_order, is_k_full, stages);
201
+ int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0;
202
int sh_zp_size = 0;
203
if (has_zp) {
204
if (is_zp_float)
@@ -217,7 +218,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
218
bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
219
int prob_m, int prob_n, int prob_k, int num_bits,
220
int group_size, bool has_act_order, bool is_k_full,
220
- int has_zp, int is_zp_float, int max_shared_mem) {
221
+ int has_zp, bool is_zp_float, bool is_a_8bit, int stages,
222
+ int max_shared_mem) {
223
// Sanity
224
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
225
th_config.num_threads == -1) {
@@ -242,7 +244,7 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
244
// Check that pipeline fits into cache
245
int cache_size = get_kernel_cache_size(
246
th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
245
- has_act_order, is_k_full, has_zp, is_zp_float);
247
+ has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages);
248
return cache_size <= max_shared_mem;
249
}
250
@@ -251,7 +253,7 @@ MarlinFuncPtr get_marlin_kernel(
253
const vllm::ScalarType c_type, const vllm::ScalarType s_type,
254
int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
255
bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
254
- int threads, bool is_zp_float) {
256
+ int threads, bool is_zp_float, int stages) {
257
int num_bits = b_type.size_bits();
258
auto kernel = MarlinDefault;
259
@@ -265,7 +267,8 @@ exec_config_t determine_exec_config(
267
const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
268
int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8,
269
int num_bits, int group_size, bool has_act_order, bool is_k_full,
268
- bool has_zp, bool is_zp_float, int max_shared_mem, int sms) {
270
+ bool has_zp, bool is_zp_float, int is_a_8bit, int stages,
271
+ int max_shared_mem, int sms) {
272
exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
273
thread_config_t* thread_configs = thread_m_blocks > 1
274
? large_batch_thread_configs
@@ -280,13 +283,15 @@ exec_config_t determine_exec_config(
283
284
if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
285
num_bits, group_size, has_act_order, is_k_full, has_zp,
283
- is_zp_float, max_shared_mem - 512)) {
286
+ is_zp_float, is_a_8bit, stages,
287
+ max_shared_mem - 512)) {
288
continue;
289
}
290
287
- int cache_size = get_kernel_cache_size(
288
- th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits,
289
- group_size, has_act_order, is_k_full, has_zp, is_zp_float);
291
+ int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m,
292
+ prob_n, prob_k, num_bits, group_size,
293
+ has_act_order, is_k_full, has_zp,
294
+ is_zp_float, is_a_8bit, stages);
295
296
int group_blocks = 0;
297
if (!has_act_order) {
@@ -297,14 +302,10 @@ exec_config_t determine_exec_config(
302
get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
303
th_config.thread_n / 16, th_config.thread_k / 16,
304
m_block_size_8, has_act_order, has_zp, group_blocks,
300
- th_config.num_threads, is_zp_float);
305
+ th_config.num_threads, is_zp_float, stages);
306
307
if (kernel == MarlinDefault) continue;
308
304
- // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16);
305
- // int n_tiles = prob_n / th_config.thread_n;
306
- // int k_tiles = prob_k / th_config.thread_k;
307
-
309
return {1, th_config};
310
}
311
@@ -321,6 +322,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
322
int group_size, int dev, cudaStream_t stream, int thread_k_init,
323
int thread_n_init, int sms, bool use_atomic_add,
324
bool use_fp32_reduce, bool is_zp_float) {
325
+ bool is_a_8bit = a_type.size_bits() == 8;
326
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
327
", ", prob_n, ", ", prob_k, "]");
328
@@ -389,8 +391,14 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
391
dev);
392
cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
393
dev);
392
- TORCH_CHECK(major_capability * 10 + minor_capability >= 80,
393
- "marlin kernel only support Ampere or newer GPUs.");
394
+ TORCH_CHECK(major_capability * 10 + minor_capability >= 75,
395
+ "marlin kernel only support Turing or newer GPUs.");
396
+ int stages = 4;
397
+ if (major_capability == 7 && minor_capability == 5) {
398
+ stages = 2;
399
+ TORCH_CHECK(a_type == vllm::kFloat16 || a_type == vllm::kS8,
400
+ "Turing only support FP16 or INT8 activation.");
401
+ }
402
if (a_type == vllm::kFE4M3fn) {
403
TORCH_CHECK(
404
major_capability * 10 + minor_capability == 89 ||
@@ -431,7 +439,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
439
exec_cfg = determine_exec_config(
440
a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k,
441
thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order,
434
- is_k_full, has_zp, is_zp_float, max_shared_mem, sms);
442
+ is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem,
443
+ sms);
444
thread_tfg = exec_cfg.tb_cfg;
445
if (thread_tfg.thread_n != -1) {
446
if (prob_n / thread_tfg.thread_n *
@@ -440,7 +449,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
449
if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split,
450
prob_n, prob_k, num_bits, group_size,
451
has_act_order, is_k_full, has_zp, is_zp_float,
443
- max_shared_mem_new)) {
452
+ is_a_8bit, stages, max_shared_mem_new)) {
453
thread_tfg = {128, 64, 128};
454
exec_cfg = {1, thread_tfg};
455
}
@@ -466,7 +475,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
475
TORCH_CHECK(
476
is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n,
477
prob_k, num_bits, group_size, has_act_order, is_k_full,
469
- has_zp, is_zp_float, max_shared_mem_new),
478
+ has_zp, is_zp_float, is_a_8bit, stages,
479
+ max_shared_mem_new),
480
"Invalid thread config: thread_m_blocks = ", thread_m_blocks,
481
", thread_k = ", thread_tfg.thread_k,
482
", thread_n = ", thread_tfg.thread_n,
@@ -475,12 +485,12 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
485
", prob_m_split = ", prob_m_split, ", group_size = ", group_size,
486
", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
487
", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
478
- ", max_shared_mem_new = ", max_shared_mem_new);
488
+ ", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new);
489
490
auto kernel = get_marlin_kernel(
491
a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
492
thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
483
- num_threads, is_zp_float);
493
+ num_threads, is_zp_float, stages);
494
495
if (kernel == MarlinDefault) {
496
TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
csrc/quantization/gptq_marlin/marlin.cuh
+60
-14
@@ -1,17 +1,19 @@
1
#pragma once
2
3
-#include <torch/all.h>
3
+#ifndef _marlin_cuh
4
+ #define _marlin_cuh
5
+ #include <torch/all.h>
6
5
-#include <ATen/cuda/CUDAContext.h>
6
-#include <c10/cuda/CUDAGuard.h>
7
-#include <cuda.h>
8
-#include <cuda_fp16.h>
9
-#include <cuda_runtime.h>
10
-#include <iostream>
7
+ #include <ATen/cuda/CUDAContext.h>
8
+ #include <c10/cuda/CUDAGuard.h>
9
+ #include <cuda.h>
10
+ #include <cuda_fp16.h>
11
+ #include <cuda_runtime.h>
12
+ #include <iostream>
13
12
-#ifndef MARLIN_NAMESPACE_NAME
13
- #define MARLIN_NAMESPACE_NAME marlin
14
-#endif
14
+ #ifndef MARLIN_NAMESPACE_NAME
15
+ #define MARLIN_NAMESPACE_NAME marlin
16
+ #endif
17
18
namespace MARLIN_NAMESPACE_NAME {
19
@@ -51,9 +53,51 @@ using I4 = Vec<int, 4>;
53
54
constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
55
54
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
55
-// No support for async
56
-#else
56
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
57
+
58
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
59
+ bool pred = true) {
60
+ if (pred) {
61
+ reinterpret_cast<int32_t*>(smem_ptr)[0] =
62
+ reinterpret_cast<const int32_t*>(glob_ptr)[0];
63
+ }
64
+}
65
+
66
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
67
+ bool pred = true) {
68
+ if (pred) {
69
+ reinterpret_cast<int64_t*>(smem_ptr)[0] =
70
+ reinterpret_cast<const int64_t*>(glob_ptr)[0];
71
+ }
72
+}
73
+
74
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
75
+ bool pred = true) {
76
+ if (pred) {
77
+ reinterpret_cast<int4*>(smem_ptr)[0] =
78
+ reinterpret_cast<const int4*>(glob_ptr)[0];
79
+ }
80
+}
81
+
82
+__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
83
+ bool pred = true) {
84
+ if (pred) {
85
+ reinterpret_cast<int4*>(smem_ptr)[0] =
86
+ reinterpret_cast<const int4*>(glob_ptr)[0];
87
+ }
88
+}
89
+
90
+__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {
91
+ reinterpret_cast<int4*>(smem_ptr)[0] =
92
+ reinterpret_cast<const int4*>(glob_ptr)[0];
93
+}
94
+
95
+__device__ inline void cp_async_fence() {}
96
+
97
+template <int n>
98
+__device__ inline void cp_async_wait() {}
99
+
100
+ #else
101
102
__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
103
bool pred = true) {
@@ -126,6 +170,8 @@ __device__ inline void cp_async_wait() {
170
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
171
}
172
129
-#endif
173
+ #endif
174
175
} // namespace MARLIN_NAMESPACE_NAME
176
+
177
+#endif
\ No newline at end of file
csrc/quantization/gptq_marlin/marlin_mma.h
new
+269
@@ -0,0 +1,269 @@
1
+
2
+#include "marlin_dtypes.cuh"
3
+
4
+namespace MARLIN_NAMESPACE_NAME {
5
+
6
+// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
7
+// output/accumulation.
8
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
9
+__device__ inline void mma(
10
+ const typename MarlinScalarType<type_id>::FragA& a_frag,
11
+ const typename MarlinScalarType<type_id>::FragB& frag_b,
12
+ typename MarlinScalarType<type_id>::FragC& frag_c, int idx = 0) {
13
+ const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
14
+ const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
15
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
16
+ if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
17
+ static_assert(!use_fp16_accum);
18
+ }
19
+
20
+ if constexpr (k_size == 16) {
21
+ if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
22
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
23
+ float* c = reinterpret_cast<float*>(&frag_c);
24
+ asm volatile(
25
+ "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
26
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
27
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
28
+ : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
29
+ "f"(c[3]));
30
+ asm volatile(
31
+ "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
32
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
33
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
34
+ : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
35
+ "f"(c[3]));
36
+#else
37
+ float* c = reinterpret_cast<float*>(&frag_c);
38
+ asm volatile(
39
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
40
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
41
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
42
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
43
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
44
+#endif
45
+ } else if constexpr (std::is_same<scalar_t, half>::value &&
46
+ use_fp16_accum) {
47
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
48
+ uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
49
+ asm volatile(
50
+ "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
51
+ "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
52
+ : "=r"(c[0]), "=r"(c[1])
53
+ : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
54
+ asm volatile(
55
+ "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
56
+ "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
57
+ : "=r"(c[0]), "=r"(c[1])
58
+ : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
59
+#else
60
+ uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
61
+ asm volatile(
62
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
63
+ "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
64
+ : "=r"(c[0]), "=r"(c[1])
65
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
66
+ "r"(c[0]), "r"(c[1]));
67
+#endif
68
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
69
+ float* c = reinterpret_cast<float*>(&frag_c);
70
+ asm volatile(
71
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
72
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
73
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
74
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
75
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
76
+ } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
77
+ float* c = reinterpret_cast<float*>(&frag_c);
78
+ asm volatile(
79
+ "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
80
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
81
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
82
+ : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]),
83
+ "f"(c[1]), "f"(c[2]), "f"(c[3]));
84
+ } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
85
+ int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
86
+ asm volatile(
87
+ "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
88
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
89
+ : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
90
+ : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]),
91
+ "r"(c[1]), "r"(c[2]), "r"(c[3]));
92
+ }
93
+ } else if (k_size == 32) {
94
+ if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
95
+ float* c = reinterpret_cast<float*>(&frag_c);
96
+ asm volatile(
97
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
98
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
99
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
100
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
101
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
102
+ } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
103
+ int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
104
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
105
+ asm volatile(
106
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
107
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
108
+ : "=r"(c[0]), "=r"(c[1])
109
+ : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1]));
110
+ asm volatile(
111
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
112
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
113
+ : "=r"(c[2]), "=r"(c[3])
114
+ : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3]));
115
+ asm volatile(
116
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
117
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
118
+ : "=r"(c[0]), "=r"(c[1])
119
+ : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1]));
120
+ asm volatile(
121
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
122
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
123
+ : "=r"(c[2]), "=r"(c[3])
124
+ : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3]));
125
+#else
126
+ asm volatile(
127
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
128
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
129
+ : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
130
+ : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
131
+ "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
132
+#endif
133
+ }
134
+ }
135
+}
136
+
137
+template <vllm::ScalarTypeId type_id, bool use_fp16_accum, int k_size = 16>
138
+__device__ inline void mma_trans(
139
+ const typename MarlinScalarType<type_id>::FragA& a_frag,
140
+ const typename MarlinScalarType<type_id>::FragB& frag_b,
141
+ const typename MarlinScalarType<type_id>::FragB& frag_b2,
142
+ typename MarlinScalarType<type_id>::FragC& frag_c) {
143
+ const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
144
+ const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
145
+ const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
146
+ float* c = reinterpret_cast<float*>(&frag_c);
147
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
148
+ if constexpr (!std::is_same<scalar_t, half>::value || k_size != 16) {
149
+ static_assert(!use_fp16_accum);
150
+ }
151
+
152
+ if constexpr (k_size == 16) {
153
+ if constexpr (std::is_same<scalar_t, half>::value && !use_fp16_accum) {
154
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
155
+ float* c = reinterpret_cast<float*>(&frag_c);
156
+ asm volatile(
157
+ "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
158
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
159
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
160
+ : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
161
+ "f"(c[3]));
162
+ asm volatile(
163
+ "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 "
164
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
165
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
166
+ : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
167
+ "f"(c[3]));
168
+#else
169
+ float* c = reinterpret_cast<float*>(&frag_c);
170
+ asm volatile(
171
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
172
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
173
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
174
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
175
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
176
+#endif
177
+ } else if constexpr (std::is_same<scalar_t, half>::value &&
178
+ use_fp16_accum) {
179
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
180
+ uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
181
+ asm volatile(
182
+ "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
183
+ "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
184
+ : "=r"(c[0]), "=r"(c[1])
185
+ : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
186
+ asm volatile(
187
+ "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 "
188
+ "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n"
189
+ : "=r"(c[0]), "=r"(c[1])
190
+ : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
191
+#else
192
+ uint32_t* c = reinterpret_cast<uint32_t*>(&frag_c);
193
+ asm volatile(
194
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
195
+ "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
196
+ : "=r"(c[0]), "=r"(c[1])
197
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
198
+ "r"(c[0]), "r"(c[1]));
199
+#endif
200
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
201
+ float* c = reinterpret_cast<float*>(&frag_c);
202
+ asm volatile(
203
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
204
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
205
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
206
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
207
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
208
+ } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
209
+ float* c = reinterpret_cast<float*>(&frag_c);
210
+ asm volatile(
211
+ "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
212
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
213
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
214
+ : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
215
+ "f"(c[3]));
216
+ } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
217
+ int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
218
+ asm volatile(
219
+ "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
220
+ "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
221
+ : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
222
+ : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
223
+ "r"(c[3]));
224
+ }
225
+ } else {
226
+ if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
227
+ float* c = reinterpret_cast<float*>(&frag_c);
228
+ asm volatile(
229
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
230
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
231
+ : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
232
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
233
+ "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
234
+ } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
235
+ int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
236
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
237
+ asm volatile(
238
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
239
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
240
+ : "=r"(c[0]), "=r"(c[1])
241
+ : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]));
242
+ asm volatile(
243
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
244
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
245
+ : "=r"(c[2]), "=r"(c[3])
246
+ : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3]));
247
+ asm volatile(
248
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
249
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
250
+ : "=r"(c[0]), "=r"(c[1])
251
+ : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1]));
252
+ asm volatile(
253
+ "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite "
254
+ "{%0,%1}, {%2}, {%3}, {%4,%5};\n"
255
+ : "=r"(c[2]), "=r"(c[3])
256
+ : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3]));
257
+#else
258
+ asm volatile(
259
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
260
+ "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
261
+ : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
262
+ : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
263
+ "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
264
+#endif
265
+ }
266
+ }
267
+}
268
+
269
+} // namespace MARLIN_NAMESPACE_NAME
\ No newline at end of file
csrc/quantization/gptq_marlin/marlin_template.h
+40
-144
@@ -26,6 +26,7 @@
26
#include "marlin.cuh"
27
#include "marlin_dtypes.cuh"
28
#include "dequant.h"
29
+#include "marlin_mma.h"
30
#include "core/scalar_type.hpp"
31
32
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
@@ -35,7 +36,7 @@
36
37
namespace MARLIN_NAMESPACE_NAME {
38
38
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
39
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750
40
41
template <typename scalar_t, // compute dtype, half or nv_float16
42
const vllm::ScalarTypeId b_type_id, // weight MarlinScalarType id
@@ -75,137 +76,6 @@ __global__ void Marlin(
76
77
#else
78
78
-// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
79
-// output/accumulation.
80
-template <vllm::ScalarTypeId type_id, int k_size = 16>
81
-__device__ inline void mma(
82
- const typename MarlinScalarType<type_id>::FragA& a_frag,
83
- const typename MarlinScalarType<type_id>::FragB& frag_b,
84
- typename MarlinScalarType<type_id>::FragC& frag_c, int idx = 0) {
85
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
86
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
87
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
88
- if constexpr (k_size == 16) {
89
- if constexpr (std::is_same<scalar_t, half>::value) {
90
- float* c = reinterpret_cast<float*>(&frag_c);
91
- asm volatile(
92
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
93
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
94
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
95
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
96
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
97
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
98
- float* c = reinterpret_cast<float*>(&frag_c);
99
- asm volatile(
100
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
101
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
102
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
103
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
104
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
105
- } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
106
- float* c = reinterpret_cast<float*>(&frag_c);
107
- asm volatile(
108
- "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
109
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
110
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
111
- : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]),
112
- "f"(c[1]), "f"(c[2]), "f"(c[3]));
113
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
114
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
115
- asm volatile(
116
- "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
117
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
118
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
119
- : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]),
120
- "r"(c[1]), "r"(c[2]), "r"(c[3]));
121
- }
122
- } else if (k_size == 32) {
123
- if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
124
- float* c = reinterpret_cast<float*>(&frag_c);
125
- asm volatile(
126
- "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
127
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
128
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
129
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
130
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
131
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
132
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
133
- asm volatile(
134
- "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
135
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
136
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
137
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
138
- "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
139
- }
140
- }
141
-}
142
-
143
-template <vllm::ScalarTypeId type_id, int k_size = 16>
144
-__device__ inline void mma_trans(
145
- const typename MarlinScalarType<type_id>::FragA& a_frag,
146
- const typename MarlinScalarType<type_id>::FragB& frag_b,
147
- const typename MarlinScalarType<type_id>::FragB& frag_b2,
148
- typename MarlinScalarType<type_id>::FragC& frag_c) {
149
- const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
150
- const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
151
- const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
152
- float* c = reinterpret_cast<float*>(&frag_c);
153
- using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
154
- if constexpr (k_size == 16) {
155
- if constexpr (std::is_same<scalar_t, half>::value) {
156
- float* c = reinterpret_cast<float*>(&frag_c);
157
- asm volatile(
158
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
159
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
160
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
161
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
162
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
163
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
164
- float* c = reinterpret_cast<float*>(&frag_c);
165
- asm volatile(
166
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
167
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
168
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
169
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
170
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
171
- } else if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
172
- float* c = reinterpret_cast<float*>(&frag_c);
173
- asm volatile(
174
- "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 "
175
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
176
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
177
- : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]),
178
- "f"(c[3]));
179
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
180
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
181
- asm volatile(
182
- "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite "
183
- "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n"
184
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
185
- : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]),
186
- "r"(c[3]));
187
- }
188
- } else {
189
- if constexpr (std::is_same<scalar_t, __nv_fp8_e4m3>::value) {
190
- float* c = reinterpret_cast<float*>(&frag_c);
191
- asm volatile(
192
- "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
193
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
194
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
195
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
196
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
197
- } else if constexpr (std::is_same<scalar_t, int8_t>::value) {
198
- int32_t* c = reinterpret_cast<int32_t*>(&frag_c);
199
- asm volatile(
200
- "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite "
201
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
202
- : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3])
203
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
204
- "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]));
205
- }
206
- }
207
-}
208
-
79
// Instruction for loading a full 16x16 matrix fragment of operand A from shared
80
// memory, directly in tensor core layout.
81
template <int count, vllm::ScalarTypeId type_id>
@@ -415,6 +285,17 @@ __global__ void Marlin(
285
if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
286
#endif
287
288
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
289
+ // Turing TensorCore only supports fp16 and int8
290
+ if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id())
291
+ return;
292
+ #endif
293
+
294
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
295
+ constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
296
+ #else
297
+ constexpr bool use_fp16_accum = false;
298
+ #endif
299
using Adtype = MarlinScalarType<a_type_id>;
300
using Cdtype = MarlinScalarType<c_type_id>;
301
const int4* A = A0;
@@ -873,10 +754,6 @@ __global__ void Marlin(
754
constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
755
: (stages * s_sh_stage);
756
int4* sh_s = sh_zp + (stages * zp_sh_stage);
876
- // shared memory reused by reduction should be smaller than
877
- // shared memory used by weight.
878
- static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
879
- stages * b_sh_stage);
757
int4* sh_a = sh_s + sh_s_size;
758
759
// Register storage for double buffer of shared memory reads.
@@ -1395,11 +1272,13 @@ __global__ void Marlin(
1272
#pragma unroll
1273
for (int i = 0; i < thread_m_blocks; i++) {
1274
if constexpr (m_block_size_8) {
1398
- mma_trans<a_type_id>(frag_a[k2][i], frag_b0, frag_b1,
1399
- frag_c[i][j][0]);
1275
+ mma_trans<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, frag_b1,
1276
+ frag_c[i][j][0]);
1277
} else {
1401
- mma<a_type_id>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1402
- mma<a_type_id>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
1278
+ mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0,
1279
+ frag_c[i][j][0]);
1280
+ mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b1,
1281
+ frag_c[i][j][1]);
1282
}
1283
}
1284
}
@@ -1433,10 +1312,12 @@ __global__ void Marlin(
1312
1313
#pragma unroll
1314
for (int i = 0; i < thread_m_blocks; i++) {
1436
- mma<a_type_id, 32>(frag_a[k2][i], frag_b[0],
1437
- (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
1438
- mma<a_type_id, 32>(frag_a[k2][i], frag_b[1],
1439
- (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
1315
+ mma<a_type_id, false, 32>(
1316
+ frag_a[k2][i], frag_b[0],
1317
+ (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]);
1318
+ mma<a_type_id, false, 32>(
1319
+ frag_a[k2][i], frag_b[1],
1320
+ (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]);
1321
}
1322
1323
if constexpr (group_blocks != -1) {
@@ -1956,6 +1837,21 @@ __global__ void Marlin(
1837
// While this pattern may not be the most readable, other ways of writing
1838
// the loop seemed to noticeably worse performance after compilation.
1839
if (slice_iters == 0) {
1840
+ // convert fp16 accum to fp32 for reduction
1841
+ if constexpr (use_fp16_accum) {
1842
+ #pragma unroll
1843
+ for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) {
1844
+ float* frag_c_part_float = reinterpret_cast<float*>(frag_c) + i * 4;
1845
+ scalar_t* frag_c_part_half =
1846
+ reinterpret_cast<scalar_t*>(frag_c_part_float);
1847
+
1848
+ #pragma unroll
1849
+ for (int i = 3; i >= 0; i--) {
1850
+ frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]);
1851
+ }
1852
+ }
1853
+ }
1854
+
1855
if constexpr (is_a_8bit) {
1856
float frag_a_s[2 * thread_m_blocks];
1857
vllm/model_executor/layers/quantization/awq_marlin.py
+1
-1
@@ -121,7 +121,7 @@ class AWQMarlinConfig(QuantizationConfig):
121
122
@classmethod
123
def get_min_capability(cls) -> int:
124
- return 80
124
+ return 75
125
126
@classmethod
127
def get_config_filenames(cls) -> list[str]:
vllm/model_executor/layers/quantization/fp8.py
+1
-1
@@ -253,7 +253,7 @@ class Fp8Config(QuantizationConfig):
253
254
@classmethod
255
def get_min_capability(cls) -> int:
256
- return 80
256
+ return 75
257
258
@classmethod
259
def get_config_filenames(cls) -> list[str]:
vllm/model_executor/layers/quantization/gptq_marlin.py
+1
-1
@@ -181,7 +181,7 @@ class GPTQMarlinConfig(QuantizationConfig):
181
182
@classmethod
183
def get_min_capability(cls) -> int:
184
- return 80
184
+ return 75
185
186
@classmethod
187
def get_config_filenames(cls) -> list[str]:
vllm/model_executor/layers/quantization/modelopt.py
+1
-1
@@ -871,7 +871,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
871
872
@classmethod
873
def get_min_capability(cls) -> int:
874
- return 80
874
+ return 75
875
876
@classmethod
877
def override_quantization_method(