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 12, 08:30 UTC
[Kernel][Quantization] add w4a8 support for marlin kernel (#24722)
Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin@redhat.com>
Jinzhen Lin committed
Nov 29, 2025 at 23:19 UTC
1656ad37045579999a5a9ef3b940f945cd92bb4e
46 files changed
+4363
-2232
CMakeLists.txt
+78
-24
@@ -354,8 +354,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
354
# Only build Marlin kernels if we are building for at least some compatible archs.
355
# Keep building Marlin for 9.0 as there are some group sizes and shapes that
356
# are not supported by Machete yet.
357
- # 9.0 for latest bf16 atomicAdd PTX
358
- cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
357
+
358
+ # marlin arches for fp16 output
359
+ cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
360
+ # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
361
+ cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
362
+ # marlin arches for fp8 input
363
+ # - sm80 doesn't support fp8 computation
364
+ # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
365
+ # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
366
+ cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
367
+
368
if (MARLIN_ARCHS)
369
370
#
@@ -365,16 +374,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
374
set(MARLIN_GEN_SCRIPT
375
${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
376
file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
377
+ list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
378
+ set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
379
369
- message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH}")
370
- message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH}")
380
+ message(STATUS "Marlin generation script hash: ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
381
+ message(STATUS "Last run Marlin generate script hash: $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
382
372
- if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH}
373
- OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH})
383
+ if (NOT DEFINED CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
384
+ OR NOT $CACHE{MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
385
execute_process(
386
COMMAND ${CMAKE_COMMAND} -E env
387
PYTHONPATH=$PYTHONPATH
377
- ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT}
388
+ ${Python_EXECUTABLE} ${MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
389
RESULT_VARIABLE marlin_generation_result
390
OUTPUT_VARIABLE marlin_generation_result
391
OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log
@@ -387,15 +398,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
398
"\nCheck the log for details: "
399
"${CMAKE_CURRENT_BINARY_DIR}/marlin_generation.log")
400
else()
390
- set(MARLIN_GEN_SCRIPT_HASH ${MARLIN_GEN_SCRIPT_HASH}
391
- CACHE STRING "Last run Marlin generate script hash" FORCE)
401
+ set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
402
+ CACHE STRING "Last run Marlin generate script hash and arch" FORCE)
403
message(STATUS "Marlin generation completed successfully.")
404
endif()
405
else()
406
message(STATUS "Marlin generation script has not changed, skipping generation.")
407
endif()
408
398
- file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/kernel_*.cu")
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}")
@@ -403,12 +414,34 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
414
set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
415
PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
416
endif()
406
-
417
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
418
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")
426
+ endif()
427
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
428
+
429
+ if (MARLIN_FP8_ARCHS)
430
+ file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
431
+ set_gencode_flags_for_srcs(
432
+ SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
433
+ CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
434
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
435
+ set_source_files_properties(${MARLIN_TEMPLATE_FP8_KERNEL_SRC}
436
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
437
+ endif()
438
+ list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_FP8_KERNEL_SRC})
439
+ endif()
440
+
441
set(MARLIN_SRCS
442
"csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
443
"csrc/quantization/gptq_marlin/gptq_marlin.cu"
444
+ "csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu"
445
"csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
446
"csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
447
set_gencode_flags_for_srcs(
@@ -941,8 +974,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
974
CUDA_ARCHS "${CUDA_ARCHS}")
975
976
list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
944
- # 9.0 for latest bf16 atomicAdd PTX
945
- cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
977
+ # moe marlin arches
978
+ # note that we always set `use_atomic_add=False` for moe marlin now,
979
+ # so we don't need 9.0 for bf16 atomicAdd PTX
980
+ cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
981
+ # moe marlin arches for fp8 input
982
+ # - sm80 doesn't support fp8 computation
983
+ # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
984
+ # so we only enable fp8 computation for SM89 (e.g. RTX 40x0) and 12.0 (e.g. RTX 50x0)
985
+ cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
986
if (MARLIN_MOE_ARCHS)
987
988
#
@@ -952,16 +992,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
992
set(MOE_MARLIN_GEN_SCRIPT
993
${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py)
994
file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH)
995
+ list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
996
+ set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MOE_MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
997
956
- message(STATUS "Marlin MOE generation script hash: ${MOE_MARLIN_GEN_SCRIPT_HASH}")
957
- message(STATUS "Last run Marlin MOE generate script hash: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}")
998
+ message(STATUS "Marlin MOE generation script hash with arch: ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
999
+ message(STATUS "Last run Marlin MOE generate script hash with arch: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}")
1000
959
- if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}
960
- OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH})
1001
+ if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
1002
+ OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH})
1003
execute_process(
1004
COMMAND ${CMAKE_COMMAND} -E env
1005
PYTHONPATH=$PYTHONPATH
964
- ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT}
1006
+ ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT} ${CUDA_ARCHS_STR}
1007
RESULT_VARIABLE moe_marlin_generation_result
1008
OUTPUT_VARIABLE moe_marlin_generation_output
1009
OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log
@@ -974,7 +1016,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1016
"\nCheck the log for details: "
1017
"${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log")
1018
else()
977
- set(MOE_MARLIN_GEN_SCRIPT_HASH ${MOE_MARLIN_GEN_SCRIPT_HASH}
1019
+ set(MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH ${MOE_MARLIN_GEN_SCRIPT_HASH_AND_ARCH}
1020
CACHE STRING "Last run Marlin MOE generate script hash" FORCE)
1021
message(STATUS "Marlin MOE generation completed successfully.")
1022
endif()
@@ -982,16 +1024,28 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
1024
message(STATUS "Marlin MOE generation script has not changed, skipping generation.")
1025
endif()
1026
985
- file(GLOB MOE_WNAA16_MARLIN_SRC "csrc/moe/marlin_moe_wna16/*.cu")
1027
+ file(GLOB MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/sm80_kernel_*.cu")
1028
+ list(APPEND MARLIN_MOE_SRC "csrc/moe/marlin_moe_wna16/ops.cu")
1029
set_gencode_flags_for_srcs(
987
- SRCS "${MOE_WNAA16_MARLIN_SRC}"
1030
+ SRCS "${MARLIN_MOE_SRC}"
1031
CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
1032
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
990
- set_source_files_properties(${MOE_WNAA16_MARLIN_SRC}
1033
+ set_source_files_properties(${MARLIN_MOE_SRC}
1034
PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1035
endif()
993
-
994
- list(APPEND VLLM_MOE_EXT_SRC ${MOE_WNAA16_MARLIN_SRC})
1036
+ list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
1037
+
1038
+ if (MARLIN_MOE_FP8_ARCHS)
1039
+ file(GLOB MARLIN_MOE_FP8_SRC "csrc/moe/marlin_moe_wna16/sm89_kernel_*.cu")
1040
+ set_gencode_flags_for_srcs(
1041
+ SRCS "${MARLIN_MOE_FP8_SRC}"
1042
+ CUDA_ARCHS "${MARLIN_MOE_FP8_ARCHS}")
1043
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
1044
+ set_source_files_properties(${MARLIN_MOE_FP8_SRC}
1045
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
1046
+ endif()
1047
+ list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_FP8_SRC})
1048
+ endif()
1049
1050
message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
1051
else()
benchmarks/kernels/benchmark_machete.py
+1
@@ -237,6 +237,7 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
237
b_q_weight=w_q,
238
b_bias=None,
239
b_scales=w_s,
240
+ a_scales=None,
241
global_scale=None,
242
b_zeros=w_zp,
243
g_idx=g_idx,
benchmarks/kernels/benchmark_marlin.py
+2
-2
@@ -263,7 +263,7 @@ def bench_run(
263
264
results.append(
265
benchmark.Timer(
266
- stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
266
+ stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
267
globals=globals,
268
label=label,
269
sub_label=sub_label,
@@ -273,7 +273,7 @@ def bench_run(
273
274
results.append(
275
benchmark.Timer(
276
- stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
276
+ stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
277
globals=globals,
278
label=label,
279
sub_label=sub_label,
csrc/moe/marlin_moe_wna16/.gitignore
+2
-1
@@ -1 +1,2 @@
1
-kernel_*.cu
\ No newline at end of file
1
+sm*_kernel_*.cu
2
+kernel_selector.h
csrc/moe/marlin_moe_wna16/generate_kernels.py
+227
-79
@@ -4,134 +4,282 @@ import glob
4
import itertools
5
import os
6
import subprocess
7
+import sys
8
9
import jinja2
10
10
-FILE_HEAD = """
11
-// auto generated by generate.py
11
+ARCHS = []
12
+SUPPORT_FP8 = False
13
+for arch in sys.argv[1].split(","):
14
+ arch = arch[: arch.index(".") + 2].replace(".", "")
15
+ arch = int(arch)
16
+ # only SM89 and SM120 fully support
17
+ # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
18
+ # SM90 and SM100 can use this PTX, but it’s simulated
19
+ # with FP16 MMA, so it cannot achieve any acceleration.
20
+ if arch in [89, 120]:
21
+ SUPPORT_FP8 = True
22
+
23
+FILE_HEAD_COMMENT = """
24
+// auto generated by generate_kernels.py
25
// clang-format off
26
+""".lstrip()
27
28
+FILE_HEAD = (
29
+ FILE_HEAD_COMMENT
30
+ + """
31
#include "kernel.h"
32
#include "marlin_template.h"
33
34
namespace MARLIN_NAMESPACE_NAME {
18
-""".strip()
35
+"""
36
+)
37
38
TEMPLATE = (
39
"template __global__ void Marlin<"
22
- "{{scalar_t}}, "
23
- "{{w_type_id}}, "
40
+ "{{a_type_id}}, "
41
+ "{{b_type_id}}, "
42
+ "{{c_type_id}}, "
43
"{{s_type_id}}, "
44
"{{threads}}, "
45
"{{thread_m_blocks}}, "
46
"{{thread_n_blocks}}, "
47
"{{thread_k_blocks}}, "
29
- "{{'true' if m_block_size_8 else 'false'}}, "
48
+ "{{m_block_size_8}}, "
49
"{{stages}}, "
50
"{{group_blocks}}, "
32
- "{{'true' if is_zp_float else 'false'}}>"
51
+ "{{is_zp_float}}>"
52
"( MARLIN_KERNEL_PARAMS );"
53
)
54
36
-# int8 with zero point case (vllm::kU8) is also supported,
37
-# we don't add it to reduce wheel size.
38
-SCALAR_TYPES = [
39
- "vllm::kU4",
40
- "vllm::kU4B8",
41
- "vllm::kU8B128",
42
- "vllm::kFE4M3fn",
43
- "vllm::kFE2M1f",
44
-]
55
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
56
57
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
48
-# group_blocks:
49
-# = 0 : act order case
50
-# = -1 : channelwise quantization
51
-# > 0 : group_size=16*group_blocks
52
-GROUP_BLOCKS = [0, -1, 1, 2, 4, 8]
53
-DTYPES = ["fp16", "bf16"]
58
+
59
+QUANT_CONFIGS = [
60
+ # AWQ-INT4
61
+ {
62
+ "b_type": "kU4",
63
+ "thread_configs": THREAD_CONFIGS,
64
+ "thread_m_blocks": THREAD_M_BLOCKS,
65
+ "group_blocks": [-1, 2, 4, 8],
66
+ },
67
+ # GPTQ-INT4
68
+ {
69
+ "b_type": "kU4B8",
70
+ "thread_configs": THREAD_CONFIGS,
71
+ "thread_m_blocks": THREAD_M_BLOCKS,
72
+ "group_blocks": [-1, 0, 2, 4, 8],
73
+ },
74
+ # AWQ-INT8
75
+ {
76
+ "b_type": "kU8B128",
77
+ "thread_configs": THREAD_CONFIGS,
78
+ "thread_m_blocks": THREAD_M_BLOCKS,
79
+ "group_blocks": [-1, 0, 2, 4, 8],
80
+ },
81
+ # FP8
82
+ {
83
+ "b_type": "kFE4M3fn",
84
+ "thread_configs": THREAD_CONFIGS,
85
+ "thread_m_blocks": THREAD_M_BLOCKS,
86
+ "group_blocks": [-1, 8],
87
+ },
88
+ # NVFP4
89
+ {
90
+ "b_type": "kFE2M1f",
91
+ "s_type": "kFE4M3fn",
92
+ "thread_configs": THREAD_CONFIGS,
93
+ "thread_m_blocks": THREAD_M_BLOCKS,
94
+ "group_blocks": [1],
95
+ },
96
+ # MXFP4
97
+ {
98
+ "a_type": ["kBFloat16"],
99
+ "b_type": "kFE2M1f",
100
+ "s_type": "kFE8M0fnu",
101
+ "thread_configs": THREAD_CONFIGS,
102
+ "thread_m_blocks": THREAD_M_BLOCKS,
103
+ "group_blocks": [2],
104
+ },
105
+ # AWQ-INT4 with INT8 activation
106
+ {
107
+ "a_type": ["kS8"],
108
+ "b_type": "kU4",
109
+ "thread_configs": THREAD_CONFIGS,
110
+ "thread_m_blocks": [1, 2, 3, 4],
111
+ "group_blocks": [-1, 2, 4, 8],
112
+ },
113
+ # GPTQ-INT4 with INT8 activation
114
+ {
115
+ "a_type": ["kS8"],
116
+ "b_type": "kU4B8",
117
+ "thread_configs": THREAD_CONFIGS,
118
+ "thread_m_blocks": [1, 2, 3, 4],
119
+ "group_blocks": [-1, 2, 4, 8],
120
+ },
121
+ # GPTQ-INT4 with FP8 activation
122
+ {
123
+ "a_type": ["kFE4M3fn"],
124
+ "b_type": "kU4B8",
125
+ "thread_configs": THREAD_CONFIGS,
126
+ "thread_m_blocks": [1, 2, 3, 4],
127
+ "group_blocks": [-1, 2, 4, 8],
128
+ },
129
+ # AWQ-INT4 with FP8 activation
130
+ {
131
+ "a_type": ["kFE4M3fn"],
132
+ "b_type": "kU4",
133
+ "thread_configs": THREAD_CONFIGS,
134
+ "thread_m_blocks": [1, 2, 3, 4],
135
+ "group_blocks": [-1, 2, 4, 8],
136
+ },
137
+ # MXFP4 with FP8 activation
138
+ {
139
+ "a_type": ["kFE4M3fn"],
140
+ "b_type": "kFE2M1f",
141
+ "c_type": ["kBFloat16"],
142
+ "s_type": "kFE8M0fnu",
143
+ "thread_configs": THREAD_CONFIGS,
144
+ "thread_m_blocks": [1, 2, 3, 4],
145
+ "group_blocks": [2],
146
+ },
147
+]
148
149
150
def remove_old_kernels():
57
- for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
151
+ for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
152
subprocess.call(["rm", "-f", filename])
153
154
+ filename = os.path.dirname(__file__) + "/kernel_selector.h"
155
+ subprocess.call(["rm", "-f", filename])
156
+
157
158
def generate_new_kernels():
62
- for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
63
- all_template_str_list = []
159
+ result_dict = {}
160
65
- for group_blocks, m_blocks, thread_configs in itertools.product(
66
- GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
67
- ):
68
- # act order case only support gptq-int4 and gptq-int8
69
- if group_blocks == 0 and scalar_type not in [
70
- "vllm::kU4B8",
71
- "vllm::kU8B128",
72
- ]:
73
- continue
74
- if thread_configs[2] == 256:
75
- # for small batch (m_blocks == 1), we only need (128, 128, 256)
76
- # for large batch (m_blocks > 1), we only need (64, 256, 256)
77
- if m_blocks <= 1 and thread_configs[0] != 128:
78
- continue
79
- if m_blocks > 1 and thread_configs[0] != 64:
80
- continue
81
-
82
- # we only support channelwise quantization and group_size == 128
83
- # for fp8
84
- if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
85
- continue
86
- # nvfp4 only supports group_size == 16
87
- # mxfp4 only supports group_size == 32
88
- if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
161
+ for quant_config in QUANT_CONFIGS:
162
+ c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
163
+ a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
164
+ b_type = quant_config["b_type"]
165
+ all_group_blocks = quant_config["group_blocks"]
166
+ all_m_blocks = quant_config["thread_m_blocks"]
167
+ all_thread_configs = quant_config["thread_configs"]
168
+
169
+ for a_type, c_type in itertools.product(a_types, c_types):
170
+ if not SUPPORT_FP8 and a_type == "kFE4M3fn":
171
continue
90
- # other quantization methods don't support group_size = 16
91
- if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
172
+ if "16" in a_type and "16" in c_type and a_type != c_type:
173
continue
174
+ s_type = quant_config.get("s_type", c_type)
175
+ if (a_type, b_type, c_type) not in result_dict:
176
+ result_dict[(a_type, b_type, c_type)] = []
177
+
178
+ for group_blocks, m_blocks, thread_configs in itertools.product(
179
+ all_group_blocks, all_m_blocks, all_thread_configs
180
+ ):
181
+ thread_k, thread_n, threads = thread_configs
182
+
183
+ if threads == 256:
184
+ # for small batch (m_blocks == 1),
185
+ # we only need (128, 128, 256)
186
+ # for large batch (m_blocks > 1),
187
+ # we only need (64, 256, 256)
188
+ if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
189
+ continue
190
+ if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
191
+ continue
192
94
- k_blocks = thread_configs[0] // 16
95
- n_blocks = thread_configs[1] // 16
96
- threads = thread_configs[2]
193
+ config = {
194
+ "threads": threads,
195
+ "s_type": s_type,
196
+ "thread_m_blocks": max(m_blocks, 1),
197
+ "thread_k_blocks": thread_k // 16,
198
+ "thread_n_blocks": thread_n // 16,
199
+ "m_block_size_8": "true" if m_blocks == 0.5 else "false",
200
+ "stages": "pipe_stages",
201
+ "group_blocks": group_blocks,
202
+ "is_zp_float": "false",
203
+ }
204
98
- c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
205
+ result_dict[(a_type, b_type, c_type)].append(config)
206
100
- if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
101
- s_type = "vllm::kFE4M3fn"
102
- elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
103
- s_type = "vllm::kFE8M0fnu"
104
- if dtype == "fp16":
105
- # we cannot safely dequantize e8m0 to fp16, so skip this
106
- continue
107
- elif dtype == "fp16":
108
- s_type = "vllm::kFloat16"
109
- elif dtype == "bf16":
110
- s_type = "vllm::kBFloat16"
207
+ kernel_selector_str = FILE_HEAD_COMMENT
208
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(
113
- scalar_t=c_dtype,
114
- w_type_id=scalar_type + ".id()",
115
- s_type_id=s_type + ".id()",
116
- threads=threads,
117
- thread_m_blocks=max(m_blocks, 1),
118
- thread_n_blocks=n_blocks,
119
- thread_k_blocks=k_blocks,
120
- m_block_size_8=m_blocks == 0.5,
121
- stages="pipe_stages",
122
- group_blocks=group_blocks,
123
- is_zp_float=False,
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
)
125
-
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(
252
+ a_type_id=f"vllm::{a_type}.id()",
253
+ b_type_id=f"vllm::{b_type}.id()",
254
+ c_type_id=f"vllm::{c_type}.id()",
255
+ s_type_id=f"vllm::{s_type}.id()",
256
+ **config,
257
+ )
258
+ + "\n"
259
+ )
260
+
261
file_content = FILE_HEAD + "\n\n"
262
file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
130
- filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"
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"
267
+
268
+ filename = filename.lower()
269
270
with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
271
f.write(file_content)
272
273
+ if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
274
+ kernel_selector_str += (
275
+ "else if (a_type == vllm::kFE4M3fn)\n"
276
+ " TORCH_CHECK(false, "
277
+ '"marlin kernel with fp8 activation is not built.");'
278
+ )
279
+
280
+ with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
281
+ f.write(kernel_selector_str)
282
+
283
284
if __name__ == "__main__":
285
remove_old_kernels()
csrc/moe/marlin_moe_wna16/kernel.h
+7
-5
@@ -11,8 +11,9 @@
11
const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
const int4 *__restrict__ b_bias_ptr, \
14
+ const float *__restrict__ a_scales_ptr, \
15
const int4 *__restrict__ scales_ptr, \
15
- const uint16_t *__restrict__ scale2_ptr, \
16
+ const uint16_t *__restrict__ global_scale_ptr, \
17
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
const int32_t *__restrict__ sorted_token_ids_ptr, \
19
const int32_t *__restrict__ expert_ids_ptr, \
@@ -20,12 +21,13 @@
21
const float *__restrict__ topk_weights_ptr, int top_k, \
22
bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
23
int prob_n, int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
23
- bool use_fp32_reduce, int max_shared_mem
24
+ bool use_fp32_reduce
25
26
namespace MARLIN_NAMESPACE_NAME {
26
-template <typename scalar_t, // compute dtype, half or nv_float16
27
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
28
- const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
27
+template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
28
+ const vllm::ScalarTypeId b_type_id, // B ScalarType id
29
+ const vllm::ScalarTypeId c_type_id, // C ScalarType id
30
+ const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id
31
const int threads, // number of threads in a threadblock
32
const int thread_m_blocks, // number of 16x16 blocks in the m
33
// dimension (batchsize) of the
csrc/moe/marlin_moe_wna16/marlin_template.h
+892
-552
@@ -38,7 +38,7 @@ namespace MARLIN_NAMESPACE_NAME {
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
41
+ const vllm::ScalarTypeId b_type_id, // weight MarlinScalarType 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
@@ -49,9 +49,10 @@ 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
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?
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
@@ -76,8 +77,8 @@ __global__ void Marlin(
77
int prob_k, // reduction dimension k
78
int* locks, // extra global storage for barrier synchronization
79
bool use_atomic_add, // whether to use atomic add to reduce
79
- bool use_fp32_reduce, // whether to use fp32 global reduce
80
- int max_shared_mem) {}
80
+ bool use_fp32_reduce // whether to use fp32 global reduce
81
+) {}
82
83
} // namespace MARLIN_NAMESPACE_NAME
84
@@ -85,65 +86,148 @@ __global__ void Marlin(
86
87
// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
88
// output/accumulation.
88
-template <typename scalar_t>
89
-__device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
90
- const typename ScalarType<scalar_t>::FragB& frag_b,
91
- typename ScalarType<scalar_t>::FragC& frag_c) {
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);
94
- float* c = reinterpret_cast<float*>(&frag_c);
95
- if constexpr (std::is_same<scalar_t, half>::value) {
96
- asm volatile(
97
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.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, nv_bfloat16>::value) {
103
- asm volatile(
104
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
105
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
106
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
107
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
108
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
109
- } else {
110
- STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
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
114
-template <typename scalar_t>
152
+template <vllm::ScalarTypeId type_id, int k_size = 16>
153
__device__ inline void mma_trans(
116
- const typename ScalarType<scalar_t>::FragA& a_frag,
117
- const typename ScalarType<scalar_t>::FragB& frag_b,
118
- const typename ScalarType<scalar_t>::FragB& frag_b2,
119
- typename ScalarType<scalar_t>::FragC& frag_c) {
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);
124
- if constexpr (std::is_same<scalar_t, half>::value) {
125
- asm volatile(
126
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.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"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
130
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
131
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
132
- asm volatile(
133
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
134
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
135
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
136
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
137
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
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 {
139
- STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
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
227
// Instruction for loading a full 16x16 matrix fragment of operand A from shared
228
// memory, directly in tensor core layout.
145
-template <int count, typename scalar_t>
146
-__device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
229
+template <int count, vllm::ScalarTypeId type_id>
230
+__device__ inline void ldsm(typename MarlinScalarType<type_id>::FragA& frag_a,
231
const void* smem_ptr) {
232
uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
233
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
@@ -167,47 +251,54 @@ __device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
251
252
// Multiply dequantized values by the corresponding quantization scale; used
253
// only for grouped quantization.
170
-template <typename scalar_t>
171
-__device__ inline void scale(typename ScalarType<scalar_t>::FragB& frag_b,
172
- typename ScalarType<scalar_t>::FragS& frag_s,
254
+template <vllm::ScalarTypeId type_id>
255
+__device__ inline void scale(typename MarlinScalarType<type_id>::FragB& frag_b,
256
+ typename MarlinScalarType<type_id>::FragS& frag_s,
257
int i) {
174
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
175
- scalar_t2 s =
176
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
258
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
259
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
260
+ scalar_t2 s = MarlinScalarType<type_id>::num2num2(
261
+ reinterpret_cast<scalar_t*>(&frag_s)[i]);
262
frag_b[0] = __hmul2(frag_b[0], s);
263
frag_b[1] = __hmul2(frag_b[1], s);
264
}
265
181
-template <typename scalar_t>
266
+template <vllm::ScalarTypeId type_id>
267
__device__ inline void scale_and_sub(
183
- typename ScalarType<scalar_t>::FragB& frag_b, scalar_t s, scalar_t zp) {
184
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
185
- scalar_t2 s2 = ScalarType<scalar_t>::num2num2(s);
186
- scalar_t2 zp2 = ScalarType<scalar_t>::num2num2(zp);
268
+ typename MarlinScalarType<type_id>::FragB& frag_b,
269
+ typename MarlinScalarType<type_id>::scalar_t s,
270
+ typename MarlinScalarType<type_id>::scalar_t zp) {
271
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
272
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
273
+ scalar_t2 s2 = MarlinScalarType<type_id>::num2num2(s);
274
+ scalar_t2 zp2 = MarlinScalarType<type_id>::num2num2(zp);
275
frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
276
frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
277
}
278
191
-template <typename scalar_t>
192
-__device__ inline void sub_zp(typename ScalarType<scalar_t>::FragB& frag_b,
193
- typename ScalarType<scalar_t>::scalar_t2& frag_zp,
194
- int i) {
195
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
196
- scalar_t2 zp =
197
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
279
+template <vllm::ScalarTypeId type_id>
280
+__device__ inline void sub_zp(
281
+ typename MarlinScalarType<type_id>::FragB& frag_b,
282
+ typename MarlinScalarType<type_id>::scalar_t2& frag_zp, int i) {
283
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
284
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
285
+ scalar_t2 zp = MarlinScalarType<type_id>::num2num2(
286
+ reinterpret_cast<scalar_t*>(&frag_zp)[i]);
287
frag_b[0] = __hsub2(frag_b[0], zp);
288
frag_b[1] = __hsub2(frag_b[1], zp);
289
}
290
291
// Same as above, but for act_order (each K is multiplied individually)
203
-template <typename scalar_t>
204
-__device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
205
- typename ScalarType<scalar_t>::FragS& frag_s_1,
206
- typename ScalarType<scalar_t>::FragS& frag_s_2,
207
- typename ScalarType<scalar_t>::FragS& frag_s_3,
208
- typename ScalarType<scalar_t>::FragS& frag_s_4,
209
- int i) {
210
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
292
+template <vllm::ScalarTypeId type_id>
293
+__device__ inline void scale4(
294
+ typename MarlinScalarType<type_id>::FragB& frag_b,
295
+ typename MarlinScalarType<type_id>::FragS& frag_s_1,
296
+ typename MarlinScalarType<type_id>::FragS& frag_s_2,
297
+ typename MarlinScalarType<type_id>::FragS& frag_s_3,
298
+ typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
299
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
300
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
301
+
302
scalar_t2 s_val_1_2;
303
s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
304
s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
@@ -221,12 +312,13 @@ __device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
312
}
313
314
// Given 2 floats multiply by 2 scales (halves)
224
-template <typename scalar_t>
225
-__device__ inline void scale_float(float* c,
226
- typename ScalarType<scalar_t>::FragS& s) {
315
+template <vllm::ScalarTypeId type_id>
316
+__device__ inline void scale_float(
317
+ float* c, typename MarlinScalarType<type_id>::FragS& s) {
318
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
319
scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
228
- c[0] = __fmul_rn(c[0], ScalarType<scalar_t>::num2float(s_ptr[0]));
229
- c[1] = __fmul_rn(c[1], ScalarType<scalar_t>::num2float(s_ptr[1]));
320
+ c[0] = __fmul_rn(c[0], MarlinScalarType<type_id>::num2float(s_ptr[0]));
321
+ c[1] = __fmul_rn(c[1], MarlinScalarType<type_id>::num2float(s_ptr[1]));
322
}
323
324
// Wait until barrier reaches `count`, then lock for current threadblock.
@@ -278,9 +370,10 @@ __device__ inline void wait_negative_and_add(int* lock) {
370
__syncthreads();
371
}
372
281
-template <typename scalar_t, // compute dtype, half or nv_float16
282
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
283
- const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
373
+template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
374
+ const vllm::ScalarTypeId b_type_id, // B ScalarType id
375
+ const vllm::ScalarTypeId c_type_id, // C ScalarType id
376
+ const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id
377
const int threads, // number of threads in a threadblock
378
const int thread_m_blocks, // number of 16x16 blocks in the m
379
// dimension (batchsize) of the
@@ -301,13 +394,18 @@ __global__ void Marlin(
394
int4* __restrict__ C, // fp16 output buffer of shape mxn
395
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
396
const int4* __restrict__ b_bias_ptr,
304
- const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
305
- // (k/groupsize)xn
306
- const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
307
- // only)
308
- const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
309
- // (k/groupsize)x(n/pack_factor)
310
- const int* __restrict__ g_idx, // int32 group indices of shape k
397
+ // float scales of input matrix, only used when is_a_8bit == true.
398
+ // shape (m,)
399
+ const float* __restrict__ a_scales_ptr,
400
+ // fp16 quantization scales. shape (k/groupsize, n)
401
+ const int4* __restrict__ scales_ptr,
402
+ // fp16 global scale (for nvfp4// only)
403
+ const uint16_t* __restrict__ global_scale_ptr,
404
+ // 4bit packed zero-points of shape
405
+ // (k/groupsize, n/pack_factor)
406
+ const int4* __restrict__ zp_ptr,
407
+ // int32 group indices of shape k
408
+ const int* __restrict__ g_idx,
409
const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
410
const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
411
const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
@@ -321,9 +419,9 @@ __global__ void Marlin(
419
int prob_k, // reduction dimension k
420
int* locks, // extra global storage for barrier synchronization
421
bool has_bias,
324
- bool use_atomic_add, // whether to use atomic add to reduce
325
- bool use_fp32_reduce, // whether to use fp32 global reduce
326
- int max_shared_mem) {
422
+ bool use_atomic_add, // whether to use atomic add to reduce
423
+ bool use_fp32_reduce // whether to use fp32 global reduce
424
+) {
425
// Each threadblock processes one "stripe" of the B matrix with (roughly) the
426
// same size, which might involve multiple column "slices" (of width 16 *
427
// `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
@@ -335,18 +433,37 @@ __global__ void Marlin(
433
// ensures good utilization of all SMs for many kinds of shape and GPU
434
// configurations, while requiring as few slow global cross-threadblock
435
// reductions as possible.
338
- using Dtype = ScalarType<scalar_t>;
339
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
340
- using FragA = typename ScalarType<scalar_t>::FragA;
341
- using FragB = typename ScalarType<scalar_t>::FragB;
342
- using FragC = typename ScalarType<scalar_t>::FragC;
343
- using FragS = typename ScalarType<scalar_t>::FragS;
344
- using FragZP = typename ScalarType<scalar_t>::FragZP;
436
+
437
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890
438
+ // FP8 computation is only supported for Ada Lovelace or newer architectures.
439
+ if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
440
+ #endif
441
+
442
+ int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
443
+ constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
444
+
445
+ using Adtype = MarlinScalarType<a_type_id>;
446
+ using Cdtype = MarlinScalarType<c_type_id>;
447
+
448
+ using scalar_t = typename MarlinScalarType<a_type_id>::scalar_t;
449
+ using scalar_t2 = typename MarlinScalarType<a_type_id>::scalar_t2;
450
+ using scalar_32bit_t = typename MarlinScalarType<a_type_id>::scalar_32bit_t;
451
+
452
+ using c_scalar_t = typename MarlinScalarType<c_type_id>::scalar_t;
453
+ using c_scalar_t2 = typename MarlinScalarType<c_type_id>::scalar_t2;
454
+
455
+ using FragA = typename MarlinScalarType<a_type_id>::FragA;
456
+ using FragB = typename MarlinScalarType<a_type_id>::FragB;
457
+ using FragC = typename MarlinScalarType<a_type_id>::FragC;
458
+ using FragS = typename MarlinScalarType<c_type_id>::FragS;
459
+ using FragZP = typename MarlinScalarType<c_type_id>::FragZP;
460
461
extern __shared__ int4 sh[];
347
- static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
462
+ static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id);
463
+ static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id);
464
+ static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id);
465
static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
349
- if constexpr (w_type == vllm::kFE2M1f) {
466
+ if constexpr (b_type == vllm::kFE2M1f) {
467
static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
468
s_type == vllm::kFE8M0fnu && group_blocks == 2);
469
} else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
@@ -355,34 +472,37 @@ __global__ void Marlin(
472
static_assert(s_type == vllm::kFloat16);
473
}
474
358
- constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
359
- constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
360
- w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
475
+ constexpr bool is_a_8bit = a_type.size_bits() == 8;
476
+ if constexpr (!is_a_8bit) {
477
+ static_assert(std::is_same<scalar_t, c_scalar_t>::value);
478
+ }
479
+ constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8;
480
+ constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 ||
481
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
482
+ b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
483
// see comments of dequant.h for more details
484
constexpr bool dequant_skip_flop =
363
- w_type == vllm::kFE4M3fn ||
364
- w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
485
+ is_a_8bit || b_type == vllm::kFE4M3fn ||
486
+ b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
487
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
366
- has_zp && !is_zp_float && !(w_type == vllm::kU8);
488
+ has_zp && !is_zp_float && !(b_type == vllm::kU8);
489
368
- scalar_t2 global_scale;
490
+ c_scalar_t2 global_scale;
491
492
constexpr bool has_act_order = group_blocks == 0;
493
372
- constexpr int pack_factor = 32 / w_type.size_bits();
494
+ constexpr int pack_factor = 32 / b_type.size_bits();
495
static_assert(thread_m_blocks == 1 || !m_block_size_8);
374
- constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
496
const int group_size =
497
(!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
498
const int scales_expert_stride =
378
- prob_n * prob_k / group_size / (w_type == vllm::kFE2M1f ? 16 : 8);
499
+ prob_n * prob_k / group_size / (b_type == vllm::kFE2M1f ? 16 : 8);
500
const int zp_expert_stride =
501
is_zp_float ? prob_n * prob_k / group_size / 8
502
: prob_n * prob_k / group_size / (pack_factor * 4);
503
const int b_bias_expert_stride = prob_n / 8;
504
505
// parallel: num valid moe blocks
385
- int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
506
int parallel = num_tokens_past_padded / moe_block_size;
507
int num_valid_blocks = parallel;
508
if (is_ep) {
@@ -395,7 +515,23 @@ __global__ void Marlin(
515
516
int k_tiles = prob_k / 16 / thread_k_blocks;
517
int n_tiles = prob_n / 16 / thread_n_blocks;
398
- int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x);
518
+
519
+ int global_mn_tiles = parallel * n_tiles;
520
+ int part2_mn_tiles = global_mn_tiles;
521
+ int part1_mn_iters = 0;
522
+ bool in_part2 = false;
523
+
524
+ // we use DP + two-tile SK here
525
+ // part1: DP
526
+ // part2: two-tile SK
527
+ // see https://github.com/vllm-project/vllm/pull/24722 for more details
528
+ if (global_mn_tiles > gridDim.x) {
529
+ part2_mn_tiles = global_mn_tiles % gridDim.x;
530
+ if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;
531
+ part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x;
532
+ }
533
+
534
+ int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
535
536
if constexpr (!has_act_order && group_blocks != -1) {
537
if (group_blocks >= thread_k_blocks) {
@@ -407,14 +543,15 @@ __global__ void Marlin(
543
}
544
}
545
410
- int slice_row = (iters * blockIdx.x) % k_tiles;
411
- int slice_col_par = (iters * blockIdx.x) / k_tiles;
412
- int slice_col = slice_col_par;
413
- int slice_iters; // number of threadblock tiles in the current slice
414
- int slice_count =
415
- 0; // total number of active threadblocks in the current slice
416
- int slice_idx; // index of threadblock in current slice; numbered bottom to
417
- // top
546
+ int slice_row = 0;
547
+ int slice_col_par = blockIdx.x;
548
+ int slice_col;
549
+ int slice_iters =
550
+ k_tiles; // number of threadblock tiles in the current slice
551
+ // total number of active threadblocks in the current slice
552
+ int slice_count = 1;
553
+ // index of threadblock in current slice; numbered bottom to top
554
+ int slice_idx = 0;
555
556
int par_id = 0;
557
int block_id = -1;
@@ -422,87 +559,89 @@ __global__ void Marlin(
559
int old_expert_id = 0;
560
int64_t B_expert_off = 0;
561
425
- int4* sh_block_sorted_ids_int4 = sh;
562
+ float* sh_a_s = reinterpret_cast<float*>(sh);
563
+ int4* sh_block_sorted_ids_int4 = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0);
564
int4* sh_rd_block_sorted_ids_int4 =
565
sh_block_sorted_ids_int4 + moe_block_size / 4;
566
int4* sh_block_topk_weights_int4 =
567
sh_rd_block_sorted_ids_int4 + moe_block_size / 4;
568
// sh_block_topk_weights_int4 only need (moe_block_size / 4);
569
// but we pad to align to 256 bytes
432
- int4* sh_new =
433
- sh_block_topk_weights_int4 + moe_block_size / 2 + moe_block_size;
570
+ int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 2;
571
int32_t* sh_block_sorted_ids =
572
reinterpret_cast<int*>(sh_block_sorted_ids_int4);
573
int32_t* sh_rd_block_sorted_ids =
574
reinterpret_cast<int*>(sh_rd_block_sorted_ids_int4);
438
- scalar_t2* sh_block_topk_weights =
439
- reinterpret_cast<scalar_t2*>(sh_block_topk_weights_int4);
575
+ c_scalar_t2* sh_block_topk_weights =
576
+ reinterpret_cast<c_scalar_t2*>(sh_block_topk_weights_int4);
577
578
int32_t block_num_valid_tokens = 0;
579
int32_t locks_off = 0;
580
581
// We can easily implement parallel problem execution by just remapping
582
// indices and advancing global pointers
446
- if (slice_col_par >= n_tiles) {
447
- slice_col = slice_col_par % n_tiles;
448
- par_id = slice_col_par / n_tiles;
449
- }
450
- if (parallel * n_tiles >= gridDim.x) {
451
- // when parallel * n_tiles >= sms
583
+ if (part2_mn_tiles >= gridDim.x) {
584
+ // when part2_mn_tiles >= sms
585
// then there are at most $sms$ conflict tile blocks
586
locks_off = blockIdx.x;
587
} else {
588
locks_off = (iters * blockIdx.x) / k_tiles - 1;
589
}
590
591
+ int prob_m_top_k = prob_m * top_k;
592
// read moe block data given block_id
593
// block_sorted_ids / block_num_valid_tokens / block_topk_weights
594
auto read_moe_block_data = [&](int block_id) {
595
block_num_valid_tokens = moe_block_size;
596
+
597
+ cp_async4_pred(sh_block_sorted_ids_int4 + threadIdx.x,
598
+ reinterpret_cast<const int4*>(sorted_token_ids_ptr) +
599
+ (block_id * moe_block_size / 4 + threadIdx.x),
600
+ threadIdx.x < moe_block_size / 4);
601
+
602
+ cp_async_fence();
603
+ cp_async_wait<0>();
604
+
605
+ __syncthreads();
606
+
607
+ if (threadIdx.x >= threads - 32) {
608
+ constexpr int size_per_thread = div_ceil(moe_block_size, 32);
609
+ int lane_id = threadIdx.x - (threads - 32);
610
+
611
+ int local_count = 0;
612
#pragma unroll
463
- for (int i = 0; i < moe_block_size / 4; i++) {
464
- int4 sorted_token_ids_int4 = reinterpret_cast<const int4*>(
465
- sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
466
- int* sorted_token_ids = reinterpret_cast<int*>(&sorted_token_ids_int4);
467
- #pragma unroll
468
- for (int j = 0; j < 4; j++) {
469
- if (sorted_token_ids[j] >= prob_m * top_k) {
470
- block_num_valid_tokens = i * 4 + j;
471
- break;
613
+ for (int i = 0; i < size_per_thread; i++) {
614
+ int j = lane_id * size_per_thread + i;
615
+ if (j < moe_block_size) {
616
+ int idx = sh_block_sorted_ids[j];
617
+ if (idx < prob_m_top_k) local_count++;
618
}
619
}
474
- if (block_num_valid_tokens != moe_block_size) break;
475
- }
620
477
- __syncthreads();
478
- int tid4 = threadIdx.x / 4;
479
- if (threadIdx.x % 4 == 0 && threadIdx.x < block_num_valid_tokens) {
480
- sh_block_sorted_ids_int4[tid4] = reinterpret_cast<const int4*>(
481
- sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4];
621
+ block_num_valid_tokens = __reduce_add_sync(0xffffffff, local_count);
622
483
- #pragma unroll
484
- for (int i = 0; i < 4; i++)
485
- sh_rd_block_sorted_ids[tid4 * 4 + i] =
486
- sh_block_sorted_ids[tid4 * 4 + i] / top_k;
623
+ if (lane_id == 0)
624
+ reinterpret_cast<int*>(sh_new)[0] = block_num_valid_tokens;
625
+ }
626
+
627
+ if (threadIdx.x < moe_block_size) {
628
+ int idx = sh_block_sorted_ids[threadIdx.x];
629
+ sh_rd_block_sorted_ids[threadIdx.x] = idx / top_k;
630
631
if (mul_topk_weights) {
489
- #pragma unroll
490
- for (int i = 0; i < 4; i++) {
491
- int idx = tid4 * 4 + i;
492
- if (idx < block_num_valid_tokens) {
493
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
494
- sh_block_topk_weights[idx] =
495
- __hmul2(global_scale,
496
- Dtype::num2num2(Dtype::float2num(
497
- topk_weights_ptr[sh_block_sorted_ids[idx]])));
498
- } else {
499
- sh_block_topk_weights[idx] = Dtype::num2num2(
500
- Dtype::float2num(topk_weights_ptr[sh_block_sorted_ids[idx]]));
501
- }
502
- }
632
+ idx = idx < prob_m_top_k ? idx : 0;
633
+ c_scalar_t2 topk_weight_val =
634
+ Cdtype::num2num2(Cdtype::float2num(topk_weights_ptr[idx]));
635
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
636
+ topk_weight_val = __hmul2(topk_weight_val, global_scale);
637
}
638
+ sh_block_topk_weights[threadIdx.x] = topk_weight_val;
639
}
640
}
641
+
642
+ __syncthreads();
643
+
644
+ block_num_valid_tokens = reinterpret_cast<int*>(sh_new)[0];
645
__syncthreads();
646
};
647
@@ -513,9 +652,8 @@ __global__ void Marlin(
652
653
old_expert_id = expert_id;
654
if (num_invalid_blocks > 0) {
516
- int skip_count = block_id == -1 ? par_id : 0;
517
- block_id++;
518
- for (int i = block_id; i < num_tokens_past_padded / moe_block_size; i++) {
655
+ int skip_count = par_id;
656
+ for (int i = 0; i < num_tokens_past_padded / moe_block_size; i++) {
657
expert_id = expert_ids_ptr[i];
658
if (expert_id != -1) {
659
if (skip_count == 0) {
@@ -530,9 +668,9 @@ __global__ void Marlin(
668
expert_id = expert_ids_ptr[block_id];
669
}
670
533
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
534
- uint16_t val = scale2_ptr[expert_id];
535
- global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
671
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
672
+ uint16_t val = global_scale_ptr[expert_id];
673
+ global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
674
}
675
676
B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
@@ -552,10 +690,11 @@ __global__ void Marlin(
690
691
// Compute all information about the current slice which is required for
692
// synchronization.
555
- auto init_slice = [&](bool first_init = false) {
693
+ bool first_init = true;
694
+ auto init_part2_slice = [&]() {
695
slice_iters =
696
iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
558
- if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
697
+ if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0;
698
if (slice_iters == 0) return;
699
if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
700
slice_count = 1;
@@ -573,7 +712,7 @@ __global__ void Marlin(
712
if (col_off > 0) slice_idx--;
713
}
714
}
576
- if (parallel * n_tiles >= gridDim.x) {
715
+ if (part2_mn_tiles >= gridDim.x) {
716
if (slice_count > 1 && slice_idx == slice_count - 1) {
717
locks_off++;
718
}
@@ -607,25 +746,61 @@ __global__ void Marlin(
746
par_id++;
747
update_next_moe_block_data();
748
}
749
+ if (is_a_8bit && (first_init || slice_col == 0)) {
750
+ __syncthreads();
751
+ cp_async1_ca_pred(&sh_a_s[threadIdx.x],
752
+ &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
753
+ threadIdx.x < block_num_valid_tokens);
754
+ }
755
};
756
612
- update_next_moe_block_data();
613
- init_slice(true);
757
+ auto init_part1_slice = [&]() {
758
+ if (part1_mn_iters) {
759
+ part1_mn_iters--;
760
+ par_id = slice_col_par / n_tiles;
761
+ slice_col = slice_col_par % n_tiles;
762
+ slice_iters = k_tiles;
763
+ update_next_moe_block_data();
764
+ if (is_a_8bit) {
765
+ __syncthreads();
766
+ cp_async1_ca_pred(&sh_a_s[threadIdx.x],
767
+ &a_scales_ptr[sh_rd_block_sorted_ids[threadIdx.x]],
768
+ threadIdx.x < block_num_valid_tokens);
769
+ }
770
+ }
771
+ };
772
+
773
+ auto init_slice = [&]() {
774
+ if (!in_part2 && !part1_mn_iters) {
775
+ in_part2 = true;
776
+ slice_col_par = (iters * blockIdx.x) / k_tiles;
777
+ slice_row = (iters * blockIdx.x) % k_tiles;
778
+ slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles;
779
+ par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles;
780
+ update_next_moe_block_data();
781
+ }
782
+ if (!in_part2) {
783
+ init_part1_slice();
784
+ } else {
785
+ init_part2_slice();
786
+ first_init = false;
787
+ }
788
+ };
789
+
790
+ init_slice();
791
792
// A sizes/strides
793
794
// stride of the A matrix in global memory
618
- int a_gl_stride = prob_k / 8;
795
+ int a_gl_stride = prob_k / (is_a_8bit ? 16 : 8);
796
// stride of an A matrix tile in shared memory
620
- constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
797
+ constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
798
// delta between subsequent A tiles in global memory
622
- constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;
799
+ constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
800
// between subsequent accesses within a tile
801
int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
802
// between shared memory writes
803
constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
627
- // between shared memory tile reads
628
- constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));
804
// within a shared memory tile
805
constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
806
// overall size of a tile
@@ -634,24 +809,25 @@ __global__ void Marlin(
809
constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
810
811
// B sizes/strides
637
- int b_gl_stride = 16 * prob_n / (pack_factor * 4);
638
- constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
639
- constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;
812
+ int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4));
813
+ constexpr int b_sh_stride =
814
+ ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4);
815
+ constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2;
816
constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
817
642
- int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
643
- int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
818
+ int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
819
constexpr int b_sh_wr_delta = threads * b_thread_vecs;
645
- constexpr int b_sh_rd_delta = threads * b_thread_vecs;
646
- constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
820
+ constexpr int b_sh_stage =
821
+ b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
822
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
823
824
// Scale sizes/strides without act_order
650
- int s_gl_stride = prob_n / 8;
651
- constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
825
+ int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8);
826
+ constexpr int s_sh_stride =
827
+ 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8);
828
constexpr int s_tb_groups =
829
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
654
- ? thread_k_blocks / group_blocks / (w_type == vllm::kFE2M1f ? 2 : 1)
830
+ ? thread_k_blocks / group_blocks
831
: 1;
832
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
833
int s_gl_rd_delta = s_gl_stride;
@@ -664,7 +840,8 @@ __global__ void Marlin(
840
constexpr int act_s_max_num_groups = 32;
841
int act_s_col_stride = 1;
842
int act_s_col_warp_stride = act_s_col_stride * 8;
667
- int tb_n_warps = thread_n_blocks / 4;
843
+
844
+ constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
845
int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
846
847
// Zero-points sizes/strides
@@ -679,7 +856,6 @@ __global__ void Marlin(
856
// Global A read index of current thread.
857
int a_gl_rd_row = threadIdx.x / a_gl_rd_delta_o;
858
int a_gl_rd_col = a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
682
-
859
// Shared write index of current thread.
860
int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
861
(threadIdx.x % a_gl_rd_delta_o);
@@ -687,17 +863,22 @@ __global__ void Marlin(
863
int a_sh_rd =
864
a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
865
(threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
690
- a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
866
+ a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters;
867
+
868
+ int b_gl_rd;
869
+ if (threads <= b_sh_stride) {
870
+ b_gl_rd = threadIdx.x;
871
+ } else {
872
+ b_gl_rd =
873
+ b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
874
+ }
875
692
- int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
693
- (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
694
- b_gl_rd += b_sh_stride * slice_col;
876
+ b_gl_rd += B_expert_off + b_sh_stride * slice_col;
877
b_gl_rd += b_gl_rd_delta_o * slice_row;
696
- auto b_sh_wr = threadIdx.x * b_thread_vecs;
878
auto b_sh_rd = threadIdx.x * b_thread_vecs;
879
+ b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
880
881
// For act_order
700
- constexpr int k_iter_size = tb_k / b_sh_wr_iters;
882
int slice_k_start = tb_k * slice_row;
883
int slice_k_finish = slice_k_start + tb_k * slice_iters;
884
int slice_k_start_shared_fetch = slice_k_start;
@@ -708,58 +889,54 @@ __global__ void Marlin(
889
if constexpr (!has_act_order) {
890
if constexpr (group_blocks == -1) {
891
s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
711
- } else {
712
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) /
713
- (w_type == vllm::kFE2M1f ? 2 : 1) +
892
+ } else if constexpr (group_blocks >= thread_k_blocks) {
893
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
894
s_sh_stride * slice_col + threadIdx.x;
895
+ } else {
896
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
897
+ threadIdx.x / s_sh_stride) +
898
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
899
}
900
}
901
auto s_sh_wr = threadIdx.x;
718
- bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
902
+ bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
903
904
// Zero-points
905
int zp_gl_rd;
906
if constexpr (has_zp) {
907
if constexpr (group_blocks == -1) {
908
zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
725
- } else {
909
+ } else if constexpr (group_blocks >= thread_k_blocks) {
910
zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
911
zp_sh_stride * slice_col + threadIdx.x;
912
+ } else {
913
+ zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
914
+ threadIdx.x / zp_sh_stride) +
915
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
916
}
917
}
918
auto zp_sh_wr = threadIdx.x;
731
- bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
919
+ bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage;
920
921
// We use a different scale layout for grouped and column-wise quantization as
922
// we scale a `half2` tile in column-major layout in the former and in
923
// row-major in the latter case.
924
int s_sh_rd;
737
- if constexpr (group_blocks != -1 && w_type == vllm::kFE2M1f) {
738
- auto warp_id = threadIdx.x / 32;
739
- int n_warps = thread_n_blocks / 4;
740
- int warp_row = warp_id / n_warps;
741
-
742
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
743
- (threadIdx.x % 32) / 4;
744
- s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2;
745
-
925
+ if constexpr (is_a_8bit) {
926
+ s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4);
927
} else if constexpr (group_blocks != -1)
747
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
748
- (threadIdx.x % 32) / 4;
928
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
929
else if constexpr (group_blocks == -1 &&
930
(m_block_size_8 || (has_zp && !dequant_skip_flop)))
751
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
752
- (threadIdx.x % 32) / 8;
931
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
932
else
754
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
755
- (threadIdx.x % 32) % 4;
933
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4;
934
935
int bias_sh_rd;
936
if constexpr (m_block_size_8) {
759
- bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
760
- (threadIdx.x % 32) / 8;
937
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
938
} else {
762
- bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
939
+ bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) +
940
(threadIdx.x % 32) % 4;
941
}
942
@@ -775,12 +952,16 @@ __global__ void Marlin(
952
if constexpr (has_zp) {
953
if constexpr (is_zp_float) {
954
if constexpr (group_blocks != -1) {
778
- zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
779
- (threadIdx.x % 32) / 4;
955
+ zp_sh_rd =
956
+ 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
957
}
958
+ } else if (is_a_8bit) {
959
+ zp_sh_rd = num_ints_per_thread * num_col_threads *
960
+ ((threadIdx.x / 32) % tb_n_warps / 2) +
961
+ num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
962
} else {
963
zp_sh_rd = num_ints_per_thread * num_col_threads *
783
- ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
964
+ ((threadIdx.x / 32) % tb_n_warps) +
965
num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
966
}
967
}
@@ -807,18 +988,13 @@ __global__ void Marlin(
988
for (int i = 0; i < b_sh_wr_iters; i++) {
989
#pragma unroll
990
for (int j = 0; j < thread_m_blocks; j++)
810
- a_sh_rd_trans[i][j] =
811
- transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
991
+ a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd);
992
}
993
994
// Since B-accesses have non-constant stride they have to be computed at
995
// runtime; we break dependencies between subsequent accesses with a tile by
996
// maintining multiple pointers (we have enough registers), a tiny
997
// optimization.
818
- const int4* B_ptr[b_sh_wr_iters];
819
- #pragma unroll
820
- for (int i = 0; i < b_sh_wr_iters; i++)
821
- B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
998
999
// Shared memory storage for global fetch pipelines.
1000
constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
@@ -847,19 +1023,12 @@ __global__ void Marlin(
1023
static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
1024
stages * b_sh_stage);
1025
int4* sh_a = sh_s + sh_s_size;
850
- constexpr int shm_size_used = moe_block_size +
851
- stages * (g_idx_stage + zp_sh_stage) +
852
- sh_s_size + sh_b_red_bias_size;
853
-
854
- // all remaining shared memory is used to cache A (input)
855
- // sh_a_max_row is at least ` stages * 16 * thread_m_blocks `
856
- int sh_a_max_row =
857
- ((max_shared_mem - 1024) / 16 - shm_size_used) / (thread_k_blocks * 2);
1026
1027
// Register storage for double buffer of shared memory reads.
1028
FragA frag_a[2][thread_m_blocks];
1029
I4 frag_b_quant[2][b_thread_vecs];
862
- FragC frag_c[thread_m_blocks][4][2];
1030
+ FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
1031
+ FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
1032
FragS frag_s[2][4]; // No act-order
1033
FragS frag_bias[2][4];
1034
FragS act_frag_s[2][4][4]; // For act-order
@@ -867,6 +1036,24 @@ __global__ void Marlin(
1036
FragZP frag_zp; // Zero-points in fp16
1037
FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
1038
1039
+ if constexpr (is_a_8bit && group_blocks != -1) {
1040
+ #pragma unroll
1041
+ for (int j = 0; j < 2; j++) {
1042
+ #pragma unroll
1043
+ for (int i = 0; i < thread_m_blocks; i++) {
1044
+ #pragma unroll
1045
+ for (int g = 0; g < 4; g++) {
1046
+ frag_c_tmp[i][j][0][g] = 0.0f;
1047
+ }
1048
+
1049
+ #pragma unroll
1050
+ for (int g = 0; g < 4; g++) {
1051
+ frag_c_tmp[i][j][1][g] = 0.0f;
1052
+ }
1053
+ }
1054
+ }
1055
+ }
1056
+
1057
// Zero accumulators.
1058
auto zero_accums = [&]() {
1059
#pragma unroll
@@ -910,43 +1097,36 @@ __global__ void Marlin(
1097
}
1098
}
1099
};
913
-
1100
// Asynchronously fetch the next A, B and s tile from global to the next
1101
// shared memory pipeline location.
916
- bool should_load_a = true;
917
- int max_num_stage_groups =
918
- ((sh_a_max_row - moe_block_size) / moe_block_size + 1) / stages;
919
- max_num_stage_groups = max(max_num_stage_groups, 1);
920
- auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true,
921
- int pipe_a = 0) {
1102
+ auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
1103
if (pred) {
923
- if (should_load_a) {
924
- int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
1104
+ int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
1105
#pragma unroll
926
- for (int i = 0; i < a_sh_wr_iters; i++) {
927
- int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
928
- int64_t sorted_row = 0;
929
- if (!m_block_size_8 || row < 8)
930
- sorted_row = sh_rd_block_sorted_ids[row];
931
- int64_t true_idx =
932
- sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
933
- cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
934
- row < block_num_valid_tokens);
935
- }
1106
+ for (int i = 0; i < a_sh_wr_iters; i++) {
1107
+ int row = a_gl_rd_delta_i / a_gl_stride * i + a_gl_rd_row;
1108
+ int64_t sorted_row = 0;
1109
+ if (!m_block_size_8 || row < 8)
1110
+ sorted_row = sh_rd_block_sorted_ids[row];
1111
+ int64_t true_idx =
1112
+ sorted_row * a_gl_stride + a_gl_rd_col + a_gl_rd_delta_o * a_off;
1113
+ cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
1114
+ row < block_num_valid_tokens);
1115
}
1116
1117
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
1118
#pragma unroll
940
- for (int i = 0; i < b_sh_wr_iters; i++) {
941
- #pragma unroll
942
- for (int j = 0; j < b_thread_vecs; j++) {
943
- cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j],
944
- B_ptr[i] + j + B_expert_off);
945
- }
1119
+ for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) {
1120
+ constexpr int count = div_ceil(b_sh_stride, threads);
1121
+ int b_gl_idx =
1122
+ b_gl_rd + (i % count) * threads +
1123
+ b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride);
1124
947
- B_ptr[i] += b_gl_rd_delta_o;
1125
+ cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]);
1126
}
1127
1128
+ b_gl_rd += b_gl_rd_delta_o;
1129
+
1130
if constexpr (has_act_order) {
1131
// Fetch g_idx thread-block portion
1132
int full_pipe = a_off;
@@ -966,44 +1146,24 @@ __global__ void Marlin(
1146
if constexpr (group_blocks != -1) {
1147
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1148
969
- if constexpr (group_blocks >= thread_k_blocks) {
970
- // Only fetch scales if this tile starts a new group
971
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
972
- if (s_sh_wr_pred) {
973
- cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
974
- }
975
- s_gl_rd += s_gl_rd_delta;
976
- }
977
- } else {
978
- for (int i = 0; i < s_tb_groups; i++) {
979
- if (s_sh_wr_pred) {
980
- cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
981
- &scales_ptr[s_gl_rd]);
982
- }
983
- s_gl_rd += s_gl_rd_delta;
1149
+ // Only fetch scales if this tile starts a new group
1150
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1151
+ if (s_sh_wr_pred) {
1152
+ cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
1153
}
1154
+ s_gl_rd += s_gl_rd_delta * s_tb_groups;
1155
}
1156
}
1157
1158
if constexpr (has_zp && group_blocks != -1) {
1159
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1160
991
- if constexpr (group_blocks >= thread_k_blocks) {
992
- // Only fetch zero-points if this tile starts a new group
993
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
994
- if (zp_sh_wr_pred) {
995
- cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
996
- }
997
- zp_gl_rd += zp_gl_rd_delta;
998
- }
999
- } else {
1000
- for (int i = 0; i < zp_tb_groups; i++) {
1001
- if (zp_sh_wr_pred) {
1002
- cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
1003
- &zp_ptr[zp_gl_rd]);
1004
- }
1005
- zp_gl_rd += zp_gl_rd_delta;
1161
+ // Only fetch zero points if this tile starts a new group
1162
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1163
+ if (zp_sh_wr_pred) {
1164
+ cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
1165
}
1166
+ zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
1167
}
1168
}
1169
}
@@ -1037,18 +1197,18 @@ __global__ void Marlin(
1197
1198
// Load the next sub-tile from the current location in the shared memory pipe
1199
// into the current register buffer.
1040
- auto fetch_to_registers = [&](int k, int pipe, int pipe_a = 0) {
1041
- int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe_a;
1200
+ auto fetch_to_registers = [&](int k, int pipe) {
1201
+ int4* sh_a_stage = sh_a + moe_block_size * a_sh_stride * pipe;
1202
#pragma unroll
1203
for (int i = 0; i < thread_m_blocks; i++)
1044
- ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
1204
+ ldsm<m_block_size_8 ? 2 : 4, a_type_id>(
1205
frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
1206
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
1207
1208
#pragma unroll
1209
for (int i = 0; i < b_thread_vecs; i++) {
1210
frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
1051
- &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
1211
+ &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]);
1212
}
1213
};
1214
@@ -1072,53 +1232,54 @@ __global__ void Marlin(
1232
1233
auto fetch_scales_to_registers = [&](int k, int full_pipe) {
1234
int pipe = full_pipe % stages;
1235
+ using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
1236
+ using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
1237
+ constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
1238
1239
if constexpr (!has_act_order) {
1240
// No act-order case
1241
if constexpr (group_blocks == -1) {
1242
// load only when starting a new slice
1080
- if (k == 0 && full_pipe == 0) {
1243
+ if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
1244
reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
1245
reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1246
}
1247
} else if constexpr (group_blocks != -1) {
1248
if constexpr (group_blocks >= thread_k_blocks) {
1086
- if (k % b_sh_wr_iters == 0) {
1087
- int4* sh_s_stage =
1088
- sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
1089
- (pipe / (group_blocks / thread_k_blocks)));
1090
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1091
- } else {
1092
- reinterpret_cast<int4*>(&frag_s[1])[0] =
1093
- reinterpret_cast<int4*>(&frag_s[0])[0];
1249
+ constexpr int g = group_blocks / thread_k_blocks;
1250
+ if (pipe % g == 0) {
1251
+ if (k % b_sh_wr_iters == 0) {
1252
+ int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
1253
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1254
+ } else {
1255
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1256
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1257
+ }
1258
}
1095
- } else {
1259
+ } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1260
auto warp_id = threadIdx.x / 32;
1097
- int n_warps = thread_n_blocks / 4;
1098
-
1099
- int warp_row = warp_id / n_warps;
1261
+ int warp_row = warp_id / tb_n_warps;
1262
1101
- int cur_k = warp_row * 16;
1102
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1103
-
1104
- int k_blocks = cur_k / 16;
1105
- int cur_group_id =
1106
- k_blocks / (group_blocks * (w_type == vllm::kFE2M1f ? 2 : 1));
1263
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1264
+ int cur_group_id = k_blocks / group_blocks2;
1265
1266
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1267
1110
- if constexpr (w_type_id != vllm::kFE2M1f.id()) {
1268
+ if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1269
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1270
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1113
- } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) {
1271
+ } else {
1272
reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1273
reinterpret_cast<int2*>(
1274
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1275
+ }
1276
+ } else if (group_blocks >= b_sh_wr_iters) {
1277
+ if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1278
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1279
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1280
} else {
1118
- reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1119
- reinterpret_cast<int2*>(
1120
- sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) +
1121
- k % 2];
1281
+ reinterpret_cast<int2*>(&frag_s[1])[0] =
1282
+ reinterpret_cast<int2*>(&frag_s[0])[0];
1283
}
1284
}
1285
}
@@ -1139,18 +1300,15 @@ __global__ void Marlin(
1300
cur_k = 0;
1301
1302
// Progress to current iteration
1142
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1303
+ cur_k += k % b_sh_wr_iters;
1304
1305
// Determine "position" inside the thread-block (based on warp and
1306
// thread-id)
1307
auto warp_id = threadIdx.x / 32;
1147
- int n_warps =
1148
- thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
1149
-
1150
- int warp_row = warp_id / n_warps;
1151
- int warp_col = warp_id % n_warps;
1308
+ int warp_row = warp_id / tb_n_warps;
1309
+ int warp_col = warp_id % tb_n_warps;
1310
1153
- cur_k += warp_row * 16;
1311
+ cur_k += warp_row * 16 * b_sh_wr_iters;
1312
1313
auto th_id = threadIdx.x % 32;
1314
cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
@@ -1205,18 +1363,16 @@ __global__ void Marlin(
1363
1364
if constexpr (group_blocks == -1) {
1365
// load only when starting a new slice
1208
- if (k == 0 && full_pipe == 0) {
1366
+ if (k == 0 && full_pipe == 0 || is_a_8bit) {
1367
#pragma unroll
1368
for (int i = 0; i < num_ints_per_thread; i++) {
1369
frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
1370
}
1371
}
1214
-
1372
} else if constexpr (group_blocks >= thread_k_blocks) {
1216
- if (k % b_sh_wr_iters == 0) {
1217
- int4* sh_zp_stage =
1218
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1219
- (pipe / (group_blocks / thread_k_blocks)));
1373
+ constexpr int g = group_blocks / thread_k_blocks;
1374
+ if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) {
1375
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
1376
#pragma unroll
1377
for (int i = 0; i < num_ints_per_thread; i++) {
1378
frag_qzp[k % 2][i] =
@@ -1225,21 +1381,11 @@ __global__ void Marlin(
1381
}
1382
} else {
1383
auto warp_id = threadIdx.x / 32;
1228
- int n_warps = thread_n_blocks / 4;
1384
1230
- int warp_row = warp_id / n_warps;
1385
+ int warp_row = warp_id / tb_n_warps;
1386
1232
- int cur_k = warp_row * 16;
1233
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1234
-
1235
- int k_blocks = cur_k / 16;
1236
- int cur_group_id = 0;
1237
-
1238
- // Suppress bogus and persistent divide-by-zero warning
1239
- #pragma nv_diagnostic push
1240
- #pragma nv_diag_suppress divide_by_zero
1241
- cur_group_id = k_blocks / group_blocks;
1242
- #pragma nv_diagnostic pop
1387
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1388
+ int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1);
1389
1390
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1391
@@ -1258,29 +1404,18 @@ __global__ void Marlin(
1404
1405
if constexpr (group_blocks != -1) {
1406
if constexpr (group_blocks >= thread_k_blocks) {
1261
- if (k % b_sh_wr_iters == 0) {
1262
- int4* sh_zp_stage =
1263
- sh_zp +
1264
- zp_sh_stage * ((group_blocks / thread_k_blocks) *
1265
- (pipe / (group_blocks / thread_k_blocks)));
1407
+ constexpr int g = group_blocks / thread_k_blocks;
1408
+ if (pipe % g == 0 && k % b_sh_wr_iters == 0) {
1409
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
1410
reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1411
sh_zp_stage[zp_sh_rd];
1412
}
1269
- } else {
1413
+ } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1414
auto warp_id = threadIdx.x / 32;
1271
- int n_warps = thread_n_blocks / 4;
1272
-
1273
- int warp_row = warp_id / n_warps;
1415
1275
- int cur_k = warp_row * 16;
1276
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1277
-
1278
- int k_blocks = cur_k / 16;
1279
- // Suppress bogus and persistent divide-by-zero warning
1280
- #pragma nv_diagnostic push
1281
- #pragma nv_diag_suppress divide_by_zero
1416
+ int warp_row = warp_id / tb_n_warps;
1417
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1418
int cur_group_id = k_blocks / group_blocks;
1283
- #pragma nv_diagnostic pop
1419
1420
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1421
@@ -1291,33 +1426,46 @@ __global__ void Marlin(
1426
}
1427
};
1428
1294
- auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1295
- dequant<scalar_t2, w_type_id, dequant_skip_flop>(q, frag_b_ptr);
1429
+ auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) {
1430
+ if constexpr (a_type.size_bits() != b_type.size_bits()) {
1431
+ if constexpr (is_a_8bit && has_zp) {
1432
+ sub_zp_and_dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(
1433
+ q, frag_b_ptr, zp);
1434
+ } else {
1435
+ dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(q, frag_b_ptr);
1436
+ }
1437
+ }
1438
};
1439
1440
// Execute the actual tensor core matmul of a sub-tile.
1441
bool is_first_matmul_in_slice = true;
1300
- auto matmul = [&](int k) {
1442
+ auto matmul = [&](int k, int pipe) {
1443
+ if (is_a_8bit) return;
1444
int k2 = k % 2;
1445
+ constexpr int g =
1446
+ group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
1447
const bool is_new_zp =
1303
- ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) ||
1448
+ (group_blocks == 0) ||
1449
+ ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
1450
+ (pipe % g == 0) ||
1451
(group_blocks == -1 && is_first_matmul_in_slice);
1452
if constexpr (has_zp && !is_zp_float) {
1453
if (is_new_zp) {
1454
if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
1455
int zp_quant_0, zp_quant_1;
1456
1310
- if constexpr (w_type.size_bits() == 4) {
1457
+ if constexpr (b_type.size_bits() == 4) {
1458
zp_quant_0 = frag_qzp[k2][0];
1459
zp_quant_1 = zp_quant_0 >> 8;
1460
} else {
1314
- static_assert(w_type.size_bits() == 8);
1461
+ static_assert(b_type.size_bits() == 8);
1462
zp_quant_0 = frag_qzp[k2][0];
1463
zp_quant_1 = frag_qzp[k2][1];
1464
}
1465
1319
- dequant_data(zp_quant_0, reinterpret_cast<scalar_t2*>(&frag_zp));
1320
- dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1466
+ dequant_data(zp_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_zp));
1467
+ dequant_data(zp_quant_1,
1468
+ reinterpret_cast<scalar_32bit_t*>(&frag_zp) + 2);
1469
}
1470
}
1471
if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
@@ -1327,14 +1475,14 @@ __global__ void Marlin(
1475
}
1476
}
1477
1330
- if constexpr (w_type == vllm::kFE2M1f) {
1478
+ if constexpr (b_type == vllm::kFE2M1f) {
1479
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1480
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1481
1334
- dequant_fp8_scales<scalar_t2, s_type_id>(
1335
- s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1336
- dequant_fp8_scales<scalar_t2, s_type_id>(
1337
- s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1482
+ dequant_fp8_scales<c_scalar_t2, s_type_id>(
1483
+ s_quant_0, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]));
1484
+ dequant_fp8_scales<c_scalar_t2, s_type_id>(
1485
+ s_quant_1, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]) + 2);
1486
}
1487
1488
// We have the m dimension as the inner loop in order to encourage overlapping
@@ -1345,61 +1493,168 @@ __global__ void Marlin(
1493
FragB frag_b1;
1494
int b_quant_0, b_quant_1;
1495
1348
- if constexpr (w_type_id == vllm::kFE2M1f.id()) {
1496
+ if constexpr (b_type_id == vllm::kFE2M1f.id()) {
1497
b_quant_1 = frag_b_quant[k2][0][j];
1498
b_quant_0 = b_quant_1 << 8;
1351
- } else if constexpr (w_type.size_bits() == 4) {
1499
+ } else if constexpr (b_type.size_bits() == 4) {
1500
b_quant_0 = frag_b_quant[k2][0][j];
1501
b_quant_1 = b_quant_0 >> 8;
1502
} else {
1355
- static_assert(w_type.size_bits() == 8);
1503
+ static_assert(b_type.size_bits() == 8);
1504
int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
1505
b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
1506
b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1507
}
1508
1361
- dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1362
- dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1509
+ dequant_data(b_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_b0));
1510
+ dequant_data(b_quant_1, reinterpret_cast<scalar_32bit_t*>(&frag_b1));
1511
1364
- if constexpr (dequant_skip_flop && has_zp && !is_zp_float) {
1365
- sub_zp<scalar_t>(frag_b0, frag_zp[j], 0);
1366
- sub_zp<scalar_t>(frag_b1, frag_zp[j], 1);
1512
+ if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) {
1513
+ sub_zp<a_type_id>(frag_b0, frag_zp[j], 0);
1514
+ sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
1515
}
1516
1517
// Apply scale to frag_b0
1370
- if constexpr (has_act_order) {
1518
+ if constexpr (has_act_order && !is_a_8bit) {
1519
static_assert(group_blocks != -1);
1372
- scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1373
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1374
- scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1375
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1520
+ scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1521
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1522
+ scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1523
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1524
} else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1377
- group_blocks == -1) {
1525
+ group_blocks == -1 && !is_a_8bit) {
1526
int idx = (threadIdx.x / 4) % 2;
1379
- scalar_t2 s2 = Dtype::nums2num2(
1527
+ scalar_t2 s2 = Adtype::nums2num2(
1528
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
1529
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
1530
if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1383
- scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1384
- scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1385
- } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) {
1531
+ scale_and_sub<a_type_id>(frag_b0, s2.x, frag_zp[j].x);
1532
+ scale_and_sub<a_type_id>(frag_b1, s2.y, frag_zp[j].y);
1533
+ } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 &&
1534
+ !is_a_8bit) {
1535
if (is_new_zp)
1536
frag_zp[j] = __hmul2(frag_zp[j],
1537
*reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1389
- scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1390
- scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1391
- } else if constexpr (group_blocks != -1) {
1392
- scale<scalar_t>(frag_b0, frag_s[k2][j], 0);
1393
- scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
1538
+ scale_and_sub<a_type_id>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1539
+ scale_and_sub<a_type_id>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1540
+ } else if constexpr (group_blocks != -1 && !is_a_8bit) {
1541
+ scale<a_type_id>(frag_b0, frag_s[k2][j], 0);
1542
+ scale<a_type_id>(frag_b1, frag_s[k2][j], 1);
1543
}
1544
1545
#pragma unroll
1546
for (int i = 0; i < thread_m_blocks; i++) {
1547
if constexpr (m_block_size_8) {
1399
- mma_trans<scalar_t>(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]);
1548
+ mma_trans<a_type_id>(frag_a[k2][i], frag_b0, frag_b1,
1549
+ frag_c[i][j][0]);
1550
} else {
1401
- mma<scalar_t>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1402
- mma<scalar_t>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
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]);
1553
+ }
1554
+ }
1555
+ }
1556
+ };
1557
+
1558
+ auto matmul_a8 = [&](int k) {
1559
+ int k2 = k % 2;
1560
+ #pragma unroll
1561
+ for (int j = 0; j < 2; j++) {
1562
+ FragB frag_b[2];
1563
+
1564
+ if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) {
1565
+ dequant_data(frag_b_quant[k2][0][j * 2],
1566
+ reinterpret_cast<scalar_32bit_t*>(&frag_b));
1567
+ dequant_data(frag_b_quant[k2][0][j * 2 + 1],
1568
+ reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2);
1569
+ } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) {
1570
+ int off = (threadIdx.x / 32) % 2 * 2 + j;
1571
+ int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF;
1572
+ dequant_data(frag_b_quant[k2][0][j * 2],
1573
+ reinterpret_cast<scalar_32bit_t*>(&frag_b), zp);
1574
+ zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF;
1575
+ dequant_data(frag_b_quant[k2][0][j * 2 + 1],
1576
+ reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2, zp);
1577
+ } else {
1578
+ reinterpret_cast<int2*>(&frag_b)[0] =
1579
+ reinterpret_cast<int2*>(&frag_b_quant[k2][j])[0];
1580
+ reinterpret_cast<int2*>(&frag_b)[1] =
1581
+ reinterpret_cast<int2*>(&frag_b_quant[k2][j])[1];
1582
+ }
1583
+
1584
+ #pragma unroll
1585
+ 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]);
1590
+ }
1591
+
1592
+ if constexpr (group_blocks != -1) {
1593
+ if (group_blocks == 2 || k == 1) {
1594
+ if constexpr (a_type == vllm::kS8) {
1595
+ int2 s_vals[2];
1596
+ s_vals[0] = {
1597
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[0],
1598
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[1]};
1599
+ s_vals[1] = {
1600
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[0],
1601
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[1]};
1602
+
1603
+ #pragma unroll
1604
+ for (int i = 0; i < thread_m_blocks; i++) {
1605
+ #pragma unroll
1606
+ for (int g = 0; g < 4; g++) {
1607
+ int scale = reinterpret_cast<int*>(&s_vals[0])[g % 2];
1608
+ *reinterpret_cast<int32_t*>(&frag_c[i][j][0][g]) +=
1609
+ *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][0][g]) *
1610
+ scale;
1611
+ frag_c_tmp[i][j][0][g] = 0.0f;
1612
+ }
1613
+
1614
+ #pragma unroll
1615
+ for (int g = 0; g < 4; g++) {
1616
+ int scale = reinterpret_cast<int*>(&s_vals[1])[g % 2];
1617
+ *reinterpret_cast<int32_t*>(&frag_c[i][j][1][g]) +=
1618
+ *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][1][g]) *
1619
+ scale;
1620
+ frag_c_tmp[i][j][1][g] = 0.0f;
1621
+ }
1622
+ }
1623
+ } else {
1624
+ float2 s_vals[2];
1625
+ if constexpr (s_type_id != vllm::kFE8M0fnu.id()) {
1626
+ static_assert(a_type.size_bits() == 16 ||
1627
+ s_type.size_bits() == 16);
1628
+ s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]);
1629
+ s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]);
1630
+ } else {
1631
+ int32_t* s_vals_int = reinterpret_cast<int32_t*>(&s_vals[0]);
1632
+ int32_t s_vals_e8m0 =
1633
+ *reinterpret_cast<int32_t*>(&frag_s[k2][j][0]);
1634
+
1635
+ s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23;
1636
+ s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15;
1637
+ s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7;
1638
+ s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1;
1639
+ }
1640
+
1641
+ #pragma unroll
1642
+ for (int i = 0; i < thread_m_blocks; i++) {
1643
+ #pragma unroll
1644
+ for (int g = 0; g < 4; g++) {
1645
+ float scale = reinterpret_cast<float*>(&s_vals[0])[g % 2];
1646
+ frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale;
1647
+ frag_c_tmp[i][j][0][g] = 0.0f;
1648
+ }
1649
+
1650
+ #pragma unroll
1651
+ for (int g = 0; g < 4; g++) {
1652
+ float scale = reinterpret_cast<float*>(&s_vals[1])[g % 2];
1653
+ frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale;
1654
+ frag_c_tmp[i][j][1][g] = 0.0f;
1655
+ }
1656
+ }
1657
+ }
1658
}
1659
}
1660
}
@@ -1413,7 +1668,8 @@ __global__ void Marlin(
1668
constexpr int red_off = threads / b_sh_stride_threads / 2;
1669
if (red_off >= 1) {
1670
auto red_idx = threadIdx.x / b_sh_stride_threads;
1416
- constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1671
+ constexpr int red_sh_stride =
1672
+ b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2;
1673
constexpr int red_sh_delta = b_sh_stride_threads;
1674
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
1675
(threadIdx.x % b_sh_stride_threads);
@@ -1428,7 +1684,8 @@ __global__ void Marlin(
1684
for (int i = red_off; i > 0; i /= 2) {
1685
if (i <= red_idx && red_idx < 2 * i) {
1686
#pragma unroll
1431
- for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) {
1687
+ for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2;
1688
+ j += (m_block_size_8 ? 2 : 1)) {
1689
int red_sh_wr =
1690
red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
1691
if (i < red_off) {
@@ -1437,24 +1694,26 @@ __global__ void Marlin(
1694
float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
1695
#pragma unroll
1696
for (int k = 0; k < 4; k++)
1440
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
1697
+ reinterpret_cast<FragC*>(
1698
+ frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] +=
1699
c_rd[k] + c_wr[k];
1700
}
1443
- sh_red[red_sh_wr] =
1444
- reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
1701
+ sh_red[red_sh_wr] = reinterpret_cast<int4*>(
1702
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j];
1703
}
1704
}
1705
__syncthreads();
1706
}
1707
if (red_idx == 0) {
1708
#pragma unroll
1451
- for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) {
1709
+ for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2;
1710
+ i += (m_block_size_8 ? 2 : 1)) {
1711
float* c_rd =
1712
reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
1713
#pragma unroll
1714
for (int j = 0; j < 4; j++)
1456
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
1457
- c_rd[j];
1715
+ reinterpret_cast<FragC*>(
1716
+ frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j];
1717
}
1718
}
1719
__syncthreads();
@@ -1470,13 +1729,13 @@ __global__ void Marlin(
1729
// We are very careful here to reduce directly in the output buffer to
1730
// maximize L2 cache utilization in this step. To do this, we write out
1731
// results in FP16 (but still reduce with FP32 compute).
1473
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1732
+ constexpr int active_threads = 32 * tb_n_warps;
1733
bool is_th_active = threadIdx.x < active_threads;
1734
if (!is_th_active) {
1735
return;
1736
}
1737
1479
- int c_gl_stride = prob_n / 8;
1738
+ int c_gl_stride = prob_n / 8 * (is_a_8bit ? 2 : 1);
1739
int c_gl_wr_delta_o = 8 * c_gl_stride;
1740
int c_gl_wr_delta_i = 4 * (active_threads / 32);
1741
int c_gl_wr;
@@ -1487,7 +1746,7 @@ __global__ void Marlin(
1746
} else {
1747
c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
1748
4 * (threadIdx.x / 32) + threadIdx.x % 4;
1490
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
1749
+ c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1);
1750
}
1751
constexpr int c_sh_wr_delta = active_threads;
1752
int c_sh_wr = threadIdx.x;
@@ -1506,7 +1765,13 @@ __global__ void Marlin(
1765
if (c_idx / c_gl_stride < block_num_valid_tokens) {
1766
int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
1767
int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
1509
- sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx];
1768
+ if constexpr (is_a_8bit) {
1769
+ int2* sh_red_int2 = reinterpret_cast<int2*>(sh_red);
1770
+ int2* c_int2 = reinterpret_cast<int2*>(C);
1771
+ sh_red_int2[c_sh_wr + c_sh_wr_delta * i] = c_int2[true_idx];
1772
+ } else {
1773
+ sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx];
1774
+ }
1775
}
1776
}
1777
}
@@ -1514,29 +1779,37 @@ __global__ void Marlin(
1779
#pragma unroll
1780
for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
1781
if (!first) {
1517
- int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
1782
+ c_scalar_t* c_red_f16;
1783
+ if constexpr (is_a_8bit) {
1784
+ int2 tmp =
1785
+ reinterpret_cast<int2*>(sh_red)[c_sh_wr + i * c_sh_wr_delta];
1786
+ c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
1787
+ } else {
1788
+ int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta];
1789
+ c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
1790
+ }
1791
#pragma unroll
1519
- for (int j = 0; j < 2 * 4; j++) {
1792
+ for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
1793
int delta = 0;
1794
if constexpr (m_block_size_8) {
1795
delta = j % 2 == 1 ? -2 : 0;
1796
}
1797
reinterpret_cast<float*>(
1525
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] +=
1526
- Dtype::num2float(reinterpret_cast<scalar_t*>(&c_red)[j]);
1798
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
1799
+ delta] += Cdtype::num2float(c_red_f16[j]);
1800
}
1801
}
1802
if (!last) {
1530
- int4 c;
1803
+ c_scalar_t c_f16[is_a_8bit ? 4 : 8];
1804
#pragma unroll
1532
- for (int j = 0; j < 2 * 4; j++) {
1805
+ for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
1806
int delta = 0;
1807
if constexpr (m_block_size_8) {
1808
delta = j % 2 == 1 ? -2 : 0;
1809
}
1537
- reinterpret_cast<scalar_t*>(&c)[j] =
1538
- Dtype::float2num(reinterpret_cast<float*>(
1539
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]);
1810
+ c_f16[j] = Cdtype::float2num(reinterpret_cast<float*>(
1811
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) +
1812
+ delta]);
1813
}
1814
1815
int c_idx;
@@ -1549,7 +1822,12 @@ __global__ void Marlin(
1822
if (c_idx / c_gl_stride < block_num_valid_tokens) {
1823
int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
1824
int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
1552
- C[true_idx] = c;
1825
+ if constexpr (is_a_8bit) {
1826
+ int2* c_int2 = reinterpret_cast<int2*>(C);
1827
+ c_int2[true_idx] = *reinterpret_cast<int2*>(c_f16);
1828
+ } else {
1829
+ C[true_idx] = *reinterpret_cast<int4*>(c_f16);
1830
+ }
1831
}
1832
}
1833
}
@@ -1563,10 +1841,10 @@ __global__ void Marlin(
1841
1842
constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
1843
1566
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1844
+ constexpr int active_threads = 32 * tb_n_warps;
1845
bool is_th_active = threadIdx.x < active_threads;
1846
1569
- constexpr int num_floats = thread_m_blocks * 4 * 2 * 4;
1847
+ constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4;
1848
constexpr int th_size = num_floats * sizeof(float) / 16;
1849
1850
int c_cur_offset = locks_off * c_size;
@@ -1634,7 +1912,7 @@ __global__ void Marlin(
1912
} else {
1913
c_sh_wr =
1914
(4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
1637
- c_sh_wr += 32 * (threadIdx.x / 32);
1915
+ c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32);
1916
}
1917
1918
int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
@@ -1643,49 +1921,49 @@ __global__ void Marlin(
1921
// We first reorder in shared memory to guarantee the most efficient final
1922
// global write patterns
1923
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1646
- scalar_t2 res =
1647
- Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1924
+ c_scalar_t2 res =
1925
+ Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
1926
1927
// For per-column quantization we finally apply the scale here (only for
1928
// 4-bit)
1651
- if constexpr (!has_act_order && group_blocks == -1 &&
1652
- w_type.size_bits() == 4 &&
1929
+ if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
1930
+ b_type.size_bits() == 4 &&
1931
(has_zp && dequant_skip_flop || !has_zp)) {
1654
- scalar_t2 tmp_scale = s[0];
1932
+ c_scalar_t2 tmp_scale = s[0];
1933
if constexpr (m_block_size_8) {
1656
- tmp_scale = Dtype::num2num2(
1934
+ tmp_scale = Cdtype::num2num2(
1935
reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
1936
}
1937
res = __hmul2(res, tmp_scale);
1938
}
1939
1662
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1940
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1941
if (!mul_topk_weights) {
1942
res = __hmul2(res, global_scale);
1943
}
1944
}
1945
if (has_bias && last) {
1668
- scalar_t2 tmp_bias = b_bias[0];
1946
+ c_scalar_t2 tmp_bias = b_bias[0];
1947
if constexpr (m_block_size_8) {
1670
- tmp_bias = Dtype::num2num2(
1948
+ tmp_bias = Cdtype::num2num2(
1949
reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
1950
}
1951
res = __hadd2(res, tmp_bias);
1952
}
1953
1954
if constexpr (m_block_size_8) {
1677
- ((scalar_t*)sh_red)[idx] = res.x;
1678
- ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1955
+ ((c_scalar_t*)sh_red)[idx] = res.x;
1956
+ ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1957
} else {
1680
- ((scalar_t2*)sh_red)[idx] = res;
1958
+ ((c_scalar_t2*)sh_red)[idx] = res;
1959
}
1960
};
1961
1684
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1962
+ if (threadIdx.x / 32 < tb_n_warps) {
1963
#pragma unroll
1964
for (int i = 0; i < thread_m_blocks; i++) {
1965
#pragma unroll
1688
- for (int j = 0; j < 4; j++) {
1966
+ for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) {
1967
if constexpr (m_block_size_8) {
1968
int wr = c_sh_wr + 16 * j;
1969
write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
@@ -1723,24 +2001,26 @@ __global__ void Marlin(
2001
if (row < block_num_valid_tokens) {
2002
int64_t sorted_row = sh_block_sorted_ids[row];
2003
int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride;
1726
- scalar_t2 topk_weight_score;
2004
+ c_scalar_t2 topk_weight_score;
2005
if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row];
2006
if (use_atomic_add && slice_count > 1 || mul_topk_weights) {
1729
- scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[true_idx]);
1730
- scalar_t2* sh_red_half2 =
1731
- reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
2007
+ c_scalar_t2* C_half2 = reinterpret_cast<c_scalar_t2*>(&C[true_idx]);
2008
+ c_scalar_t2* sh_red_half2 =
2009
+ reinterpret_cast<c_scalar_t2*>(&sh_red[c_sh_rd]);
2010
+ if (mul_topk_weights) {
2011
#pragma unroll
1733
- for (int a = 0; a < 4; a++) {
1734
- scalar_t2 res = sh_red_half2[a];
1735
- if (mul_topk_weights) {
1736
- res = __hmul2(res, topk_weight_score);
2012
+ for (int a = 0; a < 4; a++) {
2013
+ sh_red_half2[a] = __hmul2(sh_red_half2[a], topk_weight_score);
2014
}
2015
+ }
2016
1739
- if (use_atomic_add && slice_count > 1) {
1740
- atomicAdd(&C_half2[a], res);
1741
- } else {
1742
- C_half2[a] = res;
1743
- };
2017
+ if (use_atomic_add && slice_count > 1) {
2018
+ #pragma unroll
2019
+ for (int a = 0; a < 4; a++) {
2020
+ atomicAdd(&C_half2[a], sh_red_half2[a]);
2021
+ }
2022
+ } else {
2023
+ C[true_idx] = *reinterpret_cast<int4*>(sh_red_half2);
2024
}
2025
} else {
2026
C[true_idx] = sh_red[c_sh_rd];
@@ -1774,7 +2054,7 @@ __global__ void Marlin(
2054
}
2055
}
2056
}
1777
- fetch_to_shared(i, i, i < slice_iters, i);
2057
+ fetch_to_shared(i, i, i < slice_iters);
2058
}
2059
2060
zero_accums();
@@ -1799,73 +2079,100 @@ __global__ void Marlin(
2079
// have even length meaning that the next iteration will always start at
2080
// index 0.
2081
1802
- for (int stage_group_id = 0; stage_group_id < max_num_stage_groups;
1803
- stage_group_id++) {
2082
#pragma unroll
1805
- for (int pipe = 0; pipe < stages;) {
2083
+ for (int pipe = 0; pipe < stages;) {
2084
#pragma unroll
1807
- for (int k = 0; k < b_sh_wr_iters; k++) {
1808
- int idx =
1809
- (pipe >= stages && stage_group_id == max_num_stage_groups - 1)
1810
- ? (pipe - stages)
1811
- : (pipe + stage_group_id * stages);
1812
- fetch_to_registers(k + 1, pipe % stages, idx);
1813
- fetch_scales_to_registers(k + 1, pipe);
1814
- fetch_zp_to_registers(k + 1, pipe);
1815
- if (k == b_sh_wr_iters - 2) {
1816
- int idx = (pipe >= 1 && stage_group_id == max_num_stage_groups - 1)
1817
- ? (pipe - 1)
1818
- : (pipe + (stage_group_id + 1) * stages - 1);
1819
- fetch_to_shared((pipe + stages - 1) % stages, pipe,
1820
- slice_iters >= stages, idx);
1821
- pipe++;
1822
- wait_for_stage();
1823
- init_same_group(pipe % stages);
1824
- }
1825
- matmul(k);
2085
+ for (int k = 0; k < b_sh_wr_iters; k++) {
2086
+ fetch_to_registers(k + 1, pipe % stages);
2087
+ fetch_scales_to_registers(k + 1, pipe);
2088
+ fetch_zp_to_registers(k + 1, pipe);
2089
+ if (k == b_sh_wr_iters - 2) {
2090
+ fetch_to_shared((pipe + stages - 1) % stages, pipe,
2091
+ slice_iters >= stages);
2092
+ pipe++;
2093
+ wait_for_stage();
2094
+ init_same_group(pipe % stages);
2095
}
1827
- slice_iters--;
1828
- if (slice_iters == 0) {
1829
- break;
1830
- }
1831
- }
1832
-
1833
- a_gl_rd_col += a_gl_rd_delta_o * stages;
2096
1835
- if constexpr (has_act_order) {
1836
- slice_k_start += tb_k * stages;
1837
-
1838
- if (slice_k_start < prob_k) {
1839
- slice_k_start_shared_fetch += tb_k * stages;
1840
- int first_group_id = g_idx[slice_k_start];
1841
- int last_g_idx = slice_k_start + stages * tb_k * 2;
1842
- if (last_g_idx >= prob_k) {
1843
- last_g_idx = prob_k - 1;
1844
- }
1845
- int last_group_id = g_idx[last_g_idx];
1846
- if (last_group_id >= sh_first_group_id + sh_num_groups) {
1847
- fetch_act_order_scales_to_shared(false, first_group_id,
1848
- last_group_id);
1849
- __syncthreads();
1850
- }
2097
+ if constexpr (!is_a_8bit) {
2098
+ matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0));
2099
+ } else {
2100
+ static_assert(group_blocks != 0 && group_blocks != 1);
2101
+ matmul_a8(k);
2102
}
2103
}
2104
+ slice_iters--;
2105
if (slice_iters == 0) {
2106
break;
2107
}
2108
}
2109
2110
+ a_gl_rd_col += a_gl_rd_delta_o * stages;
2111
+
2112
+ if constexpr (has_act_order) {
2113
+ slice_k_start += tb_k * stages;
2114
+
2115
+ if (slice_k_start < prob_k) {
2116
+ slice_k_start_shared_fetch += tb_k * stages;
2117
+ int first_group_id = g_idx[slice_k_start];
2118
+ int last_g_idx = slice_k_start + stages * tb_k * 2;
2119
+ if (last_g_idx >= prob_k) {
2120
+ last_g_idx = prob_k - 1;
2121
+ }
2122
+ int last_group_id = g_idx[last_g_idx];
2123
+ if (last_group_id >= sh_first_group_id + sh_num_groups) {
2124
+ fetch_act_order_scales_to_shared(false, first_group_id,
2125
+ last_group_id);
2126
+ __syncthreads();
2127
+ }
2128
+ }
2129
+ }
2130
+
2131
// Process results and, if necessary, proceed to the next column slice.
2132
// While this pattern may not be the most readable, other ways of writing
2133
// the loop seemed to noticeably worse performance after compilation.
2134
if (slice_iters == 0) {
2135
+ if constexpr (is_a_8bit) {
2136
+ float frag_a_s[2 * thread_m_blocks];
2137
+
2138
+ for (int i = 0; i < 2 * thread_m_blocks; i++)
2139
+ frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4];
2140
+
2141
+ #pragma unroll
2142
+ for (int j = 0; j < 2; j++) {
2143
+ #pragma unroll
2144
+ for (int i = 0; i < thread_m_blocks; i++) {
2145
+ #pragma unroll
2146
+ for (int g = 0; g < 4; g++) {
2147
+ float c_val = frag_c[i][j][0][g];
2148
+
2149
+ if constexpr (a_type == vllm::kS8) {
2150
+ c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
2151
+ }
2152
+ float s_val = frag_a_s[i * 2 + g / 2];
2153
+ frag_c[i][j][0][g] = c_val * s_val;
2154
+ }
2155
+ #pragma unroll
2156
+ for (int g = 0; g < 4; g++) {
2157
+ float c_val = frag_c[i][j][1][g];
2158
+
2159
+ if constexpr (a_type == vllm::kS8) {
2160
+ c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
2161
+ }
2162
+ float s_val = frag_a_s[i * 2 + g / 2];
2163
+ frag_c[i][j][1][g] = c_val * s_val;
2164
+ }
2165
+ }
2166
+ }
2167
+ }
2168
+
2169
cp_async_wait<0>();
2170
bool last = slice_idx == slice_count - 1;
2171
// For per-column scales, we only fetch them here in the final step before
2172
// write-out
2173
if constexpr (!has_act_order && group_blocks == -1 &&
2174
(has_zp && dequant_skip_flop || !has_zp)) {
1868
- if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
2175
+ if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
2176
if (s_sh_wr_pred) {
2177
cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
2178
}
@@ -1883,20 +2190,27 @@ __global__ void Marlin(
2190
}
2191
2192
if constexpr (!has_act_order && group_blocks == -1 &&
1886
- (has_zp && dequant_skip_flop || !has_zp)) {
1887
- if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
2193
+ (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
2194
+ if constexpr (is_a_8bit) {
2195
+ cp_async_wait<0>();
2196
+ __syncthreads();
2197
+ if (threadIdx.x / 32 < tb_n_warps) {
2198
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
2199
+ }
2200
+ } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) {
2201
cp_async_wait<0>();
2202
__syncthreads();
1890
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
2203
+ if (threadIdx.x / 32 < tb_n_warps) {
2204
reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
2205
reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
2206
if constexpr (m_block_size_8) {
2207
int idx = (threadIdx.x / 4) % 2;
1895
- scalar_t2* frag_s_half2 = reinterpret_cast<scalar_t2*>(frag_s);
2208
+ c_scalar_t2* frag_s_half2 =
2209
+ reinterpret_cast<c_scalar_t2*>(frag_s);
2210
#pragma unroll
2211
for (int i = 0; i < 8; i++) {
1898
- frag_s_half2[i] = Dtype::num2num2(
1899
- reinterpret_cast<scalar_t*>(&frag_s_half2[i])[idx]);
2212
+ frag_s_half2[i] = Cdtype::num2num2(
2213
+ reinterpret_cast<c_scalar_t*>(&frag_s_half2[i])[idx]);
2214
}
2215
}
2216
}
@@ -1906,26 +2220,48 @@ __global__ void Marlin(
2220
// For 8-bit channelwise, we apply the scale before the global reduction
2221
// that converts the fp32 results to fp16 (so that we avoid possible
2222
// overflow in fp16)
1909
- if constexpr (!has_act_order && group_blocks == -1 &&
1910
- w_type.size_bits() == 8 &&
1911
- (has_zp && dequant_skip_flop || !has_zp)) {
1912
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
2223
+ if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
2224
+ #pragma unroll
2225
+ for (int j = 0; j < 2; j++) {
2226
+ float2 aa[2];
2227
+ aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]);
2228
+ aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]);
2229
+
2230
+ #pragma unroll
2231
+ for (int i = 0; i < thread_m_blocks; i++) {
2232
+ #pragma unroll
2233
+ for (int g = 0; g < 4; g++) {
2234
+ float scale = reinterpret_cast<float*>(&aa[0])[g % 2];
2235
+ frag_c[i][j][0][g] *= scale;
2236
+ }
2237
+
2238
+ #pragma unroll
2239
+ for (int g = 0; g < 4; g++) {
2240
+ float scale = reinterpret_cast<float*>(&aa[1])[g % 2];
2241
+ frag_c[i][j][1][g] *= scale;
2242
+ }
2243
+ }
2244
+ }
2245
+ } else if (!has_act_order && group_blocks == -1 &&
2246
+ b_type.size_bits() == 8 &&
2247
+ (has_zp && dequant_skip_flop || !has_zp)) {
2248
+ if (threadIdx.x / 32 < tb_n_warps) {
2249
#pragma unroll
2250
for (int i = 0; i < thread_m_blocks; i++) {
2251
#pragma unroll
2252
for (int j = 0; j < 4; j++) {
1917
- scale_float<scalar_t>(
2253
+ scale_float<c_type_id>(
2254
reinterpret_cast<float*>(&frag_c[i][j][0][0]),
2255
frag_s[j / 2][2 * (j % 2) + 0]);
1920
- scale_float<scalar_t>(
2256
+ scale_float<c_type_id>(
2257
reinterpret_cast<float*>(&frag_c[i][j][0][2]),
2258
frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
2259
2260
if constexpr (!m_block_size_8) {
1925
- scale_float<scalar_t>(
2261
+ scale_float<c_type_id>(
2262
reinterpret_cast<float*>(&frag_c[i][j][1][0]),
2263
frag_s[j / 2][2 * (j % 2) + 1]);
1928
- scale_float<scalar_t>(
2264
+ scale_float<c_type_id>(
2265
reinterpret_cast<float*>(&frag_c[i][j][1][2]),
2266
frag_s[j / 2][2 * (j % 2) + 1]);
2267
}
@@ -1949,7 +2285,8 @@ __global__ void Marlin(
2285
cp_async_wait<0>();
2286
__syncthreads();
2287
reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
1952
- reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
2288
+ if constexpr (!is_a_8bit)
2289
+ reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
2290
__syncthreads();
2291
}
2292
@@ -1958,37 +2295,22 @@ __global__ void Marlin(
2295
if (last || use_atomic_add)
2296
// only the last block in a slice actually writes the result
2297
write_result(last);
1961
- int old_slice_row = slice_row;
2298
slice_row = 0;
1963
- slice_col_par++;
1964
- slice_col++;
1965
- is_first_matmul_in_slice = true;
1966
- init_slice();
1967
-
1968
- // Should we load A matrix in next slice?
1969
- // `slice_col == 0`: when move to a new moe block
1970
- // `old_slice_row > 0`:
1971
- // when the last slice is not starting from k_index == 0
1972
- // (only happen when it is the first slice of a threadblock)
1973
- // `prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups`:
1974
- // when the required shared memory size is larger than
1975
- // the remaining shared memory
1976
- if (slice_col == 0 || old_slice_row ||
1977
- prob_k > thread_k_blocks * 16 * stages * max_num_stage_groups) {
1978
- should_load_a = true;
2299
+ if (!in_part2) {
2300
+ slice_col_par += gridDim.x;
2301
} else {
1980
- should_load_a = false;
2302
+ slice_col_par++;
2303
+ slice_col++;
2304
}
2305
+ is_first_matmul_in_slice = true;
2306
+ init_slice();
2307
2308
if (slice_iters) {
1984
- a_gl_rd_col = (threadIdx.x % a_gl_rd_delta_o);
1985
- #pragma unroll
1986
- for (int i = 0; i < b_sh_wr_iters; i++)
1987
- B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
1988
- if (slice_col == 0) {
1989
- #pragma unroll
1990
- for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1991
- }
2309
+ a_gl_rd_col =
2310
+ a_gl_rd_delta_o * slice_row + threadIdx.x % a_gl_rd_delta_o;
2311
+ b_gl_rd = B_expert_off + b_gl_stride * (threadIdx.x / b_sh_stride) +
2312
+ (threadIdx.x % b_sh_stride);
2313
+ b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
2314
2315
bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
2316
// Update slice k/n for scales loading
@@ -1998,8 +2320,26 @@ __global__ void Marlin(
2320
slice_k_start_shared_fetch = slice_k_start;
2321
slice_n_offset = act_s_col_tb_stride * slice_col;
2322
} else {
2001
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
2002
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2323
+ if constexpr (group_blocks == -1) {
2324
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
2325
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2326
+ } else if constexpr (group_blocks >= thread_k_blocks) {
2327
+ s_gl_rd =
2328
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2329
+ s_sh_stride * slice_col + threadIdx.x;
2330
+ zp_gl_rd =
2331
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2332
+ zp_sh_stride * slice_col + threadIdx.x;
2333
+ } else {
2334
+ s_gl_rd =
2335
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2336
+ threadIdx.x / s_sh_stride) +
2337
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
2338
+ zp_gl_rd =
2339
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2340
+ threadIdx.x / zp_sh_stride) +
2341
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
2342
+ }
2343
}
2344
start_pipes();
2345
}
csrc/moe/marlin_moe_wna16/ops.cu
+228
-364
@@ -37,39 +37,6 @@ __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
-template <int moe_block_size>
43
-__global__ void permute_cols_kernel(
44
- int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
45
- int4* __restrict__ out_int4_ptr,
46
- const int32_t* __restrict__ sorted_token_ids_ptr,
47
- const int32_t* __restrict__ expert_ids_ptr,
48
- const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
49
- int size_k, int top_k) {};
50
-
51
-} // namespace marlin
52
-
53
-torch::Tensor moe_wna16_marlin_gemm(
54
- torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
55
- torch::Tensor& b_q_weight,
56
- std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
57
- std::optional<torch::Tensor> const& b_zeros_or_none,
58
- std::optional<torch::Tensor> const& g_idx_or_none,
59
- std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
60
- torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
61
- torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
62
- int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
63
- vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
64
- int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
65
- bool is_zp_float) {
66
- TORCH_CHECK_NOT_IMPLEMENTED(false,
67
- "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
68
- return torch::empty({1, 1});
69
-}
70
-
71
-#else
72
-
40
// For a given "a" of size [M,K] performs a permutation of the K columns based
41
// on the given "perm" indices.
42
template <int moe_block_size>
@@ -207,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,
210
- int is_zp_float) {
177
+ int is_zp_float, bool is_a_8bit) {
178
int pack_factor = 32 / num_bits;
179
180
// Get B size
@@ -217,8 +184,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
184
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)
220
- int sh_block_meta_size = tb_m * 4;
221
- int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
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;
190
int sh_red_size = tb_m * (tb_n + 8) * 2;
191
int sh_bias_size = tb_n * 2;
@@ -250,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,
253
- int max_shared_mem) {
220
+ int max_shared_mem, bool is_a_8bit) {
221
// Sanity
222
if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
223
th_config.num_threads == -1) {
@@ -273,188 +240,34 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
240
}
241
242
// Check that pipeline fits into cache
276
- int cache_size = get_kernel_cache_size(
277
- th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
278
- num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
279
- return cache_size + 512 <= max_shared_mem;
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);
247
+ return cache_size <= max_shared_mem;
248
}
249
282
- #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
283
- M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
284
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
285
- thread_n_blocks == THREAD_N_BLOCKS && \
286
- thread_k_blocks == THREAD_K_BLOCKS && \
287
- m_block_size_8 == M_BLOCK_SIZE_8 && \
288
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
289
- is_zp_float == IS_ZP_FLOAT) { \
290
- constexpr auto S_TYPE = \
291
- W_TYPE == vllm::kFE2M1f \
292
- ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu) \
293
- : (std::is_same<scalar_t, half>::value ? vllm::kFloat16 \
294
- : vllm::kBFloat16); \
295
- kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS, \
296
- THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
297
- M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
298
- }
299
-
300
- // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
301
- // this is the most common cases
302
- // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
303
- // FZP: cases for float-zero-point (is_zp_float = true)
304
- // ACT: cases for act order case (group_blocks == 0)
305
- // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
306
- #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
307
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
308
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
309
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \
310
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
311
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
312
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
313
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
314
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
315
-
316
- #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
317
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
318
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
319
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
320
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
321
- \
322
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
323
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
324
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
325
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
326
- \
327
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
328
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
329
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
330
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
331
-
332
- #define COMMON_GET_IF(W_TYPE) \
333
- COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \
334
- COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \
335
- COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
336
- COMMON_GET_IF_M234(W_TYPE, 8, 4, 128)
337
-
338
- #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
339
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
340
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
341
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
342
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
343
-
344
- #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
345
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
346
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
347
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
348
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
349
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
350
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
351
-
352
- #define BIGGROUP_GET_IF(W_TYPE) \
353
- BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
354
- BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
355
- BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
356
- BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
357
-
358
- #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
359
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
360
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
361
-
362
- #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
363
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
364
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
365
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
366
-
367
- #define NVFP4_GET_IF(W_TYPE) \
368
- NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
369
- NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
370
- NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
371
- NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
372
-
373
- #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
374
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
375
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
376
-
377
- #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
378
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
379
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
380
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
381
-
382
- #define MXFP4_GET_IF(W_TYPE) \
383
- MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
384
- MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
385
- MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
386
- MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
387
-
388
- // We currently have 4-bit models only with group_blocks == 4
389
- #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
390
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
391
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
392
-
393
- #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
394
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
395
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
396
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
397
-
398
- #define FZP_GET_IF(W_TYPE) \
399
- FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \
400
- FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \
401
- FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
402
- FZP_GET_IF_M234(W_TYPE, 8, 4, 128)
403
-
404
- // We currently have 4-bit models only with group_blocks == 4
405
- #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
406
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
407
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
408
-
409
- #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
410
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
411
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
412
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
413
-
414
- #define ACT_GET_IF(W_TYPE) \
415
- ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \
416
- ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \
417
- ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
418
- ACT_GET_IF_M234(W_TYPE, 8, 4, 128)
419
-
420
-template <typename scalar_t>
421
-MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
422
- int thread_m_blocks, int thread_n_blocks,
423
- int thread_k_blocks, bool m_block_size_8,
424
- bool has_act_order, bool has_zp,
425
- int group_blocks, int num_threads,
426
- bool is_zp_float) {
427
- int num_bits = q_type.size_bits();
250
+MarlinFuncPtr get_marlin_kernel(
251
+ const vllm::ScalarType a_type, const vllm::ScalarType b_type,
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) {
256
+ int num_bits = b_type.size_bits();
257
auto kernel = MarlinDefault;
429
- if (false) {
430
- }
431
-
432
- COMMON_GET_IF(vllm::kU4)
433
- COMMON_GET_IF(vllm::kU4B8)
434
- COMMON_GET_IF(vllm::kU8B128)
258
436
- NVFP4_GET_IF(vllm::kFE2M1f)
437
-
438
- BIGGROUP_GET_IF(vllm::kFE4M3fn)
439
-
440
- ACT_GET_IF(vllm::kU4B8)
441
- ACT_GET_IF(vllm::kU8B128)
442
- if (std::is_same<scalar_t, nv_bfloat16>::value) {
443
- if (false) {
444
- }
445
- MXFP4_GET_IF(vllm::kFE2M1f)
446
- }
259
+#include "kernel_selector.h"
260
261
return kernel;
262
}
263
451
-template <typename scalar_t>
452
-exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
453
- int prob_n, int prob_k, int thread_m_blocks,
454
- bool m_block_size_8, int num_bits,
455
- int group_size, bool has_act_order,
456
- bool is_k_full, bool has_zp,
457
- bool is_zp_float, int max_shared_mem) {
264
+exec_config_t determine_exec_config(
265
+ const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
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) {
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
@@ -471,73 +284,69 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
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,
474
- is_k_full, has_zp, is_zp_float, max_shared_mem)) {
287
+ is_k_full, has_zp, is_zp_float, max_shared_mem - 512,
288
+ is_a_8bit)) {
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,
480
- num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
294
+ num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
295
+ is_a_8bit);
296
297
int group_blocks = 0;
298
if (!has_act_order) {
299
group_blocks = group_size == -1 ? -1 : (group_size / 16);
300
}
301
487
- auto kernel = get_marlin_kernel<scalar_t>(
488
- q_type, thread_m_blocks, th_config.thread_n / 16,
489
- th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
490
- group_blocks, th_config.num_threads, is_zp_float);
302
+ auto kernel =
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);
307
308
if (kernel == MarlinDefault) continue;
309
494
- if (thread_m_blocks > 1) {
495
- exec_cfg = {1, th_config};
496
- break;
497
- } else {
498
- cudaFuncAttributes attr;
499
- cudaFuncGetAttributes(&attr, kernel);
500
- int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
501
- int allow_count = min(device_max_reg_size / reg_size,
502
- max_shared_mem / (cache_size + 1024));
310
+ cudaFuncAttributes attr;
311
+ cudaFuncGetAttributes(&attr, kernel);
312
+ int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
313
+ int allow_count = min(device_max_reg_size / reg_size,
314
+ max_shared_mem / (cache_size + 1536));
315
+ if (thread_m_blocks == 1)
316
allow_count = max(min(allow_count, 4), 1);
504
- if (allow_count > count) {
505
- count = allow_count;
506
- exec_cfg = {count, th_config};
507
- };
317
+ else
318
+ allow_count = max(min(allow_count, 2), 1);
319
+
320
+ if (prob_n / th_config.thread_n * prob_m * top_k * 4 < sms * allow_count) {
321
+ allow_count =
322
+ max(prob_n / th_config.thread_n * prob_m * top_k * 4 / sms, 1);
323
}
324
+
325
+ if (allow_count > count) {
326
+ count = allow_count;
327
+ exec_cfg = {count, th_config};
328
+ };
329
}
330
331
return exec_cfg;
332
}
333
514
-template <typename scalar_t>
334
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
516
- void* s, void* s2, void* zp, void* g_idx, void* perm,
517
- void* a_tmp, void* sorted_token_ids, void* expert_ids,
518
- void* num_tokens_past_padded, void* topk_weights,
519
- int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
520
- int prob_m, int prob_n, int prob_k, void* workspace,
521
- vllm::ScalarType const& q_type, bool has_bias,
522
- bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
523
- int group_size, int dev, cudaStream_t stream, int thread_k,
524
- int thread_n, int sms, bool use_atomic_add, bool use_fp32_reduce,
525
- bool is_zp_float) {
335
+ void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
336
+ void* perm, void* a_tmp, void* sorted_token_ids,
337
+ void* expert_ids, void* num_tokens_past_padded,
338
+ void* topk_weights, int moe_block_size, int num_experts,
339
+ int top_k, bool mul_topk_weights, bool is_ep, int prob_m,
340
+ int prob_n, int prob_k, void* workspace,
341
+ vllm::ScalarType const& a_type, vllm::ScalarType const& b_type,
342
+ vllm::ScalarType const& c_type, vllm::ScalarType const& s_type,
343
+ bool has_bias, bool has_act_order, bool is_k_full, bool has_zp,
344
+ int num_groups, int group_size, int dev, cudaStream_t stream,
345
+ int thread_k, int thread_n, int sms, int blocks_per_sm,
346
+ bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) {
347
int thread_m_blocks = div_ceil(moe_block_size, 16);
348
bool m_block_size_8 = moe_block_size == 8;
528
-
529
- if (has_zp) {
530
- TORCH_CHECK(
531
- q_type == vllm::kU4 || q_type == vllm::kU8,
532
- "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
533
- } else {
534
- TORCH_CHECK(
535
- q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
536
- q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
537
- "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
538
- "has_zp = False. Got = ",
539
- q_type.str());
540
- }
349
+ bool is_a_8bit = a_type.size_bits() == 8;
350
351
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
352
", ", prob_n, ", ", prob_k, "]");
@@ -563,14 +372,15 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
372
}
373
}
374
566
- int num_bits = q_type.size_bits();
375
+ int num_bits = b_type.size_bits();
376
const int4* A_ptr = (const int4*)A;
377
const int4* B_ptr = (const int4*)B;
378
int4* C_ptr = (int4*)C;
379
int4* C_tmp_ptr = (int4*)C_tmp;
380
const int4* bias_ptr = (const int4*)b_bias;
572
- const int4* s_ptr = (const int4*)s;
573
- const uint16_t* s2_ptr = (const uint16_t*)s2;
381
+ const float* a_s_ptr = (const float*)a_s;
382
+ const int4* b_s_ptr = (const int4*)b_s;
383
+ const uint16_t* g_s_ptr = (const uint16_t*)g_s;
384
const int4* zp_ptr = (const int4*)zp;
385
const int* g_idx_ptr = (const int*)g_idx;
386
const int* perm_ptr = (const int*)perm;
@@ -618,22 +428,41 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
428
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
429
TORCH_CHECK(max_shared_mem > 0);
430
431
+ int major_capability, minor_capability;
432
+ cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
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.");
438
+ if (a_type == vllm::kFE4M3fn) {
439
+ TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
440
+ "FP8 only support Ada Lovelace or newer GPUs.");
441
+ TORCH_CHECK(
442
+ major_capability * 10 + minor_capability == 89 ||
443
+ major_capability * 10 + minor_capability == 120,
444
+ "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
445
+ "Marlin W4A16 on other devices).");
446
+ }
447
+
448
// Set thread config
449
exec_config_t exec_cfg;
450
thread_config_t thread_tfg;
451
if (thread_k != -1 && thread_n != -1) {
625
- thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
626
- exec_cfg = exec_config_t{1, thread_tfg};
452
+ thread_tfg = thread_config_t{thread_k, thread_n, thread_k * thread_n / 64};
453
+ if (blocks_per_sm == -1) blocks_per_sm = 1;
454
+ exec_cfg = exec_config_t{blocks_per_sm, thread_tfg};
455
TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
456
" is not divisible by thread_n = ", thread_n);
457
TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
458
" is not divisible by thread_k = ", thread_k);
459
} else {
460
// Auto config
633
- exec_cfg = determine_exec_config<scalar_t>(
634
- q_type, prob_m, prob_n, prob_k, thread_m_blocks, m_block_size_8,
635
- num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
636
- max_shared_mem);
461
+ exec_cfg = determine_exec_config(
462
+ a_type, b_type, c_type, s_type, prob_m, prob_n, prob_k, num_experts,
463
+ 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);
466
thread_tfg = exec_cfg.tb_cfg;
467
}
468
@@ -647,22 +476,29 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
476
int thread_k_blocks = thread_k / 16;
477
int thread_n_blocks = thread_n / 16;
478
650
- TORCH_CHECK(
651
- is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
652
- prob_n, prob_k, num_bits, group_size, has_act_order,
653
- is_k_full, has_zp, is_zp_float, max_shared_mem),
654
- "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
655
- ", thread_k = ", thread_tfg.thread_k,
656
- ", thread_n = ", thread_tfg.thread_n,
657
- ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, ", ",
658
- prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
659
- ", group_size = ", group_size, ", has_act_order = ", has_act_order,
660
- ", is_k_full = ", is_k_full, ", has_zp = ", has_zp,
661
- ", is_zp_float = ", is_zp_float, ", max_shared_mem = ", max_shared_mem);
662
-
663
- auto kernel = get_marlin_kernel<scalar_t>(
664
- q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8,
665
- has_act_order, has_zp, group_blocks, num_threads, is_zp_float);
479
+ TORCH_CHECK(is_valid_config(thread_tfg, m_block_size_8, thread_m_blocks,
480
+ prob_m, prob_n, prob_k, num_bits, group_size,
481
+ has_act_order, is_k_full, has_zp, is_zp_float,
482
+ max_shared_mem, is_a_8bit),
483
+ "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
484
+ ", thread_k = ", thread_tfg.thread_k,
485
+ ", thread_n = ", thread_tfg.thread_n,
486
+ ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
487
+ prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
488
+ ", group_size = ", group_size,
489
+ ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
490
+ ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
491
+ ", max_shared_mem = ", max_shared_mem);
492
+
493
+ int sh_cache_size =
494
+ get_kernel_cache_size(thread_tfg, m_block_size_8, thread_m_blocks, prob_m,
495
+ prob_n, prob_k, num_bits, group_size, has_act_order,
496
+ is_k_full, has_zp, is_zp_float, is_a_8bit);
497
+
498
+ auto kernel = get_marlin_kernel(
499
+ a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
500
+ thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
501
+ num_threads, is_zp_float);
502
503
if (kernel == MarlinDefault) {
504
TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
@@ -679,19 +515,20 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
515
// avoid ">>>" being formatted to "> > >"
516
// clang-format off
517
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
682
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr,
518
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, g_idx_ptr,
519
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
520
topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
685
- prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce, max_shared_mem);
521
+ prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce);
522
// clang-format on
523
}
524
525
} // namespace MARLIN_NAMESPACE_NAME
526
527
torch::Tensor moe_wna16_marlin_gemm(
692
- torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
528
+ torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
529
torch::Tensor& b_q_weight,
530
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
531
+ std::optional<torch::Tensor> const& a_scales_or_none,
532
std::optional<torch::Tensor> const& global_scale_or_none,
533
std::optional<torch::Tensor> const& b_zeros_or_none,
534
std::optional<torch::Tensor> const& g_idx_or_none,
@@ -699,11 +536,70 @@ torch::Tensor moe_wna16_marlin_gemm(
536
torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
537
torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
538
int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
702
- vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
539
+ vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
540
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
704
- bool is_zp_float) {
705
- vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
706
- int pack_factor = 32 / b_q_type.size_bits();
541
+ bool is_zp_float, int64_t thread_k, int64_t thread_n,
542
+ int64_t blocks_per_sm) {
543
+ vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
544
+
545
+ auto c_dtype = a.dtype();
546
+ if (a.scalar_type() == at::ScalarType::Half) {
547
+ a_type_id = vllm::kFloat16.id();
548
+ c_type_id = vllm::kFloat16.id();
549
+ } else if (a.scalar_type() == at::ScalarType::BFloat16) {
550
+ a_type_id = vllm::kBFloat16.id();
551
+ c_type_id = vllm::kBFloat16.id();
552
+ } else {
553
+ c_dtype = b_scales.dtype();
554
+ if (b_scales.scalar_type() == at::ScalarType::Half) {
555
+ c_type_id = vllm::kFloat16.id();
556
+ } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
557
+ c_type_id = vllm::kBFloat16.id();
558
+ } else {
559
+ c_type_id = vllm::kBFloat16.id();
560
+
561
+ TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
562
+ torch::Tensor c = c_or_none.value();
563
+ c_dtype = c.dtype();
564
+
565
+ if (c.scalar_type() == at::ScalarType::Half) {
566
+ c_type_id = vllm::kFloat16.id();
567
+ } else if (c.scalar_type() == at::ScalarType::BFloat16) {
568
+ c_type_id = vllm::kBFloat16.id();
569
+ } else {
570
+ TORCH_CHECK(false, "unsupported c dtype");
571
+ }
572
+ }
573
+
574
+ if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
575
+ a_type_id = vllm::kFE4M3fn.id();
576
+ } else if (a.scalar_type() == at::ScalarType::Char) {
577
+ a_type_id = vllm::kS8.id();
578
+ } else {
579
+ TORCH_CHECK(false, "unsupported `a` scalar_type");
580
+ }
581
+ }
582
+
583
+ s_type_id = c_type_id;
584
+ if (b_type_id == vllm::kFE2M1f.id()) {
585
+ if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
586
+ s_type_id = vllm::kFE4M3fn.id();
587
+ } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
588
+ s_type_id = vllm::kFE8M0fnu.id();
589
+ } else {
590
+ TORCH_CHECK(false,
591
+ "When b_type = float4_e2m1f, b_scale scalar type must be",
592
+ "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
593
+ }
594
+ }
595
+
596
+ vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
597
+ vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
598
+ vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
599
+ vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
600
+
601
+ int pack_factor = 32 / b_type.size_bits();
602
+ int num_experts = b_q_weight.size(0);
603
604
if (moe_block_size != 8) {
605
TORCH_CHECK(moe_block_size % 16 == 0,
@@ -745,19 +641,27 @@ torch::Tensor moe_wna16_marlin_gemm(
641
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
642
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
643
748
- // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
749
- // auto -1)
750
- int thread_k = -1;
751
- // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
752
- // auto -1)
753
- int thread_n = -1;
644
+ torch::Tensor a_scales;
645
+ auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
646
+ auto options_fp32 =
647
+ torch::TensorOptions().dtype(at::kFloat).device(a.device());
648
+
649
+ if (a_scales_or_none.has_value()) {
650
+ a_scales = a_scales_or_none.value();
651
+ TORCH_CHECK(a_type.size_bits() == 8,
652
+ "a_scales can only be used for 8bit activation.");
653
+ } else {
654
+ a_scales = torch::empty({0}, options_fp32);
655
+ TORCH_CHECK(a_type.size_bits() != 8,
656
+ "the a_scales parameter must be passed for 8bit activation.");
657
+ }
658
+
659
// sms: number of SMs to use for the kernel
660
int sms = -1;
661
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
662
663
// Alloc buffers
664
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
760
- auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
665
torch::Tensor c;
666
if (c_or_none.has_value()) {
667
c = c_or_none.value();
@@ -774,8 +678,6 @@ torch::Tensor moe_wna16_marlin_gemm(
678
679
// Alloc C tmp buffer that is going to be used for the global reduce
680
torch::Tensor c_tmp;
777
- auto options_fp32 =
778
- torch::TensorOptions().dtype(at::kFloat).device(a.device());
681
if (use_fp32_reduce && !use_atomic_add) {
682
// max num of threadblocks is sms * 4
683
long max_c_tmp_size = min(
@@ -846,11 +748,11 @@ torch::Tensor moe_wna16_marlin_gemm(
748
torch::Tensor global_scale;
749
if (global_scale_or_none.has_value()) {
750
global_scale = global_scale_or_none.value();
849
- TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
751
+ TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
752
"global_scale can only be used for nvfp4 format.");
753
} else {
754
global_scale = torch::empty({0}, options);
853
- TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
755
+ TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
756
"the global_scale parameter must be passed for nvfp4 format.");
757
}
758
@@ -877,15 +779,15 @@ torch::Tensor moe_wna16_marlin_gemm(
779
bool has_zp = b_zeros.size(-1) > 0;
780
if (has_zp) {
781
TORCH_CHECK(
880
- b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
881
- "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
782
+ b_type == vllm::kU4 || b_type == vllm::kU8,
783
+ "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
784
} else {
883
- TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
884
- b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
885
- "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
886
- "float4_e2m1f when "
887
- "has_zp = False. Got = ",
888
- b_q_type.str());
785
+ TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
786
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
787
+ b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
788
+ "b_type must be uint4b8, uint8b128, int4, int8, "
789
+ "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
790
+ b_type.str());
791
}
792
793
if (has_zp && is_zp_float) {
@@ -929,71 +831,33 @@ torch::Tensor moe_wna16_marlin_gemm(
831
" is below min_workspace_size = ", min_workspace_size);
832
833
int dev = a.get_device();
932
- if (a.scalar_type() == at::ScalarType::Half) {
933
- void* scales_ptr;
934
- if (b_q_type == vllm::kFE2M1f) {
935
- if (group_size == 16)
936
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
937
- else if (group_size == 32)
938
- scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
939
- else
940
- TORCH_CHECK(false,
941
- "float4_e2m1f only supports group_size == 16 (NVFP4) ",
942
- "and group_size == 32 (MXFP4)");
943
- } else {
944
- scales_ptr = b_scales.data_ptr<at::Half>();
945
- }
946
-
947
- MARLIN_NAMESPACE_NAME::marlin_mm<half>(
948
- a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
949
- c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
950
- global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
951
- perm.data_ptr(), a_tmp.data_ptr<at::Half>(),
952
- sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
953
- num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
954
- moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
955
- workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
956
- has_zp, num_groups, group_size, dev,
957
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
958
- use_atomic_add, use_fp32_reduce, is_zp_float);
959
- } else if (a.scalar_type() == at::ScalarType::BFloat16) {
960
- void* scales_ptr;
961
- if (b_q_type == vllm::kFE2M1f) {
962
- if (group_size == 16)
963
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
964
- else if (group_size == 32)
965
- scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
966
- else
967
- TORCH_CHECK(false,
968
- "float4_e2m1f only supports group_size == 16 (NVFP4) ",
969
- "and group_size == 32 (MXFP4)");
970
- } else {
971
- scales_ptr = b_scales.data_ptr<at::BFloat16>();
972
- }
834
974
- MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
975
- a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
976
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
977
- b_bias.data_ptr<at::BFloat16>(), scales_ptr,
978
- global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
979
- g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
980
- sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
981
- num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
982
- moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
983
- workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
984
- has_zp, num_groups, group_size, dev,
985
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
986
- use_atomic_add, use_fp32_reduce, is_zp_float);
987
- } else {
988
- TORCH_CHECK(false,
989
- "moe_wna16_marlin_gemm only supports bfloat16 and float16");
835
+ TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
836
+ "scalar type of a_scales must be float");
837
+ TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
838
+ "scalar type of global_scale must be the same with c");
839
+ if (a_type.size_bits() == 16) {
840
+ TORCH_CHECK(
841
+ a.scalar_type() == c.scalar_type(),
842
+ "scalar type of a must be the same with c for 16 bit activation");
843
}
844
845
+ MARLIN_NAMESPACE_NAME::marlin_mm(
846
+ a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
847
+ b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
848
+ global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
849
+ perm.data_ptr(), a_tmp.data_ptr(), sorted_token_ids.data_ptr(),
850
+ expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
851
+ topk_weights.data_ptr(), moe_block_size, num_experts, top_k,
852
+ mul_topk_weights, is_ep, size_m, size_n, size_k, workspace.data_ptr(),
853
+ a_type, b_type, c_type, s_type, has_bias, has_act_order, is_k_full,
854
+ has_zp, num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
855
+ thread_k, thread_n, sms, blocks_per_sm, use_atomic_add, use_fp32_reduce,
856
+ is_zp_float);
857
+
858
return c;
859
}
860
995
-#endif
996
-
861
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
862
m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
863
}
csrc/moe/torch_bindings.cpp
+5
-3
@@ -63,16 +63,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
63
m.def(
64
"moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
65
"Tensor! b_q_weight, Tensor? b_bias_or_none,"
66
- "Tensor! b_scales, Tensor? global_scale, Tensor? "
66
+ "Tensor! b_scales, Tensor? a_scales, Tensor? global_scale, Tensor? "
67
"b_zeros_or_none,"
68
"Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
69
"Tensor sorted_token_ids,"
70
"Tensor! expert_ids, Tensor! num_tokens_past_padded,"
71
"Tensor! topk_weights, int moe_block_size, int top_k, "
72
- "bool mul_topk_weights, bool is_ep, int b_q_type_id,"
72
+ "bool mul_topk_weights, bool is_ep, int b_type_id,"
73
"int size_m, int size_n, int size_k,"
74
"bool is_full_k, bool use_atomic_add,"
75
- "bool use_fp32_reduce, bool is_zp_float) -> Tensor");
75
+ "bool use_fp32_reduce, bool is_zp_float,"
76
+ "int thread_k, int thread_n, int blocks_per_sm) -> Tensor");
77
+
78
m.def(
79
"marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
80
"Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu
+4
-4
@@ -437,10 +437,10 @@ struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
437
for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
438
#pragma unroll
439
for (int k_idx = 0; k_idx < 2; ++k_idx) {
440
- FType low16 =
441
- ScalarType<FType>::float2num(C_frag[m_idx][n_idx][k_idx * 2]);
442
- FType high16 =
443
- ScalarType<FType>::float2num(C_frag[m_idx][n_idx][k_idx * 2 + 1]);
440
+ FType low16 = MarlinScalarType2<FType>::float2num(
441
+ C_frag[m_idx][n_idx][k_idx * 2]);
442
+ FType high16 = MarlinScalarType2<FType>::float2num(
443
+ C_frag[m_idx][n_idx][k_idx * 2 + 1]);
444
uint32_t tmp = (reinterpret_cast<uint32_t&>(low16) & 0xffff) |
445
(reinterpret_cast<uint32_t&>(high16) << 16);
446
int sts_offset =
csrc/quantization/gptq_allspark/allspark_utils.cuh
+3
-3
@@ -8,7 +8,7 @@
8
#include <cuda_bf16.h>
9
#include <iostream>
10
#include "../gptq_marlin/marlin_dtypes.cuh"
11
-using marlin::ScalarType;
11
+using marlin::MarlinScalarType2;
12
13
namespace allspark {
14
@@ -72,10 +72,10 @@ __global__ void f16_gemm_splitk_reduce_kernel(const FType* C_split, FType* C,
72
73
int n_mat = N_MATRIX > 0 ? N_MATRIX : (int)n_matrix;
74
for (int i = 0; i < n_mat; ++i) {
75
- sum += ScalarType<FType>::num2float(C_split[idx + i * matrix_size]);
75
+ sum += MarlinScalarType2<FType>::num2float(C_split[idx + i * matrix_size]);
76
}
77
78
- C[idx] = ScalarType<FType>::float2num(sum);
78
+ C[idx] = MarlinScalarType2<FType>::float2num(sum);
79
}
80
81
template <typename FType>
csrc/quantization/gptq_marlin/.gitignore
+2
-1
@@ -1 +1,2 @@
1
-kernel_*.cu
\ No newline at end of file
1
+sm*_kernel_*.cu
2
+kernel_selector.h
csrc/quantization/gptq_marlin/awq_marlin_repack.cu
+70
-34
@@ -4,14 +4,16 @@
4
5
namespace marlin {
6
7
-template <int const num_threads, int const num_bits>
7
+template <int const num_threads, int const num_bits, bool is_a_8bit>
8
__global__ void awq_marlin_repack_kernel(
9
uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr,
10
int size_k, int size_n) {
11
constexpr int pack_factor = 32 / num_bits;
12
13
- int k_tiles = size_k / tile_k_size;
14
- int n_tiles = size_n / tile_n_size;
13
+ constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
14
+ constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
15
+ int k_tiles = size_k / target_tile_k_size;
16
+ int n_tiles = size_n / target_tile_n_size;
17
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
18
19
auto start_k_tile = blockIdx.x * block_k_tiles;
@@ -33,10 +35,10 @@ __global__ void awq_marlin_repack_kernel(
35
36
extern __shared__ int4 sh[];
37
36
- constexpr int tile_n_ints = tile_n_size / pack_factor;
38
+ constexpr int tile_n_ints = target_tile_n_size / pack_factor;
39
40
constexpr int stage_n_threads = tile_n_ints / 4;
39
- constexpr int stage_k_threads = tile_k_size;
41
+ constexpr int stage_k_threads = target_tile_k_size;
42
constexpr int stage_size = stage_k_threads * stage_n_threads;
43
44
auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
@@ -45,7 +47,7 @@ __global__ void awq_marlin_repack_kernel(
47
return;
48
}
49
48
- int first_n = n_tile_id * tile_n_size;
50
+ int first_n = n_tile_id * target_tile_n_size;
51
int first_n_packed = first_n / pack_factor;
52
53
int4* sh_ptr = sh + stage_size * pipe;
@@ -54,7 +56,7 @@ __global__ void awq_marlin_repack_kernel(
56
auto k_id = threadIdx.x / stage_n_threads;
57
auto n_id = threadIdx.x % stage_n_threads;
58
57
- int first_k = k_tile_id * tile_k_size;
59
+ int first_k = k_tile_id * target_tile_k_size;
60
61
cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
62
reinterpret_cast<int4 const*>(
@@ -78,11 +80,11 @@ __global__ void awq_marlin_repack_kernel(
80
}
81
82
int tc_col = th_id / 4;
81
- int tc_row = (th_id % 4) * 2;
83
+ int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
84
85
constexpr int tc_offsets[4] = {0, 1, 8, 9};
86
85
- int cur_n = warp_id * 16 + tc_col;
87
+ int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
88
int cur_n_packed = cur_n / pack_factor;
89
int cur_n_pos = cur_n % pack_factor;
90
@@ -105,23 +107,50 @@ __global__ void awq_marlin_repack_kernel(
107
uint32_t vals[8];
108
#pragma unroll
109
for (int i = 0; i < 4; i++) {
108
- int cur_elem = tc_row + tc_offsets[i];
109
-
110
- int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem];
111
- int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) +
112
- sh_stride * cur_elem];
113
-
114
- vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
115
- vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
110
+ if constexpr (is_a_8bit) {
111
+ int cur_elem = tc_row + i;
112
+
113
+ int packed_src_0 =
114
+ sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) +
115
+ sh_stride * cur_elem];
116
+ int packed_src_1 =
117
+ sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) +
118
+ sh_stride * (cur_elem + 16)];
119
+
120
+ vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
121
+ vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
122
+ } else {
123
+ int cur_elem = tc_row + tc_offsets[i];
124
+
125
+ int packed_src_0 =
126
+ sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem];
127
+ int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) +
128
+ sh_stride * cur_elem];
129
+
130
+ vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask;
131
+ vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask;
132
+ }
133
}
134
118
- constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
135
+ constexpr int tile_size =
136
+ target_tile_k_size * target_tile_n_size / pack_factor;
137
int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
138
139
// Result of:
140
// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
123
- if constexpr (num_bits == 4) {
124
- constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
141
+ if constexpr (!is_a_8bit && num_bits == 4) {
142
+ int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
143
+
144
+ uint32_t res = 0;
145
+#pragma unroll
146
+ for (int i = 0; i < 8; i++) {
147
+ res |= vals[pack_idx[i]] << (i * 4);
148
+ }
149
+
150
+ out_ptr[out_offset + th_id * 4 + warp_id] = res;
151
+
152
+ } else if constexpr (is_a_8bit && num_bits == 4) {
153
+ int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
154
155
uint32_t res = 0;
156
#pragma unroll
@@ -138,8 +167,9 @@ __global__ void awq_marlin_repack_kernel(
167
uint32_t res2 = 0;
168
#pragma unroll
169
for (int i = 0; i < 4; i++) {
141
- res1 |= vals[pack_idx[i]] << (i * 8);
142
- res2 |= vals[4 + pack_idx[i]] << (i * 8);
170
+ const int ii = is_a_8bit ? i : pack_idx[i];
171
+ res1 |= vals[ii] << (i * 8);
172
+ res2 |= vals[4 + ii] << (i * 8);
173
}
174
175
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
@@ -176,18 +206,21 @@ __global__ void awq_marlin_repack_kernel(
206
207
} // namespace marlin
208
179
-#define CALL_IF(NUM_BITS) \
180
- else if (num_bits == NUM_BITS) { \
181
- cudaFuncSetAttribute( \
182
- marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS>, \
183
- cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
184
- marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS> \
185
- <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>( \
186
- b_q_weight_ptr, out_ptr, size_k, size_n); \
209
+#define CALL_IF(NUM_BITS, IS_A_8BIT) \
210
+ else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \
211
+ cudaFuncSetAttribute( \
212
+ marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
213
+ IS_A_8BIT>, \
214
+ cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
215
+ marlin::awq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
216
+ IS_A_8BIT> \
217
+ <<<blocks, marlin::repack_threads, max_shared_mem, stream>>>( \
218
+ b_q_weight_ptr, out_ptr, size_k, size_n); \
219
}
220
221
torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
190
- int64_t size_n, int64_t num_bits) {
222
+ int64_t size_n, int64_t num_bits,
223
+ bool is_a_8bit) {
224
// Verify compatibility with marlin tile of 16x64
225
TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
226
" is not divisible by tile_k_size = ", marlin::tile_k_size);
@@ -238,10 +271,13 @@ torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k,
271
272
if (false) {
273
}
241
- CALL_IF(4)
242
- CALL_IF(8)
274
+ CALL_IF(4, false)
275
+ CALL_IF(8, false)
276
+ CALL_IF(4, true)
277
+ CALL_IF(8, true)
278
else {
244
- TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits);
279
+ TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
280
+ ", is_a_8bit = ", is_a_8bit);
281
}
282
283
return out;
csrc/quantization/gptq_marlin/dequant.h
+87
@@ -470,6 +470,50 @@ __device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
470
frag_b[0] = __hmul2(frag_b[0], bias_reg);
471
}
472
473
+template <>
474
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>(
475
+ int q, __nv_fp8x4_e4m3* frag_b) {
476
+ // Constants for FP4 (E2M1) and FP16 formats
477
+ constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4;
478
+ constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT;
479
+ constexpr int MASK = 0x70707070;
480
+
481
+ // Extract and shift FP4 values to FP16 format
482
+ int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
483
+ q <<= 4;
484
+ int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT);
485
+
486
+ // Note1: reverse indexing is intentional because weights are permuted
487
+ // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of
488
+ // `frag_b[1]` to fit the layout of tensorcore
489
+ frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
490
+ frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
491
+}
492
+
493
+template <>
494
+__device__ inline void dequant<int32_t, vllm::kU4B8.id(), true>(
495
+ int q, int32_t* frag_b) {
496
+ constexpr int repeated_zp = 0x08080808;
497
+ constexpr int MASK = 0x80808080;
498
+
499
+ frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
500
+ q >>= 4;
501
+ frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
502
+}
503
+
504
+template <>
505
+__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>(
506
+ int q, __nv_fp8x4_e4m3* frag_b) {
507
+ int s = q & 0x08080808;
508
+ int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
509
+ q >>= 4;
510
+ s = q & 0x08080808;
511
+ int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3);
512
+
513
+ frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
514
+ frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
515
+}
516
+
517
template <typename scalar_t2, vllm::ScalarTypeId s_type_id>
518
__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
519
@@ -515,6 +559,49 @@ __device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE8M0fnu.id()>(
559
// Note: reverse indexing is intentional because weights are permuted
560
frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
561
frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
562
+};
563
+
564
+// subtract zero point in quanted format and then dequant
565
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
566
+ bool skip_flop = false>
567
+__device__ inline void sub_zp_and_dequant(int q, scalar_t2* frag_b, int zp);
568
+
569
+template <>
570
+__device__ inline void sub_zp_and_dequant<int32_t, vllm::kU4.id(), true>(
571
+ int q, int32_t* frag_b, int zp) {
572
+ // INT4 with zp -> INT8
573
+ // see https://github.com/vllm-project/vllm/pull/24722
574
+ int repeated_zp = 0x01010101 * zp;
575
+ int MASK = 0x80808080;
576
+
577
+ frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
578
+ q >>= 4;
579
+ frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK;
580
+}
581
+
582
+template <>
583
+__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(),
584
+ true>(int q, __nv_fp8x4_e4m3* frag_b,
585
+ int zp) {
586
+ // INT4 with zp -> FP8
587
+ // see https://github.com/vllm-project/vllm/pull/24722
588
+ uint32_t u_q = *reinterpret_cast<uint32_t*>(&q);
589
+ uint32_t u_zp = *reinterpret_cast<uint32_t*>(&zp);
590
+ uint32_t u_zp1 = u_zp + 1;
591
+ uint32_t repeated_zp = 0x01010101 * u_zp;
592
+
593
+ uint32_t q0, s;
594
+ q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
595
+ s = (q0 + repeated_zp) & 0x80808080;
596
+ uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
597
+
598
+ u_q >>= 4;
599
+ q0 = (u_q & 0x0F0F0F0F) | 0x70707070;
600
+ s = (q0 + repeated_zp) & 0x80808080;
601
+ uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s;
602
+
603
+ frag_b[0] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out1);
604
+ frag_b[1] = *reinterpret_cast<const __nv_fp8x4_e4m3*>(&Out2);
605
}
606
607
#endif
csrc/quantization/gptq_marlin/generate_kernels.py
+242
-91
@@ -4,141 +4,292 @@ import glob
4
import itertools
5
import os
6
import subprocess
7
+import sys
8
9
import jinja2
10
10
-FILE_HEAD = """
11
-// auto generated by generate.py
11
+ARCHS = []
12
+SUPPORT_FP8 = False
13
+for arch in sys.argv[1].split(","):
14
+ arch = arch[: arch.index(".") + 2].replace(".", "")
15
+ arch = int(arch)
16
+ # only SM89 and SM120 fully support
17
+ # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
18
+ # SM90 and SM100 can use this PTX, but it’s simulated
19
+ # with FP16 MMA, so it cannot achieve any acceleration.
20
+ if arch in [89, 120]:
21
+ SUPPORT_FP8 = True
22
+
23
+FILE_HEAD_COMMENT = """
24
+// auto generated by generate_kernels.py
25
// clang-format off
26
+""".lstrip()
27
28
+FILE_HEAD = (
29
+ FILE_HEAD_COMMENT
30
+ + """
31
#include "kernel.h"
32
#include "marlin_template.h"
33
34
namespace MARLIN_NAMESPACE_NAME {
18
-""".strip()
35
+"""
36
+)
37
38
TEMPLATE = (
39
"template __global__ void Marlin<"
22
- "{{scalar_t}}, "
23
- "{{w_type_id}}, "
40
+ "{{a_type_id}}, "
41
+ "{{b_type_id}}, "
42
+ "{{c_type_id}}, "
43
"{{s_type_id}}, "
44
"{{threads}}, "
45
"{{thread_m_blocks}}, "
46
"{{thread_n_blocks}}, "
47
"{{thread_k_blocks}}, "
29
- "{{'true' if m_block_size_8 else 'false'}}, "
48
+ "{{m_block_size_8}}, "
49
"{{stages}}, "
50
"{{group_blocks}}, "
32
- "{{'true' if is_zp_float else 'false'}}>"
51
+ "{{is_zp_float}}>"
52
"( MARLIN_KERNEL_PARAMS );"
53
)
54
36
-# int8 with zero point case (vllm::kU8) is also supported,
37
-# we don't add it to reduce wheel size.
38
-SCALAR_TYPES = [
39
- "vllm::kU4",
40
- "vllm::kU4B8",
41
- "vllm::kU8B128",
42
- "vllm::kFE4M3fn",
43
- "vllm::kFE2M1f",
44
-]
55
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]
56
57
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
48
-# group_blocks:
49
-# = 0 : act order case
50
-# = -1 : channelwise quantization
51
-# > 0 : group_size=16*group_blocks
52
-GROUP_BLOCKS = [0, 1, -1, 2, 4, 8]
53
-DTYPES = ["fp16", "bf16"]
58
+
59
+QUANT_CONFIGS = [
60
+ # AWQ-INT4
61
+ {
62
+ "b_type": "kU4",
63
+ "thread_configs": THREAD_CONFIGS,
64
+ "thread_m_blocks": THREAD_M_BLOCKS,
65
+ "group_blocks": [-1, 2, 4, 8],
66
+ },
67
+ # HQQ
68
+ {
69
+ "a_type": ["kFloat16"],
70
+ "b_type": "kU4",
71
+ "thread_configs": THREAD_CONFIGS,
72
+ "thread_m_blocks": THREAD_M_BLOCKS,
73
+ "group_blocks": [4],
74
+ "is_zp_float": True,
75
+ },
76
+ # GPTQ-INT4
77
+ {
78
+ "b_type": "kU4B8",
79
+ "thread_configs": THREAD_CONFIGS,
80
+ "thread_m_blocks": THREAD_M_BLOCKS,
81
+ "group_blocks": [-1, 0, 2, 4, 8],
82
+ },
83
+ # GPTQ-INT8
84
+ {
85
+ "b_type": "kU8B128",
86
+ "thread_configs": THREAD_CONFIGS,
87
+ "thread_m_blocks": THREAD_M_BLOCKS,
88
+ "group_blocks": [-1, 0, 2, 4, 8],
89
+ },
90
+ # FP8
91
+ {
92
+ "b_type": "kFE4M3fn",
93
+ "thread_configs": THREAD_CONFIGS,
94
+ "thread_m_blocks": THREAD_M_BLOCKS,
95
+ "group_blocks": [-1, 8],
96
+ },
97
+ # NVFP4
98
+ {
99
+ "b_type": "kFE2M1f",
100
+ "s_type": "kFE4M3fn",
101
+ "thread_configs": THREAD_CONFIGS,
102
+ "thread_m_blocks": THREAD_M_BLOCKS,
103
+ "group_blocks": [1],
104
+ },
105
+ # MXFP4
106
+ {
107
+ "a_type": ["kBFloat16"],
108
+ "b_type": "kFE2M1f",
109
+ "s_type": "kFE8M0fnu",
110
+ "thread_configs": THREAD_CONFIGS,
111
+ "thread_m_blocks": THREAD_M_BLOCKS,
112
+ "group_blocks": [2],
113
+ },
114
+ # AWQ-INT4 with INT8 activation
115
+ {
116
+ "a_type": ["kS8"],
117
+ "b_type": "kU4",
118
+ "thread_configs": THREAD_CONFIGS,
119
+ "thread_m_blocks": [1, 2, 3, 4],
120
+ "group_blocks": [-1, 2, 4, 8],
121
+ },
122
+ # GPTQ-INT4 with INT8 activation
123
+ {
124
+ "a_type": ["kS8"],
125
+ "b_type": "kU4B8",
126
+ "thread_configs": THREAD_CONFIGS,
127
+ "thread_m_blocks": [1, 2, 3, 4],
128
+ "group_blocks": [-1, 2, 4, 8],
129
+ },
130
+ # GPTQ-INT4 with FP8 activation
131
+ {
132
+ "a_type": ["kFE4M3fn"],
133
+ "b_type": "kU4B8",
134
+ "thread_configs": THREAD_CONFIGS,
135
+ "thread_m_blocks": [1, 2, 3, 4],
136
+ "group_blocks": [-1, 2, 4, 8],
137
+ },
138
+ # AWQ-INT4 with FP8 activation
139
+ {
140
+ "a_type": ["kFE4M3fn"],
141
+ "b_type": "kU4",
142
+ "thread_configs": THREAD_CONFIGS,
143
+ "thread_m_blocks": [1, 2, 3, 4],
144
+ "group_blocks": [-1, 2, 4, 8],
145
+ },
146
+ # MXFP4 with FP8 activation
147
+ {
148
+ "a_type": ["kFE4M3fn"],
149
+ "b_type": "kFE2M1f",
150
+ "c_type": ["kBFloat16"],
151
+ "s_type": "kFE8M0fnu",
152
+ "thread_configs": THREAD_CONFIGS,
153
+ "thread_m_blocks": [1, 2, 3, 4],
154
+ "group_blocks": [2],
155
+ },
156
+]
157
158
159
def remove_old_kernels():
57
- for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
160
+ for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
161
subprocess.call(["rm", "-f", filename])
162
163
+ filename = os.path.dirname(__file__) + "/kernel_selector.h"
164
+ subprocess.call(["rm", "-f", filename])
165
+
166
167
def generate_new_kernels():
62
- for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
63
- all_template_str_list = []
168
+ result_dict = {}
169
65
- for group_blocks, m_blocks, thread_configs in itertools.product(
66
- GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS
67
- ):
68
- # act order case only support gptq-int4 and gptq-int8
69
- if group_blocks == 0 and scalar_type not in [
70
- "vllm::kU4B8",
71
- "vllm::kU8B128",
72
- ]:
73
- continue
74
- if thread_configs[2] == 256:
75
- # for small batch (m_blocks == 1), we only need (128, 128, 256)
76
- # for large batch (m_blocks > 1), we only need (64, 256, 256)
77
- if m_blocks <= 1 and thread_configs[0] != 128:
78
- continue
79
- if m_blocks > 1 and thread_configs[0] != 64:
80
- continue
81
-
82
- # we only support channelwise quantization and group_size == 128
83
- # for fp8
84
- if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
85
- continue
86
- # nvfp4 only supports group_size == 16
87
- # mxfp4 only supports group_size == 32
88
- if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
170
+ for quant_config in QUANT_CONFIGS:
171
+ c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
172
+ a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
173
+ b_type = quant_config["b_type"]
174
+ is_zp_float = quant_config.get("is_zp_float", False)
175
+ all_group_blocks = quant_config["group_blocks"]
176
+ all_m_blocks = quant_config["thread_m_blocks"]
177
+ all_thread_configs = quant_config["thread_configs"]
178
+
179
+ for a_type, c_type in itertools.product(a_types, c_types):
180
+ if not SUPPORT_FP8 and a_type == "kFE4M3fn":
181
continue
90
- # other quantization methods don't support group_size = 16
91
- if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
182
+ if "16" in a_type and "16" in c_type and a_type != c_type:
183
continue
184
+ s_type = quant_config.get("s_type", c_type)
185
+ if (a_type, b_type, c_type) not in result_dict:
186
+ result_dict[(a_type, b_type, c_type)] = []
187
94
- k_blocks = thread_configs[0] // 16
95
- n_blocks = thread_configs[1] // 16
96
- threads = thread_configs[2]
97
-
98
- c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
99
-
100
- is_zp_float_list = [False]
101
- if dtype == "fp16" and scalar_type == "vllm::kU4" and group_blocks == 4:
102
- # HQQ (is_zp_float = true) only supports
103
- # 4bit quantization and fp16
104
- is_zp_float_list.append(True)
105
-
106
- if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
107
- s_type = "vllm::kFE4M3fn"
108
- elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
109
- s_type = "vllm::kFE8M0fnu"
110
- if dtype == "fp16":
111
- # we cannot safely dequantize e8m0 to fp16, so skip this
112
- continue
113
- elif dtype == "fp16":
114
- s_type = "vllm::kFloat16"
115
- elif dtype == "bf16":
116
- s_type = "vllm::kBFloat16"
117
-
118
- for is_zp_float in is_zp_float_list:
119
- template_str = jinja2.Template(TEMPLATE).render(
120
- scalar_t=c_dtype,
121
- w_type_id=scalar_type + ".id()",
122
- s_type_id=s_type + ".id()",
123
- threads=threads,
124
- thread_m_blocks=max(m_blocks, 1),
125
- thread_n_blocks=n_blocks,
126
- thread_k_blocks=k_blocks,
127
- m_block_size_8=m_blocks == 0.5,
128
- stages="pipe_stages",
129
- group_blocks=group_blocks,
130
- is_zp_float=is_zp_float,
131
- )
188
+ for group_blocks, m_blocks, thread_configs in itertools.product(
189
+ all_group_blocks, all_m_blocks, all_thread_configs
190
+ ):
191
+ thread_k, thread_n, threads = thread_configs
192
+
193
+ if threads == 256:
194
+ # for small batch (m_blocks == 1),
195
+ # we only need (128, 128, 256)
196
+ # for large batch (m_blocks > 1),
197
+ # we only need (64, 256, 256)
198
+ if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
199
+ continue
200
+ if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
201
+ continue
202
133
- all_template_str_list.append(template_str)
203
+ config = {
204
+ "threads": threads,
205
+ "s_type": s_type,
206
+ "thread_m_blocks": max(m_blocks, 1),
207
+ "thread_k_blocks": thread_k // 16,
208
+ "thread_n_blocks": thread_n // 16,
209
+ "m_block_size_8": "true" if m_blocks == 0.5 else "false",
210
+ "stages": "pipe_stages",
211
+ "group_blocks": group_blocks,
212
+ "is_zp_float": "true" if is_zp_float else "false",
213
+ }
214
+
215
+ result_dict[(a_type, b_type, c_type)].append(config)
216
+
217
+ kernel_selector_str = FILE_HEAD_COMMENT
218
+
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(
262
+ a_type_id=f"vllm::{a_type}.id()",
263
+ b_type_id=f"vllm::{b_type}.id()",
264
+ c_type_id=f"vllm::{c_type}.id()",
265
+ s_type_id=f"vllm::{s_type}.id()",
266
+ **config,
267
+ )
268
+ + "\n"
269
+ )
270
271
file_content = FILE_HEAD + "\n\n"
272
file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
137
- filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"
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"
277
+
278
+ filename = filename.lower()
279
280
with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
281
f.write(file_content)
282
283
+ if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
284
+ kernel_selector_str += (
285
+ "else if (a_type == vllm::kFE4M3fn)\n"
286
+ " TORCH_CHECK(false, "
287
+ '"marlin kernel with fp8 activation is not built.");'
288
+ )
289
+
290
+ with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
291
+ f.write(kernel_selector_str)
292
+
293
294
if __name__ == "__main__":
295
remove_old_kernels()
csrc/quantization/gptq_marlin/gptq_marlin.cu
+179
-294
@@ -53,7 +53,7 @@ torch::Tensor gptq_marlin_gemm(
53
std::optional<torch::Tensor> const& b_zeros_or_none,
54
std::optional<torch::Tensor> const& g_idx_or_none,
55
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
56
- vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
56
+ vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
57
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
58
bool is_zp_float) {
59
TORCH_CHECK_NOT_IMPLEMENTED(false,
@@ -243,204 +243,29 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
243
int cache_size = get_kernel_cache_size(
244
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);
246
- return cache_size + 512 <= max_shared_mem;
246
+ return cache_size <= max_shared_mem;
247
}
248
249
- #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
250
- M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
251
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
252
- thread_n_blocks == THREAD_N_BLOCKS && \
253
- thread_k_blocks == THREAD_K_BLOCKS && \
254
- m_block_size_8 == M_BLOCK_SIZE_8 && \
255
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
256
- is_zp_float == IS_ZP_FLOAT) { \
257
- constexpr auto S_TYPE = \
258
- W_TYPE == vllm::kFE2M1f \
259
- ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu) \
260
- : (std::is_same<scalar_t, half>::value ? vllm::kFloat16 \
261
- : vllm::kBFloat16); \
262
- kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS, \
263
- THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
264
- M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
265
- }
266
-
267
- // COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
268
- // this is the most common cases
269
- // BIGGROUP: cases for big group size (group_blocks in [-1, 8])
270
- // FZP: cases for float-zero-point (is_zp_float = true)
271
- // ACT: cases for act order case (group_blocks == 0)
272
- // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
273
- #define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
274
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
275
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
276
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \
277
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
278
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
279
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
280
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
281
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
282
-
283
- #define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
284
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
285
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
286
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
287
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
288
- \
289
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
290
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
291
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
292
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
293
- \
294
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
295
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
296
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \
297
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
298
-
299
- #define COMMON_GET_IF(W_TYPE) \
300
- COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \
301
- COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \
302
- COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \
303
- COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \
304
- COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \
305
- COMMON_GET_IF_M234(W_TYPE, 4, 8, 128)
306
-
307
- #define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
308
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
309
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \
310
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
311
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
312
-
313
- #define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
314
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
315
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
316
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
317
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
318
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
319
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
320
-
321
- #define BIGGROUP_GET_IF(W_TYPE) \
322
- BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
323
- BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
324
- BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \
325
- BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
326
- BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \
327
- BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
328
-
329
- #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
330
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
331
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
332
-
333
- #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
334
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
335
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
336
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
337
-
338
- #define NVFP4_GET_IF(W_TYPE) \
339
- NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
340
- NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
341
- NVFP4_GET_IF_M1(W_TYPE, 4, 8, 128) \
342
- NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
343
- NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128) \
344
- NVFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
345
-
346
- #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
347
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
348
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
349
-
350
- #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
351
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
352
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
353
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
354
-
355
- #define MXFP4_GET_IF(W_TYPE) \
356
- MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
357
- MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
358
- MXFP4_GET_IF_M1(W_TYPE, 4, 8, 128) \
359
- MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
360
- MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128) \
361
- MXFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
362
-
363
- // We currently have 4-bit models only with group_blocks == 4
364
- #define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
365
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
366
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
367
-
368
- #define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
369
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
370
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \
371
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true)
372
-
373
- #define FZP_GET_IF(W_TYPE) \
374
- FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \
375
- FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \
376
- FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \
377
- FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \
378
- FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \
379
- FZP_GET_IF_M234(W_TYPE, 4, 8, 128)
380
-
381
- // We currently have 4-bit models only with group_blocks == 4
382
- #define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
383
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \
384
- _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
385
-
386
- #define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
387
- _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
388
- _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \
389
- _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false)
390
-
391
- #define ACT_GET_IF(W_TYPE) \
392
- ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \
393
- ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \
394
- ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \
395
- ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \
396
- ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \
397
- ACT_GET_IF_M234(W_TYPE, 4, 8, 128)
398
-
399
-template <typename scalar_t>
400
-MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
401
- int thread_m_blocks, int thread_n_blocks,
402
- int thread_k_blocks, bool m_block_size_8,
403
- bool has_act_order, bool has_zp,
404
- int group_blocks, int num_threads,
405
- bool is_zp_float) {
406
- int num_bits = q_type.size_bits();
249
+MarlinFuncPtr get_marlin_kernel(
250
+ const vllm::ScalarType a_type, const vllm::ScalarType b_type,
251
+ const vllm::ScalarType c_type, const vllm::ScalarType s_type,
252
+ int thread_m_blocks, int thread_n_blocks, int thread_k_blocks,
253
+ bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks,
254
+ int threads, bool is_zp_float) {
255
+ int num_bits = b_type.size_bits();
256
auto kernel = MarlinDefault;
408
- if (false) {
409
- }
410
-
411
- COMMON_GET_IF(vllm::kU4)
412
- COMMON_GET_IF(vllm::kU4B8)
413
- COMMON_GET_IF(vllm::kU8B128)
257
415
- NVFP4_GET_IF(vllm::kFE2M1f)
416
-
417
- BIGGROUP_GET_IF(vllm::kFE4M3fn)
418
-
419
- ACT_GET_IF(vllm::kU4B8)
420
- ACT_GET_IF(vllm::kU8B128)
421
-
422
- if (std::is_same<scalar_t, half>::value) {
423
- if (false) {
424
- }
425
- FZP_GET_IF(vllm::kU4)
426
- }
427
- if (std::is_same<scalar_t, nv_bfloat16>::value) {
428
- if (false) {
429
- }
430
- MXFP4_GET_IF(vllm::kFE2M1f)
431
- }
258
+ #include "kernel_selector.h"
259
260
return kernel;
261
}
262
436
-template <typename scalar_t>
437
-exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
438
- int prob_n, int prob_k, int thread_m_blocks,
439
- bool m_block_size_8, int num_bits,
440
- int group_size, bool has_act_order,
441
- bool is_k_full, bool has_zp,
442
- bool is_zp_float, int max_shared_mem,
443
- int sms) {
263
+exec_config_t determine_exec_config(
264
+ const vllm::ScalarType& a_type, const vllm::ScalarType& b_type,
265
+ const vllm::ScalarType& c_type, const vllm::ScalarType& s_type, int prob_m,
266
+ int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8,
267
+ 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) {
269
exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
270
thread_config_t* thread_configs = thread_m_blocks > 1
271
? large_batch_thread_configs
@@ -455,7 +280,7 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
280
281
if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
282
num_bits, group_size, has_act_order, is_k_full, has_zp,
458
- is_zp_float, max_shared_mem)) {
283
+ is_zp_float, max_shared_mem - 512)) {
284
continue;
285
}
286
@@ -468,10 +293,11 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
293
group_blocks = group_size == -1 ? -1 : group_size / 16;
294
}
295
471
- auto kernel = get_marlin_kernel<scalar_t>(
472
- q_type, thread_m_blocks, th_config.thread_n / 16,
473
- th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
474
- group_blocks, th_config.num_threads, is_zp_float);
296
+ auto kernel =
297
+ get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks,
298
+ th_config.thread_n / 16, th_config.thread_k / 16,
299
+ m_block_size_8, has_act_order, has_zp, group_blocks,
300
+ th_config.num_threads, is_zp_float);
301
302
if (kernel == MarlinDefault) continue;
303
@@ -485,28 +311,16 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
311
return exec_cfg;
312
}
313
488
-template <typename scalar_t>
314
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
490
- void* s, void* s2, void* zp, void* g_idx, void* perm,
491
- void* a_tmp, int prob_m, int prob_n, int prob_k, int lda,
492
- void* workspace, vllm::ScalarType const& q_type, bool has_bias,
315
+ void* a_s, void* b_s, void* g_s, void* zp, void* g_idx,
316
+ void* perm, void* a_tmp, int prob_m, int prob_n, int prob_k,
317
+ int lda, void* workspace, vllm::ScalarType const& a_type,
318
+ vllm::ScalarType const& b_type, vllm::ScalarType const& c_type,
319
+ vllm::ScalarType const& s_type, bool has_bias,
320
bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
321
int group_size, int dev, cudaStream_t stream, int thread_k_init,
322
int thread_n_init, int sms, bool use_atomic_add,
323
bool use_fp32_reduce, bool is_zp_float) {
497
- if (has_zp) {
498
- TORCH_CHECK(
499
- q_type == vllm::kU4 || q_type == vllm::kU8,
500
- "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
501
- } else {
502
- TORCH_CHECK(
503
- q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
504
- q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
505
- "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
506
- "has_zp = False. Got = ",
507
- q_type.str());
508
- }
509
-
324
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
325
", ", prob_n, ", ", prob_k, "]");
326
@@ -531,19 +345,21 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
345
}
346
}
347
534
- int num_bits = q_type.size_bits();
348
+ int num_bits = b_type.size_bits();
349
const int4* A_ptr = (const int4*)A;
350
const int4* B_ptr = (const int4*)B;
351
int4* C_ptr = (int4*)C;
352
int4* C_tmp_ptr = (int4*)C_tmp;
353
+
354
const int4* bias_ptr = (const int4*)b_bias;
540
- const int4* s_ptr = (const int4*)s;
541
- const uint16_t* s2_ptr = (const uint16_t*)s2;
355
+ const float* a_s_ptr = (const float*)a_s;
356
+ const int4* b_s_ptr = (const int4*)b_s;
357
+ const uint16_t* g_s_ptr = (const uint16_t*)g_s;
358
+
359
const int4* zp_ptr = (const int4*)zp;
360
const int* g_idx_ptr = (const int*)g_idx;
361
const int* perm_ptr = (const int*)perm;
362
int4* a_tmp_ptr = (int4*)a_tmp;
546
-
363
int* locks = (int*)workspace;
364
365
if (has_act_order) {
@@ -568,6 +384,21 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
384
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
385
TORCH_CHECK(max_shared_mem > 0);
386
387
+ int major_capability, minor_capability;
388
+ cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
389
+ dev);
390
+ cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
391
+ dev);
392
+ TORCH_CHECK(major_capability * 10 + minor_capability >= 80,
393
+ "marlin kernel only support Ampere or newer GPUs.");
394
+ if (a_type == vllm::kFE4M3fn) {
395
+ TORCH_CHECK(
396
+ major_capability * 10 + minor_capability == 89 ||
397
+ major_capability * 10 + minor_capability == 120,
398
+ "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
399
+ "Marlin W4A16 on other devices).");
400
+ }
401
+
402
int max_par = 16;
403
if (prob_n <= 4096) max_par = 16 * 8;
404
int max_shared_mem_new = max_shared_mem;
@@ -583,7 +414,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
414
int thread_n = thread_n_init;
415
416
int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks);
586
- int m_block_size_8 = prob_m_split <= 8;
417
+ int m_block_size_8 = prob_m_split <= 8 && a_type.size_bits() == 16;
418
419
// Set thread config
420
exec_config_t exec_cfg;
@@ -597,11 +428,25 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
428
" is not divisible by thread_k = ", thread_k);
429
} else {
430
// Auto config
600
- exec_cfg = determine_exec_config<scalar_t>(
601
- q_type, prob_m_split, prob_n, prob_k, thread_m_blocks, m_block_size_8,
602
- num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
603
- max_shared_mem, sms);
431
+ exec_cfg = determine_exec_config(
432
+ a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k,
433
+ 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);
435
thread_tfg = exec_cfg.tb_cfg;
436
+ if (thread_tfg.thread_n != -1) {
437
+ if (prob_n / thread_tfg.thread_n *
438
+ div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <=
439
+ sms) {
440
+ if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split,
441
+ prob_n, prob_k, num_bits, group_size,
442
+ has_act_order, is_k_full, has_zp, is_zp_float,
443
+ max_shared_mem_new)) {
444
+ thread_tfg = {128, 64, 128};
445
+ exec_cfg = {1, thread_tfg};
446
+ }
447
+ }
448
+ }
449
+
450
if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) {
451
max_thread_m_blocks--;
452
continue;
@@ -632,10 +477,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
477
", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
478
", max_shared_mem_new = ", max_shared_mem_new);
479
635
- auto kernel = get_marlin_kernel<scalar_t>(
636
- q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks,
637
- m_block_size_8, has_act_order, has_zp, group_blocks, num_threads,
638
- is_zp_float);
480
+ auto kernel = get_marlin_kernel(
481
+ a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks,
482
+ thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks,
483
+ num_threads, is_zp_float);
484
485
if (kernel == MarlinDefault) {
486
TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
@@ -657,13 +502,15 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
502
// avoid ">>>" being formatted to "> > >"
503
// clang-format off
504
kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
660
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr,
505
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr,
506
g_idx_ptr, num_groups,
507
prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add,
508
use_fp32_reduce, max_shared_mem_new);
509
// clang-format on
510
666
- A_ptr += prob_m_split * (lda / 8);
511
+ bool is_a_8bit = a_type.size_bits() == 8;
512
+ A_ptr += prob_m_split * (lda / (is_a_8bit ? 16 : 8));
513
+ a_s_ptr += prob_m_split;
514
C_ptr += prob_m_split * (prob_n / 8);
515
rest_m -= prob_m_split;
516
}
@@ -675,15 +522,73 @@ torch::Tensor gptq_marlin_gemm(
522
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
523
torch::Tensor& b_q_weight,
524
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
525
+ std::optional<torch::Tensor> const& a_scales_or_none,
526
std::optional<torch::Tensor> const& global_scale_or_none,
527
std::optional<torch::Tensor> const& b_zeros_or_none,
528
std::optional<torch::Tensor> const& g_idx_or_none,
529
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
682
- vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
530
+ vllm::ScalarTypeId const& b_type_id, int64_t size_m, int64_t size_n,
531
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
532
bool is_zp_float) {
685
- vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
686
- int pack_factor = 32 / b_q_type.size_bits();
533
+ vllm::ScalarTypeId a_type_id, c_type_id, s_type_id;
534
+
535
+ auto c_dtype = a.dtype();
536
+ if (a.scalar_type() == at::ScalarType::Half) {
537
+ a_type_id = vllm::kFloat16.id();
538
+ c_type_id = vllm::kFloat16.id();
539
+ } else if (a.scalar_type() == at::ScalarType::BFloat16) {
540
+ a_type_id = vllm::kBFloat16.id();
541
+ c_type_id = vllm::kBFloat16.id();
542
+ } else {
543
+ c_dtype = b_scales.dtype();
544
+ if (b_scales.scalar_type() == at::ScalarType::Half) {
545
+ c_type_id = vllm::kFloat16.id();
546
+ } else if (b_scales.scalar_type() == at::ScalarType::BFloat16) {
547
+ c_type_id = vllm::kBFloat16.id();
548
+ } else {
549
+ c_type_id = vllm::kBFloat16.id();
550
+
551
+ TORCH_CHECK(c_or_none.has_value(), "c must be passed for W4A8-FP4");
552
+ torch::Tensor c = c_or_none.value();
553
+ c_dtype = c.dtype();
554
+
555
+ if (c.scalar_type() == at::ScalarType::Half) {
556
+ c_type_id = vllm::kFloat16.id();
557
+ } else if (c.scalar_type() == at::ScalarType::BFloat16) {
558
+ c_type_id = vllm::kBFloat16.id();
559
+ } else {
560
+ TORCH_CHECK(false, "unsupported c dtype");
561
+ }
562
+ }
563
+
564
+ if (a.scalar_type() == at::ScalarType::Float8_e4m3fn) {
565
+ a_type_id = vllm::kFE4M3fn.id();
566
+ } else if (a.scalar_type() == at::ScalarType::Char) {
567
+ a_type_id = vllm::kS8.id();
568
+ } else {
569
+ TORCH_CHECK(false, "unsupported `a` scalar_type");
570
+ }
571
+ }
572
+
573
+ s_type_id = c_type_id;
574
+ if (b_type_id == vllm::kFE2M1f.id()) {
575
+ if (b_scales.scalar_type() == at::ScalarType::Float8_e4m3fn) {
576
+ s_type_id = vllm::kFE4M3fn.id();
577
+ } else if (b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
578
+ s_type_id = vllm::kFE8M0fnu.id();
579
+ } else {
580
+ TORCH_CHECK(false,
581
+ "When b_type = float4_e2m1f, b_scale scalar type must be",
582
+ "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
583
+ }
584
+ }
585
+
586
+ vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
587
+ vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id);
588
+ vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id);
589
+ vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id);
590
+
591
+ int pack_factor = 32 / b_type.size_bits();
592
593
// Verify A
594
TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
@@ -721,6 +626,21 @@ torch::Tensor gptq_marlin_gemm(
626
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
627
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
628
629
+ torch::Tensor a_scales;
630
+ auto options = torch::TensorOptions().dtype(c_dtype).device(a.device());
631
+ auto options_fp32 =
632
+ torch::TensorOptions().dtype(at::kFloat).device(a.device());
633
+
634
+ if (a_scales_or_none.has_value()) {
635
+ a_scales = a_scales_or_none.value();
636
+ TORCH_CHECK(a_type.size_bits() == 8,
637
+ "a_scales can only be used for 8bit activation.");
638
+ } else {
639
+ a_scales = torch::empty({0}, options_fp32);
640
+ TORCH_CHECK(a_type.size_bits() != 8,
641
+ "the a_scales parameter must be passed for 8bit activation.");
642
+ }
643
+
644
// thread_k: `k` size of a thread_tile in `weights` (can usually be left as
645
// auto -1)
646
int thread_k = -1;
@@ -733,7 +653,6 @@ torch::Tensor gptq_marlin_gemm(
653
654
// Alloc buffers
655
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
736
- auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
656
torch::Tensor c;
657
if (c_or_none.has_value()) {
658
c = c_or_none.value();
@@ -750,8 +669,6 @@ torch::Tensor gptq_marlin_gemm(
669
670
// Alloc C tmp buffer that is going to be used for the global reduce
671
torch::Tensor c_tmp;
753
- auto options_fp32 =
754
- torch::TensorOptions().dtype(at::kFloat).device(a.device());
672
if (use_fp32_reduce) {
673
int max_m_block_size = (size_m + 16 - 1) / 16 * 16;
674
max_m_block_size = min(max_m_block_size, 64);
@@ -821,11 +738,11 @@ torch::Tensor gptq_marlin_gemm(
738
torch::Tensor global_scale;
739
if (global_scale_or_none.has_value()) {
740
global_scale = global_scale_or_none.value();
824
- TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
741
+ TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
742
"global_scale can only be used for nvfp4 format.");
743
} else {
744
global_scale = torch::empty({0}, options);
828
- TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
745
+ TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
746
"the global_scale parameter must be passed for nvfp4 format.");
747
}
748
@@ -852,15 +769,15 @@ torch::Tensor gptq_marlin_gemm(
769
bool has_zp = b_zeros.size(-1) > 0;
770
if (has_zp) {
771
TORCH_CHECK(
855
- b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
856
- "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
772
+ b_type == vllm::kU4 || b_type == vllm::kU8,
773
+ "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str());
774
} else {
858
- TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
859
- b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
860
- "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
861
- "float4_e2m1f when "
862
- "has_zp = False. Got = ",
863
- b_q_type.str());
775
+ TORCH_CHECK(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 ||
776
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
777
+ b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f,
778
+ "b_type must be uint4b8, uint8b128, int4, int8, "
779
+ "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ",
780
+ b_type.str());
781
}
782
783
if (has_zp && is_zp_float) {
@@ -902,59 +819,27 @@ torch::Tensor gptq_marlin_gemm(
819
" is below min_workspace_size = ", min_workspace_size);
820
821
int dev = a.get_device();
905
- if (a.scalar_type() == at::ScalarType::Half) {
906
- void* scales_ptr;
907
- if (b_q_type == vllm::kFE2M1f) {
908
- if (group_size == 16)
909
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
910
- else if (group_size == 32)
911
- scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
912
- else
913
- TORCH_CHECK(false,
914
- "float4_e2m1f only supports group_size == 16 (NVFP4) ",
915
- "and group_size == 32 (MXFP4)");
916
- } else {
917
- scales_ptr = b_scales.data_ptr<at::Half>();
918
- }
822
920
- marlin::marlin_mm<half>(
921
- a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
922
- c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
923
- global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
924
- perm.data_ptr(), a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k,
925
- a.stride(0), workspace.data_ptr(), b_q_type, has_bias, has_act_order,
926
- is_k_full, has_zp, num_groups, group_size, dev,
927
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
928
- use_atomic_add, use_fp32_reduce, is_zp_float);
929
- } else if (a.scalar_type() == at::ScalarType::BFloat16) {
930
- void* scales_ptr;
931
- if (b_q_type == vllm::kFE2M1f) {
932
- if (group_size == 16)
933
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
934
- else if (group_size == 32)
935
- scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
936
- else
937
- TORCH_CHECK(false,
938
- "float4_e2m1f only supports group_size == 16 (NVFP4) ",
939
- "and group_size == 32 (MXFP4)");
940
- } else {
941
- scales_ptr = b_scales.data_ptr<at::BFloat16>();
942
- }
943
-
944
- marlin::marlin_mm<nv_bfloat16>(
945
- a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
946
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
947
- b_bias.data_ptr<at::BFloat16>(), scales_ptr,
948
- global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
949
- g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
950
- size_m, size_n, size_k, a.stride(0), workspace.data_ptr(), b_q_type,
951
- has_bias, has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
952
- at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
953
- use_atomic_add, use_fp32_reduce, is_zp_float);
954
- } else {
955
- TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16");
823
+ TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
824
+ "scalar type of a_scales must be float");
825
+ TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
826
+ "scalar type of global_scale must be the same with c");
827
+ if (a_type.size_bits() == 16) {
828
+ TORCH_CHECK(
829
+ a.scalar_type() == c.scalar_type(),
830
+ "scalar type of a must be the same with c for 16 bit activation");
831
}
832
833
+ marlin::marlin_mm(
834
+ a.data_ptr(), b_q_weight.data_ptr(), c.data_ptr(), c_tmp.data_ptr(),
835
+ b_bias.data_ptr(), a_scales.data_ptr(), b_scales.data_ptr(),
836
+ global_scale.data_ptr(), b_zeros.data_ptr(), g_idx.data_ptr(),
837
+ perm.data_ptr(), a_tmp.data_ptr(), size_m, size_n, size_k, a.stride(0),
838
+ workspace.data_ptr(), a_type, b_type, c_type, s_type, has_bias,
839
+ has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
840
+ at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
841
+ use_atomic_add, use_fp32_reduce, is_zp_float);
842
+
843
return c;
844
}
845
csrc/quantization/gptq_marlin/gptq_marlin_repack.cu
+63
-32
@@ -4,15 +4,18 @@
4
5
namespace marlin {
6
7
-template <int const num_threads, int const num_bits, bool const has_perm>
7
+template <int const num_threads, int const num_bits, bool const has_perm,
8
+ bool is_a_8bit>
9
__global__ void gptq_marlin_repack_kernel(
10
uint32_t const* __restrict__ b_q_weight_ptr,
11
uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
12
int size_k, int size_n) {
13
constexpr int pack_factor = 32 / num_bits;
14
14
- int k_tiles = size_k / tile_k_size;
15
- int n_tiles = size_n / tile_n_size;
15
+ constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1);
16
+ constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1);
17
+ int k_tiles = size_k / target_tile_k_size;
18
+ int n_tiles = size_n / target_tile_n_size;
19
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
20
21
auto start_k_tile = blockIdx.x * block_k_tiles;
@@ -34,7 +37,7 @@ __global__ void gptq_marlin_repack_kernel(
37
38
extern __shared__ int4 sh[];
39
37
- constexpr int perm_size = tile_k_size / 4;
40
+ constexpr int perm_size = target_tile_k_size / 4;
41
42
int4* sh_perm_ptr = sh;
43
int4* sh_pipe_ptr = sh_perm_ptr;
@@ -42,14 +45,14 @@ __global__ void gptq_marlin_repack_kernel(
45
sh_pipe_ptr += perm_size;
46
}
47
45
- constexpr int tile_ints = tile_k_size / pack_factor;
48
+ constexpr int tile_ints = target_tile_k_size / pack_factor;
49
47
- constexpr int stage_n_threads = tile_n_size / 4;
48
- constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
50
+ constexpr int stage_n_threads = target_tile_n_size / 4;
51
+ constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints;
52
constexpr int stage_size = stage_k_threads * stage_n_threads;
53
54
auto load_perm_to_shared = [&](int k_tile_id) {
52
- int first_k_int4 = (k_tile_id * tile_k_size) / 4;
55
+ int first_k_int4 = (k_tile_id * target_tile_k_size) / 4;
56
57
int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
58
@@ -65,7 +68,7 @@ __global__ void gptq_marlin_repack_kernel(
68
return;
69
}
70
68
- int first_n = n_tile_id * tile_n_size;
71
+ int first_n = n_tile_id * target_tile_n_size;
72
73
int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
74
@@ -91,7 +94,7 @@ __global__ void gptq_marlin_repack_kernel(
94
auto k_id = threadIdx.x / stage_n_threads;
95
auto n_id = threadIdx.x % stage_n_threads;
96
94
- int first_k = k_tile_id * tile_k_size;
97
+ int first_k = k_tile_id * target_tile_k_size;
98
int first_k_packed = first_k / pack_factor;
99
100
cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
@@ -117,13 +120,13 @@ __global__ void gptq_marlin_repack_kernel(
120
}
121
122
int tc_col = th_id / 4;
120
- int tc_row = (th_id % 4) * 2;
123
+ int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2);
124
125
constexpr int tc_offsets[4] = {0, 1, 8, 9};
126
124
- int cur_n = warp_id * 16 + tc_col;
127
+ int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col;
128
126
- constexpr int sh_stride = 64;
129
+ constexpr int sh_stride = target_tile_n_size;
130
constexpr uint32_t mask = (1 << num_bits) - 1;
131
132
int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
@@ -134,6 +137,7 @@ __global__ void gptq_marlin_repack_kernel(
137
uint32_t vals[8];
138
139
if constexpr (has_perm) {
140
+ static_assert(!is_a_8bit);
141
for (int i = 0; i < 4; i++) {
142
int k_idx = tc_row + tc_offsets[i];
143
@@ -156,28 +160,49 @@ __global__ void gptq_marlin_repack_kernel(
160
161
#pragma unroll
162
for (int i = 0; i < tile_ints; i++) {
159
- b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
160
- b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
163
+ if constexpr (is_a_8bit) {
164
+ b1_vals[i] =
165
+ sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8];
166
+ } else {
167
+ b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
168
+ b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
169
+ }
170
}
171
172
#pragma unroll
173
for (int i = 0; i < 4; i++) {
165
- int cur_elem = tc_row + tc_offsets[i];
174
+ int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]);
175
int cur_int = cur_elem / pack_factor;
176
int cur_pos = cur_elem % pack_factor;
177
178
vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
170
- vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
179
+ if constexpr (is_a_8bit)
180
+ vals[4 + i] =
181
+ (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask;
182
+ else
183
+ vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
184
}
185
}
186
174
- constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
187
+ constexpr int tile_size =
188
+ target_tile_k_size * target_tile_n_size / pack_factor;
189
int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
190
191
// Result of:
192
// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
179
- if constexpr (num_bits == 4) {
180
- constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
193
+ if constexpr (!is_a_8bit && num_bits == 4) {
194
+ int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
195
+
196
+ uint32_t res = 0;
197
+#pragma unroll
198
+ for (int i = 0; i < 8; i++) {
199
+ res |= vals[pack_idx[i]] << (i * 4);
200
+ }
201
+
202
+ out_ptr[out_offset + th_id * 4 + warp_id] = res;
203
+
204
+ } else if constexpr (is_a_8bit && num_bits == 4) {
205
+ int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
206
207
uint32_t res = 0;
208
#pragma unroll
@@ -194,8 +219,9 @@ __global__ void gptq_marlin_repack_kernel(
219
uint32_t res2 = 0;
220
#pragma unroll
221
for (int i = 0; i < 4; i++) {
197
- res1 |= vals[pack_idx[i]] << (i * 8);
198
- res2 |= vals[4 + pack_idx[i]] << (i * 8);
222
+ const int ii = is_a_8bit ? i : pack_idx[i];
223
+ res1 |= vals[ii] << (i * 8);
224
+ res2 |= vals[4 + ii] << (i * 8);
225
}
226
227
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
@@ -236,21 +262,22 @@ __global__ void gptq_marlin_repack_kernel(
262
263
} // namespace marlin
264
239
-#define CALL_IF(NUM_BITS, HAS_PERM) \
240
- else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \
265
+#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \
266
+ else if (num_bits == NUM_BITS && has_perm == HAS_PERM && \
267
+ is_a_8bit == IS_A_8BIT) { \
268
cudaFuncSetAttribute( \
269
marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
243
- HAS_PERM>, \
270
+ HAS_PERM, IS_A_8BIT>, \
271
cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
272
marlin::gptq_marlin_repack_kernel<marlin::repack_threads, NUM_BITS, \
246
- HAS_PERM> \
273
+ HAS_PERM, IS_A_8BIT> \
274
<<<blocks, marlin::repack_threads, max_shared_mem, stream>>>( \
275
b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \
276
}
277
278
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
279
int64_t size_k, int64_t size_n,
253
- int64_t num_bits) {
280
+ int64_t num_bits, bool is_a_8bit) {
281
// Verify compatibility with marlin tile of 16x64
282
TORCH_CHECK(size_k % marlin::tile_k_size == 0, "size_k = ", size_k,
283
" is not divisible by tile_k_size = ", marlin::tile_k_size);
@@ -309,13 +336,17 @@ torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
336
337
if (false) {
338
}
312
- CALL_IF(4, false)
313
- CALL_IF(4, true)
314
- CALL_IF(8, false)
315
- CALL_IF(8, true)
339
+ CALL_IF(4, false, false)
340
+ CALL_IF(4, true, false)
341
+ CALL_IF(8, false, false)
342
+ CALL_IF(8, true, false)
343
+
344
+ CALL_IF(4, false, true)
345
+ CALL_IF(8, false, true)
346
+
347
else {
348
TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
318
- ", has_perm = ", has_perm);
349
+ ", has_perm = ", has_perm, ", is_a_8bit = ", is_a_8bit);
350
}
351
352
return out;
csrc/quantization/gptq_marlin/kernel.h
+6
-4
@@ -11,17 +11,19 @@
11
const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
const int4 *__restrict__ b_bias_ptr, \
14
+ const float *__restrict__ a_scales_ptr, \
15
const int4 *__restrict__ scales_ptr, \
15
- const uint16_t *__restrict__ scale2_ptr, \
16
+ const uint16_t *__restrict__ global_scale_ptr, \
17
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
19
bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \
20
int max_shared_mem
21
22
namespace MARLIN_NAMESPACE_NAME {
22
-template <typename scalar_t, // compute dtype, half or nv_float16
23
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
24
- const vllm::ScalarTypeId s_type_id, // weight ScalarType id
23
+template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
24
+ const vllm::ScalarTypeId b_type_id, // B ScalarType id
25
+ const vllm::ScalarTypeId c_type_id, // C ScalarType id
26
+ const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id
27
const int threads, // number of threads in a threadblock
28
const int thread_m_blocks, // number of 16x16 blocks in the m
29
// dimension (batchsize) of the
csrc/quantization/gptq_marlin/marlin.cuh
+39
@@ -55,6 +55,45 @@ constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; }
55
// No support for async
56
#else
57
58
+__device__ inline void cp_async1_ca_pred(void* smem_ptr, const void* glob_ptr,
59
+ bool pred = true) {
60
+ const int BYTES = 4;
61
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
62
+ asm volatile(
63
+ "{\n"
64
+ " .reg .pred p;\n"
65
+ " setp.ne.b32 p, %0, 0;\n"
66
+ " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
67
+ "}\n" ::"r"((int)pred),
68
+ "r"(smem), "l"(glob_ptr), "n"(BYTES));
69
+}
70
+
71
+__device__ inline void cp_async2_ca_pred(void* smem_ptr, const void* glob_ptr,
72
+ bool pred = true) {
73
+ const int BYTES = 8;
74
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
75
+ asm volatile(
76
+ "{\n"
77
+ " .reg .pred p;\n"
78
+ " setp.ne.b32 p, %0, 0;\n"
79
+ " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
80
+ "}\n" ::"r"((int)pred),
81
+ "r"(smem), "l"(glob_ptr), "n"(BYTES));
82
+}
83
+
84
+__device__ inline void cp_async4_ca_pred(void* smem_ptr, const void* glob_ptr,
85
+ bool pred = true) {
86
+ const int BYTES = 16;
87
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
88
+ asm volatile(
89
+ "{\n"
90
+ " .reg .pred p;\n"
91
+ " setp.ne.b32 p, %0, 0;\n"
92
+ " @p cp.async.ca.shared.global [%1], [%2], %3;\n"
93
+ "}\n" ::"r"((int)pred),
94
+ "r"(smem), "l"(glob_ptr), "n"(BYTES));
95
+}
96
+
97
__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,
98
bool pred = true) {
99
const int BYTES = 16;
csrc/quantization/gptq_marlin/marlin_dtypes.cuh
+70
-4
@@ -2,8 +2,10 @@
2
#ifndef _data_types_cuh
3
#define _data_types_cuh
4
#include "marlin.cuh"
5
+#include "core/scalar_type.hpp"
6
#include <cuda_fp16.h>
7
#include <cuda_bf16.h>
8
+#include <cuda_fp8.h>
9
10
#ifndef MARLIN_NAMESPACE_NAME
11
#define MARLIN_NAMESPACE_NAME marlin
@@ -11,14 +13,16 @@
13
14
namespace MARLIN_NAMESPACE_NAME {
15
14
-template <typename scalar_t>
15
-class ScalarType {};
16
+template <long scalar_type_id>
17
+class MarlinScalarType {};
18
19
template <>
18
-class ScalarType<half> {
20
+class MarlinScalarType<vllm::kFloat16.id()> {
21
public:
22
using scalar_t = half;
23
using scalar_t2 = half2;
24
+ using scalar_t4 = half2;
25
+ using scalar_32bit_t = half2;
26
27
// Matrix fragments for tensor core instructions; their precise layout is
28
// documented here:
@@ -27,6 +31,7 @@ class ScalarType<half> {
31
using FragB = Vec<half2, 2>;
32
using FragC = Vec<float, 4>;
33
using FragS = Vec<half2, 1>;
34
+ using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
35
using FragZP = Vec<half2, 4>;
36
37
static __device__ float inline num2float(const half x) {
@@ -44,18 +49,25 @@ class ScalarType<half> {
49
static __host__ __device__ half inline float2num(const float x) {
50
return __float2half(x);
51
}
52
+
53
+ static __host__ __device__ float2 inline num22float2(const half2 x) {
54
+ return __half22float2(x);
55
+ }
56
};
57
58
template <>
50
-class ScalarType<nv_bfloat16> {
59
+class MarlinScalarType<vllm::kBFloat16.id()> {
60
public:
61
using scalar_t = nv_bfloat16;
62
using scalar_t2 = nv_bfloat162;
63
+ using scalar_t4 = nv_bfloat162;
64
+ using scalar_32bit_t = nv_bfloat162;
65
66
using FragA = Vec<nv_bfloat162, 4>;
67
using FragB = Vec<nv_bfloat162, 2>;
68
using FragC = Vec<float, 4>;
69
using FragS = Vec<nv_bfloat162, 1>;
70
+ using FragS0 = Vec<__nv_fp8x2_e4m3, 1>;
71
using FragZP = Vec<nv_bfloat162, 4>;
72
73
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
@@ -75,9 +87,63 @@ class ScalarType<nv_bfloat16> {
87
static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
88
return __float2bfloat16(x);
89
}
90
+
91
+ static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
92
+ return __bfloat1622float2(x);
93
+ }
94
#endif
95
};
96
97
+template <>
98
+class MarlinScalarType<vllm::kFE4M3fn.id()> {
99
+ public:
100
+ using scalar_t = __nv_fp8_e4m3;
101
+ using scalar_t2 = __nv_fp8x2_e4m3;
102
+ using scalar_t4 = __nv_fp8x4_e4m3;
103
+ using scalar_32bit_t = __nv_fp8x4_e4m3;
104
+
105
+ using FragA = Vec<__nv_fp8x4_e4m3, 4>;
106
+ using FragB = Vec<__nv_fp8x4_e4m3, 2>;
107
+ using FragC = Vec<float, 4>;
108
+ using FragZP = Vec<__nv_fp8x2_e4m3, 4>;
109
+
110
+ static __host__ __device__
111
+ float2 inline num22float2(const __nv_fp8x2_e4m3 x) {
112
+ return (float2)x;
113
+ }
114
+};
115
+
116
+template <>
117
+class MarlinScalarType<vllm::kS8.id()> {
118
+ public:
119
+ using scalar_t = int8_t;
120
+ using scalar_t2 = int16_t;
121
+ using scalar_t4 = int32_t;
122
+ using scalar_32bit_t = int32_t;
123
+
124
+ using FragA = Vec<int32_t, 4>;
125
+ using FragB = Vec<int32_t, 2>;
126
+ using FragC = Vec<float, 4>;
127
+ using FragZP = Vec<int16_t, 4>;
128
+};
129
+
130
+template <typename scalar_t>
131
+class MarlinScalarType2 {};
132
+
133
+template <>
134
+class MarlinScalarType2<half> : public MarlinScalarType<vllm::kFloat16.id()> {};
135
+
136
+template <>
137
+class MarlinScalarType2<nv_bfloat16>
138
+ : public MarlinScalarType<vllm::kBFloat16.id()> {};
139
+
140
+template <>
141
+class MarlinScalarType2<__nv_fp8_e4m3>
142
+ : public MarlinScalarType<vllm::kFE4M3fn.id()> {};
143
+
144
+template <>
145
+class MarlinScalarType2<int8_t> : public MarlinScalarType<vllm::kS8.id()> {};
146
+
147
} // namespace MARLIN_NAMESPACE_NAME
148
149
#endif
csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu
new
+106
@@ -0,0 +1,106 @@
1
+
2
+
3
+#include "marlin.cuh"
4
+
5
+#include "core/registration.h"
6
+
7
+// for only non-zp format (like gptq)
8
+__global__ void marlin_int4_fp8_preprocess_kernel_without_zp(
9
+ // qweight: (size_k * size_n // 8,)
10
+ const int32_t* __restrict__ qweight,
11
+ // output: same shape with qweight
12
+ int32_t* __restrict__ output) {
13
+ int32_t val = qweight[blockIdx.x * 32 + threadIdx.x];
14
+ int32_t new_val = 0;
15
+
16
+#pragma unroll
17
+ for (int32_t i = 0; i < 8; i++) {
18
+ int32_t single_val = val & 0xF;
19
+ single_val = single_val >= 8 ? single_val - 8 : 15 - single_val;
20
+ new_val |= single_val << (i * 4);
21
+ val >>= 4;
22
+ }
23
+
24
+ output[blockIdx.x * 32 + threadIdx.x] = new_val;
25
+}
26
+
27
+// for awq format only (with zp and with awq weight layout)
28
+__global__ void marlin_int4_fp8_preprocess_kernel_awq(
29
+ // AWQ qweight: (size_k, size_n // 8)
30
+ const int32_t* __restrict__ qweight,
31
+ // output: same shape with qweight
32
+ int32_t* __restrict__ output,
33
+ // AWQ zeros: (size_k // group_size, size_n // 8)
34
+ const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
35
+ int32_t group_size) {
36
+ int32_t val =
37
+ qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
38
+ int32_t zero =
39
+ qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
40
+ blockIdx.y];
41
+ int32_t new_val = 0;
42
+
43
+#pragma unroll
44
+ for (int32_t i = 0; i < 8; i++) {
45
+ int32_t single_val = val & 0xF;
46
+ int32_t single_zero = zero & 0xF;
47
+
48
+ single_val =
49
+ single_val >= single_zero ? single_val - single_zero : 15 - single_val;
50
+ new_val |= single_val << (i * 4);
51
+ val >>= 4;
52
+ zero >>= 4;
53
+ }
54
+
55
+ output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
56
+}
57
+
58
+torch::Tensor marlin_int4_fp8_preprocess(
59
+ torch::Tensor& qweight, std::optional<torch::Tensor> qzeros_or_none,
60
+ bool inplace) {
61
+ TORCH_CHECK(qweight.device().is_cuda(), "qweight is not on GPU");
62
+ TORCH_CHECK(qweight.scalar_type() == at::ScalarType::Int,
63
+ "qweight.dtype != torch.int32");
64
+
65
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
66
+
67
+ torch::Tensor output = inplace ? qweight : torch::empty_like(qweight);
68
+
69
+ if (!qzeros_or_none.has_value()) {
70
+ TORCH_CHECK(qweight.numel() * 8 % 256 == 0,
71
+ "qweight.numel() * 8 % 256 != 0");
72
+
73
+ int blocks = qweight.numel() * 8 / 256;
74
+ marlin_int4_fp8_preprocess_kernel_without_zp<<<blocks, 32>>>(
75
+ (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr());
76
+ } else {
77
+ int32_t size_k = qweight.size(0);
78
+ int32_t size_n = qweight.size(1) * 8;
79
+ torch::Tensor qzeros = qzeros_or_none.value();
80
+
81
+ TORCH_CHECK(size_k % 32 == 0, "size_k % 32 != 0");
82
+ TORCH_CHECK(qzeros.device().is_cuda(), "qzeros is not on GPU");
83
+ TORCH_CHECK(qzeros.scalar_type() == at::ScalarType::Int,
84
+ "qweight.dtype != torch.int32");
85
+ TORCH_CHECK(device_of(qweight) == device_of(qzeros),
86
+ "qzeros is not on the same device with qweight");
87
+
88
+ int32_t group_size = qweight.size(0) / qzeros.size(0);
89
+ TORCH_CHECK(qweight.size(1) == qzeros.size(1),
90
+ "qweight.size(1) != qzeros.size(1)");
91
+ TORCH_CHECK(qweight.size(0) % qzeros.size(0) == 0,
92
+ "qweight.size(0) % qzeros.size(0) != 0");
93
+ TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
94
+
95
+ dim3 blocks(size_k / 32, size_n / 8);
96
+ marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32>>>(
97
+ (const int32_t*)qweight.data_ptr(), (int32_t*)output.data_ptr(),
98
+ (const int32_t*)qzeros.data_ptr(), size_n, size_k, group_size);
99
+ }
100
+
101
+ return output;
102
+}
103
+
104
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
105
+ m.impl("marlin_int4_fp8_preprocess", &marlin_int4_fp8_preprocess);
106
+}
csrc/quantization/gptq_marlin/marlin_template.h
+785
-416
@@ -38,7 +38,7 @@ namespace MARLIN_NAMESPACE_NAME {
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
41
+ const vllm::ScalarTypeId b_type_id, // weight MarlinScalarType id
42
const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
43
const int threads, // number of threads in a threadblock
44
const int thread_m_blocks, // number of 16x16 blocks in the m
@@ -77,65 +77,139 @@ __global__ void Marlin(
77
78
// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
79
// output/accumulation.
80
-template <typename scalar_t>
81
-__device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
82
- const typename ScalarType<scalar_t>::FragB& frag_b,
83
- typename ScalarType<scalar_t>::FragC& frag_c) {
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);
86
- float* c = reinterpret_cast<float*>(&frag_c);
87
- if constexpr (std::is_same<scalar_t, half>::value) {
88
- asm volatile(
89
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
90
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
91
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
92
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
93
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
94
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
95
- asm volatile(
96
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
97
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
98
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
99
- : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
100
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
101
- } else {
102
- STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
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
106
-template <typename scalar_t>
143
+template <vllm::ScalarTypeId type_id, int k_size = 16>
144
__device__ inline void mma_trans(
108
- const typename ScalarType<scalar_t>::FragA& a_frag,
109
- const typename ScalarType<scalar_t>::FragB& frag_b,
110
- const typename ScalarType<scalar_t>::FragB& frag_b2,
111
- typename ScalarType<scalar_t>::FragC& frag_c) {
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);
116
- if constexpr (std::is_same<scalar_t, half>::value) {
117
- asm volatile(
118
- "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
119
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
120
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
121
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
122
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
123
- } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
124
- asm volatile(
125
- "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
126
- "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
127
- : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
128
- : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
129
- "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
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 {
131
- STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
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
209
// Instruction for loading a full 16x16 matrix fragment of operand A from shared
210
// memory, directly in tensor core layout.
137
-template <int count, typename scalar_t>
138
-__device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
211
+template <int count, vllm::ScalarTypeId type_id>
212
+__device__ inline void ldsm(typename MarlinScalarType<type_id>::FragA& frag_a,
213
const void* smem_ptr) {
214
uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
215
uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
@@ -159,47 +233,54 @@ __device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
233
234
// Multiply dequantized values by the corresponding quantization scale; used
235
// only for grouped quantization.
162
-template <typename scalar_t>
163
-__device__ inline void scale(typename ScalarType<scalar_t>::FragB& frag_b,
164
- typename ScalarType<scalar_t>::FragS& frag_s,
236
+template <vllm::ScalarTypeId type_id>
237
+__device__ inline void scale(typename MarlinScalarType<type_id>::FragB& frag_b,
238
+ typename MarlinScalarType<type_id>::FragS& frag_s,
239
int i) {
166
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
167
- scalar_t2 s =
168
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
240
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
241
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
242
+ scalar_t2 s = MarlinScalarType<type_id>::num2num2(
243
+ reinterpret_cast<scalar_t*>(&frag_s)[i]);
244
frag_b[0] = __hmul2(frag_b[0], s);
245
frag_b[1] = __hmul2(frag_b[1], s);
246
}
247
173
-template <typename scalar_t>
248
+template <vllm::ScalarTypeId type_id>
249
__device__ inline void scale_and_sub(
175
- typename ScalarType<scalar_t>::FragB& frag_b, scalar_t s, scalar_t zp) {
176
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
177
- scalar_t2 s2 = ScalarType<scalar_t>::num2num2(s);
178
- scalar_t2 zp2 = ScalarType<scalar_t>::num2num2(zp);
250
+ typename MarlinScalarType<type_id>::FragB& frag_b,
251
+ typename MarlinScalarType<type_id>::scalar_t s,
252
+ typename MarlinScalarType<type_id>::scalar_t zp) {
253
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
254
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
255
+ scalar_t2 s2 = MarlinScalarType<type_id>::num2num2(s);
256
+ scalar_t2 zp2 = MarlinScalarType<type_id>::num2num2(zp);
257
frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
258
frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
259
}
260
183
-template <typename scalar_t>
184
-__device__ inline void sub_zp(typename ScalarType<scalar_t>::FragB& frag_b,
185
- typename ScalarType<scalar_t>::scalar_t2& frag_zp,
186
- int i) {
187
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
188
- scalar_t2 zp =
189
- ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
261
+template <vllm::ScalarTypeId type_id>
262
+__device__ inline void sub_zp(
263
+ typename MarlinScalarType<type_id>::FragB& frag_b,
264
+ typename MarlinScalarType<type_id>::scalar_t2& frag_zp, int i) {
265
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
266
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
267
+ scalar_t2 zp = MarlinScalarType<type_id>::num2num2(
268
+ reinterpret_cast<scalar_t*>(&frag_zp)[i]);
269
frag_b[0] = __hsub2(frag_b[0], zp);
270
frag_b[1] = __hsub2(frag_b[1], zp);
271
}
272
273
// Same as above, but for act_order (each K is multiplied individually)
195
-template <typename scalar_t>
196
-__device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
197
- typename ScalarType<scalar_t>::FragS& frag_s_1,
198
- typename ScalarType<scalar_t>::FragS& frag_s_2,
199
- typename ScalarType<scalar_t>::FragS& frag_s_3,
200
- typename ScalarType<scalar_t>::FragS& frag_s_4,
201
- int i) {
202
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
274
+template <vllm::ScalarTypeId type_id>
275
+__device__ inline void scale4(
276
+ typename MarlinScalarType<type_id>::FragB& frag_b,
277
+ typename MarlinScalarType<type_id>::FragS& frag_s_1,
278
+ typename MarlinScalarType<type_id>::FragS& frag_s_2,
279
+ typename MarlinScalarType<type_id>::FragS& frag_s_3,
280
+ typename MarlinScalarType<type_id>::FragS& frag_s_4, int i) {
281
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
282
+ using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2;
283
+
284
scalar_t2 s_val_1_2;
285
s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
286
s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
@@ -213,12 +294,13 @@ __device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
294
}
295
296
// Given 2 floats multiply by 2 scales (halves)
216
-template <typename scalar_t>
217
-__device__ inline void scale_float(float* c,
218
- typename ScalarType<scalar_t>::FragS& s) {
297
+template <vllm::ScalarTypeId type_id>
298
+__device__ inline void scale_float(
299
+ float* c, typename MarlinScalarType<type_id>::FragS& s) {
300
+ using scalar_t = typename MarlinScalarType<type_id>::scalar_t;
301
scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
220
- c[0] = __fmul_rn(c[0], ScalarType<scalar_t>::num2float(s_ptr[0]));
221
- c[1] = __fmul_rn(c[1], ScalarType<scalar_t>::num2float(s_ptr[1]));
302
+ c[0] = __fmul_rn(c[0], MarlinScalarType<type_id>::num2float(s_ptr[0]));
303
+ c[1] = __fmul_rn(c[1], MarlinScalarType<type_id>::num2float(s_ptr[1]));
304
}
305
306
// Wait until barrier reaches `count`, then lock for current threadblock.
@@ -270,9 +352,10 @@ __device__ inline void wait_negative_and_add(int* lock) {
352
__syncthreads();
353
}
354
273
-template <typename scalar_t, // compute dtype, half or nv_float16
274
- const vllm::ScalarTypeId w_type_id, // weight ScalarType id
275
- const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
355
+template <const vllm::ScalarTypeId a_type_id, // A ScalarType id
356
+ const vllm::ScalarTypeId b_type_id, // B ScalarType id
357
+ const vllm::ScalarTypeId c_type_id, // C ScalarType id
358
+ const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id
359
const int threads, // number of threads in a threadblock
360
const int thread_m_blocks, // number of 16x16 blocks in the m
361
// dimension (batchsize) of the
@@ -288,18 +371,23 @@ template <typename scalar_t, // compute dtype, half or nv_float16
371
const bool is_zp_float // is zero point of float16 type?
372
>
373
__global__ void Marlin(
291
- const int4* __restrict__ A, // fp16 input matrix of shape mxk
292
- const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
293
- int4* __restrict__ C, // fp16 output buffer of shape mxn
294
- int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
374
+ const int4* __restrict__ A0, // fp16 input matrix of shape mxk
375
+ const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
376
+ int4* __restrict__ C0, // fp16 output buffer of shape mxn
377
+ int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
378
const int4* __restrict__ b_bias_ptr,
296
- const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
297
- // (k/groupsize)xn
298
- const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
299
- // only)
300
- const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
301
- // (k/groupsize)x(n/pack_factor)
302
- const int* __restrict__ g_idx, // int32 group indices of shape k
379
+ // float scales of input matrix, only used when is_a_8bit == true.
380
+ // shape (m,)
381
+ const float* __restrict__ a_scales_ptr,
382
+ // fp16 quantization scales. shape (k/groupsize, n)
383
+ const int4* __restrict__ scales_ptr,
384
+ // fp16 global scale (for nvfp4// only)
385
+ const uint16_t* __restrict__ global_scale_ptr,
386
+ // 4bit packed zero-points of shape
387
+ // (k/groupsize, n/pack_factor)
388
+ const int4* __restrict__ zp_ptr,
389
+ // int32 group indices of shape k
390
+ const int* __restrict__ g_idx,
391
int num_groups, // number of scale groups per output channel
392
int prob_m, // batch dimension m
393
int prob_n, // output dimension n
@@ -321,17 +409,35 @@ __global__ void Marlin(
409
// ensures good utilization of all SMs for many kinds of shape and GPU
410
// configurations, while requiring as few slow global cross-threadblock
411
// reductions as possible.
324
- using Dtype = ScalarType<scalar_t>;
325
- using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
326
- using FragA = typename ScalarType<scalar_t>::FragA;
327
- using FragB = typename ScalarType<scalar_t>::FragB;
328
- using FragC = typename ScalarType<scalar_t>::FragC;
329
- using FragS = typename ScalarType<scalar_t>::FragS;
330
- using FragZP = typename ScalarType<scalar_t>::FragZP;
331
-
332
- static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
412
+
413
+ #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890
414
+ // FP8 computation is only supported for Ada Lovelace or newer architectures.
415
+ if constexpr (a_type_id == vllm::kFE4M3fn.id()) return;
416
+ #endif
417
+
418
+ using Adtype = MarlinScalarType<a_type_id>;
419
+ using Cdtype = MarlinScalarType<c_type_id>;
420
+ const int4* A = A0;
421
+ int4* C = C0;
422
+
423
+ using scalar_t = typename MarlinScalarType<a_type_id>::scalar_t;
424
+ using scalar_t2 = typename MarlinScalarType<a_type_id>::scalar_t2;
425
+ using scalar_32bit_t = typename MarlinScalarType<a_type_id>::scalar_32bit_t;
426
+
427
+ using c_scalar_t = typename MarlinScalarType<c_type_id>::scalar_t;
428
+ using c_scalar_t2 = typename MarlinScalarType<c_type_id>::scalar_t2;
429
+
430
+ using FragA = typename MarlinScalarType<a_type_id>::FragA;
431
+ using FragB = typename MarlinScalarType<a_type_id>::FragB;
432
+ using FragC = typename MarlinScalarType<a_type_id>::FragC;
433
+ using FragS = typename MarlinScalarType<c_type_id>::FragS;
434
+ using FragZP = typename MarlinScalarType<c_type_id>::FragZP;
435
+
436
+ static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id);
437
+ static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id);
438
+ static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id);
439
static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
334
- if constexpr (w_type == vllm::kFE2M1f) {
440
+ if constexpr (b_type == vllm::kFE2M1f) {
441
static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
442
s_type == vllm::kFE8M0fnu && group_blocks == 2);
443
} else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
@@ -340,27 +446,35 @@ __global__ void Marlin(
446
static_assert(s_type == vllm::kFloat16);
447
}
448
343
- constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
344
- constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
345
- w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
449
+ constexpr bool is_a_8bit = a_type.size_bits() == 8;
450
+ if constexpr (!is_a_8bit) {
451
+ static_assert(std::is_same<scalar_t, c_scalar_t>::value);
452
+ }
453
+ constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8;
454
+ constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 ||
455
+ b_type == vllm::kS4 || b_type == vllm::kS8 ||
456
+ b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
457
// see comments of dequant.h for more details
458
constexpr bool dequant_skip_flop =
348
- w_type == vllm::kFE4M3fn ||
349
- w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
459
+ is_a_8bit || b_type == vllm::kFE4M3fn ||
460
+ b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
461
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
351
- has_zp && !is_zp_float && !(w_type == vllm::kU8);
462
+ has_zp && !is_zp_float && !(b_type == vllm::kU8);
463
+
464
+ c_scalar_t2 global_scale;
465
353
- scalar_t2 global_scale;
354
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
355
- // NVFP4 format requires global scale
356
- uint16_t val = scale2_ptr[0];
357
- global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
466
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
467
+ uint16_t val = global_scale_ptr[0];
468
+ global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
469
}
470
471
constexpr bool has_act_order = group_blocks == 0;
472
constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
473
363
- constexpr int pack_factor = 32 / w_type.size_bits();
474
+ extern __shared__ int4 sh[];
475
+ float* sh_a_s = reinterpret_cast<float*>(sh);
476
+ int4* sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0);
477
+ constexpr int pack_factor = 32 / b_type.size_bits();
478
static_assert(thread_m_blocks == 1 || !m_block_size_8);
479
480
// For larger GEMMs we run multiple batchsize 64 versions in parallel for a
@@ -373,7 +487,19 @@ __global__ void Marlin(
487
488
int k_tiles = prob_k / 16 / thread_k_blocks;
489
int n_tiles = prob_n / 16 / thread_n_blocks;
376
- int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x);
490
+
491
+ int global_mn_tiles = parallel * n_tiles;
492
+ int part2_mn_tiles = global_mn_tiles;
493
+ int part1_mn_iters = 0;
494
+ bool in_part2 = false;
495
+
496
+ if (global_mn_tiles > gridDim.x) {
497
+ part2_mn_tiles = global_mn_tiles % gridDim.x;
498
+ if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x;
499
+ part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x;
500
+ }
501
+
502
+ int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x);
503
504
if constexpr (!has_act_order && group_blocks != -1) {
505
if (group_blocks >= thread_k_blocks) {
@@ -385,28 +511,21 @@ __global__ void Marlin(
511
}
512
}
513
388
- int slice_row = (iters * blockIdx.x) % k_tiles;
389
- int slice_col_par = (iters * blockIdx.x) / k_tiles;
390
- int slice_col = slice_col_par;
391
- int slice_iters; // number of threadblock tiles in the current slice
392
- int slice_count =
393
- 0; // total number of active threadblocks in the current slice
394
- int slice_idx; // index of threadblock in current slice; numbered bottom to
395
- // top
514
+ int slice_row = 0;
515
+ int slice_col_par = blockIdx.x;
516
+ int slice_col;
517
+ int slice_iters =
518
+ k_tiles; // number of threadblock tiles in the current slice
519
+ // total number of active threadblocks in the current slice
520
+ int slice_count = 1;
521
+ // index of threadblock in current slice; numbered bottom to top
522
+ int slice_idx = 0;
523
524
int par_id = 0;
525
int locks_off = 0;
526
400
- // We can easily implement parallel problem execution by just remapping
401
- // indices and advancing global pointers
402
- if (slice_col_par >= n_tiles) {
403
- A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8;
404
- C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
405
- slice_col = slice_col_par % n_tiles;
406
- par_id = slice_col_par / n_tiles;
407
- }
408
- if (parallel * n_tiles >= gridDim.x) {
409
- // when parallel * n_tiles >= sms
527
+ if (part2_mn_tiles >= gridDim.x) {
528
+ // when part2_mn_tiles >= sms
529
// then there are at most $sms$ conflict tile blocks
530
locks_off = blockIdx.x;
531
} else {
@@ -415,10 +534,11 @@ __global__ void Marlin(
534
535
// Compute all information about the current slice which is required for
536
// synchronization.
418
- auto init_slice = [&](bool first_init = false) {
537
+ bool first_init = true;
538
+ auto init_part2_slice = [&]() {
539
slice_iters =
540
iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
421
- if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
541
+ if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0;
542
if (slice_iters == 0) return;
543
if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
544
slice_count = 1;
@@ -436,7 +556,7 @@ __global__ void Marlin(
556
if (col_off > 0) slice_idx--;
557
}
558
}
439
- if (parallel * n_tiles >= gridDim.x) {
559
+ if (part2_mn_tiles >= gridDim.x) {
560
if (slice_count > 1 && slice_idx == slice_count - 1) {
561
locks_off++;
562
}
@@ -466,28 +586,68 @@ __global__ void Marlin(
586
}
587
588
if (slice_col == n_tiles) {
469
- A += 16 * thread_m_blocks * lda / 8;
589
+ A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8);
590
C += 16 * thread_m_blocks * prob_n / 8;
591
slice_col = 0;
592
par_id++;
593
}
594
+ if (is_a_8bit && (first_init || slice_col == 0)) {
595
+ __syncthreads();
596
+ int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x;
597
+ cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd],
598
+ threadIdx.x < prob_m);
599
+ }
600
};
475
- init_slice(true);
601
+
602
+ auto init_part1_slice = [&]() {
603
+ if (part1_mn_iters) {
604
+ part1_mn_iters--;
605
+ par_id = slice_col_par / n_tiles;
606
+ slice_col = slice_col_par % n_tiles;
607
+ slice_iters = k_tiles;
608
+ A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda;
609
+ C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n;
610
+ if (is_a_8bit) {
611
+ __syncthreads();
612
+ int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x;
613
+ cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd],
614
+ threadIdx.x < prob_m);
615
+ }
616
+ }
617
+ };
618
+
619
+ auto init_slice = [&]() {
620
+ if (!in_part2 && !part1_mn_iters) {
621
+ in_part2 = true;
622
+ slice_col_par = (iters * blockIdx.x) / k_tiles;
623
+ slice_row = (iters * blockIdx.x) % k_tiles;
624
+ slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles;
625
+ par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles;
626
+ A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda;
627
+ C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n;
628
+ }
629
+ if (!in_part2) {
630
+ init_part1_slice();
631
+ } else {
632
+ init_part2_slice();
633
+ first_init = false;
634
+ }
635
+ };
636
+
637
+ init_slice();
638
639
// A sizes/strides
640
641
// stride of the A matrix in global memory
480
- int a_gl_stride = lda / 8;
642
+ int a_gl_stride = lda / (is_a_8bit ? 16 : 8);
643
// stride of an A matrix tile in shared memory
482
- constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
644
+ constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
645
// delta between subsequent A tiles in global memory
484
- constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;
646
+ constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8);
647
// between subsequent accesses within a tile
648
int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
649
// between shared memory writes
650
constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
489
- // between shared memory tile reads
490
- constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));
651
// within a shared memory tile
652
constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
653
// overall size of a tile
@@ -496,24 +656,25 @@ __global__ void Marlin(
656
constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
657
658
// B sizes/strides
499
- int b_gl_stride = 16 * prob_n / (pack_factor * 4);
500
- constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
501
- constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;
659
+ int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4));
660
+ constexpr int b_sh_stride =
661
+ ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4);
662
+ constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2;
663
constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
664
504
- int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
505
- int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
665
+ int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
666
constexpr int b_sh_wr_delta = threads * b_thread_vecs;
507
- constexpr int b_sh_rd_delta = threads * b_thread_vecs;
508
- constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
667
+ constexpr int b_sh_stage =
668
+ b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1);
669
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
670
671
// Scale sizes/strides without act_order
512
- int s_gl_stride = prob_n / 8;
513
- constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
672
+ int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8);
673
+ constexpr int s_sh_stride =
674
+ 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8);
675
constexpr int s_tb_groups =
676
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
516
- ? thread_k_blocks / group_blocks / (w_type == vllm::kFE2M1f ? 2 : 1)
677
+ ? thread_k_blocks / group_blocks
678
: 1;
679
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
680
int s_gl_rd_delta = s_gl_stride;
@@ -527,7 +688,7 @@ __global__ void Marlin(
688
int act_s_col_stride = 1;
689
int act_s_col_warp_stride = act_s_col_stride * 8;
690
530
- int tb_n_warps = thread_n_blocks / 4;
691
+ constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4);
692
int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
693
694
// Zero-points sizes/strides
@@ -550,17 +711,22 @@ __global__ void Marlin(
711
int a_sh_rd =
712
a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
713
(threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
553
- a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
714
+ a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters;
715
+
716
+ int b_gl_rd;
717
+ if (threads <= b_sh_stride) {
718
+ b_gl_rd = threadIdx.x;
719
+ } else {
720
+ b_gl_rd =
721
+ b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride);
722
+ }
723
555
- int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
556
- (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;
559
- auto b_sh_wr = threadIdx.x * b_thread_vecs;
726
auto b_sh_rd = threadIdx.x * b_thread_vecs;
727
+ b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1));
728
729
// For act_order
563
- constexpr int k_iter_size = tb_k / b_sh_wr_iters;
730
int slice_k_start = tb_k * slice_row;
731
int slice_k_finish = slice_k_start + tb_k * slice_iters;
732
int slice_k_start_shared_fetch = slice_k_start;
@@ -571,58 +737,54 @@ __global__ void Marlin(
737
if constexpr (!has_act_order) {
738
if constexpr (group_blocks == -1) {
739
s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
574
- } else {
575
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) /
576
- (w_type == vllm::kFE2M1f ? 2 : 1) +
740
+ } else if constexpr (group_blocks >= thread_k_blocks) {
741
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
742
s_sh_stride * slice_col + threadIdx.x;
743
+ } else {
744
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
745
+ threadIdx.x / s_sh_stride) +
746
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
747
}
748
}
749
auto s_sh_wr = threadIdx.x;
581
- bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
750
+ bool s_sh_wr_pred = threadIdx.x < s_sh_stage;
751
752
// Zero-points
753
int zp_gl_rd;
754
if constexpr (has_zp) {
755
if constexpr (group_blocks == -1) {
756
zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
588
- } else {
757
+ } else if constexpr (group_blocks >= thread_k_blocks) {
758
zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
759
zp_sh_stride * slice_col + threadIdx.x;
760
+ } else {
761
+ zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
762
+ threadIdx.x / zp_sh_stride) +
763
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
764
}
765
}
766
auto zp_sh_wr = threadIdx.x;
594
- bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
767
+ bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage;
768
769
// We use a different scale layout for grouped and column-wise quantization as
770
// we scale a `half2` tile in column-major layout in the former and in
771
// row-major in the latter case.
772
int s_sh_rd;
600
- if constexpr (group_blocks != -1 && w_type == vllm::kFE2M1f) {
601
- auto warp_id = threadIdx.x / 32;
602
- int n_warps = thread_n_blocks / 4;
603
- int warp_row = warp_id / n_warps;
604
-
605
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
606
- (threadIdx.x % 32) / 4;
607
- s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2;
608
-
773
+ if constexpr (is_a_8bit) {
774
+ s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4);
775
} else if constexpr (group_blocks != -1)
610
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
611
- (threadIdx.x % 32) / 4;
776
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
777
else if constexpr (group_blocks == -1 &&
778
(m_block_size_8 || (has_zp && !dequant_skip_flop)))
614
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
615
- (threadIdx.x % 32) / 8;
779
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
780
else
617
- s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
618
- (threadIdx.x % 32) % 4;
781
+ s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4;
782
783
int bias_sh_rd;
784
if constexpr (m_block_size_8) {
622
- bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
623
- (threadIdx.x % 32) / 8;
785
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8;
786
} else {
625
- bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
787
+ bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) +
788
(threadIdx.x % 32) % 4;
789
}
790
@@ -638,12 +800,16 @@ __global__ void Marlin(
800
if constexpr (has_zp) {
801
if constexpr (is_zp_float) {
802
if constexpr (group_blocks != -1) {
641
- zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
642
- (threadIdx.x % 32) / 4;
803
+ zp_sh_rd =
804
+ 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4;
805
}
806
+ } else if (is_a_8bit) {
807
+ zp_sh_rd = num_ints_per_thread * num_col_threads *
808
+ ((threadIdx.x / 32) % tb_n_warps / 2) +
809
+ num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
810
} else {
811
zp_sh_rd = num_ints_per_thread * num_col_threads *
646
- ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
812
+ ((threadIdx.x / 32) % tb_n_warps) +
813
num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
814
}
815
}
@@ -678,26 +844,19 @@ __global__ void Marlin(
844
for (int i = 0; i < b_sh_wr_iters; i++) {
845
#pragma unroll
846
for (int j = 0; j < thread_m_blocks; j++)
681
- a_sh_rd_trans[i][j] =
682
- transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
847
+ a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd);
848
}
849
850
// Since B-accesses have non-constant stride they have to be computed at
851
// runtime; we break dependencies between subsequent accesses with a tile by
852
// maintining multiple pointers (we have enough registers), a tiny
853
// optimization.
689
- const int4* B_ptr[b_sh_wr_iters];
690
- #pragma unroll
691
- for (int i = 0; i < b_sh_wr_iters; i++)
692
- B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
854
694
- extern __shared__ int4 sh[];
855
// Shared memory storage for global fetch pipelines.
856
constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks;
857
constexpr int sh_b_size = stages * b_sh_stage;
698
- int4* sh_b = sh;
699
- int4* sh_red = sh;
700
-
858
+ int4* sh_b = sh_new;
859
+ int4* sh_red = sh_new;
860
constexpr int sh_size_b_red_min =
861
(sh_red_size < sh_b_size ? sh_red_size : sh_b_size);
862
constexpr int sh_size_b_red_max =
@@ -708,8 +867,8 @@ __global__ void Marlin(
867
? sh_size_b_red_max
868
: (sh_size_b_red_min + sh_bias_size);
869
711
- int4* sh_bias = sh + sh_size_b_red_min;
712
- int4* sh_g_idx = sh + sh_b_red_bias_size;
870
+ int4* sh_bias = sh_new + sh_size_b_red_min;
871
+ int4* sh_g_idx = sh_new + sh_b_red_bias_size;
872
int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
873
constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
874
: (stages * s_sh_stage);
@@ -723,7 +882,8 @@ __global__ void Marlin(
882
// Register storage for double buffer of shared memory reads.
883
FragA frag_a[2][thread_m_blocks];
884
I4 frag_b_quant[2][b_thread_vecs];
726
- FragC frag_c[thread_m_blocks][4][2];
885
+ FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2];
886
+ FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2];
887
FragS frag_s[2][4]; // No act-order
888
FragS frag_bias[2][4];
889
FragS act_frag_s[2][4][4]; // For act-order
@@ -731,6 +891,24 @@ __global__ void Marlin(
891
FragZP frag_zp; // Zero-points in fp16
892
FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
893
894
+ if constexpr (is_a_8bit) {
895
+ #pragma unroll
896
+ for (int j = 0; j < 2; j++) {
897
+ #pragma unroll
898
+ for (int i = 0; i < thread_m_blocks; i++) {
899
+ #pragma unroll
900
+ for (int g = 0; g < 4; g++) {
901
+ frag_c_tmp[i][j][0][g] = 0.0f;
902
+ }
903
+
904
+ #pragma unroll
905
+ for (int g = 0; g < 4; g++) {
906
+ frag_c_tmp[i][j][1][g] = 0.0f;
907
+ }
908
+ }
909
+ }
910
+ }
911
+
912
// Zero accumulators.
913
auto zero_accums = [&]() {
914
#pragma unroll
@@ -788,15 +966,17 @@ __global__ void Marlin(
966
}
967
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
968
#pragma unroll
791
- for (int i = 0; i < b_sh_wr_iters; i++) {
792
- #pragma unroll
793
- for (int j = 0; j < b_thread_vecs; j++) {
794
- cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j);
795
- }
969
+ for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) {
970
+ constexpr int count = div_ceil(b_sh_stride, threads);
971
+ int b_gl_idx =
972
+ b_gl_rd + (i % count) * threads +
973
+ b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride);
974
797
- B_ptr[i] += b_gl_rd_delta_o;
975
+ cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]);
976
}
977
978
+ b_gl_rd += b_gl_rd_delta_o;
979
+
980
if constexpr (has_act_order) {
981
// Fetch g_idx thread-block portion
982
int full_pipe = a_off;
@@ -816,44 +996,24 @@ __global__ void Marlin(
996
if constexpr (group_blocks != -1) {
997
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
998
819
- if constexpr (group_blocks >= thread_k_blocks) {
820
- // Only fetch scales if this tile starts a new group
821
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
822
- if (s_sh_wr_pred) {
823
- cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
824
- }
825
- s_gl_rd += s_gl_rd_delta;
826
- }
827
- } else {
828
- for (int i = 0; i < s_tb_groups; i++) {
829
- if (s_sh_wr_pred) {
830
- cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
831
- &scales_ptr[s_gl_rd]);
832
- }
833
- s_gl_rd += s_gl_rd_delta;
999
+ // Only fetch scales if this tile starts a new group
1000
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1001
+ if (s_sh_wr_pred) {
1002
+ cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
1003
}
1004
+ s_gl_rd += s_gl_rd_delta * s_tb_groups;
1005
}
1006
}
1007
1008
if constexpr (has_zp && group_blocks != -1) {
1009
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1010
841
- if constexpr (group_blocks >= thread_k_blocks) {
842
- // Only fetch zero-points if this tile starts a new group
843
- if (pipe % (group_blocks / thread_k_blocks) == 0) {
844
- if (zp_sh_wr_pred) {
845
- cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
846
- }
847
- zp_gl_rd += zp_gl_rd_delta;
848
- }
849
- } else {
850
- for (int i = 0; i < zp_tb_groups; i++) {
851
- if (zp_sh_wr_pred) {
852
- cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
853
- &zp_ptr[zp_gl_rd]);
854
- }
855
- zp_gl_rd += zp_gl_rd_delta;
1011
+ // Only fetch zero points if this tile starts a new group
1012
+ if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) {
1013
+ if (zp_sh_wr_pred) {
1014
+ cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
1015
}
1016
+ zp_gl_rd += zp_gl_rd_delta * zp_tb_groups;
1017
}
1018
}
1019
}
@@ -891,14 +1051,14 @@ __global__ void Marlin(
1051
int4* sh_a_stage = sh_a + a_sh_stage * pipe;
1052
#pragma unroll
1053
for (int i = 0; i < thread_m_blocks; i++)
894
- ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
1054
+ ldsm<m_block_size_8 ? 2 : 4, a_type_id>(
1055
frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
1056
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
1057
1058
#pragma unroll
1059
for (int i = 0; i < b_thread_vecs; i++) {
1060
frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
901
- &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
1061
+ &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]);
1062
}
1063
};
1064
@@ -922,53 +1082,54 @@ __global__ void Marlin(
1082
1083
auto fetch_scales_to_registers = [&](int k, int full_pipe) {
1084
int pipe = full_pipe % stages;
1085
+ using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>;
1086
+ using IT0 = typename std::conditional_t<is_a_8bit, int, int2>;
1087
+ constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1);
1088
1089
if constexpr (!has_act_order) {
1090
// No act-order case
1091
if constexpr (group_blocks == -1) {
1092
// load only when starting a new slice
930
- if (k == 0 && full_pipe == 0) {
1093
+ if (k == 0 && full_pipe == 0 && dequant_skip_flop) {
1094
reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
1095
reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1096
}
1097
} else if constexpr (group_blocks != -1) {
1098
if constexpr (group_blocks >= thread_k_blocks) {
936
- if (k % b_sh_wr_iters == 0) {
937
- int4* sh_s_stage =
938
- sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
939
- (pipe / (group_blocks / thread_k_blocks)));
940
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
941
- } else {
942
- reinterpret_cast<int4*>(&frag_s[1])[0] =
943
- reinterpret_cast<int4*>(&frag_s[0])[0];
1099
+ constexpr int g = group_blocks / thread_k_blocks;
1100
+ if (pipe % g == 0) {
1101
+ if (k % b_sh_wr_iters == 0) {
1102
+ int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g));
1103
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1104
+ } else {
1105
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1106
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1107
+ }
1108
}
945
- } else {
1109
+ } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1110
auto warp_id = threadIdx.x / 32;
947
- int n_warps = thread_n_blocks / 4;
1111
+ int warp_row = warp_id / tb_n_warps;
1112
949
- int warp_row = warp_id / n_warps;
950
-
951
- int cur_k = warp_row * 16;
952
- cur_k += k_iter_size * (k % b_sh_wr_iters);
953
-
954
- int k_blocks = cur_k / 16;
955
- int cur_group_id =
956
- k_blocks / (group_blocks * (w_type == vllm::kFE2M1f ? 2 : 1));
1113
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1114
+ int cur_group_id = k_blocks / group_blocks2;
1115
1116
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1117
960
- if constexpr (w_type_id != vllm::kFE2M1f.id()) {
1118
+ if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1119
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1120
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
963
- } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) {
1121
+ } else {
1122
reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1123
reinterpret_cast<int2*>(
1124
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1125
+ }
1126
+ } else if (group_blocks >= b_sh_wr_iters) {
1127
+ if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1128
+ reinterpret_cast<int4*>(&frag_s[1])[0] =
1129
+ reinterpret_cast<int4*>(&frag_s[0])[0];
1130
} else {
968
- reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
969
- reinterpret_cast<int2*>(
970
- sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) +
971
- k % 2];
1131
+ reinterpret_cast<int2*>(&frag_s[1])[0] =
1132
+ reinterpret_cast<int2*>(&frag_s[0])[0];
1133
}
1134
}
1135
}
@@ -989,18 +1150,15 @@ __global__ void Marlin(
1150
cur_k = 0;
1151
1152
// Progress to current iteration
992
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1153
+ cur_k += k % b_sh_wr_iters;
1154
1155
// Determine "position" inside the thread-block (based on warp and
1156
// thread-id)
1157
auto warp_id = threadIdx.x / 32;
997
- int n_warps =
998
- thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
999
-
1000
- int warp_row = warp_id / n_warps;
1001
- int warp_col = warp_id % n_warps;
1158
+ int warp_row = warp_id / tb_n_warps;
1159
+ int warp_col = warp_id % tb_n_warps;
1160
1003
- cur_k += warp_row * 16;
1161
+ cur_k += warp_row * 16 * b_sh_wr_iters;
1162
1163
auto th_id = threadIdx.x % 32;
1164
cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
@@ -1055,18 +1213,16 @@ __global__ void Marlin(
1213
1214
if constexpr (group_blocks == -1) {
1215
// load only when starting a new slice
1058
- if (k == 0 && full_pipe == 0) {
1216
+ if (k == 0 && full_pipe == 0 || is_a_8bit) {
1217
#pragma unroll
1218
for (int i = 0; i < num_ints_per_thread; i++) {
1219
frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
1220
}
1221
}
1064
-
1222
} else if constexpr (group_blocks >= thread_k_blocks) {
1066
- if (k % b_sh_wr_iters == 0) {
1067
- int4* sh_zp_stage =
1068
- sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1069
- (pipe / (group_blocks / thread_k_blocks)));
1223
+ constexpr int g = group_blocks / thread_k_blocks;
1224
+ if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) {
1225
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
1226
#pragma unroll
1227
for (int i = 0; i < num_ints_per_thread; i++) {
1228
frag_qzp[k % 2][i] =
@@ -1075,21 +1231,11 @@ __global__ void Marlin(
1231
}
1232
} else {
1233
auto warp_id = threadIdx.x / 32;
1078
- int n_warps = thread_n_blocks / 4;
1079
-
1080
- int warp_row = warp_id / n_warps;
1234
1082
- int cur_k = warp_row * 16;
1083
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1235
+ int warp_row = warp_id / tb_n_warps;
1236
1085
- int k_blocks = cur_k / 16;
1086
- int cur_group_id = 0;
1087
-
1088
- // Suppress bogus and persistent divide-by-zero warning
1089
- #pragma nv_diagnostic push
1090
- #pragma nv_diag_suppress divide_by_zero
1091
- cur_group_id = k_blocks / group_blocks;
1092
- #pragma nv_diagnostic pop
1237
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1238
+ int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1);
1239
1240
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1241
@@ -1108,29 +1254,18 @@ __global__ void Marlin(
1254
1255
if constexpr (group_blocks != -1) {
1256
if constexpr (group_blocks >= thread_k_blocks) {
1111
- if (k % b_sh_wr_iters == 0) {
1112
- int4* sh_zp_stage =
1113
- sh_zp +
1114
- zp_sh_stage * ((group_blocks / thread_k_blocks) *
1115
- (pipe / (group_blocks / thread_k_blocks)));
1257
+ constexpr int g = group_blocks / thread_k_blocks;
1258
+ if (pipe % g == 0 && k % b_sh_wr_iters == 0) {
1259
+ int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g));
1260
reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1261
sh_zp_stage[zp_sh_rd];
1262
}
1119
- } else {
1263
+ } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) {
1264
auto warp_id = threadIdx.x / 32;
1121
- int n_warps = thread_n_blocks / 4;
1265
1123
- int warp_row = warp_id / n_warps;
1124
-
1125
- int cur_k = warp_row * 16;
1126
- cur_k += k_iter_size * (k % b_sh_wr_iters);
1127
-
1128
- int k_blocks = cur_k / 16;
1129
- // Suppress bogus and persistent divide-by-zero warning
1130
- #pragma nv_diagnostic push
1131
- #pragma nv_diag_suppress divide_by_zero
1266
+ int warp_row = warp_id / tb_n_warps;
1267
+ int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters;
1268
int cur_group_id = k_blocks / group_blocks;
1133
- #pragma nv_diagnostic pop
1269
1270
int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1271
@@ -1141,33 +1276,46 @@ __global__ void Marlin(
1276
}
1277
};
1278
1144
- auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1145
- dequant<scalar_t2, w_type_id, dequant_skip_flop>(q, frag_b_ptr);
1279
+ auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) {
1280
+ if constexpr (a_type.size_bits() != b_type.size_bits()) {
1281
+ if constexpr (is_a_8bit && has_zp) {
1282
+ sub_zp_and_dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(
1283
+ q, frag_b_ptr, zp);
1284
+ } else {
1285
+ dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(q, frag_b_ptr);
1286
+ }
1287
+ }
1288
};
1289
1290
// Execute the actual tensor core matmul of a sub-tile.
1291
bool is_first_matmul_in_slice = true;
1150
- auto matmul = [&](int k) {
1292
+ auto matmul = [&](int k, int pipe) {
1293
+ if (is_a_8bit) return;
1294
int k2 = k % 2;
1295
+ constexpr int g =
1296
+ group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1;
1297
const bool is_new_zp =
1153
- ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) ||
1298
+ (group_blocks == 0) ||
1299
+ ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) &&
1300
+ (pipe % g == 0) ||
1301
(group_blocks == -1 && is_first_matmul_in_slice);
1302
if constexpr (has_zp && !is_zp_float) {
1303
if (is_new_zp) {
1304
if constexpr (group_blocks == -1) is_first_matmul_in_slice = false;
1305
int zp_quant_0, zp_quant_1;
1306
1160
- if constexpr (w_type.size_bits() == 4) {
1307
+ if constexpr (b_type.size_bits() == 4) {
1308
zp_quant_0 = frag_qzp[k2][0];
1309
zp_quant_1 = zp_quant_0 >> 8;
1310
} else {
1164
- static_assert(w_type.size_bits() == 8);
1311
+ static_assert(b_type.size_bits() == 8);
1312
zp_quant_0 = frag_qzp[k2][0];
1313
zp_quant_1 = frag_qzp[k2][1];
1314
}
1315
1169
- dequant_data(zp_quant_0, reinterpret_cast<scalar_t2*>(&frag_zp));
1170
- dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1316
+ dequant_data(zp_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_zp));
1317
+ dequant_data(zp_quant_1,
1318
+ reinterpret_cast<scalar_32bit_t*>(&frag_zp) + 2);
1319
}
1320
}
1321
if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
@@ -1177,14 +1325,14 @@ __global__ void Marlin(
1325
}
1326
}
1327
1180
- if constexpr (w_type == vllm::kFE2M1f) {
1328
+ if constexpr (b_type == vllm::kFE2M1f) {
1329
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1330
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1331
1184
- dequant_fp8_scales<scalar_t2, s_type_id>(
1185
- s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1186
- dequant_fp8_scales<scalar_t2, s_type_id>(
1187
- s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1332
+ dequant_fp8_scales<c_scalar_t2, s_type_id>(
1333
+ s_quant_0, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]));
1334
+ dequant_fp8_scales<c_scalar_t2, s_type_id>(
1335
+ s_quant_1, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]) + 2);
1336
}
1337
1338
// We have the m dimension as the inner loop in order to encourage overlapping
@@ -1195,61 +1343,168 @@ __global__ void Marlin(
1343
FragB frag_b1;
1344
int b_quant_0, b_quant_1;
1345
1198
- if constexpr (w_type_id == vllm::kFE2M1f.id()) {
1346
+ if constexpr (b_type_id == vllm::kFE2M1f.id()) {
1347
b_quant_1 = frag_b_quant[k2][0][j];
1348
b_quant_0 = b_quant_1 << 8;
1201
- } else if constexpr (w_type.size_bits() == 4) {
1349
+ } else if constexpr (b_type.size_bits() == 4) {
1350
b_quant_0 = frag_b_quant[k2][0][j];
1351
b_quant_1 = b_quant_0 >> 8;
1352
} else {
1205
- static_assert(w_type.size_bits() == 8);
1353
+ static_assert(b_type.size_bits() == 8);
1354
int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
1355
b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
1356
b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1357
}
1358
1211
- dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1212
- dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1359
+ dequant_data(b_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_b0));
1360
+ dequant_data(b_quant_1, reinterpret_cast<scalar_32bit_t*>(&frag_b1));
1361
1214
- if constexpr (dequant_skip_flop && has_zp && !is_zp_float) {
1215
- sub_zp<scalar_t>(frag_b0, frag_zp[j], 0);
1216
- sub_zp<scalar_t>(frag_b1, frag_zp[j], 1);
1362
+ if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) {
1363
+ sub_zp<a_type_id>(frag_b0, frag_zp[j], 0);
1364
+ sub_zp<a_type_id>(frag_b1, frag_zp[j], 1);
1365
}
1366
1367
// Apply scale to frag_b0
1220
- if constexpr (has_act_order) {
1368
+ if constexpr (has_act_order && !is_a_8bit) {
1369
static_assert(group_blocks != -1);
1222
- scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1223
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1224
- scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1225
- act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1370
+ scale4<a_type_id>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1371
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1372
+ scale4<a_type_id>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1373
+ act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1374
} else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1227
- group_blocks == -1) {
1375
+ group_blocks == -1 && !is_a_8bit) {
1376
int idx = (threadIdx.x / 4) % 2;
1229
- scalar_t2 s2 = Dtype::nums2num2(
1377
+ scalar_t2 s2 = Adtype::nums2num2(
1378
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
1379
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
1380
if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1233
- scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1234
- scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1235
- } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) {
1381
+ scale_and_sub<a_type_id>(frag_b0, s2.x, frag_zp[j].x);
1382
+ scale_and_sub<a_type_id>(frag_b1, s2.y, frag_zp[j].y);
1383
+ } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 &&
1384
+ !is_a_8bit) {
1385
if (is_new_zp)
1386
frag_zp[j] = __hmul2(frag_zp[j],
1387
*reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1239
- scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1240
- scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1241
- } else if constexpr (group_blocks != -1) {
1242
- scale<scalar_t>(frag_b0, frag_s[k2][j], 0);
1243
- scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
1388
+ scale_and_sub<a_type_id>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x);
1389
+ scale_and_sub<a_type_id>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y);
1390
+ } else if constexpr (group_blocks != -1 && !is_a_8bit) {
1391
+ scale<a_type_id>(frag_b0, frag_s[k2][j], 0);
1392
+ scale<a_type_id>(frag_b1, frag_s[k2][j], 1);
1393
}
1394
1395
#pragma unroll
1396
for (int i = 0; i < thread_m_blocks; i++) {
1397
if constexpr (m_block_size_8) {
1249
- mma_trans<scalar_t>(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]);
1398
+ mma_trans<a_type_id>(frag_a[k2][i], frag_b0, frag_b1,
1399
+ frag_c[i][j][0]);
1400
} else {
1251
- mma<scalar_t>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1252
- mma<scalar_t>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
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]);
1403
+ }
1404
+ }
1405
+ }
1406
+ };
1407
+
1408
+ auto matmul_a8 = [&](int k) {
1409
+ int k2 = k % 2;
1410
+ #pragma unroll
1411
+ for (int j = 0; j < 2; j++) {
1412
+ FragB frag_b[2];
1413
+
1414
+ if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) {
1415
+ dequant_data(frag_b_quant[k2][0][j * 2],
1416
+ reinterpret_cast<scalar_32bit_t*>(&frag_b));
1417
+ dequant_data(frag_b_quant[k2][0][j * 2 + 1],
1418
+ reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2);
1419
+ } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) {
1420
+ int off = (threadIdx.x / 32) % 2 * 2 + j;
1421
+ int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF;
1422
+ dequant_data(frag_b_quant[k2][0][j * 2],
1423
+ reinterpret_cast<scalar_32bit_t*>(&frag_b), zp);
1424
+ zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF;
1425
+ dequant_data(frag_b_quant[k2][0][j * 2 + 1],
1426
+ reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2, zp);
1427
+ } else {
1428
+ reinterpret_cast<int2*>(&frag_b)[0] =
1429
+ reinterpret_cast<int2*>(&frag_b_quant[k2][j])[0];
1430
+ reinterpret_cast<int2*>(&frag_b)[1] =
1431
+ reinterpret_cast<int2*>(&frag_b_quant[k2][j])[1];
1432
+ }
1433
+
1434
+ #pragma unroll
1435
+ 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]);
1440
+ }
1441
+
1442
+ if constexpr (group_blocks != -1) {
1443
+ if (group_blocks == 2 || k == 1) {
1444
+ if constexpr (a_type == vllm::kS8) {
1445
+ int2 s_vals[2];
1446
+ s_vals[0] = {
1447
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[0],
1448
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[1]};
1449
+ s_vals[1] = {
1450
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[0],
1451
+ (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[1]};
1452
+
1453
+ #pragma unroll
1454
+ for (int i = 0; i < thread_m_blocks; i++) {
1455
+ #pragma unroll
1456
+ for (int g = 0; g < 4; g++) {
1457
+ int scale = reinterpret_cast<int*>(&s_vals[0])[g % 2];
1458
+ *reinterpret_cast<int32_t*>(&frag_c[i][j][0][g]) +=
1459
+ *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][0][g]) *
1460
+ scale;
1461
+ frag_c_tmp[i][j][0][g] = 0.0f;
1462
+ }
1463
+
1464
+ #pragma unroll
1465
+ for (int g = 0; g < 4; g++) {
1466
+ int scale = reinterpret_cast<int*>(&s_vals[1])[g % 2];
1467
+ *reinterpret_cast<int32_t*>(&frag_c[i][j][1][g]) +=
1468
+ *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][1][g]) *
1469
+ scale;
1470
+ frag_c_tmp[i][j][1][g] = 0.0f;
1471
+ }
1472
+ }
1473
+ } else {
1474
+ float2 s_vals[2];
1475
+ if constexpr (s_type_id != vllm::kFE8M0fnu.id()) {
1476
+ static_assert(a_type.size_bits() == 16 ||
1477
+ s_type.size_bits() == 16);
1478
+ s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]);
1479
+ s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]);
1480
+ } else {
1481
+ int32_t* s_vals_int = reinterpret_cast<int32_t*>(&s_vals[0]);
1482
+ int32_t s_vals_e8m0 =
1483
+ *reinterpret_cast<int32_t*>(&frag_s[k2][j][0]);
1484
+
1485
+ s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23;
1486
+ s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15;
1487
+ s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7;
1488
+ s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1;
1489
+ }
1490
+
1491
+ #pragma unroll
1492
+ for (int i = 0; i < thread_m_blocks; i++) {
1493
+ #pragma unroll
1494
+ for (int g = 0; g < 4; g++) {
1495
+ float scale = reinterpret_cast<float*>(&s_vals[0])[g % 2];
1496
+ frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale;
1497
+ frag_c_tmp[i][j][0][g] = 0.0f;
1498
+ }
1499
+
1500
+ #pragma unroll
1501
+ for (int g = 0; g < 4; g++) {
1502
+ float scale = reinterpret_cast<float*>(&s_vals[1])[g % 2];
1503
+ frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale;
1504
+ frag_c_tmp[i][j][1][g] = 0.0f;
1505
+ }
1506
+ }
1507
+ }
1508
}
1509
}
1510
}
@@ -1263,7 +1518,8 @@ __global__ void Marlin(
1518
constexpr int red_off = threads / b_sh_stride_threads / 2;
1519
if (red_off >= 1) {
1520
auto red_idx = threadIdx.x / b_sh_stride_threads;
1266
- constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1521
+ constexpr int red_sh_stride =
1522
+ b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2;
1523
constexpr int red_sh_delta = b_sh_stride_threads;
1524
int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
1525
(threadIdx.x % b_sh_stride_threads);
@@ -1278,7 +1534,8 @@ __global__ void Marlin(
1534
for (int i = red_off; i > 0; i /= 2) {
1535
if (i <= red_idx && red_idx < 2 * i) {
1536
#pragma unroll
1281
- for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) {
1537
+ for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2;
1538
+ j += (m_block_size_8 ? 2 : 1)) {
1539
int red_sh_wr =
1540
red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
1541
if (i < red_off) {
@@ -1287,24 +1544,26 @@ __global__ void Marlin(
1544
float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
1545
#pragma unroll
1546
for (int k = 0; k < 4; k++)
1290
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
1547
+ reinterpret_cast<FragC*>(
1548
+ frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] +=
1549
c_rd[k] + c_wr[k];
1550
}
1293
- sh_red[red_sh_wr] =
1294
- reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
1551
+ sh_red[red_sh_wr] = reinterpret_cast<int4*>(
1552
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j];
1553
}
1554
}
1555
__syncthreads();
1556
}
1557
if (red_idx == 0) {
1558
#pragma unroll
1301
- for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) {
1559
+ for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2;
1560
+ i += (m_block_size_8 ? 2 : 1)) {
1561
float* c_rd =
1562
reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
1563
#pragma unroll
1564
for (int j = 0; j < 4; j++)
1306
- reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
1307
- c_rd[j];
1565
+ reinterpret_cast<FragC*>(
1566
+ frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j];
1567
}
1568
}
1569
__syncthreads();
@@ -1320,10 +1579,10 @@ __global__ void Marlin(
1579
// We are very careful here to reduce directly in the output buffer to
1580
// maximize L2 cache utilization in this step. To do this, we write out
1581
// results in FP16 (but still reduce with FP32 compute).
1323
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1582
+ constexpr int active_threads = 32 * tb_n_warps;
1583
if (threadIdx.x < active_threads) {
1584
int c_gl_stride = prob_n / 8;
1326
- int c_gl_wr_delta_o = 8 * c_gl_stride;
1585
+ int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1);
1586
int c_gl_wr_delta_i = 4 * (active_threads / 32);
1587
int c_gl_wr;
1588
if constexpr (m_block_size_8) {
@@ -1331,9 +1590,9 @@ __global__ void Marlin(
1590
4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8;
1591
c_gl_wr += (2 * thread_n_blocks) * slice_col;
1592
} else {
1334
- c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
1593
+ c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) +
1594
4 * (threadIdx.x / 32) + threadIdx.x % 4;
1336
- c_gl_wr += (2 * thread_n_blocks) * slice_col;
1595
+ c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1);
1596
}
1597
constexpr int c_sh_wr_delta = active_threads;
1598
auto c_sh_wr = threadIdx.x;
@@ -1351,6 +1610,14 @@ __global__ void Marlin(
1610
&C[c_gl_wr + i * c_gl_stride +
1611
(threadIdx.x % 8) / 4 * c_gl_wr_delta_i],
1612
(threadIdx.x % 4) * 2 + i < prob_m);
1613
+ } else if constexpr (is_a_8bit) {
1614
+ int2* sh_red_int2 = reinterpret_cast<int2*>(sh_red);
1615
+ int2* c_int2 = reinterpret_cast<int2*>(C);
1616
+ cp_async2_ca_pred(
1617
+ &sh_red_int2[c_sh_wr + c_sh_wr_delta * i],
1618
+ &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1619
+ c_gl_wr_delta_i * (i % 2)],
1620
+ i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m);
1621
} else {
1622
cp_async4_pred(
1623
&sh_red[c_sh_wr + c_sh_wr_delta * i],
@@ -1370,36 +1637,51 @@ __global__ void Marlin(
1637
(m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m);
1638
if (mask) {
1639
if (!first) {
1373
- int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
1640
+ c_scalar_t* c_red_f16;
1641
+ if constexpr (is_a_8bit) {
1642
+ int2 tmp =
1643
+ reinterpret_cast<int2*>(sh_red)[c_sh_wr + i * c_sh_wr_delta];
1644
+ c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
1645
+ } else {
1646
+ int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta];
1647
+ c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp);
1648
+ }
1649
#pragma unroll
1375
- for (int j = 0; j < 2 * 4; j++) {
1650
+ for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
1651
int delta = 0;
1652
if constexpr (m_block_size_8) {
1653
delta = j % 2 == 1 ? -2 : 0;
1654
}
1655
reinterpret_cast<float*>(
1381
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] +=
1382
- Dtype::num2float(reinterpret_cast<scalar_t*>(&c_red)[j]);
1656
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j +
1657
+ (i % 4) + delta] += Cdtype::num2float(c_red_f16[j]);
1658
}
1659
}
1660
if (!last) {
1386
- int4 c;
1661
+ c_scalar_t c_f16[is_a_8bit ? 4 : 8];
1662
#pragma unroll
1388
- for (int j = 0; j < 2 * 4; j++) {
1663
+ for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) {
1664
int delta = 0;
1665
if constexpr (m_block_size_8) {
1666
delta = j % 2 == 1 ? -2 : 0;
1667
}
1393
- reinterpret_cast<scalar_t*>(&c)[j] =
1394
- Dtype::float2num(reinterpret_cast<float*>(
1395
- &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]);
1668
+ c_f16[j] = Cdtype::float2num(reinterpret_cast<float*>(
1669
+ &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j +
1670
+ (i % 4) + delta]);
1671
}
1397
- if constexpr (m_block_size_8)
1672
+ if constexpr (m_block_size_8) {
1673
C[c_gl_wr + i * c_gl_stride +
1399
- (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = c;
1400
- else
1674
+ (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] =
1675
+ *reinterpret_cast<int4*>(c_f16);
1676
+ } else if constexpr (is_a_8bit) {
1677
+ int2* c_int2 = reinterpret_cast<int2*>(C);
1678
+ c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1679
+ c_gl_wr_delta_i * (i % 2)] =
1680
+ *reinterpret_cast<int2*>(c_f16);
1681
+ } else {
1682
C[c_gl_wr + c_gl_wr_delta_o * (i / 2) +
1402
- c_gl_wr_delta_i * (i % 2)] = c;
1683
+ c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast<int4*>(c_f16);
1684
+ }
1685
}
1686
}
1687
}
@@ -1414,10 +1696,10 @@ __global__ void Marlin(
1696
1697
constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
1698
1417
- constexpr int active_threads = 32 * thread_n_blocks / 4;
1699
+ constexpr int active_threads = 32 * tb_n_warps;
1700
bool is_th_active = threadIdx.x < active_threads;
1701
1420
- constexpr int num_floats = thread_m_blocks * 4 * 2 * 4;
1702
+ constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4;
1703
constexpr int th_size = num_floats * sizeof(float) / 16;
1704
1705
int c_cur_offset = locks_off * c_size;
@@ -1471,7 +1753,7 @@ __global__ void Marlin(
1753
} else {
1754
c_sh_wr =
1755
(4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
1474
- c_sh_wr += 32 * (threadIdx.x / 32);
1756
+ c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32);
1757
}
1758
1759
int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
@@ -1481,47 +1763,47 @@ __global__ void Marlin(
1763
// We first reorder in shared memory to guarantee the most efficient final
1764
// global write patterns
1765
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1484
- scalar_t2 res =
1485
- Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1766
+ c_scalar_t2 res =
1767
+ Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
1768
1769
// For per-column quantization we finally apply the scale here (only for
1770
// 4-bit)
1489
- if constexpr (!has_act_order && group_blocks == -1 &&
1490
- w_type.size_bits() == 4 &&
1771
+ if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit &&
1772
+ b_type.size_bits() == 4 &&
1773
(has_zp && dequant_skip_flop || !has_zp)) {
1492
- scalar_t2 tmp_scale = s[0];
1774
+ c_scalar_t2 tmp_scale = s[0];
1775
if constexpr (m_block_size_8) {
1494
- tmp_scale = Dtype::num2num2(
1776
+ tmp_scale = Cdtype::num2num2(
1777
reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
1778
}
1779
res = __hmul2(res, tmp_scale);
1780
}
1781
1500
- if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1782
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1783
res = __hmul2(res, global_scale);
1784
}
1785
if (has_bias && last) {
1504
- scalar_t2 tmp_bias = b_bias[0];
1786
+ c_scalar_t2 tmp_bias = b_bias[0];
1787
if constexpr (m_block_size_8) {
1506
- tmp_bias = Dtype::num2num2(
1788
+ tmp_bias = Cdtype::num2num2(
1789
reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
1790
}
1791
res = __hadd2(res, tmp_bias);
1792
}
1793
1794
if constexpr (m_block_size_8) {
1513
- ((scalar_t*)sh_red)[idx] = res.x;
1514
- ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1795
+ ((c_scalar_t*)sh_red)[idx] = res.x;
1796
+ ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1797
} else {
1516
- ((scalar_t2*)sh_red)[idx] = res;
1798
+ ((c_scalar_t2*)sh_red)[idx] = res;
1799
}
1800
};
1801
1520
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
1802
+ if (threadIdx.x / 32 < tb_n_warps) {
1803
#pragma unroll
1804
for (int i = 0; i < thread_m_blocks; i++) {
1805
#pragma unroll
1524
- for (int j = 0; j < 4; j++) {
1806
+ for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) {
1807
if constexpr (m_block_size_8) {
1808
int wr = c_sh_wr + 16 * j;
1809
write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
@@ -1557,9 +1839,9 @@ __global__ void Marlin(
1839
i++) {
1840
if (c_gl_wr < c_gl_wr_end) {
1841
if (use_atomic_add && slice_count > 1) {
1560
- scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[c_gl_wr]);
1561
- scalar_t2* sh_red_half2 =
1562
- reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
1842
+ c_scalar_t2* C_half2 = reinterpret_cast<c_scalar_t2*>(&C[c_gl_wr]);
1843
+ c_scalar_t2* sh_red_half2 =
1844
+ reinterpret_cast<c_scalar_t2*>(&sh_red[c_sh_rd]);
1845
#pragma unroll
1846
for (int a = 0; a < 4; a++) {
1847
atomicAdd(&C_half2[a], sh_red_half2[a]);
@@ -1635,7 +1917,13 @@ __global__ void Marlin(
1917
wait_for_stage();
1918
init_same_group(pipe % stages);
1919
}
1638
- matmul(k);
1920
+
1921
+ if constexpr (!is_a_8bit) {
1922
+ matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0));
1923
+ } else {
1924
+ static_assert(group_blocks != 0 && group_blocks != 1);
1925
+ matmul_a8(k);
1926
+ }
1927
}
1928
slice_iters--;
1929
if (slice_iters == 0) {
@@ -1668,13 +1956,47 @@ __global__ void Marlin(
1956
// While this pattern may not be the most readable, other ways of writing
1957
// the loop seemed to noticeably worse performance after compilation.
1958
if (slice_iters == 0) {
1959
+ if constexpr (is_a_8bit) {
1960
+ float frag_a_s[2 * thread_m_blocks];
1961
+
1962
+ for (int i = 0; i < 2 * thread_m_blocks; i++)
1963
+ frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4];
1964
+
1965
+ #pragma unroll
1966
+ for (int j = 0; j < 2; j++) {
1967
+ #pragma unroll
1968
+ for (int i = 0; i < thread_m_blocks; i++) {
1969
+ #pragma unroll
1970
+ for (int g = 0; g < 4; g++) {
1971
+ float c_val = frag_c[i][j][0][g];
1972
+
1973
+ if constexpr (a_type == vllm::kS8) {
1974
+ c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
1975
+ }
1976
+ float s_val = frag_a_s[i * 2 + g / 2];
1977
+ frag_c[i][j][0][g] = c_val * s_val;
1978
+ }
1979
+ #pragma unroll
1980
+ for (int g = 0; g < 4; g++) {
1981
+ float c_val = frag_c[i][j][1][g];
1982
+
1983
+ if constexpr (a_type == vllm::kS8) {
1984
+ c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val));
1985
+ }
1986
+ float s_val = frag_a_s[i * 2 + g / 2];
1987
+ frag_c[i][j][1][g] = c_val * s_val;
1988
+ }
1989
+ }
1990
+ }
1991
+ }
1992
+
1993
cp_async_wait<0>();
1994
bool last = slice_idx == slice_count - 1;
1995
// For per-column scales, we only fetch them here in the final step before
1996
// write-out
1997
if constexpr (!has_act_order && group_blocks == -1 &&
1998
(has_zp && dequant_skip_flop || !has_zp)) {
1677
- if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1999
+ if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) {
2000
if (s_sh_wr_pred) {
2001
cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
2002
}
@@ -1692,20 +2014,27 @@ __global__ void Marlin(
2014
}
2015
2016
if constexpr (!has_act_order && group_blocks == -1 &&
1695
- (has_zp && dequant_skip_flop || !has_zp)) {
1696
- if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
2017
+ (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) {
2018
+ if constexpr (is_a_8bit) {
2019
cp_async_wait<0>();
2020
__syncthreads();
1699
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
2021
+ if (threadIdx.x / 32 < tb_n_warps) {
2022
+ reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
2023
+ }
2024
+ } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) {
2025
+ cp_async_wait<0>();
2026
+ __syncthreads();
2027
+ if (threadIdx.x / 32 < tb_n_warps) {
2028
reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
2029
reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
2030
if constexpr (m_block_size_8) {
2031
int idx = (threadIdx.x / 4) % 2;
1704
- scalar_t2* frag_s_half2 = reinterpret_cast<scalar_t2*>(frag_s);
2032
+ c_scalar_t2* frag_s_half2 =
2033
+ reinterpret_cast<c_scalar_t2*>(frag_s);
2034
#pragma unroll
2035
for (int i = 0; i < 8; i++) {
1707
- frag_s_half2[i] = Dtype::num2num2(
1708
- reinterpret_cast<scalar_t*>(&frag_s_half2[i])[idx]);
2036
+ frag_s_half2[i] = Cdtype::num2num2(
2037
+ reinterpret_cast<c_scalar_t*>(&frag_s_half2[i])[idx]);
2038
}
2039
}
2040
}
@@ -1715,26 +2044,48 @@ __global__ void Marlin(
2044
// For 8-bit channelwise, we apply the scale before the global reduction
2045
// that converts the fp32 results to fp16 (so that we avoid possible
2046
// overflow in fp16)
1718
- if constexpr (!has_act_order && group_blocks == -1 &&
1719
- w_type.size_bits() == 8 &&
1720
- (has_zp && dequant_skip_flop || !has_zp)) {
1721
- if (threadIdx.x / 32 < thread_n_blocks / 4) {
2047
+ if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) {
2048
+ #pragma unroll
2049
+ for (int j = 0; j < 2; j++) {
2050
+ float2 aa[2];
2051
+ aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]);
2052
+ aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]);
2053
+
2054
+ #pragma unroll
2055
+ for (int i = 0; i < thread_m_blocks; i++) {
2056
+ #pragma unroll
2057
+ for (int g = 0; g < 4; g++) {
2058
+ float scale = reinterpret_cast<float*>(&aa[0])[g % 2];
2059
+ frag_c[i][j][0][g] *= scale;
2060
+ }
2061
+
2062
+ #pragma unroll
2063
+ for (int g = 0; g < 4; g++) {
2064
+ float scale = reinterpret_cast<float*>(&aa[1])[g % 2];
2065
+ frag_c[i][j][1][g] *= scale;
2066
+ }
2067
+ }
2068
+ }
2069
+ } else if (!has_act_order && group_blocks == -1 &&
2070
+ b_type.size_bits() == 8 &&
2071
+ (has_zp && dequant_skip_flop || !has_zp)) {
2072
+ if (threadIdx.x / 32 < tb_n_warps) {
2073
#pragma unroll
2074
for (int i = 0; i < thread_m_blocks; i++) {
2075
#pragma unroll
2076
for (int j = 0; j < 4; j++) {
1726
- scale_float<scalar_t>(
2077
+ scale_float<c_type_id>(
2078
reinterpret_cast<float*>(&frag_c[i][j][0][0]),
2079
frag_s[j / 2][2 * (j % 2) + 0]);
1729
- scale_float<scalar_t>(
2080
+ scale_float<c_type_id>(
2081
reinterpret_cast<float*>(&frag_c[i][j][0][2]),
2082
frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
2083
2084
if constexpr (!m_block_size_8) {
1734
- scale_float<scalar_t>(
2085
+ scale_float<c_type_id>(
2086
reinterpret_cast<float*>(&frag_c[i][j][1][0]),
2087
frag_s[j / 2][2 * (j % 2) + 1]);
1737
- scale_float<scalar_t>(
2088
+ scale_float<c_type_id>(
2089
reinterpret_cast<float*>(&frag_c[i][j][1][2]),
2090
frag_s[j / 2][2 * (j % 2) + 1]);
2091
}
@@ -1758,7 +2109,8 @@ __global__ void Marlin(
2109
cp_async_wait<0>();
2110
__syncthreads();
2111
reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
1761
- reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
2112
+ if constexpr (!is_a_8bit)
2113
+ reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
2114
__syncthreads();
2115
}
2116
@@ -1768,21 +2120,22 @@ __global__ void Marlin(
2120
// only the last block in a slice actually writes the result
2121
write_result(last);
2122
slice_row = 0;
1771
- slice_col_par++;
1772
- slice_col++;
2123
+ if (!in_part2) {
2124
+ slice_col_par += gridDim.x;
2125
+ } else {
2126
+ slice_col_par++;
2127
+ slice_col++;
2128
+ }
2129
is_first_matmul_in_slice = true;
2130
init_slice();
2131
2132
if (slice_iters) {
2133
a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
2134
(threadIdx.x % a_gl_rd_delta_o);
1779
- #pragma unroll
1780
- for (int i = 0; i < b_sh_wr_iters; i++)
1781
- B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
1782
- if (slice_col == 0) {
1783
- #pragma unroll
1784
- for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1785
- }
2135
+ a_gl_rd += a_gl_rd_delta_o * slice_row;
2136
+ b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) +
2137
+ (threadIdx.x % b_sh_stride);
2138
+ b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row;
2139
2140
bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
2141
// Update slice k/n for scales loading
@@ -1791,12 +2144,28 @@ __global__ void Marlin(
2144
slice_k_finish = slice_k_start + tb_k * slice_iters;
2145
slice_k_start_shared_fetch = slice_k_start;
2146
slice_n_offset = act_s_col_tb_stride * slice_col;
1794
-
2147
} else {
1796
- s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1797
- zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2148
+ if constexpr (group_blocks == -1) {
2149
+ s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
2150
+ zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
2151
+ } else if constexpr (group_blocks >= thread_k_blocks) {
2152
+ s_gl_rd =
2153
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2154
+ s_sh_stride * slice_col + threadIdx.x;
2155
+ zp_gl_rd =
2156
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
2157
+ zp_sh_stride * slice_col + threadIdx.x;
2158
+ } else {
2159
+ s_gl_rd =
2160
+ s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2161
+ threadIdx.x / s_sh_stride) +
2162
+ s_sh_stride * slice_col + threadIdx.x % s_sh_stride;
2163
+ zp_gl_rd =
2164
+ zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks +
2165
+ threadIdx.x / zp_sh_stride) +
2166
+ zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride;
2167
+ }
2168
}
1799
-
2169
start_pipes();
2170
}
2171
}
csrc/torch_bindings.cpp
+12
-5
@@ -298,9 +298,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
298
// gptq_marlin Optimized Quantized GEMM for GPTQ.
299
ops.def(
300
"gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
301
- "Tensor? b_bias_or_none,"
302
- "Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
303
- "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
301
+ "Tensor? b_bias_or_none,Tensor b_scales, "
302
+ "Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
303
+ "Tensor? "
304
+ "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_type_id, "
305
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
306
"bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor");
307
// conditionally compiled so impl registration is in source file
@@ -308,13 +309,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
309
// gptq_marlin repack from GPTQ.
310
ops.def(
311
"gptq_marlin_repack(Tensor b_q_weight, Tensor perm, "
311
- "SymInt size_k, SymInt size_n, int num_bits) -> Tensor");
312
+ "SymInt size_k, SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
313
// conditionally compiled so impl registrations are in source file
314
315
// awq_marlin repack from AWQ.
316
ops.def(
317
"awq_marlin_repack(Tensor b_q_weight, SymInt size_k, "
317
- "SymInt size_n, int num_bits) -> Tensor");
318
+ "SymInt size_n, int num_bits, bool is_a_8bit) -> Tensor");
319
+ // conditionally compiled so impl registrations are in source file
320
+
321
+ // preprocess W-int4A-fp8 weight for marlin kernel
322
+ ops.def(
323
+ "marlin_int4_fp8_preprocess(Tensor qweight, "
324
+ "Tensor? qzeros_or_none, bool inplace) -> Tensor");
325
// conditionally compiled so impl registrations are in source file
326
327
// CUTLASS w4a8 GEMM
docs/design/moe_kernel_features.md
+1
-1
@@ -60,7 +60,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
60
61
- [`ModelOptFp8MoEMethod`][vllm.model_executor.layers.quantization.modelopt.ModelOptFp8MoEMethod]
62
- [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]
63
-- [`CompressedTensorsW4A4Nvfp4MoeMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoeMethod]
63
+- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod]
64
- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW8A8Fp8MoEMethod]
65
- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
66
- [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
tests/kernels/moe/test_moe.py
+197
-59
@@ -21,7 +21,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
21
22
import vllm.model_executor.layers.fused_moe # noqa
23
from tests.kernels.moe.utils import fused_moe
24
-from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
24
+from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
25
from vllm._aiter_ops import rocm_aiter_ops
26
from vllm.config import VllmConfig, set_current_vllm_config
27
from vllm.distributed.parallel_state import init_distributed_environment
@@ -65,6 +65,64 @@ NUM_EXPERTS = [8, 64, 192]
65
EP_SIZE = [1, 4]
66
TOP_KS = [2, 6]
67
68
+MOE_MARLIN_QUANT_TEST_CONFIGS = [
69
+ # AWQ-INT4
70
+ {"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]},
71
+ # GPTQ-INT4
72
+ {
73
+ "b_type": scalar_types.uint4b8,
74
+ "support_act_order": True,
75
+ "group_blocks": [-1, 2, 4, 8],
76
+ },
77
+ # GPTQ-INT8
78
+ {
79
+ "b_type": scalar_types.uint8b128,
80
+ "support_act_order": True,
81
+ "group_blocks": [-1, 2, 4, 8],
82
+ },
83
+ # FP8
84
+ {"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]},
85
+ # NVFP4
86
+ {"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]},
87
+ # MXFP4
88
+ {
89
+ "a_type": [scalar_types.bfloat16],
90
+ "b_type": scalar_types.float4_e2m1f,
91
+ "group_blocks": [2],
92
+ },
93
+ # AWQ-INT4 with INT8 activation
94
+ {
95
+ "a_type": [scalar_types.int8],
96
+ "b_type": scalar_types.uint4,
97
+ "group_blocks": [-1, 2, 4, 8],
98
+ },
99
+ # GPTQ-INT4 with INT8 activation
100
+ {
101
+ "a_type": [scalar_types.int8],
102
+ "b_type": scalar_types.uint4b8,
103
+ "group_blocks": [-1, 2, 4, 8],
104
+ },
105
+ # GPTQ-INT4 with FP8 activation
106
+ {
107
+ "a_type": [scalar_types.float8_e4m3fn],
108
+ "b_type": scalar_types.uint4b8,
109
+ "group_blocks": [-1, 2, 4, 8],
110
+ },
111
+ # AWQ-INT4 with FP8 activation
112
+ {
113
+ "a_type": [scalar_types.float8_e4m3fn],
114
+ "b_type": scalar_types.uint4,
115
+ "group_blocks": [-1, 2, 4, 8],
116
+ },
117
+ # MXFP4 with FP8 activation
118
+ {
119
+ "a_type": [scalar_types.float8_e4m3fn],
120
+ "b_type": scalar_types.float4_e2m1f,
121
+ "c_type": [scalar_types.bfloat16],
122
+ "group_blocks": [2],
123
+ },
124
+]
125
+
126
FUSED_MOE_MNK_FACTORS = [
127
(1, 128, 128),
128
(1, 2048, 128),
@@ -505,63 +563,74 @@ def marlin_moe_generate_valid_test_cases():
563
m_list = [1, 123, 666]
564
n_list = [128, 1024]
565
k_list = [256, 2048]
508
- e_list = [4, 12]
566
+ e_list = [5, 12]
567
topk_list = [2, 3]
568
ep_size_list = [1, 4]
511
- dtype_list = [torch.bfloat16]
512
- group_size_list = [-1, 32, 128]
569
act_order_list = [True, False]
514
- quant_type_list = [
515
- scalar_types.float4_e2m1f,
516
- scalar_types.float8_e4m3fn,
517
- scalar_types.uint4,
518
- scalar_types.uint4b8,
519
- scalar_types.uint8b128,
520
- ]
570
is_k_full_list = [True, False]
571
572
all_combinations = itertools.product(
573
+ MOE_MARLIN_QUANT_TEST_CONFIGS,
574
m_list,
575
n_list,
576
k_list,
577
e_list,
578
topk_list,
579
ep_size_list,
530
- dtype_list,
531
- group_size_list,
580
act_order_list,
533
- quant_type_list,
581
is_k_full_list,
582
)
583
584
def is_invalid(
538
- m, n, k, e, topk, ep_size, dtype, group_size, act_order, quant_type, is_k_full
585
+ a_type,
586
+ b_type,
587
+ c_type,
588
+ group_blocks,
589
+ m,
590
+ n,
591
+ k,
592
+ e,
593
+ topk,
594
+ ep_size,
595
+ act_order,
596
+ is_k_full,
597
):
540
- if quant_type == scalar_types.float8_e4m3fn and group_size not in [-1, 128]:
541
- return False
542
- if quant_type == scalar_types.float4_e2m1f:
543
- if group_size not in [16, 32]:
544
- return False
545
- if dtype == torch.float16 and group_size == 32:
546
- return False
547
- if quant_type != scalar_types.float4_e2m1f and group_size == 16:
598
+ group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
599
+ if group_size > 0 and k % group_size != 0:
600
return False
601
550
- # Filter act_order
551
- if act_order:
552
- if group_size in (-1, k, n):
553
- return False
554
- if quant_type not in [scalar_types.uint4b8]:
555
- return False
556
- elif not is_k_full:
602
+ if act_order and group_size in [-1, k, n]:
603
+ return False
604
+ if group_size in [k, n]:
605
+ return False
606
+ if not act_order and is_k_full:
607
return False
608
559
- return True
609
+ return a_type.size_bits < 16 or a_type is c_type
610
611
cases = []
612
for case in all_combinations:
563
- if is_invalid(*case):
564
- cases.append(case)
613
+ quant_test_config, m, n, k, _, _, _, act_order, *_ = case
614
+ if act_order and not quant_test_config.get("support_act_order", False):
615
+ continue
616
+
617
+ f16_types = [scalar_types.float16]
618
+ inner_combinations = itertools.product(
619
+ quant_test_config.get("a_type", f16_types),
620
+ [quant_test_config["b_type"]],
621
+ quant_test_config.get("c_type", f16_types),
622
+ quant_test_config["group_blocks"],
623
+ )
624
+
625
+ for sub_case in inner_combinations:
626
+ if (
627
+ sub_case[0] == scalar_types.float8_e4m3fn
628
+ and current_platform.get_device_capability() not in [89, 120]
629
+ ):
630
+ continue
631
+ args = sub_case + (m, n, k) + case[4:]
632
+ if is_invalid(*args):
633
+ cases.append(args)
634
return cases
635
636
@@ -571,6 +640,7 @@ class MarlinMoEWeightData:
640
qweight: torch.Tensor
641
scales: torch.Tensor
642
global_scale: torch.Tensor | None
643
+ a_scales_factor: torch.Tensor | None
644
g_idx: torch.Tensor | None
645
zeros: torch.Tensor | None
646
sort_indices: torch.Tensor | None
@@ -583,11 +653,20 @@ class MarlinMoEWeightData:
653
group_size: int,
654
act_order: bool | None = None,
655
bias: torch.Tensor | None = None,
656
+ input_type: ScalarType = None,
657
) -> "MarlinMoEWeightData":
658
assert w.ndim == 3
659
+
660
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
661
k = w.shape[-1]
662
663
+ if input_type == scalar_types.int8:
664
+ input_dtype = torch.int8
665
+ elif input_type == scalar_types.float8_e4m3fn:
666
+ input_dtype = torch.float8_e4m3fn
667
+ else:
668
+ input_dtype = w.dtype
669
+
670
w_ref_l: list[torch.Tensor] = []
671
qweight_l: list[torch.Tensor] = []
672
scales_l: list[torch.Tensor] = []
@@ -601,11 +680,13 @@ class MarlinMoEWeightData:
680
if quant_type == scalar_types.float4_e2m1f:
681
if group_size == 16:
682
w_ref, qweight, scales, global_scale = (
604
- rand_marlin_weight_nvfp4_like(w[i], group_size)
683
+ rand_marlin_weight_nvfp4_like(
684
+ w[i], group_size, input_dtype=input_dtype
685
+ )
686
)
687
else:
688
w_ref, qweight, scales = rand_marlin_weight_mxfp4_like(
608
- w[i], group_size
689
+ w[i], group_size, input_dtype=input_dtype
690
)
691
global_scale = None
692
@@ -615,13 +696,18 @@ class MarlinMoEWeightData:
696
if global_scale is not None:
697
global_scale_l.append(global_scale)
698
elif quant_type == scalar_types.float8_e4m3fn:
618
- w_ref, qweight, scales = marlin_quant_fp8_torch(w[i], group_size)
699
+ w_ref, qweight, scales = marlin_quant_fp8_torch(
700
+ w[i], group_size, input_dtype=input_dtype
701
+ )
702
w_ref_l.append(w_ref.T)
703
qweight_l.append(qweight)
704
scales_l.append(scales)
705
elif has_zp:
706
w_ref, qweight, scales, zeros = awq_marlin_quantize(
624
- w[i].transpose(1, 0), quant_type, group_size
707
+ w[i].transpose(1, 0),
708
+ quant_type,
709
+ group_size,
710
+ input_dtype=input_dtype,
711
)
712
713
w_ref_l.append(w_ref.T)
@@ -631,7 +717,12 @@ class MarlinMoEWeightData:
717
else:
718
test_perm = torch.randperm(k)
719
w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
634
- w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm
720
+ w[i].transpose(1, 0),
721
+ quant_type,
722
+ group_size,
723
+ act_order,
724
+ test_perm,
725
+ input_dtype=input_dtype,
726
)
727
728
w_ref_l.append(w_ref.T)
@@ -652,11 +743,18 @@ class MarlinMoEWeightData:
743
sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
744
marlin_bias = stack_and_dev(bias_l) if bias_l else None
745
746
+ a_scales_factor = None
747
+ if input_type == scalar_types.int8 and group_size != -1:
748
+ a_scales_factor = 1 / 4096 * scales.max().float()
749
+ scales = scales / scales.max() * 4096
750
+ scales = scales.round().to(torch.int16).view(w.dtype)
751
+
752
return MarlinMoEWeightData(
753
w_ref=w_ref,
754
qweight=qweight,
755
scales=scales,
756
global_scale=global_scale,
757
+ a_scales_factor=a_scales_factor,
758
g_idx=g_idx,
759
zeros=zeros,
760
sort_indices=sort_indices,
@@ -666,28 +764,47 @@ class MarlinMoEWeightData:
764
765
@pytest.mark.flaky(reruns=2)
766
@pytest.mark.parametrize(
669
- ("m, n, k, e, topk, ep_size, dtype, group_size,act_order, quant_type, is_k_full"),
767
+ (
768
+ "a_type, b_type, c_type, group_blocks,"
769
+ "m, n, k, e, topk, ep_size, act_order, is_k_full"
770
+ ),
771
marlin_moe_generate_valid_test_cases(),
772
)
773
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
774
def test_fused_marlin_moe(
674
- m: int,
675
- n: int,
676
- k: int,
677
- e: int,
678
- topk: int,
679
- ep_size: int,
680
- dtype: torch.dtype,
681
- group_size: int,
682
- act_order: bool,
683
- quant_type: ScalarType,
684
- is_k_full: bool,
775
+ a_type,
776
+ b_type,
777
+ c_type,
778
+ group_blocks,
779
+ m,
780
+ n,
781
+ k,
782
+ e,
783
+ topk,
784
+ ep_size,
785
+ act_order,
786
+ is_k_full,
787
):
686
- torch.cuda.manual_seed(0)
788
+ torch.cuda.manual_seed(1)
789
+ group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
790
+
791
+ if c_type == scalar_types.float16:
792
+ dtype = torch.float16
793
+ elif c_type == scalar_types.bfloat16:
794
+ dtype = torch.bfloat16
795
+ else:
796
+ raise RuntimeError("unsupported c_type")
797
+
798
+ if a_type == scalar_types.int8:
799
+ a_dtype = torch.int8
800
+ elif a_type == scalar_types.float8_e4m3fn:
801
+ a_dtype = torch.float8_e4m3fn
802
+ else:
803
+ a_dtype = dtype
804
805
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
689
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
690
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
806
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
807
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
808
809
if ep_size > 1:
810
local_e = e // ep_size
@@ -700,11 +817,19 @@ def test_fused_marlin_moe(
817
e_map = None
818
819
w1_data = MarlinMoEWeightData.make(
703
- w=w1, quant_type=quant_type, group_size=group_size, act_order=act_order
820
+ w=w1,
821
+ quant_type=b_type,
822
+ group_size=group_size,
823
+ act_order=act_order,
824
+ input_type=a_type,
825
)
826
827
w2_data = MarlinMoEWeightData.make(
707
- w=w2, quant_type=quant_type, group_size=group_size, act_order=act_order
828
+ w=w2,
829
+ quant_type=b_type,
830
+ group_size=group_size,
831
+ act_order=act_order,
832
+ input_type=a_type,
833
)
834
835
score = torch.randn((m, e), device="cuda", dtype=dtype)
@@ -712,8 +837,18 @@ def test_fused_marlin_moe(
837
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
838
839
with set_current_vllm_config(vllm_config):
715
- torch_output = torch_moe(
716
- a, w1_data.w_ref, w2_data.w_ref, score, topk, expert_map=e_map
840
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
841
+ topk_weight, topk_ids = torch.topk(score, topk)
842
+ torch_output = torch_experts(
843
+ a,
844
+ w1_data.w_ref,
845
+ w2_data.w_ref,
846
+ topk_weight=topk_weight,
847
+ topk_ids=topk_ids,
848
+ global_num_experts=e,
849
+ expert_map=e_map,
850
+ quant_dtype=a_dtype,
851
+ per_act_token_quant=True,
852
)
853
854
marlin_output = fused_marlin_moe(
@@ -733,15 +868,18 @@ def test_fused_marlin_moe(
868
global_scale2=w2_data.global_scale,
869
g_idx1=w1_data.g_idx,
870
g_idx2=w2_data.g_idx,
871
+ input_global_scale1=w1_data.a_scales_factor,
872
+ input_global_scale2=w2_data.a_scales_factor,
873
sort_indices1=w1_data.sort_indices,
874
sort_indices2=w2_data.sort_indices,
875
w1_zeros=w1_data.zeros,
876
w2_zeros=w2_data.zeros,
740
- quant_type_id=quant_type.id,
877
+ input_dtype=a_dtype,
878
+ quant_type_id=b_type.id,
879
is_k_full=is_k_full,
880
)
881
744
- torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
882
+ torch.testing.assert_close(marlin_output, torch_output, atol=4e-2, rtol=0)
883
884
885
@pytest.mark.flaky(reruns=2)
tests/kernels/quantization/test_marlin_gemm.py
+286
-111
@@ -5,6 +5,8 @@
5
Run `pytest tests/kernels/quantization/test_marlin_gemm.py`.
6
"""
7
8
+import itertools
9
+
10
import pytest
11
import torch
12
@@ -17,8 +19,10 @@ from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
19
GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES,
20
GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES,
21
)
22
+from vllm.model_executor.layers.quantization.utils.int8_utils import (
23
+ per_token_quant_int8,
24
+)
25
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
21
- MARLIN_SUPPORTED_GROUP_SIZES,
26
marlin_make_empty_g_idx,
27
marlin_make_workspace_new,
28
marlin_permute_bias,
@@ -26,7 +30,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
30
query_marlin_supported_quant_types,
31
)
32
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
29
- FP4_MARLIN_SUPPORTED_GROUP_SIZES,
33
rand_marlin_weight_mxfp4_like,
34
rand_marlin_weight_nvfp4_like,
35
)
@@ -50,6 +53,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
53
quantize_weights,
54
sort_weights,
55
)
56
+from vllm.platforms import current_platform
57
from vllm.scalar_type import scalar_types
58
59
ACT_ORDER_OPTS = [False, True]
@@ -65,6 +69,12 @@ MARLIN_24_N_CHUNKS = [512]
69
70
HQQ_SUPPORTED_GROUP_SIZES = [64]
71
72
+MARLIN_REPACK_NK_FACTORS = [
73
+ (4, 8),
74
+ (7, 5),
75
+ (13, 11),
76
+]
77
+
78
MNK_FACTORS = [
79
(1, 1, 1),
80
(1, 4, 8),
@@ -74,6 +84,64 @@ MNK_FACTORS = [
84
85
DTYPES = [torch.float16, torch.bfloat16]
86
87
+DENSE_MARLIN_QUANT_TEST_CONFIGS = [
88
+ # AWQ-INT4
89
+ {"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]},
90
+ # GPTQ-INT4
91
+ {
92
+ "b_type": scalar_types.uint4b8,
93
+ "support_act_order": True,
94
+ "group_blocks": [-1, 2, 4, 8],
95
+ },
96
+ # GPTQ-INT8
97
+ {
98
+ "b_type": scalar_types.uint8b128,
99
+ "support_act_order": True,
100
+ "group_blocks": [-1, 2, 4, 8],
101
+ },
102
+ # FP8
103
+ {"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]},
104
+ # NVFP4
105
+ {"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]},
106
+ # MXFP4
107
+ {
108
+ "a_type": [scalar_types.bfloat16],
109
+ "b_type": scalar_types.float4_e2m1f,
110
+ "group_blocks": [2],
111
+ },
112
+ # AWQ-INT4 with INT8 activation
113
+ {
114
+ "a_type": [scalar_types.int8],
115
+ "b_type": scalar_types.uint4,
116
+ "group_blocks": [-1, 2, 4, 8],
117
+ },
118
+ # GPTQ-INT4 with INT8 activation
119
+ {
120
+ "a_type": [scalar_types.int8],
121
+ "b_type": scalar_types.uint4b8,
122
+ "group_blocks": [-1, 2, 4, 8],
123
+ },
124
+ # GPTQ-INT4 with FP8 activation
125
+ {
126
+ "a_type": [scalar_types.float8_e4m3fn],
127
+ "b_type": scalar_types.uint4b8,
128
+ "group_blocks": [-1, 2, 4, 8],
129
+ },
130
+ # AWQ-INT4 with FP8 activation
131
+ {
132
+ "a_type": [scalar_types.float8_e4m3fn],
133
+ "b_type": scalar_types.uint4,
134
+ "group_blocks": [-1, 2, 4, 8],
135
+ },
136
+ # MXFP4 with FP8 activation
137
+ {
138
+ "a_type": [scalar_types.float8_e4m3fn],
139
+ "b_type": scalar_types.float4_e2m1f,
140
+ "c_type": [scalar_types.bfloat16],
141
+ "group_blocks": [2],
142
+ },
143
+]
144
+
145
146
def compute_max_diff(output, output_ref):
147
return torch.mean(torch.abs(output - output_ref)) / torch.mean(
@@ -85,6 +153,58 @@ def rand_data(shape, dtype=torch.float16):
153
return torch.randn(shape, dtype=dtype, device="cuda")
154
155
156
+@pytest.mark.skipif(
157
+ not is_quant_method_supported("gptq_marlin"),
158
+ reason="Marlin is not supported on this GPU type.",
159
+)
160
+def test_marlin_int4_fp8_preprocess_without_zp():
161
+ qweight_unpacked = torch.randint(
162
+ 0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda"
163
+ )
164
+ qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2]
165
+ qweight_packed = qweight_packed.to(torch.int8).view(torch.int32)
166
+
167
+ cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed)
168
+
169
+ torch_res = torch.where(
170
+ qweight_unpacked >= 8, qweight_unpacked - 8, 15 - qweight_unpacked
171
+ )
172
+ torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2]
173
+ torch_res = torch_res.to(torch.int8).view(torch.int32)
174
+
175
+ assert (cuda_res == torch_res).all()
176
+
177
+
178
+@pytest.mark.skipif(
179
+ not is_quant_method_supported("gptq_marlin"),
180
+ reason="Marlin is not supported on this GPU type.",
181
+)
182
+def test_marlin_int4_fp8_preprocess_awq():
183
+ group_size = 128
184
+
185
+ qweight_unpacked = torch.randint(
186
+ 0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda"
187
+ )
188
+ qzeros_unpacked = torch.randint(
189
+ 0, 16, size=(2048 // group_size, 2048), dtype=torch.int32, device="cuda"
190
+ )
191
+
192
+ qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2]
193
+ qweight_packed = qweight_packed.to(torch.int8).view(torch.int32)
194
+ qzeros_packed = qzeros_unpacked[:, ::2] * 16 + qzeros_unpacked[:, 1::2]
195
+ qzeros_packed = qzeros_packed.to(torch.int8).view(torch.int32)
196
+
197
+ cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed, qzeros_packed)
198
+
199
+ repeated_zp = qzeros_unpacked.repeat_interleave(group_size, 0)
200
+ torch_res = qweight_unpacked - repeated_zp
201
+ torch_res[torch_res < 0] = 15 - qweight_unpacked[torch_res < 0]
202
+ torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2]
203
+ torch_res = torch_res.to(torch.int8).view(torch.int32)
204
+
205
+ assert (cuda_res == torch_res).all()
206
+
207
+
208
@pytest.mark.skipif(
209
not is_quant_method_supported("gptq_marlin"),
210
reason="Marlin is not supported on this GPU type.",
@@ -92,16 +212,17 @@ def rand_data(shape, dtype=torch.float16):
212
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
213
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
214
@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(False, False))
95
-@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
215
@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
97
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
216
+@pytest.mark.parametrize("is_a_8bit", [True, False])
217
+@pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS)
218
def test_gptq_marlin_repack(
99
- k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors
219
+ k_chunk, n_chunk, quant_type, act_order, is_a_8bit, nk_factors
220
):
101
- m_factor, n_factor, k_factor = mnk_factors
221
+ n_factor, k_factor = nk_factors
222
223
size_k = k_chunk * k_factor
224
size_n = n_chunk * n_factor
225
+ group_size = 128
226
227
# Filter act_order
228
if act_order:
@@ -109,6 +230,8 @@ def test_gptq_marlin_repack(
230
return
231
if group_size == size_k:
232
return
233
+ if is_a_8bit:
234
+ return
235
236
# Normalize group_size
237
if group_size == -1:
@@ -133,23 +256,19 @@ def test_gptq_marlin_repack(
256
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
257
258
# Pack to Marlin format
136
- weight_perm = get_weight_perm(quant_type.size_bits)
259
+ weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
260
marlin_q_w_1 = marlin_weights(
138
- q_w, size_k, size_n, quant_type.size_bits, weight_perm
261
+ q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit
262
)
263
264
opcheck(
265
torch.ops._C.gptq_marlin_repack,
143
- (q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits),
266
+ (q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit),
267
)
268
269
# Run Marlin repack GPU kernel
270
marlin_q_w_2 = ops.gptq_marlin_repack(
148
- q_w_gptq,
149
- sort_indices,
150
- size_k,
151
- size_n,
152
- quant_type.size_bits,
271
+ q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit
272
)
273
torch.cuda.synchronize()
274
@@ -163,18 +282,15 @@ def test_gptq_marlin_repack(
282
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
283
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
284
@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(True))
166
-@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
167
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
168
-def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size, mnk_factors):
169
- m_factor, n_factor, k_factor = mnk_factors
285
+@pytest.mark.parametrize("is_a_8bit", [True, False])
286
+@pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS)
287
+def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors):
288
+ n_factor, k_factor = nk_factors
289
290
size_k = k_chunk * k_factor
291
size_n = n_chunk * n_factor
292
174
- # Normalize group_size
175
- if group_size == -1:
176
- group_size = size_k
177
- assert group_size <= size_k
293
+ group_size = 128
294
295
# Create input
296
b_weight = rand_data((size_k, size_n))
@@ -188,162 +304,221 @@ def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size, mnk_factors
304
q_w_awq = awq_pack(q_w, quant_type.size_bits, size_k, size_n)
305
306
# Pack to Marlin format
191
- weight_perm = get_weight_perm(quant_type.size_bits)
307
+ weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
308
marlin_q_w_1 = marlin_weights(
193
- q_w, size_k, size_n, quant_type.size_bits, weight_perm
309
+ q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit
310
)
311
312
opcheck(
197
- torch.ops._C.awq_marlin_repack, (q_w_awq, size_k, size_n, quant_type.size_bits)
313
+ torch.ops._C.awq_marlin_repack,
314
+ (q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit),
315
)
316
317
# Run Marlin repack GPU kernel
318
marlin_q_w_2 = ops.awq_marlin_repack(
202
- q_w_awq,
203
- size_k,
204
- size_n,
205
- quant_type.size_bits,
319
+ q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit
320
)
321
torch.cuda.synchronize()
322
323
torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2)
324
325
326
+def marlin_generate_valid_test_cases():
327
+ all_combinations = itertools.product(
328
+ DENSE_MARLIN_QUANT_TEST_CONFIGS,
329
+ MNK_FACTORS,
330
+ MARLIN_N_CHUNKS,
331
+ MARLIN_K_CHUNKS,
332
+ ACT_ORDER_OPTS,
333
+ K_FULL_OPTS,
334
+ USE_ATOMIC_ADD_OPTS,
335
+ USE_FP32_REDUCE_OPTS,
336
+ )
337
+
338
+ def is_invalid(
339
+ a_type,
340
+ b_type,
341
+ c_type,
342
+ group_blocks,
343
+ size_m,
344
+ size_n,
345
+ size_k,
346
+ act_order,
347
+ is_k_full,
348
+ use_atomic_add,
349
+ use_fp32_reduce,
350
+ ):
351
+ if use_atomic_add:
352
+ if use_fp32_reduce:
353
+ return False
354
+ if (
355
+ c_type == scalar_types.bfloat16
356
+ and torch.cuda.get_device_capability()[0] < 9
357
+ ):
358
+ return False
359
+
360
+ group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
361
+ if group_size > 0 and size_k % group_size != 0:
362
+ return False
363
+
364
+ if act_order and group_size in [-1, size_k]:
365
+ return False
366
+ if group_size == size_k:
367
+ return False
368
+ if not act_order and is_k_full:
369
+ return False
370
+
371
+ return a_type.size_bits < 16 or a_type is c_type
372
+
373
+ cases = []
374
+ for case in all_combinations:
375
+ quant_test_config, mnk_factors, n_chunk, k_chunk, act_order, *_ = case
376
+ size_m = mnk_factors[0]
377
+ size_n = mnk_factors[1] * n_chunk
378
+ size_k = mnk_factors[2] * k_chunk
379
+
380
+ if act_order and not quant_test_config.get("support_act_order", False):
381
+ continue
382
+
383
+ f16_types = [scalar_types.float16, scalar_types.bfloat16]
384
+ inner_combinations = itertools.product(
385
+ quant_test_config.get("a_type", f16_types),
386
+ [quant_test_config["b_type"]],
387
+ quant_test_config.get("c_type", f16_types),
388
+ quant_test_config["group_blocks"],
389
+ )
390
+
391
+ for sub_case in inner_combinations:
392
+ if (
393
+ sub_case[0] == scalar_types.float8_e4m3fn
394
+ and current_platform.get_device_capability() not in [89, 120]
395
+ ):
396
+ continue
397
+ args = sub_case + (size_m, size_n, size_k) + case[4:]
398
+ if is_invalid(*args):
399
+ cases.append(args)
400
+ return cases
401
+
402
+
403
@pytest.mark.skipif(
404
not is_quant_method_supported("gptq_marlin"),
405
reason="Marlin is not supported on this GPU type.",
406
)
216
-@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
217
-@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
218
-@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types())
407
@pytest.mark.parametrize(
220
- "group_size", set(MARLIN_SUPPORTED_GROUP_SIZES + FP4_MARLIN_SUPPORTED_GROUP_SIZES)
408
+ (
409
+ "a_type, b_type, c_type, group_blocks,"
410
+ "size_m, size_n, size_k, act_order, is_k_full,"
411
+ "use_atomic_add, use_fp32_reduce"
412
+ ),
413
+ marlin_generate_valid_test_cases(),
414
)
222
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
223
-@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
224
-@pytest.mark.parametrize("is_k_full", K_FULL_OPTS)
225
-@pytest.mark.parametrize("use_atomic_add", USE_ATOMIC_ADD_OPTS)
226
-@pytest.mark.parametrize("use_fp32_reduce", USE_FP32_REDUCE_OPTS)
227
-@pytest.mark.parametrize("dtype", DTYPES)
415
def test_gptq_marlin_gemm(
229
- k_chunk,
230
- n_chunk,
231
- quant_type,
232
- group_size,
233
- mnk_factors,
416
+ a_type,
417
+ b_type,
418
+ c_type,
419
+ group_blocks,
420
+ size_m,
421
+ size_n,
422
+ size_k,
423
act_order,
424
is_k_full,
425
use_atomic_add,
426
use_fp32_reduce,
238
- dtype,
427
):
240
- m_factor, n_factor, k_factor = mnk_factors
241
- has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
428
+ has_zp = b_type in [scalar_types.uint4, scalar_types.uint8]
429
243
- size_m = m_factor
244
- size_k = k_chunk * k_factor
245
- size_n = n_chunk * n_factor
246
-
247
- if act_order:
248
- if group_size == -1:
249
- return
250
- if group_size == size_k:
251
- return
252
- if has_zp:
253
- return
430
+ group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
431
255
- if size_k % group_size != 0:
256
- return
432
+ if c_type == scalar_types.float16:
433
+ dtype = torch.float16
434
+ elif c_type == scalar_types.bfloat16:
435
+ dtype = torch.bfloat16
436
+ else:
437
+ raise RuntimeError("unsupported c_type")
438
258
- a_input = rand_data((size_m, size_k), dtype)
259
- b_weight = rand_data((size_k, size_n), dtype)
439
+ if a_type == scalar_types.int8:
440
+ a_dtype = torch.int8
441
+ elif a_type == scalar_types.float8_e4m3fn:
442
+ a_dtype = torch.float8_e4m3fn
443
+ else:
444
+ a_dtype = dtype
445
261
- if quant_type == scalar_types.float4_e2m1f:
262
- if group_size not in [16, 32] or act_order:
263
- return
264
- if group_size == 32 and dtype == torch.float16:
265
- return
446
+ a_input = rand_data((size_m, size_k), dtype=dtype)
447
+ b_weight = rand_data((size_k, size_n), dtype=dtype)
448
449
+ if b_type == scalar_types.float4_e2m1f:
450
if group_size == 16:
451
w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_nvfp4_like(
269
- b_weight.T, group_size
452
+ b_weight.T, group_size, input_dtype=a_dtype
453
)
454
else:
455
w_ref, marlin_q_w, marlin_s = rand_marlin_weight_mxfp4_like(
273
- b_weight.T, group_size
456
+ b_weight.T, group_size, input_dtype=a_dtype
457
)
458
marlin_s2 = None
459
460
g_idx = None
461
sort_indices = None
462
marlin_zp = None
280
- elif quant_type == scalar_types.float8_e4m3fn:
281
- if group_size not in [-1, 128]:
282
- return
283
- if act_order:
284
- return
285
- w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(b_weight.T, group_size)
463
+ elif b_type == scalar_types.float8_e4m3fn:
464
+ w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(
465
+ b_weight.T, group_size, input_dtype=a_dtype
466
+ )
467
g_idx = None
468
sort_indices = None
469
marlin_zp = None
470
marlin_s2 = None
471
elif has_zp:
291
- if group_size == 16:
292
- return
472
w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
294
- b_weight, quant_type, group_size
473
+ b_weight, b_type, group_size, input_dtype=a_dtype
474
)
475
g_idx = None
476
sort_indices = None
477
marlin_s2 = None
478
else:
300
- if group_size == 16:
301
- return
479
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
303
- b_weight, quant_type, group_size, act_order
480
+ b_weight, b_type, group_size, act_order, input_dtype=a_dtype
481
)
482
+
483
marlin_zp = None
484
marlin_s2 = None
485
486
workspace = marlin_make_workspace_new(w_ref.device)
487
310
- opcheck(
311
- torch.ops._C.gptq_marlin_gemm,
312
- (
313
- a_input,
314
- None,
315
- marlin_q_w,
316
- None,
317
- marlin_s,
318
- marlin_s2,
319
- marlin_zp,
320
- g_idx,
321
- sort_indices,
322
- workspace,
323
- quant_type.id,
324
- a_input.shape[0],
325
- b_weight.shape[1],
326
- a_input.shape[1],
327
- is_k_full,
328
- use_atomic_add,
329
- use_fp32_reduce,
330
- False,
331
- ),
332
- test_utils=DEFAULT_OPCHECK_TEST_UTILS,
333
- )
488
+ if a_type == scalar_types.int8:
489
+ a_input, a_scales = per_token_quant_int8(a_input)
490
+ a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1)
491
+ a_input_ref = a_input_ref.to(dtype)
492
+
493
+ if group_size != -1:
494
+ a_scales = a_scales / 4096 * marlin_s.max()
495
+ a_scales = a_scales.float()
496
+ marlin_s = marlin_s / marlin_s.max() * 4096
497
+ marlin_s = marlin_s.round().to(torch.int16).view(dtype)
498
+ elif a_type == scalar_types.float8_e4m3fn:
499
+ a_input, a_scales = ops.scaled_fp8_quant(a_input, use_per_token_if_dynamic=True)
500
+ a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1)
501
+ a_input_ref = a_input_ref.to(dtype)
502
+ else:
503
+ assert a_type.size_bits == 16
504
+ a_input_ref = a_input
505
+ a_scales = None
506
+
507
+ output = torch.empty((size_m, size_n), dtype=dtype, device=a_input.device)
508
509
output = ops.gptq_marlin_gemm(
510
a_input,
337
- None,
511
+ output,
512
marlin_q_w,
513
None,
514
marlin_s,
515
+ a_scales,
516
marlin_s2,
517
marlin_zp,
518
g_idx,
519
sort_indices,
520
workspace,
346
- quant_type,
521
+ b_type,
522
a_input.shape[0],
523
b_weight.shape[1],
524
a_input.shape[1],
@@ -352,12 +527,9 @@ def test_gptq_marlin_gemm(
527
use_fp32_reduce=use_fp32_reduce,
528
is_zp_float=False,
529
)
355
- output_ref = torch.matmul(a_input, w_ref)
356
-
357
- torch.cuda.synchronize()
530
+ output_ref = torch.matmul(a_input_ref, w_ref)
531
532
max_diff = compute_max_diff(output, output_ref)
360
-
533
assert max_diff < 0.04
534
535
@@ -507,6 +679,7 @@ def test_hqq_marlin_gemm(
679
None,
680
marlin_s,
681
None,
682
+ None,
683
marlin_zp,
684
g_idx,
685
g_idx_sort_indices,
@@ -559,6 +732,7 @@ def test_marlin_gemm_subset_input():
732
None,
733
marlin_s,
734
None,
735
+ None,
736
marlin_zp,
737
g_idx,
738
sort_indices,
@@ -607,6 +781,7 @@ def test_marlin_gemm_with_bias(size_m):
781
marlin_bias,
782
marlin_s,
783
None,
784
+ None,
785
marlin_zp,
786
g_idx,
787
sort_indices,
tests/kernels/utils.py
+18
@@ -846,6 +846,13 @@ def torch_experts(
846
or (expert_map is not None and global_num_experts == expert_map.shape[0])
847
)
848
849
+ if quant_dtype in [torch.float16, torch.bfloat16]:
850
+ quant_dtype = None
851
+ quant_input_only = quant_dtype is not None and w1_scale is None and w2_scale is None
852
+ if quant_input_only:
853
+ assert a1_scale is None and a2_scale is None
854
+ assert per_act_token_quant
855
+
856
M, K = a.shape
857
topk = topk_ids.shape[1]
858
@@ -863,6 +870,9 @@ def torch_experts(
870
a, a1_scale, quant_dtype, per_act_token_quant, block_shape
871
)
872
873
+ if quant_input_only:
874
+ a = (a.float() * a_scale.view(-1, 1)).to(w1.dtype)
875
+
876
num_experts = w1.shape[0]
877
878
topk_ids = topk_ids.view(-1)
@@ -882,6 +892,14 @@ def torch_experts(
892
out[mask] = tmp2 @ w2[i].transpose(0, 1)
893
if b_bias2 is not None:
894
out[mask] = out[mask] + b_bias2[i].view(1, -1).to(tmp1.dtype)
895
+ elif quant_input_only:
896
+ tmp1 = a[mask] @ w1[i].transpose(0, 1)
897
+ tmp2 = SiluAndMul()(tmp1)
898
+ tmp2, tmp2_scale = moe_kernel_quantize_input(
899
+ tmp2, None, quant_dtype, per_act_token_quant
900
+ )
901
+ tmp2 = (tmp2.float() * tmp2_scale.view(-1, 1)).to(w2.dtype)
902
+ out[mask] = tmp2 @ w2[i].transpose(0, 1)
903
elif block_shape is not None:
904
# block quantized
905
assert (
vllm/_custom_ops.py
+48
-10
@@ -554,6 +554,7 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
554
b_q_weight: torch.Tensor,
555
b_bias: torch.Tensor | None,
556
b_scales: torch.Tensor,
557
+ a_scales: torch.Tensor | None,
558
global_scale: torch.Tensor | None,
559
b_zeros: torch.Tensor | None,
560
g_idx: torch.Tensor | None,
@@ -568,7 +569,10 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
569
use_fp32_reduce: bool = False,
570
is_zp_float: bool = False,
571
) -> torch.Tensor:
571
- return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
572
+ dtype = a.dtype
573
+ if dtype not in [torch.half, torch.bfloat16]:
574
+ dtype = b_scales.dtype
575
+ return torch.empty((size_m, size_n), device=a.device, dtype=dtype)
576
577
@register_fake("_C::awq_dequantize")
578
def _awq_dequantize_fake(
@@ -1167,8 +1171,11 @@ def gptq_marlin_repack(
1171
size_k: int,
1172
size_n: int,
1173
num_bits: int,
1174
+ is_a_8bit: bool = False,
1175
) -> torch.Tensor:
1171
- return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n, num_bits)
1176
+ return torch.ops._C.gptq_marlin_repack(
1177
+ b_q_weight, perm, size_k, size_n, num_bits, is_a_8bit
1178
+ )
1179
1180
1181
if hasattr(torch.ops._C, "gptq_marlin_repack"):
@@ -1180,6 +1187,7 @@ if hasattr(torch.ops._C, "gptq_marlin_repack"):
1187
size_k: torch.SymInt,
1188
size_n: torch.SymInt,
1189
num_bits: int,
1190
+ is_a_8bit: bool = False,
1191
) -> torch.Tensor:
1192
pack_factor = 32 // num_bits
1193
marlin_tile_size = 16
@@ -1192,9 +1200,15 @@ if hasattr(torch.ops._C, "gptq_marlin_repack"):
1200
1201
# awq_marlin
1202
def awq_marlin_repack(
1195
- b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int
1203
+ b_q_weight: torch.Tensor,
1204
+ size_k: int,
1205
+ size_n: int,
1206
+ num_bits: int,
1207
+ is_a_8bit: bool = False,
1208
) -> torch.Tensor:
1197
- return torch.ops._C.awq_marlin_repack(b_q_weight, size_k, size_n, num_bits)
1209
+ return torch.ops._C.awq_marlin_repack(
1210
+ b_q_weight, size_k, size_n, num_bits, is_a_8bit
1211
+ )
1212
1213
1214
if hasattr(torch.ops._C, "awq_marlin_repack"):
@@ -1205,6 +1219,7 @@ if hasattr(torch.ops._C, "awq_marlin_repack"):
1219
size_k: torch.SymInt,
1220
size_n: torch.SymInt,
1221
num_bits: int,
1222
+ is_a_8bit: bool = False,
1223
) -> torch.Tensor:
1224
pack_factor = 32 // num_bits
1225
marlin_tile_size = 16
@@ -1221,6 +1236,7 @@ def gptq_marlin_moe_repack(
1236
size_k: int,
1237
size_n: int,
1238
num_bits: int,
1239
+ is_a_8bit: bool = False,
1240
) -> torch.Tensor:
1241
num_experts = b_q_weight.shape[0]
1242
assert size_k % 16 == 0
@@ -1231,7 +1247,7 @@ def gptq_marlin_moe_repack(
1247
)
1248
for e in range(num_experts):
1249
output[e] = torch.ops._C.gptq_marlin_repack(
1234
- b_q_weight[e], perm[e], size_k, size_n, num_bits
1250
+ b_q_weight[e], perm[e], size_k, size_n, num_bits, is_a_8bit
1251
)
1252
return output
1253
@@ -1242,6 +1258,7 @@ def awq_marlin_moe_repack(
1258
size_k: int,
1259
size_n: int,
1260
num_bits: int,
1261
+ is_a_8bit: bool = False,
1262
) -> torch.Tensor:
1263
num_experts = b_q_weight.shape[0]
1264
assert size_k % 16 == 0
@@ -1252,17 +1269,26 @@ def awq_marlin_moe_repack(
1269
)
1270
for e in range(num_experts):
1271
output[e] = torch.ops._C.awq_marlin_repack(
1255
- b_q_weight[e], size_k, size_n, num_bits
1272
+ b_q_weight[e], size_k, size_n, num_bits, is_a_8bit
1273
)
1274
return output
1275
1276
1277
+def marlin_int4_fp8_preprocess(
1278
+ qweight: torch.Tensor,
1279
+ qzeros_or_none: torch.Tensor | None = None,
1280
+ inplace: bool = False,
1281
+):
1282
+ return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace)
1283
+
1284
+
1285
def gptq_marlin_gemm(
1286
a: torch.Tensor,
1287
c: torch.Tensor | None,
1288
b_q_weight: torch.Tensor,
1289
b_bias: torch.Tensor | None,
1290
b_scales: torch.Tensor,
1291
+ a_scales: torch.Tensor | None,
1292
global_scale: torch.Tensor | None,
1293
b_zeros: torch.Tensor | None,
1294
g_idx: torch.Tensor | None,
@@ -1283,6 +1309,7 @@ def gptq_marlin_gemm(
1309
b_q_weight,
1310
b_bias,
1311
b_scales,
1312
+ a_scales,
1313
global_scale,
1314
b_zeros,
1315
g_idx,
@@ -1600,7 +1627,7 @@ def allspark_repack_weight(
1627
if use asymmetric quantization, has_zp = True.
1628
1629
Returns:
1603
- tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] :
1630
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] :
1631
rearranged weight, scale, and optionally zero_point.
1632
"""
1633
K = qweight.shape[0]
@@ -1683,7 +1710,7 @@ def scaled_int8_quant(
1710
symmetric: Whether to use symmetric quantization (scale only, azp ignored).
1711
1712
Returns:
1686
- tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp.
1713
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.
1714
"""
1715
output = torch.empty_like(input, dtype=torch.int8)
1716
if scale is not None:
@@ -2004,6 +2031,7 @@ def moe_wna16_marlin_gemm(
2031
b_qweight: torch.Tensor,
2032
b_bias: torch.Tensor | None,
2033
b_scales: torch.Tensor,
2034
+ a_scales: torch.Tensor | None,
2035
global_scale: torch.Tensor | None,
2036
b_qzeros: torch.Tensor | None,
2037
g_idx: torch.Tensor | None,
@@ -2025,6 +2053,9 @@ def moe_wna16_marlin_gemm(
2053
use_atomic_add: bool,
2054
use_fp32_reduce: bool,
2055
is_zp_float: bool,
2056
+ thread_k: int = -1,
2057
+ thread_n: int = -1,
2058
+ blocks_per_sm: int = -1,
2059
) -> torch.Tensor:
2060
return torch.ops._moe_C.moe_wna16_marlin_gemm(
2061
input,
@@ -2032,6 +2063,7 @@ def moe_wna16_marlin_gemm(
2063
b_qweight,
2064
b_bias,
2065
b_scales,
2066
+ a_scales,
2067
global_scale,
2068
b_qzeros,
2069
g_idx,
@@ -2053,6 +2085,9 @@ def moe_wna16_marlin_gemm(
2085
use_atomic_add,
2086
use_fp32_reduce,
2087
is_zp_float,
2088
+ thread_k,
2089
+ thread_n,
2090
+ blocks_per_sm,
2091
)
2092
2093
@@ -2088,7 +2123,10 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe")
2123
input: torch.Tensor,
2124
output: torch.Tensor | None,
2125
b_qweight: torch.Tensor,
2126
+ b_bias: torch.Tensor | None,
2127
b_scales: torch.Tensor,
2128
+ a_scales: torch.Tensor | None,
2129
+ global_scale: torch.Tensor | None,
2130
b_qzeros: torch.Tensor | None,
2131
g_idx: torch.Tensor | None,
2132
perm: torch.Tensor | None,
@@ -2109,7 +2147,7 @@ if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe")
2147
use_atomic_add: bool,
2148
use_fp32_reduce: bool,
2149
is_zp_float: bool,
2112
- ) -> torch.Tensor:
2150
+ ):
2151
return torch.empty(
2152
(size_m * top_k, size_n), dtype=input.dtype, device=input.device
2153
)
@@ -2583,7 +2621,7 @@ def onednn_scaled_int8_quant(
2621
symmetric: Whether to use symmetric quantization (scale only, azp ignored).
2622
2623
Returns:
2586
- tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp.
2624
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor | None] : Output int8 tensor, scales, and optionally azp.
2625
"""
2626
output = torch.empty_like(input, dtype=torch.int8)
2627
token_num = input.numel() // input.shape[-1]
vllm/envs.py
+5
@@ -145,6 +145,7 @@ if TYPE_CHECKING:
145
VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
146
VLLM_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict"
147
VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
148
+ VLLM_MARLIN_INPUT_DTYPE: Literal["int8", "fp8"] | None = None
149
VLLM_MXFP4_USE_MARLIN: bool | None = None
150
VLLM_V1_USE_OUTLINES_CACHE: bool = False
151
VLLM_TPU_BUCKET_PADDING_GAP: int = 0
@@ -1122,6 +1123,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1123
"VLLM_MXFP4_USE_MARLIN": lambda: maybe_convert_bool(
1124
os.environ.get("VLLM_MXFP4_USE_MARLIN", None)
1125
),
1126
+ # The activation dtype for marlin kernel
1127
+ "VLLM_MARLIN_INPUT_DTYPE": env_with_choices(
1128
+ "VLLM_MARLIN_INPUT_DTYPE", None, ["int8", "fp8"]
1129
+ ),
1130
# Whether to turn on the outlines cache for V1
1131
# This cache is unbounded and on disk, so it's not safe to use in
1132
# an environment with potentially malicious users.
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+38
-9
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.fused_moe.utils import _resize_cache, disable_in
24
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
25
marlin_make_workspace_new,
26
marlin_moe_intermediate_size,
27
- maybe_warn_marlin_atomic_add,
27
+ marlin_quant_input,
28
)
29
from vllm.scalar_type import ScalarType, scalar_types
30
@@ -65,6 +65,8 @@ def _fused_marlin_moe(
65
activation_func: Callable[
66
[str, torch.Tensor, torch.Tensor], None
67
] = default_activation_func,
68
+ input_global_scale1: torch.Tensor | None = None,
69
+ input_global_scale2: torch.Tensor | None = None,
70
global_scale1: torch.Tensor | None = None,
71
global_scale2: torch.Tensor | None = None,
72
g_idx1: torch.Tensor | None = None,
@@ -77,6 +79,7 @@ def _fused_marlin_moe(
79
intermediate_cache13: torch.Tensor | None = None,
80
intermediate_cache2: torch.Tensor | None = None,
81
output: torch.Tensor | None = None,
82
+ input_dtype: torch.dtype | None = None,
83
is_k_full: bool = True,
84
) -> torch.Tensor:
85
assert hidden_states.ndim == 2
@@ -106,18 +109,22 @@ def _fused_marlin_moe(
109
110
intermediate_cache2 = _resize_cache(intermediate_cache2, (M * num_topk, N))
111
109
- maybe_warn_marlin_atomic_add(hidden_states.device, hidden_states.dtype)
110
- use_atomic_add = (
111
- hidden_states.dtype == torch.half
112
- or torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
113
- )
112
+ a_scales1 = None
113
+ gate_up_input = hidden_states
114
+ if input_dtype == torch.int8:
115
+ gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype)
116
+ if input_global_scale1 is not None:
117
+ a_scales1 = a_scales1 * input_global_scale1
118
+ elif input_dtype == torch.float8_e4m3fn:
119
+ gate_up_input, a_scales1 = marlin_quant_input(hidden_states, input_dtype)
120
121
intermediate_cache1 = ops.moe_wna16_marlin_gemm(
116
- hidden_states,
122
+ gate_up_input,
123
intermediate_cache1,
124
w1,
125
bias1,
126
w1_scale,
127
+ a_scales1,
128
global_scale1,
129
w1_zeros,
130
g_idx1,
@@ -136,7 +143,7 @@ def _fused_marlin_moe(
143
size_n=2 * N,
144
size_k=K,
145
is_k_full=is_k_full,
139
- use_atomic_add=use_atomic_add,
146
+ use_atomic_add=False,
147
use_fp32_reduce=True,
148
is_zp_float=False,
149
)
@@ -151,12 +158,25 @@ def _fused_marlin_moe(
158
if expert_map is not None:
159
output.zero_()
160
161
+ a_scales2 = None
162
+ if input_dtype == torch.int8:
163
+ intermediate_cache2, a_scales2 = marlin_quant_input(
164
+ intermediate_cache2, input_dtype
165
+ )
166
+ if input_global_scale2 is not None:
167
+ a_scales2 = a_scales2 * input_global_scale2
168
+ elif input_dtype == torch.float8_e4m3fn:
169
+ intermediate_cache2, a_scales2 = marlin_quant_input(
170
+ intermediate_cache2, input_dtype
171
+ )
172
+
173
output = ops.moe_wna16_marlin_gemm(
174
intermediate_cache2,
175
output,
176
w2,
177
bias2,
178
w2_scale,
179
+ a_scales2,
180
global_scale2,
181
w2_zeros,
182
g_idx2,
@@ -175,7 +195,7 @@ def _fused_marlin_moe(
195
size_n=K,
196
size_k=N,
197
is_k_full=is_k_full,
178
- use_atomic_add=use_atomic_add,
198
+ use_atomic_add=False,
199
use_fp32_reduce=True,
200
is_zp_float=False,
201
)
@@ -203,6 +223,8 @@ def fused_marlin_moe(
223
] = default_activation_func,
224
moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
225
expert_map: torch.Tensor | None = None,
226
+ input_global_scale1: torch.Tensor | None = None,
227
+ input_global_scale2: torch.Tensor | None = None,
228
global_scale1: torch.Tensor | None = None,
229
global_scale2: torch.Tensor | None = None,
230
g_idx1: torch.Tensor | None = None,
@@ -216,6 +238,7 @@ def fused_marlin_moe(
238
intermediate_cache2: torch.Tensor | None = None,
239
is_k_full: bool = True,
240
output: torch.Tensor | None = None,
241
+ input_dtype: torch.dtype | None = None,
242
inplace: bool = False,
243
) -> torch.Tensor:
244
"""
@@ -287,6 +310,9 @@ def fused_marlin_moe(
310
if M * topk / E / block_size_m < 0.9:
311
break
312
313
+ if input_dtype is not None and input_dtype.itemsize == 1:
314
+ block_size_m = max(block_size_m, 16)
315
+
316
if global_num_experts == -1:
317
global_num_experts = E
318
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
@@ -313,6 +339,8 @@ def fused_marlin_moe(
339
num_tokens_post_padded=num_tokens_post_padded,
340
activation=activation,
341
activation_func=activation_func,
342
+ input_global_scale1=input_global_scale1,
343
+ input_global_scale2=input_global_scale2,
344
global_scale1=global_scale1,
345
global_scale2=global_scale2,
346
g_idx1=g_idx1,
@@ -325,6 +353,7 @@ def fused_marlin_moe(
353
intermediate_cache13=intermediate_cache13,
354
intermediate_cache2=intermediate_cache2,
355
output=None,
356
+ input_dtype=input_dtype,
357
is_k_full=is_k_full,
358
).view(-1, topk, K)
359
vllm/model_executor/layers/quantization/auto_round.py
+2
-2
@@ -266,7 +266,7 @@ class AutoRoundConfig(QuantizationConfig):
266
from vllm.model_executor.layers.quantization.awq_marlin import (
267
AWQMarlinConfig,
268
AWQMarlinLinearMethod,
269
- AWQMoEMethod,
269
+ AWQMarlinMoEMethod,
270
)
271
272
quant_args_marlin = AWQMarlinConfig(
@@ -291,7 +291,7 @@ class AutoRoundConfig(QuantizationConfig):
291
292
if isinstance(layer, FusedMoE):
293
if use_marlin:
294
- return AWQMoEMethod(quant_args_marlin, layer.moe_config)
294
+ return AWQMarlinMoEMethod(quant_args_marlin, layer.moe)
295
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
296
297
config = {
vllm/model_executor/layers/quantization/awq.py
+2
-2
@@ -106,7 +106,7 @@ class AWQConfig(QuantizationConfig):
106
return AWQLinearMethod(self)
107
elif isinstance(layer, FusedMoE):
108
# Lazy import to avoid circular import.
109
- from .awq_marlin import AWQMarlinConfig, AWQMoEMethod
109
+ from .awq_marlin import AWQMarlinConfig, AWQMarlinMoEMethod
110
from .moe_wna16 import MoeWNA16Config
111
from .utils.marlin_utils import check_moe_marlin_supports_layer
112
@@ -136,7 +136,7 @@ class AWQConfig(QuantizationConfig):
136
awq_marlin_config = AWQMarlinConfig.from_config(
137
marlin_compatible_config_dict
138
)
139
- return AWQMoEMethod(awq_marlin_config, layer.moe_config)
139
+ return AWQMarlinMoEMethod(awq_marlin_config, layer.moe_config)
140
return None
141
142
def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
vllm/model_executor/layers/quantization/awq_marlin.py
+77
-5
@@ -40,6 +40,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
40
check_marlin_supported,
41
check_marlin_supports_layer,
42
check_moe_marlin_supports_layer,
43
+ get_marlin_input_dtype,
44
+ marlin_act_int8_process_scales,
45
marlin_make_empty_g_idx,
46
marlin_make_workspace_new,
47
marlin_moe_permute_scales,
@@ -69,7 +71,6 @@ class AWQMarlinConfig(QuantizationConfig):
71
# num_bits -> type
72
TYPE_MAP = {
73
4: scalar_types.uint4,
72
- 8: scalar_types.uint8,
74
}
75
76
def __init__(
@@ -193,7 +194,9 @@ class AWQMarlinConfig(QuantizationConfig):
194
return AWQConfig.from_config(self.full_config).get_quant_method(
195
layer, prefix
196
)
196
- return AWQMarlinLinearMethod(self)
197
+ quant_method = AWQMarlinLinearMethod(self)
198
+ quant_method.input_dtype = get_marlin_input_dtype(prefix)
199
+ return quant_method
200
elif isinstance(layer, FusedMoE):
201
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
202
@@ -211,7 +214,9 @@ class AWQMarlinConfig(QuantizationConfig):
214
return MoeWNA16Config.from_config(self.full_config).get_quant_method(
215
layer, prefix
216
)
214
- return AWQMoEMethod(self, layer.moe_config)
217
+ moe_quant_method = AWQMarlinMoEMethod(self, layer.moe_config)
218
+ moe_quant_method.input_dtype = get_marlin_input_dtype(prefix)
219
+ return moe_quant_method
220
return None
221
222
@classmethod
@@ -270,6 +275,8 @@ class AWQMarlinLinearMethod(LinearMethodBase):
275
276
def __init__(self, quant_config: AWQMarlinConfig) -> None:
277
self.quant_config = quant_config
278
+ self.quant_type = scalar_types.uint4
279
+ self.input_dtype = None
280
281
def create_weights(
282
self,
@@ -312,6 +319,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):
319
)
320
321
num_groups = input_size_per_partition // group_size
322
+ layer.num_groups = num_groups
323
324
qzeros = PackedvLLMParameter(
325
data=torch.empty(
@@ -358,12 +366,19 @@ class AWQMarlinLinearMethod(LinearMethodBase):
366
# Allocate marlin workspace
367
layer.workspace = marlin_make_workspace_new(device)
368
369
+ is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
370
+
371
+ if self.input_dtype == torch.float8_e4m3fn:
372
+ ops.marlin_int4_fp8_preprocess(layer.qweight, layer.qzeros, inplace=True)
373
+ layer.scales.data = layer.scales.data * 512
374
+
375
# Repack weights from AWQ format to marlin format.
376
marlin_qweight = ops.awq_marlin_repack(
377
layer.qweight,
378
size_k=layer.input_size_per_partition,
379
size_n=layer.output_size_per_partition,
380
num_bits=self.quant_config.quant_type.size_bits,
381
+ is_a_8bit=is_a_8bit,
382
)
383
replace_parameter(layer, "qweight", marlin_qweight)
384
@@ -373,7 +388,16 @@ class AWQMarlinLinearMethod(LinearMethodBase):
388
size_k=layer.input_size_per_partition,
389
size_n=layer.output_size_per_partition,
390
group_size=self.quant_config.group_size,
391
+ is_a_8bit=is_a_8bit,
392
)
393
+ if self.input_dtype == torch.int8 and layer.num_groups > 1:
394
+ marlin_scales, input_global_scale = marlin_act_int8_process_scales(
395
+ marlin_scales
396
+ )
397
+ layer.register_parameter(
398
+ "input_global_scale", Parameter(input_global_scale, requires_grad=False)
399
+ )
400
+
401
replace_parameter(layer, "scales", marlin_scales)
402
403
# Permute zero-points from AWQ format to marlin format.
@@ -382,6 +406,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):
406
size_k=layer.num_groups,
407
size_n=layer.output_size_per_partition,
408
num_bits=self.quant_config.quant_type.size_bits,
409
+ is_a_8bit=is_a_8bit,
410
)
411
replace_parameter(layer, "qzeros", marlin_zp)
412
@@ -409,11 +434,13 @@ class AWQMarlinLinearMethod(LinearMethodBase):
434
quant_type=self.quant_config.quant_type,
435
output_size_per_partition=layer.output_size_per_partition,
436
input_size_per_partition=layer.input_size_per_partition,
437
+ input_global_scale=getattr(layer, "input_global_scale", None),
438
bias=bias,
439
+ input_dtype=self.input_dtype,
440
)
441
442
416
-class AWQMoEMethod(FusedMoEMethodBase):
443
+class AWQMarlinMoEMethod(FusedMoEMethodBase):
444
def __init__(
445
self,
446
quant_config: AWQMarlinConfig,
@@ -422,8 +449,9 @@ class AWQMoEMethod(FusedMoEMethodBase):
449
super().__init__(moe)
450
self.quant_config = quant_config
451
if self.quant_config.weight_bits != 4:
425
- raise ValueError("AWQMoEMethod only supports 4bit now.")
452
+ raise ValueError("AWQMarlinMoEMethod only supports 4bit now.")
453
self.quant_type = scalar_types.uint4
454
+ self.input_dtype = None
455
self.use_marlin = True
456
457
def create_weights(
@@ -435,6 +463,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
463
params_dtype: torch.dtype,
464
**extra_weight_attrs,
465
):
466
+ layer.input_dtype = self.input_dtype
467
extra_weight_attrs.update(
468
{
469
"is_transposed": True,
@@ -468,6 +497,8 @@ class AWQMoEMethod(FusedMoEMethodBase):
497
498
num_groups_w13 = hidden_size // self.quant_config.group_size
499
num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size
500
+ layer.num_groups_w13 = num_groups_w13
501
+ layer.num_groups_w2 = num_groups_w2
502
503
# WEIGHT_SCALES
504
# Allocate 2 scales for w1 and w3 respectively.
@@ -522,6 +553,21 @@ class AWQMoEMethod(FusedMoEMethodBase):
553
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
554
num_experts = layer.w13_qweight.shape[0]
555
device = layer.w13_qweight.device
556
+ is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
557
+
558
+ if self.input_dtype == torch.float8_e4m3fn:
559
+ ops.marlin_int4_fp8_preprocess(
560
+ layer.w13_qweight.view(-1, layer.w13_qweight.size(2)),
561
+ layer.w13_qzeros.view(-1, layer.w13_qzeros.size(2)),
562
+ inplace=True,
563
+ )
564
+ ops.marlin_int4_fp8_preprocess(
565
+ layer.w2_qweight.view(-1, layer.w2_qweight.size(2)),
566
+ layer.w2_qzeros.view(-1, layer.w2_qzeros.size(2)),
567
+ inplace=True,
568
+ )
569
+ layer.w13_scales.data = layer.w13_scales.data * 512
570
+ layer.w2_scales.data = layer.w2_scales.data * 512
571
572
layer.w13_g_idx_sort_indices = torch.nn.Parameter(
573
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
@@ -538,6 +584,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
584
size_k=layer.w13_qweight.shape[1],
585
size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor,
586
num_bits=self.quant_config.weight_bits,
587
+ is_a_8bit=is_a_8bit,
588
)
589
replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
590
@@ -547,6 +594,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
594
size_k=layer.w2_qweight.shape[1],
595
size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor,
596
num_bits=self.quant_config.weight_bits,
597
+ is_a_8bit=is_a_8bit,
598
)
599
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
600
@@ -556,7 +604,16 @@ class AWQMoEMethod(FusedMoEMethodBase):
604
size_k=layer.intermediate_size_per_partition,
605
size_n=layer.w13_scales.shape[2],
606
group_size=self.quant_config.group_size,
607
+ is_a_8bit=is_a_8bit,
608
)
609
+ if self.input_dtype == torch.int8 and layer.num_groups_w13 > 1:
610
+ marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
611
+ marlin_w13_scales
612
+ )
613
+ layer.register_parameter(
614
+ "w13_input_global_scale",
615
+ Parameter(w13_input_global_scale, requires_grad=False),
616
+ )
617
618
replace_parameter(layer, "w13_scales", marlin_w13_scales)
619
@@ -565,7 +622,17 @@ class AWQMoEMethod(FusedMoEMethodBase):
622
size_k=layer.intermediate_size_per_partition,
623
size_n=layer.w2_scales.shape[2],
624
group_size=self.quant_config.group_size,
625
+ is_a_8bit=is_a_8bit,
626
)
627
+ if self.input_dtype == torch.int8 and layer.num_groups_w2 > 1:
628
+ marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
629
+ marlin_w2_scales
630
+ )
631
+ layer.register_parameter(
632
+ "w2_input_global_scale",
633
+ Parameter(w2_input_global_scale, requires_grad=False),
634
+ )
635
+
636
replace_parameter(layer, "w2_scales", marlin_w2_scales)
637
638
marlin_w13_zp = moe_awq_to_marlin_zero_points(
@@ -573,6 +640,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
640
size_k=layer.w13_qzeros.shape[1],
641
size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor,
642
num_bits=self.quant_config.weight_bits,
643
+ is_a_8bit=is_a_8bit,
644
)
645
replace_parameter(layer, "w13_qzeros", marlin_w13_zp)
646
@@ -581,6 +649,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
649
size_k=layer.w2_qzeros.shape[1],
650
size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor,
651
num_bits=self.quant_config.weight_bits,
652
+ is_a_8bit=is_a_8bit,
653
)
654
replace_parameter(layer, "w2_qzeros", marlin_w2_zp)
655
@@ -636,6 +705,8 @@ class AWQMoEMethod(FusedMoEMethodBase):
705
router_logits,
706
topk_weights,
707
topk_ids,
708
+ input_global_scale1=getattr(layer, "w13_input_global_scale", None),
709
+ input_global_scale2=getattr(layer, "w2_input_global_scale", None),
710
quant_type_id=self.quant_type.id,
711
apply_router_weight_on_input=apply_router_weight_on_input,
712
global_num_experts=global_num_experts,
@@ -643,4 +714,5 @@ class AWQMoEMethod(FusedMoEMethodBase):
714
w1_zeros=layer.w13_qzeros,
715
w2_zeros=layer.w2_qzeros,
716
workspace=layer.workspace,
717
+ input_dtype=self.input_dtype,
718
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+9
-2
@@ -157,7 +157,9 @@ class CompressedTensorsConfig(QuantizationConfig):
157
if isinstance(layer, Attention):
158
return CompressedTensorsKVCacheMethod(self)
159
if isinstance(layer, FusedMoE):
160
- return CompressedTensorsMoEMethod.get_moe_method(self, layer, prefix)
160
+ return CompressedTensorsMoEMethod.get_moe_method(
161
+ self, layer, layer_name=prefix
162
+ )
163
return None
164
165
def _add_fused_moe_to_target_scheme_map(self):
@@ -547,6 +549,7 @@ class CompressedTensorsConfig(QuantizationConfig):
549
weight_quant: QuantizationArgs,
550
input_quant: QuantizationArgs,
551
format: str | None = None,
552
+ layer_name: str | None = None,
553
) -> "CompressedTensorsScheme":
554
# use the per-layer format if defined, otherwise, use global format
555
format = format if format is not None else self.quant_format
@@ -585,6 +588,7 @@ class CompressedTensorsConfig(QuantizationConfig):
588
symmetric=weight_quant.symmetric,
589
group_size=weight_quant.group_size,
590
actorder=weight_quant.actorder,
591
+ layer_name=layer_name,
592
)
593
594
act_quant_format = is_activation_quantization_format(format)
@@ -724,7 +728,10 @@ class CompressedTensorsConfig(QuantizationConfig):
728
else:
729
# Find the quant_scheme
730
scheme = self._get_scheme_from_parts( # type: ignore
727
- weight_quant=weight_quant, input_quant=input_quant, format=format
731
+ weight_quant=weight_quant,
732
+ input_quant=input_quant,
733
+ format=format,
734
+ layer_name=layer_name,
735
)
736
737
# Raise error if device does not support the scheme
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+87
-18
@@ -64,6 +64,8 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
64
)
65
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
66
check_moe_marlin_supports_layer,
67
+ get_marlin_input_dtype,
68
+ marlin_act_int8_process_scales,
69
marlin_make_workspace_new,
70
marlin_moe_permute_scales,
71
)
@@ -101,7 +103,7 @@ __all__ = [
103
"CompressedTensorsW8A8Int8MoEMethod",
104
"CompressedTensorsWNA16MarlinMoEMethod",
105
"CompressedTensorsWNA16MoEMethod",
104
- "CompressedTensorsW4A4Nvfp4MoeMethod",
106
+ "CompressedTensorsW4A4Nvfp4MoEMethod",
107
"CompressedTensorsW4A8Int8MoEMethod",
108
]
109
@@ -111,13 +113,13 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
113
def get_moe_method(
114
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
115
layer: torch.nn.Module,
114
- prefix: str,
116
+ layer_name: str,
117
) -> "CompressedTensorsMoEMethod":
118
# FusedMoE was made by combining multiple Linears so need to
119
# make sure quantization config for Linear can target it
120
quant_config._add_fused_moe_to_target_scheme_map()
121
unfused_names = [
120
- prefix + proj_name
122
+ layer_name + proj_name
123
for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
124
]
125
# TODO: refactor this to use expert_mapping and check all layer numbers
@@ -158,32 +160,40 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
160
"WNA16MoE is not supported with actorder=group/dynamic."
161
)
162
logger.info_once("Using CompressedTensorsWNA16MoEMethod")
161
- return CompressedTensorsWNA16MoEMethod(quant_config, layer.moe_config)
163
+ return CompressedTensorsWNA16MoEMethod(
164
+ quant_config, layer.moe_config, layer_name
165
+ )
166
else:
167
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
168
return CompressedTensorsWNA16MarlinMoEMethod(
165
- quant_config, layer.moe_config
169
+ quant_config, layer.moe_config, layer_name
170
)
171
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
168
- return CompressedTensorsW4A4Nvfp4MoeMethod(layer.moe_config)
172
+ return CompressedTensorsW4A4Nvfp4MoEMethod(layer.moe_config, layer_name)
173
elif (
174
quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
175
or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
176
or quant_config._is_fp8_w8a8(weight_quant, input_quant)
177
):
174
- return CompressedTensorsW8A8Fp8MoEMethod(quant_config, layer.moe_config)
178
+ return CompressedTensorsW8A8Fp8MoEMethod(
179
+ quant_config, layer.moe_config, layer_name
180
+ )
181
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
176
- return CompressedTensorsW8A8Int8MoEMethod(quant_config, layer.moe_config)
182
+ return CompressedTensorsW8A8Int8MoEMethod(
183
+ quant_config, layer.moe_config, layer_name
184
+ )
185
elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
178
- return CompressedTensorsW4A8Int8MoEMethod(quant_config, layer.moe_config)
186
+ return CompressedTensorsW4A8Int8MoEMethod(
187
+ quant_config, layer.moe_config, layer_name
188
+ )
189
else:
190
raise RuntimeError(
191
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
192
)
193
194
185
-class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
186
- def __init__(self, moe: FusedMoEConfig):
195
+class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
196
+ def __init__(self, moe: FusedMoEConfig, layer_name: str | None = None):
197
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
198
detect_nvfp4_moe_support,
199
)
@@ -194,17 +204,21 @@ class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
204
self.allow_flashinfer = _nvfp4.allow_flashinfer
205
self.use_marlin = _nvfp4.use_marlin
206
self.group_size = 16
207
+ self.layer_name = layer_name
208
+ self.marlin_input_dtype = (
209
+ get_marlin_input_dtype(layer_name) if self.use_marlin else None
210
+ )
211
self.flashinfer_moe_backend = None
212
if self.allow_flashinfer:
213
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
214
logger.info_once(
215
f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
202
- " for CompressedTensorsW4A4Nvfp4MoeMethod."
216
+ " for CompressedTensorsW4A4Nvfp4MoEMethod."
217
)
218
elif self.use_marlin:
205
- logger.info_once("Using Marlin for CompressedTensorsW4A4Nvfp4MoeMethod.")
219
+ logger.info_once("Using Marlin for CompressedTensorsW4A4Nvfp4MoEMethod.")
220
else:
207
- logger.info_once("Using Cutlass for CompressedTensorsW4A4Nvfp4MoeMethod.")
221
+ logger.info_once("Using Cutlass for CompressedTensorsW4A4Nvfp4MoEMethod.")
222
223
def create_weights(
224
self,
@@ -354,7 +368,7 @@ class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
368
)
369
370
if self.use_marlin:
357
- prepare_moe_fp4_layer_for_marlin(layer)
371
+ prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)
372
return
373
# w13
374
if (
@@ -538,7 +552,7 @@ class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
552
):
553
if enable_eplb:
554
raise NotImplementedError(
541
- "EPLB not supported for `CompressedTensorsW4A4MoeMethod` yet."
555
+ "EPLB not supported for `CompressedTensorsW4A4MoEMethod` yet."
556
)
557
558
return flashinfer_trtllm_fp4_moe(
@@ -576,6 +590,7 @@ class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
590
apply_router_weight_on_input=apply_router_weight_on_input,
591
global_num_experts=global_num_experts,
592
expert_map=expert_map,
593
+ input_dtype=self.marlin_input_dtype,
594
workspace=layer.workspace,
595
)
596
@@ -610,7 +625,7 @@ class CompressedTensorsW4A4Nvfp4MoeMethod(CompressedTensorsMoEMethod):
625
assert expert_map is None, (
626
"Expert Parallelism / expert_map "
627
"is currently not supported for "
613
- "CompressedTensorsW4A4Nvfp4MoeMethod."
628
+ "CompressedTensorsW4A4Nvfp4MoEMethod."
629
)
630
assert self.moe_quant_config is not None
631
@@ -637,6 +652,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
652
self,
653
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
654
moe: FusedMoEConfig,
655
+ layer_name: str | None = None,
656
):
657
super().__init__(moe)
658
self.quant_config = quant_config
@@ -690,6 +706,10 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
706
or self.is_fp8_w8a8_sm100
707
)
708
self.disable_expert_map = False
709
+ self.layer_name = layer_name
710
+ self.marlin_input_dtype = (
711
+ get_marlin_input_dtype(layer_name) if self.use_marlin else None
712
+ )
713
714
def create_weights(
715
self,
@@ -931,7 +951,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
951
layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
952
953
elif self.use_marlin:
934
- prepare_moe_fp8_layer_for_marlin(layer, False)
954
+ prepare_moe_fp8_layer_for_marlin(
955
+ layer, False, input_dtype=self.marlin_input_dtype
956
+ )
957
# Activations not quantized for marlin.
958
del layer.w13_input_scale
959
del layer.w2_input_scale
@@ -1144,6 +1166,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1166
apply_router_weight_on_input=apply_router_weight_on_input,
1167
global_num_experts=global_num_experts,
1168
expert_map=expert_map,
1169
+ input_dtype=self.marlin_input_dtype,
1170
workspace=layer.workspace,
1171
)
1172
@@ -1240,6 +1263,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1263
self,
1264
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1265
moe: FusedMoEConfig,
1266
+ layer_name: str | None = None,
1267
):
1268
super().__init__(moe)
1269
self.quant_config = quant_config
@@ -1392,6 +1416,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1416
self,
1417
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1418
moe: FusedMoEConfig,
1419
+ layer_name: str | None = None,
1420
):
1421
super().__init__(moe)
1422
self.quant_config = quant_config
@@ -1403,6 +1428,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1428
self.strategy = config.strategy
1429
self.group_size = config.group_size
1430
self.actorder = config.actorder
1431
+ self.layer_name = layer_name
1432
+ self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
1433
assert config.symmetric, "Only symmetric quantization is supported for MoE"
1434
1435
if not (
@@ -1477,6 +1504,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1504
num_groups_w2 = w2_scales_size // self.group_size
1505
num_groups_w13 = hidden_size // self.group_size
1506
1507
+ layer.num_groups_w13 = num_groups_w13
1508
+ layer.num_groups_w2 = num_groups_w2
1509
+
1510
w13_scale = torch.nn.Parameter(
1511
torch.ones(
1512
num_experts,
@@ -1560,6 +1590,17 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1590
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1591
num_experts = layer.w13_weight_g_idx.shape[0]
1592
device = layer.w13_weight_g_idx.device
1593
+ is_a_8bit = (
1594
+ self.marlin_input_dtype is not None
1595
+ and self.marlin_input_dtype.itemsize == 1
1596
+ )
1597
+
1598
+ if self.marlin_input_dtype == torch.float8_e4m3fn:
1599
+ # NOTE: for non-zp quantization format only
1600
+ ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True)
1601
+ ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True)
1602
+ layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512
1603
+ layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512
1604
1605
# when running models with grouped act order,
1606
# resort to g_idx values provided in checkpoint
@@ -1610,31 +1651,54 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1651
layer.w13_weight_packed.shape[1] * self.packed_factor,
1652
layer.w13_weight_packed.shape[2],
1653
self.num_bits,
1654
+ is_a_8bit=is_a_8bit,
1655
)
1656
replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight)
1657
+
1658
marlin_w2_qweight = ops.gptq_marlin_moe_repack(
1659
layer.w2_weight_packed,
1660
layer.w2_g_idx_sort_indices,
1661
layer.w2_weight_packed.shape[1] * self.packed_factor,
1662
layer.w2_weight_packed.shape[2],
1663
self.num_bits,
1664
+ is_a_8bit=is_a_8bit,
1665
)
1666
replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight)
1667
+
1668
# Repack scales
1669
marlin_w13_scales = marlin_moe_permute_scales(
1670
s=layer.w13_weight_scale,
1671
size_k=layer.w13_weight_packed.shape[2],
1672
size_n=layer.w13_weight_scale.shape[2],
1673
group_size=self.group_size,
1674
+ is_a_8bit=is_a_8bit,
1675
)
1676
+ if self.marlin_input_dtype == torch.int8 and layer.num_groups_w13 > 1:
1677
+ marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
1678
+ marlin_w13_scales
1679
+ )
1680
+ layer.register_parameter(
1681
+ "w13_input_global_scale",
1682
+ torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
1683
+ )
1684
replace_parameter(layer, "w13_weight_scale", marlin_w13_scales)
1685
+
1686
marlin_w2_scales = marlin_moe_permute_scales(
1687
s=layer.w2_weight_scale,
1688
size_k=layer.w2_weight_scale.shape[1]
1689
* (self.group_size if self.group_size != -1 else self.packed_factor),
1690
size_n=layer.w2_weight_scale.shape[2],
1691
group_size=self.group_size,
1692
+ is_a_8bit=is_a_8bit,
1693
)
1694
+ if self.marlin_input_dtype == torch.int8 and layer.num_groups_w2 > 1:
1695
+ marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
1696
+ marlin_w2_scales
1697
+ )
1698
+ layer.register_parameter(
1699
+ "w2_input_global_scale",
1700
+ torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
1701
+ )
1702
replace_parameter(layer, "w2_weight_scale", marlin_w2_scales)
1703
1704
layer.workspace = marlin_make_workspace_new(device, 4)
@@ -1729,6 +1793,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1793
router_logits,
1794
topk_weights,
1795
topk_ids,
1796
+ input_global_scale1=getattr(layer, "w13_input_global_scale", None),
1797
+ input_global_scale2=getattr(layer, "w2_input_global_scale", None),
1798
quant_type_id=self.quant_type.id,
1799
apply_router_weight_on_input=apply_router_weight_on_input,
1800
global_num_experts=global_num_experts,
@@ -1738,6 +1804,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1804
sort_indices1=layer.w13_g_idx_sort_indices,
1805
sort_indices2=layer.w2_g_idx_sort_indices,
1806
workspace=layer.workspace,
1807
+ input_dtype=self.marlin_input_dtype,
1808
is_k_full=self.is_k_full,
1809
)
1810
@@ -1747,6 +1814,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1814
self,
1815
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1816
moe: FusedMoEConfig,
1817
+ layer_name: str | None = None,
1818
):
1819
super().__init__(moe)
1820
self.quant_config = quant_config
@@ -1999,6 +2067,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2067
self,
2068
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
2069
moe: FusedMoEConfig,
2070
+ layer_name: str | None = None,
2071
):
2072
super().__init__(moe)
2073
self.has_bias = self.moe.has_bias
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+11
@@ -14,7 +14,11 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
14
MPLinearLayerConfig,
15
choose_mp_linear_kernel,
16
)
17
+from vllm.model_executor.layers.quantization.kernels.mixed_precision.marlin import (
18
+ MarlinLinearKernel,
19
+)
20
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
21
+ get_marlin_input_dtype,
22
marlin_repeat_scales_on_all_ranks,
23
)
24
from vllm.model_executor.parameter import (
@@ -45,12 +49,14 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
49
group_size: int | None = None,
50
symmetric: bool | None = True,
51
actorder: ActivationOrdering | None = None,
52
+ layer_name: str | None = None,
53
):
54
self.pack_factor = 32 // num_bits
55
self.strategy = strategy
56
self.symmetric = symmetric
57
self.group_size = -1 if group_size is None else group_size
58
self.has_g_idx = actorder == ActivationOrdering.GROUP
59
+ self.layer_name = layer_name
60
61
if self.group_size == -1 and self.strategy != "channel":
62
raise ValueError(
@@ -108,6 +114,11 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):
114
logger.info("Using %s for CompressedTensorsWNA16", kernel_type.__name__)
115
self._kernel_backends_being_used.add(kernel_type.__name__)
116
117
+ if isinstance(kernel_type, MarlinLinearKernel):
118
+ input_dtype = get_marlin_input_dtype(self.layer_name)
119
+ if input_dtype is not None:
120
+ mp_linear_kernel_config.act_type = input_dtype
121
+
122
# If group_size is -1, we are in channelwise case.
123
group_size = self.group_size if self.group_size != -1 else input_size
124
row_parallel = input_size != input_size_per_partition
vllm/model_executor/layers/quantization/fp8.py
+19
-4
@@ -69,6 +69,9 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
69
process_fp8_weight_tensor_strategy,
70
validate_fp8_block_shape,
71
)
72
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
73
+ get_marlin_input_dtype,
74
+)
75
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
76
apply_fp8_marlin_linear,
77
prepare_fp8_layer_for_marlin,
@@ -316,7 +319,9 @@ class Fp8Config(QuantizationConfig):
319
fused_mapping=self.packed_modules_mapping,
320
):
321
return UnquantizedLinearMethod()
319
- return Fp8LinearMethod(self)
322
+ quant_method = Fp8LinearMethod(self)
323
+ quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
324
+ return quant_method
325
elif isinstance(layer, FusedMoE):
326
if is_layer_skipped(
327
prefix=prefix,
@@ -324,7 +329,9 @@ class Fp8Config(QuantizationConfig):
329
fused_mapping=self.packed_modules_mapping,
330
):
331
return UnquantizedFusedMoEMethod(layer.moe_config)
327
- return Fp8MoEMethod(self, layer)
332
+ moe_quant_method = Fp8MoEMethod(self, layer)
333
+ moe_quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
334
+ return moe_quant_method
335
elif isinstance(layer, Attention):
336
return Fp8KVCacheMethod(self)
337
return None
@@ -375,6 +382,7 @@ class Fp8LinearMethod(LinearMethodBase):
382
383
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
384
# kernel for fast weight-only FP8 quantization
385
+ self.marlin_input_dtype = None
386
self.use_marlin = (
387
not current_platform.has_device_capability(89)
388
or envs.VLLM_TEST_FORCE_FP8_MARLIN
@@ -552,7 +560,9 @@ class Fp8LinearMethod(LinearMethodBase):
560
)
561
562
if self.use_marlin:
555
- prepare_fp8_layer_for_marlin(layer, size_k_first)
563
+ prepare_fp8_layer_for_marlin(
564
+ layer, size_k_first, input_dtype=self.marlin_input_dtype
565
+ )
566
# Activations not quantized for marlin.
567
del layer.input_scale
568
return
@@ -610,6 +620,7 @@ class Fp8LinearMethod(LinearMethodBase):
620
workspace=layer.workspace,
621
size_n=layer.output_size_per_partition,
622
size_k=layer.input_size_per_partition,
623
+ input_dtype=self.marlin_input_dtype,
624
bias=bias,
625
)
626
@@ -657,6 +668,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
668
self.block_quant, layer.moe_parallel_config
669
)
670
671
+ self.marlin_input_dtype = None
672
self.use_marlin = self.fp8_backend == Fp8MoeBackend.MARLIN
673
self.flashinfer_moe_backend: FlashinferMoeBackend | None = None
674
if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
@@ -1031,7 +1043,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1043
layer.w13_weight.data = w13_weight.data
1044
1045
if self.use_marlin:
1034
- prepare_moe_fp8_layer_for_marlin(layer, False)
1046
+ prepare_moe_fp8_layer_for_marlin(
1047
+ layer, False, input_dtype=self.marlin_input_dtype
1048
+ )
1049
# Activations not quantized for marlin.
1050
del layer.w13_input_scale
1051
del layer.w2_input_scale
@@ -1270,6 +1284,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1284
apply_router_weight_on_input=apply_router_weight_on_input,
1285
global_num_experts=global_num_experts,
1286
expert_map=expert_map,
1287
+ input_dtype=self.marlin_input_dtype,
1288
workspace=layer.workspace,
1289
)
1290
elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
vllm/model_executor/layers/quantization/gptq_marlin.py
+71
-3
@@ -41,6 +41,8 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import (
41
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
42
check_marlin_supported,
43
check_moe_marlin_supports_layer,
44
+ get_marlin_input_dtype,
45
+ marlin_act_int8_process_scales,
46
marlin_make_workspace_new,
47
marlin_moe_permute_scales,
48
marlin_permute_bias,
@@ -251,8 +253,21 @@ class GPTQMarlinConfig(QuantizationConfig):
253
return MoeWNA16Config.from_config(self.full_config).get_quant_method(
254
layer, prefix
255
)
254
- return get_moe_quant_method(self, layer, prefix, GPTQMarlinMoEMethod)
255
- return get_linear_quant_method(self, layer, prefix, GPTQMarlinLinearMethod)
256
+ moe_quant_method = get_moe_quant_method(
257
+ self, layer, prefix, GPTQMarlinMoEMethod
258
+ )
259
+ if moe_quant_method is None:
260
+ return None
261
+ moe_quant_method.input_dtype = get_marlin_input_dtype(prefix)
262
+ return moe_quant_method
263
+
264
+ quant_method = get_linear_quant_method(
265
+ self, layer, prefix, GPTQMarlinLinearMethod
266
+ )
267
+ if quant_method is None:
268
+ return None
269
+ quant_method.input_dtype = get_marlin_input_dtype(prefix)
270
+ return quant_method
271
272
@classmethod
273
def is_gptq_marlin_compatible(cls, quant_config: dict[str, Any]):
@@ -319,6 +334,8 @@ class GPTQMarlinLinearMethod(LinearMethodBase):
334
335
def __init__(self, quant_config: GPTQMarlinConfig) -> None:
336
self.quant_config = quant_config
337
+ self.input_dtype = None
338
+ self.quant_type = self.quant_config.quant_type
339
340
# Verify supported on platform.
341
verify_marlin_supported(
@@ -339,6 +356,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):
356
output_size_per_partition = sum(output_partition_sizes)
357
is_row_parallel = input_size != input_size_per_partition
358
weight_loader = extra_weight_attrs.get("weight_loader")
359
+ input_dtype = self.input_dtype
360
361
mp_linear_kernel_config = MPLinearLayerConfig(
362
full_weight_shape=(input_size, output_size),
@@ -347,7 +365,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):
365
output_size_per_partition,
366
),
367
weight_type=self.quant_config.quant_type,
350
- act_type=params_dtype,
368
+ act_type=params_dtype if input_dtype is None else input_dtype,
369
group_size=self.quant_config.group_size,
370
zero_points=False,
371
has_g_idx=self.quant_config.desc_act,
@@ -482,6 +500,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
500
self.quant_type = scalar_types.uint8b128
501
else:
502
raise ValueError("GPTQMarlinMoEMethod only supports int4 and int8 now.")
503
+ self.input_dtype = None
504
self.use_marlin = True
505
506
def create_weights(
@@ -493,6 +512,14 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
512
params_dtype: torch.dtype,
513
**extra_weight_attrs,
514
):
515
+ layer.input_dtype = self.input_dtype
516
+ is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
517
+
518
+ if is_a_8bit:
519
+ assert self.quant_type == scalar_types.uint4b8, (
520
+ "W8A8-INT8 is not supported by marlin kernel."
521
+ )
522
+
523
intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
524
525
self.is_k_full = (not self.quant_config.desc_act) or (
@@ -513,6 +540,9 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
540
scales_size2 = 1
541
strategy = FusedMoeWeightScaleSupported.CHANNEL.value
542
543
+ layer.num_groups_w13 = scales_size13
544
+ layer.num_groups_w2 = scales_size2
545
+
546
extra_weight_attrs.update({"quant_method": strategy, "is_transposed": True})
547
# Fused gate_up_proj (column parallel)
548
w13_qweight = torch.nn.Parameter(
@@ -630,6 +660,19 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
660
layer.workspace = marlin_make_workspace_new(device, 4)
661
662
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
663
+ is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
664
+
665
+ if is_a_8bit:
666
+ assert self.quant_type == scalar_types.uint4b8, (
667
+ "W8A8-INT8 is not supported by marlin kernel."
668
+ )
669
+
670
+ if self.input_dtype == torch.float8_e4m3fn:
671
+ ops.marlin_int4_fp8_preprocess(layer.w13_qweight, inplace=True)
672
+ ops.marlin_int4_fp8_preprocess(layer.w2_qweight, inplace=True)
673
+ layer.w13_scales.data = layer.w13_scales.data * 512
674
+ layer.w2_scales.data = layer.w2_scales.data * 512
675
+
676
# Process act_order
677
if self.quant_config.desc_act:
678
# Get sorting based on g_idx
@@ -678,6 +721,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
721
layer.w13_qweight.shape[1] * self.quant_config.pack_factor,
722
layer.w13_qweight.shape[2],
723
self.quant_config.quant_type.size_bits,
724
+ is_a_8bit=is_a_8bit,
725
)
726
replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
727
marlin_w2_qweight = ops.gptq_marlin_moe_repack(
@@ -686,6 +730,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
730
layer.w2_qweight.shape[1] * self.quant_config.pack_factor,
731
layer.w2_qweight.shape[2],
732
self.quant_config.quant_type.size_bits,
733
+ is_a_8bit=is_a_8bit,
734
)
735
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
736
# Repack scales
@@ -694,7 +739,17 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
739
size_k=layer.intermediate_size_per_partition,
740
size_n=layer.w13_scales.shape[2],
741
group_size=self.quant_config.group_size,
742
+ is_a_8bit=is_a_8bit,
743
)
744
+ if self.input_dtype == torch.int8 and layer.num_groups_w13 > 1:
745
+ marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
746
+ marlin_w13_scales
747
+ )
748
+ layer.register_parameter(
749
+ "w13_input_global_scale",
750
+ torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
751
+ )
752
+
753
replace_parameter(layer, "w13_scales", marlin_w13_scales)
754
marlin_w2_scales = marlin_moe_permute_scales(
755
s=layer.w2_scales,
@@ -706,7 +761,17 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
761
),
762
size_n=layer.w2_scales.shape[2],
763
group_size=self.quant_config.group_size,
764
+ is_a_8bit=is_a_8bit,
765
)
766
+ if self.input_dtype == torch.int8 and layer.num_groups_w2 > 1:
767
+ marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
768
+ marlin_w2_scales
769
+ )
770
+ layer.register_parameter(
771
+ "w2_input_global_scale",
772
+ torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
773
+ )
774
+
775
replace_parameter(layer, "w2_scales", marlin_w2_scales)
776
777
if hasattr(layer, "w13_bias") and layer.w13_bias is not None:
@@ -761,6 +826,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
826
router_logits,
827
topk_weights,
828
topk_ids,
829
+ input_global_scale1=getattr(layer, "w13_input_global_scale", None),
830
+ input_global_scale2=getattr(layer, "w2_input_global_scale", None),
831
quant_type_id=self.quant_type.id,
832
apply_router_weight_on_input=apply_router_weight_on_input,
833
global_num_experts=global_num_experts,
@@ -771,4 +838,5 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
838
sort_indices2=layer.w2_g_idx_sort_indices,
839
workspace=layer.workspace,
840
is_k_full=self.is_k_full,
841
+ input_dtype=self.input_dtype,
842
)
vllm/model_executor/layers/quantization/hqq_marlin.py
+1
@@ -351,6 +351,7 @@ class HQQMarlinMethod(LinearMethodBase):
351
bias,
352
scales,
353
None,
354
+ None,
355
zeros,
356
layer.g_idx,
357
layer.g_idx_sort_indices,
vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py
+34
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
9
MARLIN_SUPPORTED_GROUP_SIZES,
10
apply_gptq_marlin_linear,
11
check_marlin_supports_shape,
12
+ marlin_act_int8_process_scales,
13
marlin_is_k_full,
14
marlin_make_empty_g_idx,
15
marlin_make_workspace_new,
@@ -21,6 +22,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
22
)
23
from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
24
from vllm.platforms import current_platform
25
+from vllm.scalar_type import scalar_types
26
27
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
28
@@ -65,6 +67,18 @@ class MarlinLinearKernel(MPLinearKernel):
67
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
68
device = getattr(layer, self.w_q_name).device
69
c = self.config
70
+ is_a_8bit = c.act_type is not None and c.act_type.itemsize == 1
71
+
72
+ if is_a_8bit:
73
+ assert c.weight_type == scalar_types.uint4b8, (
74
+ "W8A8 is not supported by marlin kernel."
75
+ )
76
+
77
+ if c.act_type == torch.float8_e4m3fn:
78
+ ops.marlin_int4_fp8_preprocess(getattr(layer, self.w_q_name), inplace=True)
79
+ getattr(layer, self.w_s_name).data = (
80
+ getattr(layer, self.w_s_name).data * 512
81
+ )
82
83
row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
84
self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
@@ -88,6 +102,7 @@ class MarlinLinearKernel(MPLinearKernel):
102
size_k=c.partition_weight_shape[0],
103
size_n=c.partition_weight_shape[1],
104
num_bits=c.weight_type.size_bits,
105
+ is_a_8bit=is_a_8bit,
106
)
107
return x
108
@@ -99,7 +114,22 @@ class MarlinLinearKernel(MPLinearKernel):
114
size_k=c.partition_weight_shape[0],
115
size_n=c.partition_weight_shape[1],
116
group_size=c.group_size,
117
+ is_a_8bit=is_a_8bit,
118
)
119
+
120
+ if c.group_size == -1:
121
+ num_groups = 1
122
+ else:
123
+ num_groups = c.partition_weight_shape[0] // c.group_size
124
+
125
+ if c.act_type == torch.int8 and num_groups > 1:
126
+ x.data, input_global_scale = marlin_act_int8_process_scales(x.data)
127
+ layer.register_parameter(
128
+ "input_global_scale",
129
+ torch.nn.Parameter(input_global_scale, requires_grad=False),
130
+ )
131
+ else:
132
+ layer.input_global_scale = None
133
return x
134
135
if c.has_g_idx:
@@ -129,6 +159,7 @@ class MarlinLinearKernel(MPLinearKernel):
159
size_k=grouped_k,
160
size_n=c.partition_weight_shape[1],
161
num_bits=c.weight_type.size_bits,
162
+ is_a_8bit=is_a_8bit,
163
),
164
)
165
else:
@@ -150,6 +181,7 @@ class MarlinLinearKernel(MPLinearKernel):
181
182
# `process_weights_after_loading` will ensure w_zp and w_gidx are not
183
# None for marlin
184
+
185
return apply_gptq_marlin_linear(
186
input=x,
187
weight=w_q,
@@ -162,5 +194,7 @@ class MarlinLinearKernel(MPLinearKernel):
194
input_size_per_partition=c.partition_weight_shape[0],
195
output_size_per_partition=c.partition_weight_shape[1],
196
is_k_full=self.is_k_full,
197
+ input_global_scale=getattr(layer, "input_global_scale", None),
198
bias=bias,
199
+ input_dtype=c.act_type,
200
)
vllm/model_executor/layers/quantization/modelopt.py
+15
-3
@@ -55,6 +55,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
55
select_cutlass_fp8_gemm_impl,
56
swap_w13_to_w31,
57
)
58
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
59
+ get_marlin_input_dtype,
60
+)
61
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
62
apply_fp4_marlin_linear,
63
is_fp4_marlin_supported,
@@ -170,9 +173,15 @@ class ModelOptQuantConfigBase(QuantizationConfig):
173
174
# now, the layer is quantized, handle it here
175
if isinstance(layer, LinearBase):
173
- return self.LinearMethodCls(self)
176
+ quant_method = self.LinearMethodCls(self)
177
+ if getattr(quant_method, "backend", "") == "marlin":
178
+ quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
179
+ return quant_method
180
elif isinstance(layer, FusedMoE):
175
- return self.FusedMoEMethodCls(quant_config=self, layer=layer)
181
+ quant_method = self.FusedMoEMethodCls(quant_config=self, layer=layer)
182
+ if getattr(quant_method, "backend", "") == "marlin":
183
+ quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
184
+ return quant_method
185
186
return None
187
@@ -898,6 +907,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
907
908
def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
909
self.quant_config = quant_config
910
+ self.marlin_input_dtype = None
911
912
self.backend = "none"
913
if envs.VLLM_NVFP4_GEMM_BACKEND is None:
@@ -1065,6 +1075,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1075
size_n=layer.output_size_per_partition,
1076
size_k=layer.input_size_per_partition,
1077
bias=bias,
1078
+ input_dtype=self.marlin_input_dtype,
1079
)
1080
1081
output_dtype = x.dtype
@@ -1124,6 +1135,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1135
self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported
1136
self.allow_flashinfer = _nvfp4.allow_flashinfer
1137
self.use_marlin = _nvfp4.use_marlin
1138
+ self.marlin_input_dtype = None
1139
self.flashinfer_moe_backend = None
1140
if self.allow_flashinfer:
1141
self.flashinfer_moe_backend = get_flashinfer_moe_backend()
@@ -1517,7 +1529,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1529
apply_router_weight_on_input=apply_router_weight_on_input,
1530
global_num_experts=global_num_experts,
1531
expert_map=expert_map,
1520
- workspace=layer.workspace,
1532
+ input_dtype=self.marlin_input_dtype,
1533
)
1534
1535
elif self.allow_flashinfer:
vllm/model_executor/layers/quantization/mxfp4.py
+10
-2
@@ -38,6 +38,9 @@ from vllm.model_executor.layers.quantization.base_config import (
38
QuantizationConfig,
39
QuantizeMethodBase,
40
)
41
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
42
+ get_marlin_input_dtype,
43
+)
44
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
45
prepare_moe_fp4_layer_for_marlin,
46
)
@@ -205,7 +208,9 @@ class Mxfp4Config(QuantizationConfig):
208
if current_platform.is_xpu():
209
return IpexMxfp4MoEMethod(layer.moe_config)
210
else:
208
- return Mxfp4MoEMethod(layer.moe_config)
211
+ quant_method = Mxfp4MoEMethod(layer.moe_config)
212
+ quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
213
+ return quant_method
214
elif isinstance(layer, Attention):
215
# TODO: Add support for MXFP4 Attention.
216
logger.debug_once(
@@ -220,6 +225,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
225
def __init__(self, moe: FusedMoEConfig):
226
super().__init__(moe)
227
self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
228
+
229
+ self.marlin_input_dtype = None
230
self.use_marlin = self.mxfp4_backend == Mxfp4Backend.MARLIN
231
self.max_capture_size = (
232
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
@@ -385,7 +392,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
392
393
def process_weights_after_loading(self, layer):
394
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
388
- prepare_moe_fp4_layer_for_marlin(layer)
395
+ prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)
396
elif (
397
self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
398
or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
@@ -914,6 +921,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
921
global_num_experts=global_num_experts,
922
activation=activation,
923
expert_map=expert_map,
924
+ input_dtype=self.marlin_input_dtype,
925
)
926
927
assert _can_support_mxfp4(
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+112
-13
@@ -9,6 +9,11 @@ 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.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13
+from vllm.model_executor.layers.quantization.utils.int8_utils import (
14
+ per_token_quant_int8,
15
+)
16
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
17
from vllm.platforms import current_platform
18
from vllm.scalar_type import ScalarType, scalar_types
19
@@ -286,10 +291,10 @@ def get_scale_perms():
291
292
293
def marlin_permute_scales(
289
- s: torch.Tensor, size_k: int, size_n: int, group_size: int
294
+ s: torch.Tensor, size_k: int, size_n: int, group_size: int, is_a_8bit: bool = False
295
) -> torch.Tensor:
296
scale_perm, scale_perm_single = get_scale_perms()
292
- if group_size < size_k and group_size != -1:
297
+ if group_size < size_k and group_size != -1 and not is_a_8bit:
298
s = s.reshape((-1, len(scale_perm)))[:, scale_perm]
299
else:
300
s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
@@ -305,11 +310,15 @@ def marlin_permute_bias(s: torch.Tensor) -> torch.Tensor:
310
return s.reshape(*origin_shape).contiguous()
311
312
313
+def marlin_act_int8_process_scales(s: torch.Tensor):
314
+ a_scales_scale_factor = 1 / 4096 * s.max().float()
315
+ s = s / s.max() * 4096
316
+ s = s.round().to(torch.int16).view(s.dtype)
317
+ return s, a_scales_scale_factor
318
+
319
+
320
def marlin_moe_permute_scales(
309
- s: torch.Tensor,
310
- size_k: int,
311
- size_n: int,
312
- group_size: int,
321
+ s: torch.Tensor, size_k: int, size_n: int, group_size: int, is_a_8bit: bool = False
322
):
323
num_experts = s.shape[0]
324
output = torch.empty(
@@ -319,12 +328,12 @@ def marlin_moe_permute_scales(
328
)
329
330
for e in range(num_experts):
322
- output[e] = marlin_permute_scales(s[e], size_k, size_n, group_size)
331
+ output[e] = marlin_permute_scales(s[e], size_k, size_n, group_size, is_a_8bit)
332
return output
333
334
335
def marlin_zero_points(
327
- zp: torch.Tensor, size_k: int, size_n: int, num_bits: int
336
+ zp: torch.Tensor, size_k: int, size_n: int, num_bits: int, is_a_8bit: bool = False
337
) -> torch.Tensor:
338
# Permute zero-points in a similar way to scales, but do not use the
339
# "single" permutation, since zero-points are applied on every MMA
@@ -339,7 +348,8 @@ def marlin_zero_points(
348
else:
349
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
350
342
- zp = zp.reshape((-1, len(interleave)))[:, interleave].ravel()
351
+ if not is_a_8bit:
352
+ zp = zp.reshape((-1, len(interleave)))[:, interleave].ravel()
353
zp = zp.reshape((-1, size_n)).contiguous()
354
zp = pack_cols(zp, num_bits, size_k, size_n)
355
@@ -347,7 +357,11 @@ def marlin_zero_points(
357
358
359
def awq_to_marlin_zero_points(
350
- q_zp_packed: torch.Tensor, size_k: int, size_n: int, num_bits: int
360
+ q_zp_packed: torch.Tensor,
361
+ size_k: int,
362
+ size_n: int,
363
+ num_bits: int,
364
+ is_a_8bit: bool = False,
365
) -> torch.Tensor:
366
# AWQ zero-points are quantized and packed on the column dim.
367
# In addition, the values are permuted based on dequantizer.
@@ -366,12 +380,16 @@ def awq_to_marlin_zero_points(
380
q_zp = q_zp.reshape((-1, len(undo_interleave)))[:, undo_interleave].ravel()
381
q_zp = q_zp.reshape((-1, size_n)).contiguous()
382
369
- marlin_zp = marlin_zero_points(q_zp, size_k, size_n, num_bits)
383
+ marlin_zp = marlin_zero_points(q_zp, size_k, size_n, num_bits, is_a_8bit)
384
return marlin_zp
385
386
387
def moe_awq_to_marlin_zero_points(
374
- q_zp_packed: torch.Tensor, size_k: int, size_n: int, num_bits: int
388
+ q_zp_packed: torch.Tensor,
389
+ size_k: int,
390
+ size_n: int,
391
+ num_bits: int,
392
+ is_a_8bit: bool = False,
393
):
394
num_experts = q_zp_packed.shape[0]
395
output = torch.empty(
@@ -380,7 +398,9 @@ def moe_awq_to_marlin_zero_points(
398
dtype=q_zp_packed.dtype,
399
)
400
for e in range(num_experts):
383
- output[e] = awq_to_marlin_zero_points(q_zp_packed[e], size_k, size_n, num_bits)
401
+ output[e] = awq_to_marlin_zero_points(
402
+ q_zp_packed[e], size_k, size_n, num_bits, is_a_8bit
403
+ )
404
return output
405
406
@@ -432,6 +452,48 @@ def should_use_atomic_add_reduce(
452
return True
453
454
455
+_quant_fp8_method: QuantFP8 | None = None
456
+
457
+
458
+def get__quant_fp8_method() -> QuantFP8:
459
+ global _quant_fp8_method
460
+ if _quant_fp8_method is None:
461
+ _quant_fp8_method = QuantFP8(False, GroupShape.PER_TOKEN)
462
+ return _quant_fp8_method
463
+
464
+
465
+def get_marlin_input_dtype(prefix):
466
+ if envs.VLLM_MARLIN_INPUT_DTYPE is None:
467
+ return
468
+ elif envs.VLLM_MARLIN_INPUT_DTYPE.lower() == "int8":
469
+ return torch.int8
470
+ elif envs.VLLM_MARLIN_INPUT_DTYPE.lower() == "fp8":
471
+ if not current_platform.is_device_capability(
472
+ 89
473
+ ) and not current_platform.is_device_capability(120):
474
+ raise ValueError(
475
+ "Marlin W4A8-FP8 only support SM89 or SM120 device "
476
+ "(It is slower than Marlin W4A16 on other devices). "
477
+ "You can consider using W4A8-INT8 instead"
478
+ "(set VLLM_MARLIN_INPUT_DTYPE=int8)."
479
+ )
480
+
481
+ _ = get__quant_fp8_method()
482
+ return torch.float8_e4m3fn
483
+ else:
484
+ return
485
+
486
+
487
+def marlin_quant_input(x: torch.Tensor, quant_dtype: torch.dtype):
488
+ x = x.reshape(-1, x.shape[-1])
489
+ if quant_dtype == torch.int8:
490
+ return per_token_quant_int8(x)
491
+ elif quant_dtype == torch.float8_e4m3fn:
492
+ return get__quant_fp8_method()(x)
493
+ else:
494
+ raise ValueError(f"unsupported quant_dtype {quant_dtype}")
495
+
496
+
497
def apply_gptq_marlin_linear(
498
input: torch.Tensor,
499
weight: torch.Tensor,
@@ -444,8 +506,10 @@ def apply_gptq_marlin_linear(
506
output_size_per_partition: int,
507
input_size_per_partition: int,
508
is_k_full: bool,
509
+ input_global_scale: torch.Tensor | None = None,
510
bias: torch.Tensor | None = None,
511
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
512
+ input_dtype: torch.dtype | None = None,
513
) -> torch.Tensor:
514
reshaped_x = input.reshape(-1, input.shape[-1])
515
out_shape = input.shape[:-1] + (output_size_per_partition,)
@@ -458,12 +522,27 @@ def apply_gptq_marlin_linear(
522
dtype=input.dtype,
523
)
524
525
+ a_scales = None
526
+ if input_dtype == torch.int8:
527
+ assert wtype == scalar_types.uint4b8, (
528
+ "W8A8-INT8 is not supported by marlin kernel."
529
+ )
530
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
531
+ a_scales = a_scales * input_global_scale
532
+ elif input_dtype == torch.float8_e4m3fn:
533
+ assert wtype == scalar_types.uint4b8, (
534
+ "INT8 weight + FP8 activation is not supported."
535
+ )
536
+
537
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
538
+
539
output = ops.gptq_marlin_gemm(
540
reshaped_x,
541
None,
542
weight,
543
bias,
544
weight_scale,
545
+ a_scales,
546
None,
547
weight_zp,
548
g_idx,
@@ -493,8 +572,10 @@ def apply_awq_marlin_linear(
572
quant_type: ScalarType,
573
output_size_per_partition: int,
574
input_size_per_partition: int,
575
+ input_global_scale: torch.Tensor | None = None,
576
bias: torch.Tensor | None = None,
577
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
578
+ input_dtype: torch.dtype | None = None,
579
) -> torch.Tensor:
580
reshaped_x = input.reshape(-1, input.shape[-1])
581
out_shape = input.shape[:-1] + (output_size_per_partition,)
@@ -507,12 +588,20 @@ def apply_awq_marlin_linear(
588
dtype=input.dtype,
589
)
590
591
+ a_scales = None
592
+ if input_dtype == torch.int8:
593
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
594
+ a_scales = a_scales * input_global_scale
595
+ elif input_dtype == torch.float8_e4m3fn:
596
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
597
+
598
output = ops.gptq_marlin_gemm(
599
reshaped_x,
600
None,
601
weight,
602
bias,
603
weight_scale,
604
+ a_scales,
605
None,
606
weight_zp,
607
g_idx,
@@ -538,8 +627,10 @@ def apply_rtn_marlin_linear(
627
quant_type: ScalarType,
628
output_size_per_partition: int,
629
input_size_per_partition: int,
630
+ input_global_scale: torch.Tensor | None = None,
631
bias: torch.Tensor | None = None,
632
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
633
+ input_dtype: torch.dtype | None = None,
634
) -> torch.Tensor:
635
reshaped_x = input.reshape(-1, input.shape[-1])
636
out_shape = input.shape[:-1] + (output_size_per_partition,)
@@ -552,12 +643,20 @@ def apply_rtn_marlin_linear(
643
dtype=input.dtype,
644
)
645
646
+ a_scales = None
647
+ if input_dtype == torch.int8:
648
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
649
+ a_scales = a_scales * input_global_scale
650
+ elif input_dtype == torch.float8_e4m3fn:
651
+ reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)
652
+
653
output = ops.gptq_marlin_gemm(
654
reshaped_x,
655
None,
656
weight,
657
bias,
658
weight_scale,
659
+ a_scales,
660
None,
661
None,
662
None,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+88
-33
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
marlin_make_workspace_new,
12
marlin_permute_bias,
13
marlin_permute_scales,
14
+ marlin_quant_input,
15
should_use_atomic_add_reduce,
16
)
17
from vllm.platforms import current_platform
@@ -37,12 +38,6 @@ def nvfp4_marlin_process_scales(marlin_scales):
38
# convert to half first, we would convert to fp8 later
39
marlin_scales = marlin_scales.to(torch.half)
40
40
- # 8 is the number of scale number using by one thread
41
- marlin_scales = marlin_scales.view(marlin_scales.size(0) // 2, 2, -1, 8)
42
- marlin_scales = marlin_scales.permute(0, 2, 1, 3).reshape(
43
- marlin_scales.size(0) * 2, -1
44
- )
45
-
41
# fit the layout of fp8 dequantization
42
marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
43
marlin_scales.size(0), -1
@@ -62,18 +57,20 @@ def nvfp4_marlin_process_scales(marlin_scales):
57
return marlin_scales
58
59
65
-def mxfp4_marlin_process_scales(marlin_scales):
66
- # 8 is the number of scale number using by one thread
67
- marlin_scales = marlin_scales.view(marlin_scales.size(0) // 2, 2, -1, 8)
68
- marlin_scales = marlin_scales.permute(0, 2, 1, 3).reshape(
69
- marlin_scales.size(0) * 2, -1
70
- )
71
-
60
+def mxfp4_marlin_process_scales(marlin_scales, input_dtype=None):
61
# fit the layout of fp8 dequantization
73
- marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
74
- marlin_scales.size(0), -1
75
- )
62
+ if input_dtype is None or input_dtype.itemsize == 2:
63
+ marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
64
+ marlin_scales.size(0), -1
65
+ )
66
+
67
marlin_scales = marlin_scales.to(torch.float8_e8m0fnu)
68
+ if input_dtype == torch.float8_e4m3fn:
69
+ marlin_scales = marlin_scales.view(torch.uint8)
70
+ assert marlin_scales.max() <= 249
71
+ # exponent_bias (fp4->fp8) = 2 ** 3 - 2 ** 1 = 6
72
+ marlin_scales = marlin_scales + 6
73
+ marlin_scales = marlin_scales.view(torch.float8_e8m0fnu)
74
return marlin_scales
75
76
@@ -99,6 +96,7 @@ def apply_fp4_marlin_linear(
96
size_n: int,
97
size_k: int,
98
bias: torch.Tensor | None = None,
99
+ input_dtype: torch.dtype | None = None,
100
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
101
) -> torch.Tensor:
102
# For GPUs that lack FP4 hardware support, we can leverage the
@@ -111,12 +109,24 @@ def apply_fp4_marlin_linear(
109
m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype
110
)
111
112
+ inputs = reshaped_x
113
+ a_scales = None
114
+ is_nvfp4 = weight_scale_2 is not None
115
+ if input_dtype is not None and input_dtype.itemsize == 1:
116
+ if is_nvfp4:
117
+ raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.")
118
+ elif input_dtype != torch.float8_e4m3fn:
119
+ raise RuntimeError("MXFP4 weight + INT8 activation is not supported.")
120
+
121
+ inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
122
+
123
output = ops.gptq_marlin_gemm(
115
- a=reshaped_x,
124
+ a=inputs,
125
c=None,
126
b_q_weight=weight,
127
b_bias=bias,
128
b_scales=weight_scale,
129
+ a_scales=a_scales,
130
global_scale=weight_scale_2,
131
b_zeros=None,
132
g_idx=None,
@@ -133,7 +143,9 @@ def apply_fp4_marlin_linear(
143
return output.reshape(out_shape)
144
145
136
-def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
146
+def prepare_fp4_layer_for_marlin(
147
+ layer: torch.nn.Module, input_dtype: torch.dtype | None = None
148
+) -> None:
149
logger.warning_once(
150
"Your GPU does not have native support for FP4 computation but "
151
"FP4 quantization is being used. Weight-only FP4 compression will "
@@ -160,12 +172,14 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
172
perm = torch.empty(0, dtype=torch.int, device=device)
173
qweight = layer.weight.view(torch.int32).T.contiguous()
174
175
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
176
marlin_qweight = ops.gptq_marlin_repack(
177
b_q_weight=qweight,
178
perm=perm,
179
size_k=part_size_k,
180
size_n=part_size_n,
181
num_bits=4,
182
+ is_a_8bit=is_a_8bit,
183
)
184
layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
185
@@ -178,7 +192,11 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
192
193
weight_scale = weight_scale.to(param_dtype)
194
weight_scale = marlin_permute_scales(
181
- s=weight_scale, size_k=part_size_k, size_n=part_size_n, group_size=group_size
195
+ s=weight_scale,
196
+ size_k=part_size_k,
197
+ size_n=part_size_n,
198
+ group_size=group_size,
199
+ is_a_8bit=is_a_8bit,
200
)
201
202
if is_nvfp4:
@@ -189,7 +207,9 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
207
weight_scale_2 = nvfp4_marlin_process_global_scale(weight_scale_2)
208
layer.weight_scale_2 = torch.nn.Parameter(weight_scale_2, requires_grad=False)
209
else:
192
- weight_scale = mxfp4_marlin_process_scales(weight_scale)
210
+ weight_scale = mxfp4_marlin_process_scales(
211
+ weight_scale, input_dtype=input_dtype
212
+ )
213
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
214
215
if hasattr(layer, "bias") and layer.bias is not None:
@@ -200,7 +220,9 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
220
return
221
222
203
-def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
223
+def prepare_moe_fp4_layer_for_marlin(
224
+ layer: torch.nn.Module, input_dtype: torch.dtype | None = None
225
+) -> None:
226
logger.warning_once(
227
"Your GPU does not have native support for FP4 computation but "
228
"FP4 quantization is being used. Weight-only FP4 compression will "
@@ -220,6 +242,7 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
242
param_dtype = layer.params_dtype
243
layer.workspace = marlin_make_workspace_new(device, 4)
244
perm = torch.empty(0, dtype=torch.int, device=device)
245
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
246
247
# WEIGHT
248
# Repack weights to marlin format
@@ -237,7 +260,12 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
260
qweight = weight[i].view(torch.int32).T.contiguous()
261
262
marlin_qweight = ops.gptq_marlin_repack(
240
- b_q_weight=qweight, perm=perm, size_k=size_k, size_n=size_n, num_bits=4
263
+ b_q_weight=qweight,
264
+ perm=perm,
265
+ size_k=size_k,
266
+ size_n=size_n,
267
+ num_bits=4,
268
+ is_a_8bit=is_a_8bit,
269
)
270
tensor_list.append(marlin_qweight)
271
@@ -266,12 +294,18 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
294
scale = scales[i].T
295
296
marlin_scales = marlin_permute_scales(
269
- s=scale, size_k=size_k, size_n=size_n, group_size=group_size
297
+ s=scale,
298
+ size_k=size_k,
299
+ size_n=size_n,
300
+ group_size=group_size,
301
+ is_a_8bit=is_a_8bit,
302
)
303
if is_nvfp4:
304
marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
305
else:
274
- marlin_scales = mxfp4_marlin_process_scales(marlin_scales)
306
+ marlin_scales = mxfp4_marlin_process_scales(
307
+ marlin_scales, input_dtype=input_dtype
308
+ )
309
tensor_list.append(marlin_scales)
310
311
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
@@ -301,7 +335,10 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
335
setattr(layer, name, bias)
336
337
304
-def rand_marlin_weight_nvfp4_like(weight, group_size):
338
+def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
339
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
340
+
341
+ assert not is_a_8bit, "NVFP4 weight + INT8/FP8 activation is not supported."
342
assert group_size > 0
343
size_n, size_k = weight.shape
344
device = weight.device
@@ -337,10 +374,15 @@ def rand_marlin_weight_nvfp4_like(weight, group_size):
374
size_k=size_k,
375
size_n=size_n,
376
num_bits=4,
377
+ is_a_8bit=is_a_8bit,
378
)
379
380
marlin_scales = marlin_permute_scales(
343
- s=scales.T.to(weight.dtype), size_k=size_k, size_n=size_n, group_size=group_size
381
+ s=scales.T.to(weight.dtype),
382
+ size_k=size_k,
383
+ size_n=size_n,
384
+ group_size=group_size,
385
+ is_a_8bit=is_a_8bit,
386
)
387
marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
388
@@ -349,14 +391,20 @@ def rand_marlin_weight_nvfp4_like(weight, group_size):
391
return weight_ref.T, marlin_qweight, marlin_scales, global_scale
392
393
352
-def rand_marlin_weight_mxfp4_like(weight, group_size):
394
+def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None):
395
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
396
+ if is_a_8bit:
397
+ assert input_dtype == torch.float8_e4m3fn, (
398
+ "MXFP4 weight + INT8 activation is not supported."
399
+ )
400
+
401
assert group_size > 0
402
size_n, size_k = weight.shape
403
device = weight.device
404
405
scales = torch.randint(
358
- 100,
359
- 125,
406
+ 110,
407
+ 120,
408
(size_n, size_k // group_size),
409
dtype=torch.uint8,
410
device=weight.device,
@@ -380,18 +428,25 @@ def rand_marlin_weight_mxfp4_like(weight, group_size):
428
).view(size_n, size_k)
429
weight_ref = weight_ref * scales.repeat_interleave(group_size, 1).to(weight.dtype)
430
431
+ perm = torch.empty(0, dtype=torch.int, device=device)
432
+ fp4_weight = fp4_weight.view(torch.int32).T.contiguous()
433
marlin_qweight = ops.gptq_marlin_repack(
384
- b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
385
- perm=torch.empty(0, dtype=torch.int, device=device),
434
+ b_q_weight=fp4_weight,
435
+ perm=perm,
436
size_k=size_k,
437
size_n=size_n,
438
num_bits=4,
439
+ is_a_8bit=is_a_8bit,
440
)
441
442
marlin_scales = marlin_permute_scales(
392
- s=scales.T.to(weight.dtype), size_k=size_k, size_n=size_n, group_size=group_size
443
+ s=scales.T.to(weight.dtype),
444
+ size_k=size_k,
445
+ size_n=size_n,
446
+ group_size=group_size,
447
+ is_a_8bit=is_a_8bit,
448
)
449
395
- marlin_scales = mxfp4_marlin_process_scales(marlin_scales)
450
+ marlin_scales = mxfp4_marlin_process_scales(marlin_scales, input_dtype=input_dtype)
451
452
return weight_ref.T, marlin_qweight, marlin_scales.to(torch.float8_e8m0fnu)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+34
-7
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
marlin_make_workspace_new,
12
marlin_permute_bias,
13
marlin_permute_scales,
14
+ marlin_quant_input,
15
should_use_atomic_add_reduce,
16
)
17
from vllm.platforms import current_platform
@@ -45,6 +46,7 @@ def apply_fp8_marlin_linear(
46
size_n: int,
47
size_k: int,
48
bias: torch.Tensor | None,
49
+ input_dtype: torch.dtype | None = None,
50
use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
51
) -> torch.Tensor:
52
# For GPUs that lack FP8 hardware support, we can leverage the
@@ -57,12 +59,21 @@ def apply_fp8_marlin_linear(
59
m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype
60
)
61
62
+ inputs = reshaped_x
63
+ a_scales = None
64
+ if input_dtype is not None and input_dtype.itemsize == 1:
65
+ if input_dtype != torch.float8_e4m3fn:
66
+ raise RuntimeError("FP8 weight + INT8 activation is not supported.")
67
+
68
+ inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)
69
+
70
output = ops.gptq_marlin_gemm(
71
a=reshaped_x,
72
c=None,
73
b_q_weight=weight,
74
b_bias=bias,
75
b_scales=weight_scale,
76
+ a_scales=a_scales,
77
global_scale=None,
78
b_zeros=None,
79
g_idx=None,
@@ -80,7 +91,9 @@ def apply_fp8_marlin_linear(
91
92
93
def prepare_fp8_layer_for_marlin(
83
- layer: torch.nn.Module, size_k_first: bool = True
94
+ layer: torch.nn.Module,
95
+ size_k_first: bool = True,
96
+ input_dtype: torch.dtype | None = None,
97
) -> None:
98
logger.warning_once(
99
"Your GPU does not have native support for FP8 computation but "
@@ -162,7 +175,8 @@ def prepare_fp8_layer_for_marlin(
175
marlin_scales = marlin_permute_scales(
176
s=scales, size_k=part_size_k, size_n=part_size_n, group_size=group_size
177
)
165
- marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
178
+ if input_dtype != torch.float8_e4m3fn:
179
+ marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
180
layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)
181
182
if hasattr(layer, "bias") and layer.bias is not None:
@@ -172,7 +186,9 @@ def prepare_fp8_layer_for_marlin(
186
187
188
def prepare_moe_fp8_layer_for_marlin(
175
- layer: torch.nn.Module, size_k_first: bool = True
189
+ layer: torch.nn.Module,
190
+ size_k_first: bool = True,
191
+ input_dtype: torch.dtype | None = None,
192
) -> None:
193
logger.warning_once(
194
"Your GPU does not have native support for FP8 computation but "
@@ -278,7 +294,8 @@ def prepare_moe_fp8_layer_for_marlin(
294
tensor_list.append(marlin_scales)
295
296
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
281
- scales = fp8_fused_exponent_bias_into_scales(scales)
297
+ if input_dtype != torch.float8_e4m3fn:
298
+ scales = fp8_fused_exponent_bias_into_scales(scales)
299
scales = torch.nn.Parameter(scales, requires_grad=False)
300
301
setattr(layer, name + "_weight_scale", scales)
@@ -318,7 +335,11 @@ def pack_fp8_to_int32(
335
return int32_tensor.T.contiguous() if size_k_first else int32_tensor
336
337
321
-def marlin_quant_fp8_torch(weight, group_size):
338
+def marlin_quant_fp8_torch(weight, group_size, input_dtype=None):
339
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
340
+ if is_a_8bit:
341
+ assert input_dtype == torch.float8_e4m3fn
342
+
343
size_n, size_k = weight.shape
344
device = weight.device
345
@@ -334,16 +355,22 @@ def marlin_quant_fp8_torch(weight, group_size):
355
weight_ref = fp8_weight.to(weight.dtype) * repeated_scales
356
357
packed_weight = pack_fp8_to_int32(fp8_weight, False).T.contiguous()
358
+ perm = torch.empty(0, dtype=torch.int, device=device)
359
marlin_qweight = ops.gptq_marlin_repack(
360
b_q_weight=packed_weight,
339
- perm=torch.empty(0, dtype=torch.int, device=device),
361
+ perm=perm,
362
size_k=size_k,
363
size_n=size_n,
364
num_bits=8,
365
+ is_a_8bit=is_a_8bit,
366
)
367
368
marlin_scales = marlin_permute_scales(
346
- s=scales.T, size_k=size_k, size_n=size_n, group_size=group_size
369
+ s=scales.T,
370
+ size_k=size_k,
371
+ size_n=size_n,
372
+ group_size=group_size,
373
+ is_a_8bit=is_a_8bit,
374
)
375
376
marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
vllm/model_executor/layers/quantization/utils/marlin_utils_test.py
+88
-30
@@ -5,7 +5,8 @@
5
import numpy as np
6
import torch
7
8
-from vllm.scalar_type import ScalarType
8
+from vllm import _custom_ops as ops
9
+from vllm.scalar_type import ScalarType, scalar_types
10
11
from .marlin_utils import GPTQ_MARLIN_TILE, marlin_permute_scales, marlin_zero_points
12
from .quant_utils import (
@@ -29,13 +30,19 @@ class MarlinWorkspace:
30
self.scratch = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda")
31
32
32
-def marlin_permute_weights(q_w, size_k, size_n, perm, tile=GPTQ_MARLIN_TILE):
33
+def marlin_permute_weights(
34
+ q_w, size_k, size_n, perm, tile=GPTQ_MARLIN_TILE, is_a_8bit=False
35
+):
36
assert q_w.shape == (size_k, size_n)
37
assert size_k % tile == 0, f"size_k = {size_k}, tile = {tile}"
38
assert size_n % tile == 0, f"size_k = {size_n}, tile = {tile}"
39
37
- # Permute weights to 16x64 marlin tiles
38
- q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile))
40
+ if is_a_8bit:
41
+ # Permute weights to 32x32 marlin tiles
42
+ q_w = q_w.reshape((size_k // (tile * 2), tile * 2, size_n // tile, tile))
43
+ else:
44
+ # Permute weights to 16x64 marlin tiles
45
+ q_w = q_w.reshape((size_k // tile, tile, size_n // tile, tile))
46
q_w = q_w.permute((0, 2, 1, 3))
47
q_w = q_w.reshape((size_k // tile, size_n * tile))
48
@@ -44,9 +51,9 @@ def marlin_permute_weights(q_w, size_k, size_n, perm, tile=GPTQ_MARLIN_TILE):
51
return q_w
52
53
47
-def marlin_weights(q_w, size_k, size_n, num_bits, perm):
54
+def marlin_weights(q_w, size_k, size_n, num_bits, perm, is_a_8bit=False):
55
# Permute
49
- q_w = marlin_permute_weights(q_w, size_k, size_n, perm)
56
+ q_w = marlin_permute_weights(q_w, size_k, size_n, perm, is_a_8bit=is_a_8bit)
57
58
# Pack
59
pack_factor = get_pack_factor(num_bits)
@@ -63,28 +70,53 @@ def marlin_weights(q_w, size_k, size_n, num_bits, perm):
70
return q_packed
71
72
66
-def get_weight_perm(num_bits: int):
73
+def get_weight_perm(num_bits: int, is_a_8bit: bool = False):
74
perm_list: list[int] = []
68
- for i in range(32):
69
- perm1: list[int] = []
70
- col = i // 4
71
- for block in [0, 1]:
72
- for row in [
73
- 2 * (i % 4),
74
- 2 * (i % 4) + 1,
75
- 2 * (i % 4 + 4),
76
- 2 * (i % 4 + 4) + 1,
77
- ]:
78
- perm1.append(16 * row + col + 8 * block)
79
- for j in range(4):
80
- perm_list.extend([p + 256 * j for p in perm1])
75
+ if is_a_8bit:
76
+ for i in range(32):
77
+ perm1 = []
78
+ col = i // 4
79
+ for block in [0, 1]:
80
+ for row in [
81
+ 4 * (i % 4),
82
+ 4 * (i % 4) + 1,
83
+ 4 * (i % 4) + 2,
84
+ 4 * (i % 4) + 3,
85
+ 4 * (i % 4 + 4),
86
+ 4 * (i % 4 + 4) + 1,
87
+ 4 * (i % 4 + 4) + 2,
88
+ 4 * (i % 4 + 4) + 3,
89
+ ]:
90
+ perm1.append(16 * row + col + 8 * block)
91
+ for j in range(2):
92
+ perm_list.extend([p + 512 * j for p in perm1])
93
+ else:
94
+ for i in range(32):
95
+ perm1 = []
96
+ col = i // 4
97
+ for block in [0, 1]:
98
+ for row in [
99
+ 2 * (i % 4),
100
+ 2 * (i % 4) + 1,
101
+ 2 * (i % 4 + 4),
102
+ 2 * (i % 4 + 4) + 1,
103
+ ]:
104
+ perm1.append(16 * row + col + 8 * block)
105
+ for j in range(4):
106
+ perm_list.extend([p + 256 * j for p in perm1])
107
108
perm = np.array(perm_list)
109
110
if num_bits == 4:
85
- interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
111
+ if is_a_8bit: # noqa: SIM108
112
+ interleave = np.array([0, 4, 1, 5, 2, 6, 3, 7])
113
+ else:
114
+ interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
115
elif num_bits == 8:
87
- interleave = np.array([0, 2, 1, 3])
116
+ if is_a_8bit: # noqa: SIM108
117
+ interleave = np.array([0, 1, 2, 3])
118
+ else:
119
+ interleave = np.array([0, 2, 1, 3])
120
else:
121
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
122
@@ -99,7 +131,10 @@ def marlin_quantize(
131
group_size: int,
132
act_order: bool,
133
test_perm: torch.Tensor | None = None,
134
+ input_dtype: torch.dtype | None = None,
135
):
136
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
137
+
138
size_k, size_n = w.shape
139
num_bits = quant_type.size_bits
140
@@ -120,9 +155,15 @@ def marlin_quantize(
155
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
156
157
# Reformat to marlin
123
- weight_perm = get_weight_perm(num_bits)
124
- marlin_q_w = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm)
125
- marlin_s = marlin_permute_scales(s, size_k, size_n, group_size)
158
+ weight_perm = get_weight_perm(num_bits, is_a_8bit)
159
+ marlin_q_w = marlin_weights(
160
+ q_w, size_k, size_n, num_bits, weight_perm, is_a_8bit=is_a_8bit
161
+ )
162
+ marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, is_a_8bit=is_a_8bit)
163
+
164
+ if input_dtype == torch.float8_e4m3fn and quant_type == scalar_types.uint4b8:
165
+ ops.marlin_int4_fp8_preprocess(marlin_q_w, inplace=True)
166
+ marlin_s = marlin_s * 512
167
168
# Create result
169
res_list = [w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, rand_perm]
@@ -132,7 +173,13 @@ def marlin_quantize(
173
return res_list
174
175
135
-def awq_marlin_quantize(w: torch.Tensor, quant_type: ScalarType, group_size: int):
176
+def awq_marlin_quantize(
177
+ w: torch.Tensor,
178
+ quant_type: ScalarType,
179
+ group_size: int,
180
+ input_dtype: torch.dtype | None = None,
181
+):
182
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
183
size_k, size_n = w.shape
184
185
# Normalize group_size
@@ -147,11 +194,22 @@ def awq_marlin_quantize(w: torch.Tensor, quant_type: ScalarType, group_size: int
194
# Quantize with zp
195
w_ref, q_w, s, zp = quantize_weights(w, quant_type, group_size, zero_points=True)
196
197
+ if input_dtype == torch.float8_e4m3fn and quant_type == scalar_types.uint4:
198
+ repeated_zp = zp.repeat_interleave(group_size, 0)
199
+ q_w_old = q_w
200
+ q_w = q_w_old - repeated_zp
201
+ q_w[q_w < 0] = 15 - q_w_old[q_w < 0]
202
+ s = s * 512
203
+
204
# Reformat to marlin
151
- weight_perm = get_weight_perm(quant_type.size_bits)
152
- marlin_q_w = marlin_weights(q_w, size_k, size_n, quant_type.size_bits, weight_perm)
153
- marlin_s = marlin_permute_scales(s, size_k, size_n, group_size)
154
- marlin_zp = marlin_zero_points(zp, num_groups, size_n, quant_type.size_bits)
205
+ weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
206
+ marlin_q_w = marlin_weights(
207
+ q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit=is_a_8bit
208
+ )
209
+ marlin_s = marlin_permute_scales(s, size_k, size_n, group_size, is_a_8bit=is_a_8bit)
210
+ marlin_zp = marlin_zero_points(
211
+ zp, num_groups, size_n, quant_type.size_bits, is_a_8bit=is_a_8bit
212
+ )
213
214
# Create result
215
res_list = [w_ref, marlin_q_w, marlin_s, marlin_zp]