@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d06ba4ed3f
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 14:45 UTC

[Kernel] moe wna16 marlin kernel (#14447)

Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com> Co-authored-by: Michael Goin <michael@neuralmagic.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Jinzhen Lin committed Apr 15, 2025 at 11:05 UTC d06ba4ed3f9a5929eabb404842a5c02da42e960b
16 files changed +3484 -336
CMakeLists.txt
+41 -11
@@ -609,21 +609,51 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
609 list(APPEND VLLM_MOE_EXT_SRC "${VLLM_MOE_WNA16_SRC}")
610 cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
611 if (MARLIN_MOE_ARCHS)
612 - set(MARLIN_MOE_SRC
613 - "csrc/moe/marlin_kernels/marlin_moe_kernel.h"
614 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4b8.h"
615 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4b8.cu"
616 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku8b128.h"
617 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku8b128.cu"
618 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4.h"
619 - "csrc/moe/marlin_kernels/marlin_moe_kernel_ku4.cu"
620 - "csrc/moe/marlin_moe_ops.cu")
612
613 + #
614 + # For the Marlin MOE kernels we automatically generate sources for various
615 + # preselected input type pairs and schedules.
616 + # Generate sources:
617 + set(MOE_MARLIN_GEN_SCRIPT
618 + ${CMAKE_CURRENT_SOURCE_DIR}/csrc/moe/marlin_moe_wna16/generate_kernels.py)
619 + file(MD5 ${MOE_MARLIN_GEN_SCRIPT} MOE_MARLIN_GEN_SCRIPT_HASH)
620 +
621 + message(STATUS "Marlin MOE generation script hash: ${MOE_MARLIN_GEN_SCRIPT_HASH}")
622 + message(STATUS "Last run Marlin MOE generate script hash: $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}")
623 +
624 + if (NOT DEFINED CACHE{MOE_MARLIN_GEN_SCRIPT_HASH}
625 + OR NOT $CACHE{MOE_MARLIN_GEN_SCRIPT_HASH} STREQUAL ${MOE_MARLIN_GEN_SCRIPT_HASH})
626 + execute_process(
627 + COMMAND ${CMAKE_COMMAND} -E env
628 + PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$PYTHONPATH
629 + ${Python_EXECUTABLE} ${MOE_MARLIN_GEN_SCRIPT}
630 + RESULT_VARIABLE moe_marlin_generation_result
631 + OUTPUT_VARIABLE moe_marlin_generation_output
632 + OUTPUT_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log
633 + ERROR_FILE ${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log
634 + )
635 +
636 + if (NOT moe_marlin_generation_result EQUAL 0)
637 + message(FATAL_ERROR "Marlin MOE generation failed."
638 + " Result: \"${moe_marlin_generation_result}\""
639 + "\nCheck the log for details: "
640 + "${CMAKE_CURRENT_BINARY_DIR}/moe_marlin_generation.log")
641 + else()
642 + set(MOE_MARLIN_GEN_SCRIPT_HASH ${MOE_MARLIN_GEN_SCRIPT_HASH}
643 + CACHE STRING "Last run Marlin MOE generate script hash" FORCE)
644 + message(STATUS "Marlin MOE generation completed successfully.")
645 + endif()
646 + else()
647 + message(STATUS "Marlin MOE generation script has not changed, skipping generation.")
648 + endif()
649 +
650 + file(GLOB MOE_WNAA16_MARLIN_SRC "csrc/moe/marlin_moe_wna16/*.cu")
651 set_gencode_flags_for_srcs(
623 - SRCS "${MARLIN_MOE_SRC}"
652 + SRCS "${MOE_WNAA16_MARLIN_SRC}"
653 CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
654
626 - list(APPEND VLLM_MOE_EXT_SRC "${MARLIN_MOE_SRC}")
655 + list(APPEND VLLM_MOE_EXT_SRC ${MOE_WNAA16_MARLIN_SRC})
656 +
657 message(STATUS "Building Marlin MOE kernels for archs: ${MARLIN_MOE_ARCHS}")
658 else()
659 message(STATUS "Not building Marlin MOE kernels as no compatible archs found"
csrc/moe/marlin_moe_wna16/generate_kernels.py new
+103
@@ -0,0 +1,103 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import glob
3 +import itertools
4 +import os
5 +import subprocess
6 +
7 +import jinja2
8 +
9 +FILE_HEAD = """
10 +// auto generated by generate.py
11 +// clang-format off
12 +
13 +#include "kernel.h"
14 +#include "marlin_template.h"
15 +
16 +namespace MARLIN_NAMESPACE_NAME {
17 +""".strip()
18 +
19 +TEMPLATE = ("template __global__ void Marlin<"
20 + "{{scalar_t}}, "
21 + "{{w_type_id}}, "
22 + "{{threads}}, "
23 + "{{thread_m_blocks}}, "
24 + "{{thread_n_blocks}}, "
25 + "{{thread_k_blocks}}, "
26 + "{{'true' if m_block_size_8 else 'false'}}, "
27 + "{{stages}}, "
28 + "{{'true' if has_act_order else 'false'}}, "
29 + "{{'true' if has_zp else 'false'}}, "
30 + "{{group_blocks}}, "
31 + "{{'true' if is_zp_float else 'false'}}>"
32 + "( MARLIN_KERNEL_PARAMS );")
33 +
34 +# int8 with zero point case (vllm::kU8) is also supported,
35 +# we don't add it to reduce wheel size.
36 +SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128"]
37 +THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
38 +
39 +THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
40 +# group_blocks:
41 +# = 0 : act order case
42 +# = -1 : channelwise quantization
43 +# > 0 : group_size=16*group_blocks
44 +GROUP_BLOCKS = [0, -1, 2, 4, 8]
45 +DTYPES = ["fp16", "bf16"]
46 +
47 +
48 +def remove_old_kernels():
49 + for filename in glob.glob(os.path.dirname(__file__) + "/kernel_*.cu"):
50 + subprocess.call(["rm", "-f", filename])
51 +
52 +
53 +def generate_new_kernels():
54 + for scalar_type, dtype in itertools.product(SCALAR_TYPES, DTYPES):
55 + has_zp = "B" not in scalar_type
56 + all_template_str_list = []
57 +
58 + for group_blocks, m_blocks, thread_configs in itertools.product(
59 + GROUP_BLOCKS, THREAD_M_BLOCKS, THREAD_CONFIGS):
60 +
61 + has_act_order = group_blocks == 0
62 + if has_zp and has_act_order:
63 + continue
64 + if thread_configs[2] == 256:
65 + if m_blocks <= 1 and thread_configs[0] != 128:
66 + continue
67 + if m_blocks > 1 and thread_configs[0] != 64:
68 + continue
69 +
70 + k_blocks = thread_configs[0] // 16
71 + n_blocks = thread_configs[1] // 16
72 + threads = thread_configs[2]
73 +
74 + c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
75 +
76 + template_str = jinja2.Template(TEMPLATE).render(
77 + scalar_t=c_dtype,
78 + w_type_id=scalar_type + ".id()",
79 + threads=threads,
80 + thread_m_blocks=max(m_blocks, 1),
81 + thread_n_blocks=n_blocks,
82 + thread_k_blocks=k_blocks,
83 + m_block_size_8=m_blocks == 0.5,
84 + stages="pipe_stages",
85 + has_act_order=has_act_order,
86 + has_zp=has_zp,
87 + group_blocks=group_blocks,
88 + is_zp_float=False,
89 + )
90 +
91 + all_template_str_list.append(template_str)
92 +
93 + file_content = FILE_HEAD + "\n\n"
94 + file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
95 + filename = f"kernel_{dtype}_{scalar_type[6:].lower()}.cu"
96 +
97 + with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
98 + f.write(file_content)
99 +
100 +
101 +if __name__ == "__main__":
102 + remove_old_kernels()
103 + generate_new_kernels()
csrc/moe/marlin_moe_wna16/kernel.h new
+44
@@ -0,0 +1,44 @@
1 +
2 +#ifndef MARLIN_NAMESPACE_NAME
3 + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
4 +#endif
5 +
6 +#include "quantization/gptq_marlin/marlin.cuh"
7 +#include "quantization/gptq_marlin/marlin_dtypes.cuh"
8 +#include "core/scalar_type.hpp"
9 +
10 +#define MARLIN_KERNEL_PARAMS \
11 + const int4 *__restrict__ A, const int4 *__restrict__ B, \
12 + int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13 + const int4 *__restrict__ scales_ptr, const int4 *__restrict__ zp_ptr, \
14 + const int *__restrict__ g_idx, \
15 + const int32_t *__restrict__ sorted_token_ids_ptr, \
16 + const int32_t *__restrict__ expert_ids_ptr, \
17 + const int32_t *__restrict__ num_tokens_past_padded_ptr, \
18 + const float *__restrict__ topk_weights_ptr, int top_k, \
19 + bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
20 + int prob_n, int prob_k, int *locks, bool use_atomic_add, \
21 + bool use_fp32_reduce
22 +
23 +namespace MARLIN_NAMESPACE_NAME {
24 +template <typename scalar_t, // compute dtype, half or nv_float16
25 + const vllm::ScalarTypeId w_type_id, // weight ScalarType id
26 + const int threads, // number of threads in a threadblock
27 + const int thread_m_blocks, // number of 16x16 blocks in the m
28 + // dimension (batchsize) of the
29 + // threadblock
30 + const int thread_n_blocks, // same for n dimension (output)
31 + const int thread_k_blocks, // same for k dimension (reduction)
32 + const bool m_block_size_8, // whether m_block_size == 8
33 + // only works when thread_m_blocks == 1
34 + const int stages, // number of stages for the async global->shared
35 + // fetch pipeline
36 + const bool has_act_order, // whether act_order is enabled
37 + const bool has_zp, // whether zero-points are enabled
38 + const int group_blocks, // number of consecutive 16x16 blocks
39 + // with a separate quantization scale
40 + const bool is_zp_float // is zero point of float16 type?
41 + >
42 +__global__ void Marlin(MARLIN_KERNEL_PARAMS);
43 +
44 +}
csrc/moe/marlin_moe_wna16/marlin_template.h new
+1917
@@ -0,0 +1,1917 @@
1 +/*
2 + * Modified by Neural Magic
3 + * Copyright (C) Marlin.2024 Elias Frantar
4 + *
5 + * Licensed under the Apache License, Version 2.0 (the "License");
6 + * you may not use this file except in compliance with the License.
7 + * You may obtain a copy of the License at
8 + *
9 + * http://www.apache.org/licenses/LICENSE-2.0
10 + *
11 + * Unless required by applicable law or agreed to in writing, software
12 + * distributed under the License is distributed on an "AS IS" BASIS,
13 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 + * See the License for the specific language governing permissions and
15 + * limitations under the License.
16 + */
17 +
18 +/*
19 + * Adapted from https://github.com/IST-DASLab/marlin
20 + */
21 +
22 +#ifndef MARLIN_NAMESPACE_NAME
23 + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
24 +#endif
25 +
26 +#include "quantization/gptq_marlin/marlin.cuh"
27 +#include "quantization/gptq_marlin/marlin_dtypes.cuh"
28 +#include "core/scalar_type.hpp"
29 +
30 +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
31 + static_assert(std::is_same<scalar_t, half>::value || \
32 + std::is_same<scalar_t, nv_bfloat16>::value, \
33 + "only float16 and bfloat16 is supported");
34 +
35 +namespace MARLIN_NAMESPACE_NAME {
36 +
37 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
38 +
39 +template <typename scalar_t, // compute dtype, half or nv_float16
40 + const vllm::ScalarTypeId w_type_id, // weight ScalarType id
41 + const int threads, // number of threads in a threadblock
42 + const int thread_m_blocks, // number of 16x16 blocks in the m
43 + // dimension (batchsize) of the
44 + // threadblock
45 + const int thread_n_blocks, // same for n dimension (output)
46 + const int thread_k_blocks, // same for k dimension (reduction)
47 + const bool m_block_size_8, // whether m_block_size == 8
48 + // only works when thread_m_blocks == 1
49 + const int stages, // number of stages for the async global->shared
50 + // fetch pipeline
51 + const bool has_act_order, // whether act_order is enabled
52 + const bool has_zp, // whether zero-points are enabled
53 + const int group_blocks, // number of consecutive 16x16 blocks
54 + // with a separate quantization scale
55 + const bool is_zp_float // is zero point of float16 type?
56 + >
57 +__global__ void Marlin(
58 + const int4* __restrict__ A, // fp16 input matrix of shape mxk
59 + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
60 + int4* __restrict__ C, // fp16 output buffer of shape mxn
61 + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
62 + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
63 + // (k/groupsize)xn
64 + const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
65 + // (k/groupsize)x(n/pack_factor)
66 + const int* __restrict__ g_idx, // int32 group indices of shape k
67 + const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
68 + const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
69 + const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
70 + const float* __restrict__ topk_weights_ptr, // moe top weights
71 + int top_k, // num of experts per token
72 + bool mul_topk_weights, // mul topk weights or not
73 + bool is_ep, // expert parallelism
74 + int num_groups, // number of scale groups per output channel
75 + int prob_m, // batch dimension m
76 + int prob_n, // output dimension n
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
80 + bool use_fp32_reduce // whether to use fp32 global reduce
81 +) {}
82 +
83 +} // namespace MARLIN_NAMESPACE_NAME
84 +
85 +#else
86 +
87 +// m16n8k16 tensor core mma instruction with fp16 inputs and fp32
88 +// output/accumulation.
89 +template <typename scalar_t>
90 +__device__ inline void mma(const typename ScalarType<scalar_t>::FragA& a_frag,
91 + const typename ScalarType<scalar_t>::FragB& frag_b,
92 + typename ScalarType<scalar_t>::FragC& frag_c) {
93 + const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
94 + const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
95 + float* c = reinterpret_cast<float*>(&frag_c);
96 + if constexpr (std::is_same<scalar_t, half>::value) {
97 + asm volatile(
98 + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
99 + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
100 + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
101 + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
102 + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
103 + } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
104 + asm volatile(
105 + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
106 + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
107 + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
108 + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
109 + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
110 + } else {
111 + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
112 + }
113 +}
114 +
115 +template <typename scalar_t>
116 +__device__ inline void mma_trans(
117 + const typename ScalarType<scalar_t>::FragA& a_frag,
118 + const typename ScalarType<scalar_t>::FragB& frag_b,
119 + const typename ScalarType<scalar_t>::FragB& frag_b2,
120 + typename ScalarType<scalar_t>::FragC& frag_c) {
121 + const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);
122 + const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);
123 + const uint32_t* b2 = reinterpret_cast<const uint32_t*>(&frag_b2);
124 + float* c = reinterpret_cast<float*>(&frag_c);
125 + if constexpr (std::is_same<scalar_t, half>::value) {
126 + asm volatile(
127 + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
128 + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
129 + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
130 + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
131 + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
132 + } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
133 + asm volatile(
134 + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
135 + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
136 + : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])
137 + : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]),
138 + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
139 + } else {
140 + STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t);
141 + }
142 +}
143 +
144 +// Instruction for loading a full 16x16 matrix fragment of operand A from shared
145 +// memory, directly in tensor core layout.
146 +template <int count, typename scalar_t>
147 +__device__ inline void ldsm(typename ScalarType<scalar_t>::FragA& frag_a,
148 + const void* smem_ptr) {
149 + uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);
150 + uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
151 + if constexpr (count == 4) {
152 + asm volatile(
153 + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
154 + : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])
155 + : "r"(smem));
156 + } else if constexpr (count == 2) {
157 + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n"
158 + : "=r"(a[0]), "=r"(a[1])
159 + : "r"(smem));
160 + } else if constexpr (count == 1) {
161 + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n"
162 + : "=r"(a[0])
163 + : "r"(smem));
164 + } else {
165 + static_assert(count == 1 || count == 2 || count == 4, "invalid count");
166 + }
167 +}
168 +
169 +// Lookup-table based 3-input logical operation; explicitly used for
170 +// dequantization as the compiler does not seem to automatically recognize it in
171 +// all cases.
172 +template <int lut>
173 +__device__ inline int lop3(int a, int b, int c) {
174 + int res;
175 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
176 + : "=r"(res)
177 + : "r"(a), "r"(b), "r"(c), "n"(lut));
178 + return res;
179 +}
180 +
181 +// Constructs destination register by taking bytes from 2 sources (based on
182 +// mask)
183 +template <int start_byte, int mask>
184 +__device__ inline uint32_t prmt(uint32_t a) {
185 + uint32_t res;
186 + asm volatile("prmt.b32 %0, %1, %2, %3;\n"
187 + : "=r"(res)
188 + : "r"(a), "n"(start_byte), "n"(mask));
189 + return res;
190 +}
191 +
192 +template <typename scalar_t, int bit>
193 +__device__ inline typename ScalarType<scalar_t>::FragB dequant(
194 + int q, typename ScalarType<scalar_t>::FragB& frag_b);
195 +
196 +//
197 +// Efficiently dequantize 4bit values packed in an int32 value into a full
198 +// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,
199 +// with some small changes:
200 +// - FP16:
201 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287
202 +// - BF16:
203 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
204 +//
205 +template <>
206 +__device__ inline typename ScalarType<half>::FragB dequant<half, 4>(
207 + int q, typename ScalarType<half>::FragB& frag_b) {
208 + const int LO = 0x000f000f;
209 + const int HI = 0x00f000f0;
210 + const int EX = 0x64006400;
211 + // Guarantee that the `(a & b) | c` operations are LOP3s.
212 + int lo = lop3 < (0xf0 & 0xcc) | 0xaa > (q, LO, EX);
213 + int hi = lop3 < (0xf0 & 0xcc) | 0xaa > (q, HI, EX);
214 + // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point
215 + // directly into `SUB` and `ADD`.
216 + const int SUB = 0x64086408;
217 + const int MUL = 0x2c002c00;
218 + const int ADD = 0xd480d480;
219 + frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
220 + *reinterpret_cast<const half2*>(&SUB));
221 + frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),
222 + *reinterpret_cast<const half2*>(&MUL),
223 + *reinterpret_cast<const half2*>(&ADD));
224 + return frag_b;
225 +}
226 +
227 +template <>
228 +__device__ inline typename ScalarType<nv_bfloat16>::FragB
229 +dequant<nv_bfloat16, 4>(int q,
230 + typename ScalarType<nv_bfloat16>::FragB& frag_b) {
231 + static constexpr uint32_t MASK = 0x000f000f;
232 + static constexpr uint32_t EX = 0x43004300;
233 +
234 + // Guarantee that the `(a & b) | c` operations are LOP3s.
235 +
236 + int lo = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
237 + q >>= 4;
238 + int hi = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
239 +
240 + static constexpr uint32_t MUL = 0x3F803F80;
241 + static constexpr uint32_t ADD = 0xC308C308;
242 +
243 + frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
244 + *reinterpret_cast<const nv_bfloat162*>(&MUL),
245 + *reinterpret_cast<const nv_bfloat162*>(&ADD));
246 + frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
247 + *reinterpret_cast<const nv_bfloat162*>(&MUL),
248 + *reinterpret_cast<const nv_bfloat162*>(&ADD));
249 + return frag_b;
250 +}
251 +
252 +//
253 +// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or
254 +// bf16 Reference:
255 +// - FP16:
256 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85
257 +// - BF16:
258 +// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
259 +//
260 +template <>
261 +__device__ inline typename ScalarType<half>::FragB dequant<half, 8>(
262 + int q, typename ScalarType<half>::FragB& frag_b) {
263 + static constexpr uint32_t mask_for_elt_01 = 0x5250;
264 + static constexpr uint32_t mask_for_elt_23 = 0x5351;
265 + static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
266 +
267 + uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
268 + uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
269 +
270 + static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
271 +
272 + frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
273 + *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
274 + frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
275 + *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
276 + return frag_b;
277 +}
278 +
279 +template <>
280 +__device__ inline typename ScalarType<nv_bfloat16>::FragB
281 +dequant<nv_bfloat16, 8>(int q,
282 + typename ScalarType<nv_bfloat16>::FragB& frag_b) {
283 + float fp32_intermediates[4];
284 + uint32_t* fp32_intermediates_casted =
285 + reinterpret_cast<uint32_t*>(fp32_intermediates);
286 +
287 + static constexpr uint32_t fp32_base = 0x4B000000;
288 + fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
289 + fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
290 + fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
291 + fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
292 +
293 + fp32_intermediates[0] -= 8388736.f;
294 + fp32_intermediates[1] -= 8388736.f;
295 + fp32_intermediates[2] -= 8388736.f;
296 + fp32_intermediates[3] -= 8388736.f;
297 +
298 + uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(&frag_b);
299 + bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
300 + fp32_intermediates_casted[1], 0x7632);
301 + bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
302 + fp32_intermediates_casted[3], 0x7632);
303 +
304 + return frag_b;
305 +}
306 +
307 +// Multiply dequantized values by the corresponding quantization scale; used
308 +// only for grouped quantization.
309 +template <typename scalar_t>
310 +__device__ inline void scale(typename ScalarType<scalar_t>::FragB& frag_b,
311 + typename ScalarType<scalar_t>::FragS& frag_s,
312 + int i) {
313 + using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
314 + scalar_t2 s =
315 + ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_s)[i]);
316 + frag_b[0] = __hmul2(frag_b[0], s);
317 + frag_b[1] = __hmul2(frag_b[1], s);
318 +}
319 +
320 +template <typename scalar_t>
321 +__device__ inline void scale_and_sub(
322 + typename ScalarType<scalar_t>::FragB& frag_b, scalar_t s, scalar_t zp) {
323 + using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
324 + scalar_t2 s2 = ScalarType<scalar_t>::num2num2(s);
325 + scalar_t2 zp2 = ScalarType<scalar_t>::num2num2(zp);
326 + frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2));
327 + frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2));
328 +}
329 +
330 +template <typename scalar_t>
331 +__device__ inline void sub_zp(typename ScalarType<scalar_t>::FragB& frag_b,
332 + typename ScalarType<scalar_t>::scalar_t2& frag_zp,
333 + int i) {
334 + using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
335 + scalar_t2 zp =
336 + ScalarType<scalar_t>::num2num2(reinterpret_cast<scalar_t*>(&frag_zp)[i]);
337 + frag_b[0] = __hsub2(frag_b[0], zp);
338 + frag_b[1] = __hsub2(frag_b[1], zp);
339 +}
340 +
341 +// Same as above, but for act_order (each K is multiplied individually)
342 +template <typename scalar_t>
343 +__device__ inline void scale4(typename ScalarType<scalar_t>::FragB& frag_b,
344 + typename ScalarType<scalar_t>::FragS& frag_s_1,
345 + typename ScalarType<scalar_t>::FragS& frag_s_2,
346 + typename ScalarType<scalar_t>::FragS& frag_s_3,
347 + typename ScalarType<scalar_t>::FragS& frag_s_4,
348 + int i) {
349 + using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
350 + scalar_t2 s_val_1_2;
351 + s_val_1_2.x = reinterpret_cast<scalar_t*>(&frag_s_1)[i];
352 + s_val_1_2.y = reinterpret_cast<scalar_t*>(&frag_s_2)[i];
353 +
354 + scalar_t2 s_val_3_4;
355 + s_val_3_4.x = reinterpret_cast<scalar_t*>(&frag_s_3)[i];
356 + s_val_3_4.y = reinterpret_cast<scalar_t*>(&frag_s_4)[i];
357 +
358 + frag_b[0] = __hmul2(frag_b[0], s_val_1_2);
359 + frag_b[1] = __hmul2(frag_b[1], s_val_3_4);
360 +}
361 +
362 +// Given 2 floats multiply by 2 scales (halves)
363 +template <typename scalar_t>
364 +__device__ inline void scale_float(float* c,
365 + typename ScalarType<scalar_t>::FragS& s) {
366 + scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s);
367 + c[0] = __fmul_rn(c[0], ScalarType<scalar_t>::num2float(s_ptr[0]));
368 + c[1] = __fmul_rn(c[1], ScalarType<scalar_t>::num2float(s_ptr[1]));
369 +}
370 +
371 +// Wait until barrier reaches `count`, then lock for current threadblock.
372 +__device__ inline void barrier_acquire(int* lock, int count) {
373 + if (threadIdx.x == 0) {
374 + int state = -1;
375 + do
376 + // Guarantee that subsequent writes by this threadblock will be visible
377 + // globally.
378 + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
379 + : "=r"(state)
380 + : "l"(lock));
381 + while (state != count);
382 + }
383 + __syncthreads();
384 +}
385 +
386 +// Release barrier and increment visitation count.
387 +__device__ inline void barrier_release(int* lock, bool reset = false) {
388 + __syncthreads();
389 + if (threadIdx.x == 0) {
390 + if (reset) {
391 + lock[0] = 0;
392 + return;
393 + }
394 + int val = 1;
395 + // Make sure that all writes since acquiring this barrier are visible
396 + // globally, while releasing the barrier.
397 + asm volatile("fence.acq_rel.gpu;\n");
398 + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
399 + :
400 + : "l"(lock), "r"(val));
401 + }
402 +}
403 +
404 +// Wait until value of lock to be negative, and then add 1
405 +__device__ inline void wait_negative_and_add(int* lock) {
406 + if (threadIdx.x == 0) {
407 + int state = 0;
408 + do
409 + // Guarantee that subsequent writes by this threadblock will be visible
410 + // globally.
411 + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
412 + : "=r"(state)
413 + : "l"(lock));
414 + while (state >= 0);
415 + atomicAdd(lock, 1);
416 + }
417 + __syncthreads();
418 +}
419 +
420 +template <typename scalar_t, // compute dtype, half or nv_float16
421 + const vllm::ScalarTypeId w_type_id, // weight ScalarType id
422 + const int threads, // number of threads in a threadblock
423 + const int thread_m_blocks, // number of 16x16 blocks in the m
424 + // dimension (batchsize) of the
425 + // threadblock
426 + const int thread_n_blocks, // same for n dimension (output)
427 + const int thread_k_blocks, // same for k dimension (reduction)
428 + const bool m_block_size_8, // whether m_block_size == 8
429 + // only works when thread_m_blocks == 1
430 + const int stages, // number of stages for the async global->shared
431 + // fetch pipeline
432 + const bool has_act_order, // whether act_order is enabled
433 + const bool has_zp, // whether zero-points are enabled
434 + const int group_blocks, // number of consecutive 16x16 blocks
435 + // with a separate quantization scale
436 + const bool is_zp_float // is zero point of float16 type?
437 + >
438 +__global__ void Marlin(
439 + const int4* __restrict__ A, // fp16 input matrix of shape mxk
440 + const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
441 + int4* __restrict__ C, // fp16 output buffer of shape mxn
442 + int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
443 + const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
444 + // (k/groupsize)xn
445 + const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
446 + // (k/groupsize)x(n/pack_factor)
447 + const int* __restrict__ g_idx, // int32 group indices of shape k
448 + const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
449 + const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
450 + const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
451 + const float* __restrict__ topk_weights_ptr, // moe top weights
452 + int top_k, // num of experts per token
453 + bool mul_topk_weights, // mul topk weights or not
454 + bool is_ep, // expert parallelism
455 + int num_groups, // number of scale groups per output channel
456 + int prob_m, // batch dimension m
457 + int prob_n, // output dimension n
458 + int prob_k, // reduction dimension k
459 + int* locks, // extra global storage for barrier synchronization
460 + bool use_atomic_add, // whether to use atomic add to reduce
461 + bool use_fp32_reduce // whether to use fp32 global reduce
462 +) {
463 + // Each threadblock processes one "stripe" of the B matrix with (roughly) the
464 + // same size, which might involve multiple column "slices" (of width 16 *
465 + // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM
466 + // example:
467 + // 0 1 3
468 + // 0 2 3
469 + // 1 2 4
470 + // While this kind of partitioning makes things somewhat more complicated, it
471 + // ensures good utilization of all SMs for many kinds of shape and GPU
472 + // configurations, while requiring as few slow global cross-threadblock
473 + // reductions as possible.
474 + using Dtype = ScalarType<scalar_t>;
475 + using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
476 + using FragA = typename ScalarType<scalar_t>::FragA;
477 + using FragB = typename ScalarType<scalar_t>::FragB;
478 + using FragC = typename ScalarType<scalar_t>::FragC;
479 + using FragS = typename ScalarType<scalar_t>::FragS;
480 + using FragZP = typename ScalarType<scalar_t>::FragZP;
481 +
482 + extern __shared__ int4 sh[];
483 + static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
484 +
485 + constexpr int pack_factor = 32 / w_type.size_bits();
486 + static_assert(thread_m_blocks == 1 || !m_block_size_8);
487 + constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
488 + const int group_size =
489 + (!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
490 + const int scales_expert_stride = prob_n * prob_k / group_size / 8;
491 + const int zp_expert_stride =
492 + is_zp_float ? prob_n * prob_k / group_size / 8
493 + : prob_n * prob_k / group_size / (pack_factor * 4);
494 +
495 + // parallel: num valid moe blocks
496 + int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
497 + int parallel = num_tokens_past_padded / moe_block_size;
498 + int num_valid_blocks = parallel;
499 + if (is_ep) {
500 + for (int i = 0; i < parallel; i++) {
501 + if (expert_ids_ptr[i] == -1) num_valid_blocks--;
502 + }
503 + }
504 + int num_invalid_blocks = parallel - num_valid_blocks;
505 + parallel = num_valid_blocks;
506 +
507 + int k_tiles = prob_k / 16 / thread_k_blocks;
508 + int n_tiles = prob_n / 16 / thread_n_blocks;
509 + int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x);
510 +
511 + if constexpr (!has_act_order && group_blocks != -1) {
512 + if (group_blocks >= thread_k_blocks) {
513 + // Ensure that the number of tiles in each stripe is a multiple of the
514 + // groupsize; this avoids an annoying special case where a stripe starts
515 + // in the middle of group.
516 + iters = (group_blocks / thread_k_blocks) *
517 + div_ceil(iters, (group_blocks / thread_k_blocks));
518 + }
519 + }
520 +
521 + int slice_row = (iters * blockIdx.x) % k_tiles;
522 + int slice_col_par = (iters * blockIdx.x) / k_tiles;
523 + int slice_col = slice_col_par;
524 + int slice_iters; // number of threadblock tiles in the current slice
525 + int slice_count =
526 + 0; // total number of active threadblocks in the current slice
527 + int slice_idx; // index of threadblock in current slice; numbered bottom to
528 + // top
529 +
530 + int par_id = 0;
531 + int block_id = -1;
532 + int64_t expert_id = 0; // use int64 to avoid computation result overflow
533 + int old_expert_id = 0;
534 + int64_t B_expert_off = 0;
535 +
536 + int4* sh_block_sorted_ids_int4 = sh;
537 + int32_t* sh_block_sorted_ids =
538 + reinterpret_cast<int*>(sh_block_sorted_ids_int4);
539 + int4* sh_block_topk_weights_int4 =
540 + sh_block_sorted_ids_int4 + moe_block_size / 4;
541 + scalar_t2* sh_block_topk_weights =
542 + reinterpret_cast<scalar_t2*>(sh_block_topk_weights_int4);
543 + int4* sh_new = sh_block_topk_weights_int4 + moe_block_size / 4;
544 +
545 + int32_t block_num_valid_tokens = 0;
546 + int32_t locks_off = 0;
547 +
548 + // We can easily implement parallel problem execution by just remapping
549 + // indices and advancing global pointers
550 + if (slice_col_par >= n_tiles) {
551 + slice_col = slice_col_par % n_tiles;
552 + par_id = slice_col_par / n_tiles;
553 + }
554 + if (parallel * n_tiles >= gridDim.x) {
555 + // when parallel * n_tiles >= sms
556 + // then there are at most $sms$ conflict tile blocks
557 + locks_off = blockIdx.x;
558 + } else {
559 + locks_off = (iters * blockIdx.x) / k_tiles - 1;
560 + }
561 +
562 + // read moe block data given block_id
563 + // block_sorted_ids / block_num_valid_tokens / block_topk_weights
564 + auto read_moe_block_data = [&](int block_id) {
565 + block_num_valid_tokens = moe_block_size;
566 + #pragma unroll
567 + for (int i = 0; i < moe_block_size / 4; i++) {
568 + int4 sorted_token_ids_int4 = reinterpret_cast<const int4*>(
569 + sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
570 + int* sorted_token_ids = reinterpret_cast<int*>(&sorted_token_ids_int4);
571 + #pragma unroll
572 + for (int j = 0; j < 4; j++) {
573 + if (sorted_token_ids[j] >= prob_m * top_k) {
574 + block_num_valid_tokens = i * 4 + j;
575 + break;
576 + }
577 + }
578 + if (block_num_valid_tokens != moe_block_size) break;
579 + }
580 +
581 + __syncthreads();
582 + int tid4 = threadIdx.x / 4;
583 + if (threadIdx.x % 4 == 0 && threadIdx.x < block_num_valid_tokens) {
584 + sh_block_sorted_ids_int4[tid4] = reinterpret_cast<const int4*>(
585 + sorted_token_ids_ptr)[block_id * moe_block_size / 4 + tid4];
586 +
587 + if (mul_topk_weights) {
588 + #pragma unroll
589 + for (int i = 0; i < 4; i++) {
590 + sh_block_topk_weights[tid4 * 4 + i] =
591 + Dtype::num2num2(Dtype::float2num(
592 + topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]]));
593 + }
594 + }
595 + }
596 + __syncthreads();
597 + };
598 +
599 + // when move to next moe block, find the next block_id and expert_id
600 + // and then read moe block data
601 + auto update_next_moe_block_data = [&]() {
602 + if (par_id >= parallel) return;
603 +
604 + old_expert_id = expert_id;
605 + if (num_invalid_blocks > 0) {
606 + int skip_count = block_id == -1 ? par_id : 0;
607 + block_id++;
608 + for (int i = block_id; i < num_tokens_past_padded / moe_block_size; i++) {
609 + expert_id = expert_ids_ptr[i];
610 + if (expert_id != -1) {
611 + if (skip_count == 0) {
612 + block_id = i;
613 + break;
614 + };
615 + skip_count--;
616 + };
617 + }
618 + } else {
619 + block_id = par_id;
620 + expert_id = expert_ids_ptr[block_id];
621 + }
622 +
623 + B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
624 + scales_ptr += (expert_id - old_expert_id) * scales_expert_stride;
625 + if constexpr (has_zp) {
626 + zp_ptr += (expert_id - old_expert_id) * zp_expert_stride;
627 + }
628 + if constexpr (has_act_order) {
629 + g_idx += (expert_id - old_expert_id) * prob_k;
630 + }
631 +
632 + read_moe_block_data(block_id);
633 + };
634 +
635 + // Compute all information about the current slice which is required for
636 + // synchronization.
637 + auto init_slice = [&](bool first_init = false) {
638 + slice_iters =
639 + iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);
640 + if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;
641 + if (slice_iters == 0) return;
642 + if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;
643 + slice_count = 1;
644 + slice_idx = 0;
645 + int col_first = iters * div_ceil(k_tiles * slice_col_par, iters);
646 + if (col_first <= k_tiles * (slice_col_par + 1)) {
647 + int col_off = col_first - k_tiles * slice_col_par;
648 + slice_count = div_ceil(k_tiles - col_off, iters);
649 + if (col_off > 0) slice_count++;
650 + int delta_first = iters * blockIdx.x - col_first;
651 + if (delta_first < 0 || (col_off == 0 && delta_first == 0))
652 + slice_idx = slice_count - 1;
653 + else {
654 + slice_idx = slice_count - 1 - delta_first / iters;
655 + if (col_off > 0) slice_idx--;
656 + }
657 + }
658 + if (parallel * n_tiles >= gridDim.x) {
659 + if (slice_count > 1 && slice_idx == slice_count - 1) {
660 + locks_off++;
661 + }
662 + } else {
663 + locks_off++;
664 + }
665 +
666 + if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) {
667 + constexpr int threads_per_m = 16 * thread_n_blocks / 8;
668 + int m_per_thread =
669 + div_ceil(block_num_valid_tokens, threads / threads_per_m);
670 + for (int i = 0; i < m_per_thread; i++) {
671 + int row = threads / threads_per_m * i + threadIdx.x / threads_per_m;
672 + if (row < block_num_valid_tokens) {
673 + int64_t sorted_row = sh_block_sorted_ids[row];
674 + int col = slice_col * 16 * thread_n_blocks / 8 +
675 + threadIdx.x % threads_per_m;
676 + C[sorted_row * prob_n / 8 + col] = {0, 0, 0, 0};
677 + }
678 + }
679 + // After write zero to output, write a negative value to lock.
680 + // Every SM that processes the same slice would wait for
681 + // the negative value, and then atomicAdd 1 to it.
682 + // After all SMs are processed, the lock value would back to 0 again.
683 + __syncthreads();
684 + if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count;
685 + }
686 +
687 + if (slice_col == n_tiles) {
688 + slice_col = 0;
689 + par_id++;
690 + update_next_moe_block_data();
691 + }
692 + };
693 +
694 + update_next_moe_block_data();
695 + init_slice(true);
696 +
697 + // A sizes/strides
698 +
699 + // stride of the A matrix in global memory
700 + int a_gl_stride = prob_k / 8;
701 + // stride of an A matrix tile in shared memory
702 + constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
703 + // delta between subsequent A tiles in global memory
704 + constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;
705 + // between subsequent accesses within a tile
706 + int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);
707 + // between shared memory writes
708 + constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);
709 + // between shared memory tile reads
710 + constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));
711 + // within a shared memory tile
712 + constexpr int a_sh_rd_delta_i = a_sh_stride * 16;
713 + // overall size of a tile
714 + constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);
715 + // number of shared write iterations for a tile
716 + constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta);
717 +
718 + // B sizes/strides
719 + int b_gl_stride = 16 * prob_n / (pack_factor * 4);
720 + constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;
721 + constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;
722 + constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;
723 +
724 + int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;
725 + int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);
726 + constexpr int b_sh_wr_delta = threads * b_thread_vecs;
727 + constexpr int b_sh_rd_delta = threads * b_thread_vecs;
728 + constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;
729 + constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
730 +
731 + // Scale sizes/strides without act_order
732 + int s_gl_stride = prob_n / 8;
733 + constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
734 + constexpr int s_tb_groups =
735 + !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
736 + ? thread_k_blocks / group_blocks
737 + : 1;
738 + constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
739 + int s_gl_rd_delta = s_gl_stride;
740 +
741 + // Scale size/strides with act_order
742 + constexpr int tb_k = 16 * thread_k_blocks;
743 + constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;
744 + // constexpr int act_s_row_stride = 1;
745 + // int act_s_col_stride = act_s_row_stride * num_groups;
746 + int act_s_col_stride = 1;
747 + int act_s_col_warp_stride = act_s_col_stride * 8;
748 + int tb_n_warps = thread_n_blocks / 4;
749 + int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;
750 +
751 + // Zero-points sizes/strides
752 + int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4;
753 + constexpr int zp_sh_stride = is_zp_float
754 + ? 16 * thread_n_blocks / 8
755 + : ((16 * thread_n_blocks) / pack_factor) / 4;
756 + constexpr int zp_tb_groups = s_tb_groups;
757 + constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;
758 + int zp_gl_rd_delta = zp_gl_stride;
759 +
760 + // Global A read index of current thread.
761 + int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
762 + (threadIdx.x % a_gl_rd_delta_o);
763 + a_gl_rd += a_gl_rd_delta_o * slice_row;
764 + // Shared write index of current thread.
765 + int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +
766 + (threadIdx.x % a_gl_rd_delta_o);
767 + // Shared read index.
768 + int a_sh_rd =
769 + a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) +
770 + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1));
771 + a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));
772 +
773 + int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +
774 + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;
775 + b_gl_rd += b_sh_stride * slice_col;
776 + b_gl_rd += b_gl_rd_delta_o * slice_row;
777 + int b_sh_wr = threadIdx.x * b_thread_vecs;
778 + int b_sh_rd = threadIdx.x * b_thread_vecs;
779 +
780 + // For act_order
781 + constexpr int k_iter_size = tb_k / b_sh_wr_iters;
782 + int slice_k_start = tb_k * slice_row;
783 + int slice_k_finish = slice_k_start + tb_k * slice_iters;
784 + int slice_k_start_shared_fetch = slice_k_start;
785 + int slice_n_offset = act_s_col_tb_stride * slice_col;
786 +
787 + // No act_order
788 + int s_gl_rd;
789 + if constexpr (!has_act_order) {
790 + if constexpr (group_blocks == -1) {
791 + s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
792 + } else {
793 + s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
794 + s_sh_stride * slice_col + threadIdx.x;
795 + }
796 + }
797 + int s_sh_wr = threadIdx.x;
798 + bool s_sh_wr_pred = threadIdx.x < s_sh_stride;
799 +
800 + // Zero-points
801 + int zp_gl_rd;
802 + if constexpr (has_zp) {
803 + if constexpr (group_blocks == -1) {
804 + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
805 + } else {
806 + zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
807 + zp_sh_stride * slice_col + threadIdx.x;
808 + }
809 + }
810 + int zp_sh_wr = threadIdx.x;
811 + bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;
812 +
813 + // We use a different scale layout for grouped and column-wise quantization as
814 + // we scale a `half2` tile in column-major layout in the former and in
815 + // row-major in the latter case.
816 + int s_sh_rd;
817 + if constexpr (group_blocks != -1)
818 + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
819 + (threadIdx.x % 32) / 4;
820 + else if constexpr (group_blocks == -1 && (m_block_size_8 || has_zp))
821 + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
822 + (threadIdx.x % 32) / 8;
823 + else
824 + s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
825 + (threadIdx.x % 32) % 4;
826 +
827 + // Zero-points have the same read layout as the scales
828 + // (without column-wise case)
829 + constexpr int num_col_threads = 8;
830 + constexpr int num_row_threads = 4;
831 + constexpr int num_ints_per_thread = 8 / pack_factor;
832 + int zp_sh_rd;
833 + if constexpr (has_zp) {
834 + if constexpr (is_zp_float) {
835 + if constexpr (group_blocks != -1) {
836 + zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
837 + (threadIdx.x % 32) / 4;
838 + }
839 + } else {
840 + zp_sh_rd = num_ints_per_thread * num_col_threads *
841 + ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
842 + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);
843 + }
844 + }
845 +
846 + // To ensure that writing and reading A tiles to/from shared memory, the
847 + // latter in fragment format, is fully bank conflict free, we need to use a
848 + // rather fancy XOR-based layout. The key here is that neither reads nor
849 + // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the
850 + // same shared memory banks. Further, it seems (based on NSight-Compute) that
851 + // each warp must also write a consecutive memory segment?
852 + auto transform_a = [&](int i) {
853 + int row = i / a_gl_rd_delta_o;
854 + return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;
855 + };
856 + // Since the computation of this remapping is non-trivial and, due to our main
857 + // loop unrolls, all shared memory accesses are static, we simply precompute
858 + // both transformed reads and writes.
859 + int a_sh_wr_trans[a_sh_wr_iters];
860 + #pragma unroll
861 + for (int i = 0; i < a_sh_wr_iters; i++)
862 + a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);
863 + int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];
864 + #pragma unroll
865 + for (int i = 0; i < b_sh_wr_iters; i++) {
866 + #pragma unroll
867 + for (int j = 0; j < thread_m_blocks; j++)
868 + a_sh_rd_trans[i][j] =
869 + transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);
870 + }
871 +
872 + // Since B-accesses have non-constant stride they have to be computed at
873 + // runtime; we break dependencies between subsequent accesses with a tile by
874 + // maintining multiple pointers (we have enough registers), a tiny
875 + // optimization.
876 + const int4* B_ptr[b_sh_wr_iters];
877 + #pragma unroll
878 + for (int i = 0; i < b_sh_wr_iters; i++)
879 + B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;
880 +
881 + // Shared memory storage for global fetch pipelines.
882 + int4* sh_a = sh_new;
883 + int4* sh_b = sh_a + (stages * a_sh_stage);
884 + int4* sh_g_idx = sh_b + (stages * b_sh_stage);
885 + int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
886 + int4* sh_s = sh_zp + (stages * zp_sh_stage);
887 + int4* sh_red = sh_b;
888 +
889 + // Register storage for double buffer of shared memory reads.
890 + FragA frag_a[2][thread_m_blocks];
891 + I4 frag_b_quant[2][b_thread_vecs];
892 + FragC frag_c[thread_m_blocks][4][2];
893 + FragS frag_s[2][4]; // No act-order
894 + FragS act_frag_s[2][4][4]; // For act-order
895 + int frag_qzp[2][num_ints_per_thread]; // Zero-points
896 + FragZP frag_zp; // Zero-points in fp16
897 + FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ
898 +
899 + // Zero accumulators.
900 + auto zero_accums = [&]() {
901 + #pragma unroll
902 + for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)
903 + reinterpret_cast<float*>(frag_c)[i] = 0;
904 + };
905 +
906 + int sh_first_group_id = -1;
907 + int sh_num_groups = -1;
908 + constexpr int sh_max_num_groups = 32;
909 +
910 + auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id,
911 + int last_group_id) {
912 + sh_first_group_id = first_group_id;
913 + sh_num_groups = last_group_id - first_group_id + 1;
914 +
915 + if (sh_num_groups < sh_max_num_groups) {
916 + sh_num_groups = sh_max_num_groups;
917 + }
918 +
919 + if (sh_first_group_id + sh_num_groups > num_groups) {
920 + sh_num_groups = num_groups - sh_first_group_id;
921 + }
922 +
923 + int row_offset = first_group_id * s_gl_stride;
924 +
925 + if (is_async) {
926 + for (int i = 0; i < sh_num_groups; i++) {
927 + if (threadIdx.x < s_sh_stride) {
928 + cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],
929 + &scales_ptr[row_offset + (i * s_gl_stride) +
930 + slice_n_offset + threadIdx.x]);
931 + }
932 + }
933 + } else {
934 + for (int i = 0; i < sh_num_groups; i++) {
935 + if (threadIdx.x < s_sh_stride) {
936 + sh_s[(i * s_sh_stride) + threadIdx.x] =
937 + scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +
938 + threadIdx.x];
939 + }
940 + }
941 + }
942 + };
943 + // Asynchronously fetch the next A, B and s tile from global to the next
944 + // shared memory pipeline location.
945 + int a_remaining_load_count_in_slice = stages;
946 + auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {
947 + if (pred) {
948 + int4* sh_a_stage = sh_a + a_sh_stage * pipe;
949 + if (prob_k > thread_k_blocks * 16 * stages || slice_col == 0 ||
950 + a_remaining_load_count_in_slice > 0) {
951 + a_remaining_load_count_in_slice--;
952 + #pragma unroll
953 + for (int i = 0; i < a_sh_wr_iters; i++) {
954 + int a_idx = a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off;
955 + int row = a_idx / a_gl_stride;
956 + int64_t sorted_row = 0;
957 + if (!m_block_size_8 || row < 8)
958 + sorted_row = sh_block_sorted_ids[row] / top_k;
959 + int64_t true_idx = sorted_row * a_gl_stride + a_idx % a_gl_stride;
960 + cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[true_idx],
961 + row < block_num_valid_tokens);
962 + }
963 + }
964 + int4* sh_b_stage = sh_b + b_sh_stage * pipe;
965 + #pragma unroll
966 + for (int i = 0; i < b_sh_wr_iters; i++) {
967 + #pragma unroll
968 + for (int j = 0; j < b_thread_vecs; j++) {
969 + cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j],
970 + B_ptr[i] + j + B_expert_off);
971 + }
972 +
973 + B_ptr[i] += b_gl_rd_delta_o;
974 + }
975 +
976 + if constexpr (has_act_order) {
977 + // Fetch g_idx thread-block portion
978 + int full_pipe = a_off;
979 + int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;
980 + if (cur_k < prob_k && cur_k < slice_k_finish) {
981 + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
982 +
983 + int4 const* cur_g_idx_stage_ptr =
984 + reinterpret_cast<int4 const*>(&g_idx[cur_k]);
985 +
986 + if (threadIdx.x < g_idx_stage) {
987 + cp_async4_pred(&sh_g_idx_stage[threadIdx.x],
988 + &cur_g_idx_stage_ptr[threadIdx.x]);
989 + }
990 + }
991 + } else {
992 + if constexpr (group_blocks != -1) {
993 + int4* sh_s_stage = sh_s + s_sh_stage * pipe;
994 +
995 + if constexpr (group_blocks >= thread_k_blocks) {
996 + // Only fetch scales if this tile starts a new group
997 + if (pipe % (group_blocks / thread_k_blocks) == 0) {
998 + if (s_sh_wr_pred) {
999 + cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);
1000 + }
1001 + s_gl_rd += s_gl_rd_delta;
1002 + }
1003 + } else {
1004 + for (int i = 0; i < s_tb_groups; i++) {
1005 + if (s_sh_wr_pred) {
1006 + cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],
1007 + &scales_ptr[s_gl_rd]);
1008 + }
1009 + s_gl_rd += s_gl_rd_delta;
1010 + }
1011 + }
1012 + }
1013 +
1014 + if constexpr (has_zp && group_blocks != -1) {
1015 + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1016 +
1017 + if constexpr (group_blocks >= thread_k_blocks) {
1018 + // Only fetch zero-points if this tile starts a new group
1019 + if (pipe % (group_blocks / thread_k_blocks) == 0) {
1020 + if (zp_sh_wr_pred) {
1021 + cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);
1022 + }
1023 + zp_gl_rd += zp_gl_rd_delta;
1024 + }
1025 + } else {
1026 + for (int i = 0; i < zp_tb_groups; i++) {
1027 + if (zp_sh_wr_pred) {
1028 + cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],
1029 + &zp_ptr[zp_gl_rd]);
1030 + }
1031 + zp_gl_rd += zp_gl_rd_delta;
1032 + }
1033 + }
1034 + }
1035 + }
1036 + }
1037 + // Insert a fence even when we are winding down the pipeline to ensure that
1038 + // waiting is also correct at this point.
1039 + cp_async_fence();
1040 + };
1041 +
1042 + auto fetch_col_zp_to_shared = [&]() {
1043 + if (zp_sh_wr_pred) {
1044 + cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);
1045 + }
1046 + };
1047 +
1048 + auto fetch_col_scale_to_shared = [&]() {
1049 + if (s_sh_wr_pred) {
1050 + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
1051 + }
1052 + };
1053 +
1054 + // Wait until the next thread tile has been loaded to shared memory.
1055 + auto wait_for_stage = [&]() {
1056 + // We only have `stages - 2` active fetches since we are double buffering
1057 + // and can only issue the next fetch when it is guaranteed that the previous
1058 + // shared memory load is fully complete (as it may otherwise be
1059 + // overwritten).
1060 + cp_async_wait<stages - 2>();
1061 + __syncthreads();
1062 + };
1063 +
1064 + // Load the next sub-tile from the current location in the shared memory pipe
1065 + // into the current register buffer.
1066 + auto fetch_to_registers = [&](int k, int pipe) {
1067 + int4* sh_a_stage = sh_a + a_sh_stage * pipe;
1068 + #pragma unroll
1069 + for (int i = 0; i < thread_m_blocks; i++)
1070 + ldsm<m_block_size_8 ? 2 : 4, scalar_t>(
1071 + frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);
1072 + int4* sh_b_stage = sh_b + b_sh_stage * pipe;
1073 +
1074 + #pragma unroll
1075 + for (int i = 0; i < b_thread_vecs; i++) {
1076 + frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(
1077 + &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);
1078 + }
1079 + };
1080 +
1081 + bool is_same_group[stages];
1082 + int same_group_id[stages];
1083 +
1084 + auto init_same_group = [&](int pipe) {
1085 + if constexpr (!has_act_order) {
1086 + return;
1087 + }
1088 +
1089 + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1090 + int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1091 +
1092 + int group_id_1 = sh_g_idx_int_ptr[0];
1093 + int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];
1094 +
1095 + is_same_group[pipe] = group_id_1 == group_id_2;
1096 + same_group_id[pipe] = group_id_1;
1097 + };
1098 +
1099 + auto fetch_scales_to_registers = [&](int k, int full_pipe) {
1100 + int pipe = full_pipe % stages;
1101 +
1102 + if constexpr (!has_act_order) {
1103 + // No act-order case
1104 + if constexpr (group_blocks == -1) {
1105 + // load only when starting a new slice
1106 + if (k == 0 && full_pipe == 0) {
1107 + reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd];
1108 + reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1109 + }
1110 + } else if constexpr (group_blocks != -1) {
1111 + if constexpr (group_blocks >= thread_k_blocks) {
1112 + int4* sh_s_stage =
1113 + sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *
1114 + (pipe / (group_blocks / thread_k_blocks)));
1115 + reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];
1116 + } else {
1117 + int warp_id = threadIdx.x / 32;
1118 + int n_warps = thread_n_blocks / 4;
1119 +
1120 + int warp_row = warp_id / n_warps;
1121 +
1122 + int cur_k = warp_row * 16;
1123 + cur_k += k_iter_size * (k % b_sh_wr_iters);
1124 +
1125 + int k_blocks = cur_k / 16;
1126 + int cur_group_id = k_blocks / group_blocks;
1127 +
1128 + int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1129 +
1130 + reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1131 + sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1132 + }
1133 + }
1134 +
1135 + return;
1136 + }
1137 +
1138 + // Act-order case
1139 +
1140 + // Determine K of the "current" thread-block
1141 + int cur_k = slice_k_start + tb_k * full_pipe;
1142 + if (cur_k >= prob_k || cur_k >= slice_k_finish) {
1143 + return;
1144 + }
1145 +
1146 + // Reset (to current thread-block) since we read g_idx portion from the
1147 + // shared memory
1148 + cur_k = 0;
1149 +
1150 + // Progress to current iteration
1151 + cur_k += k_iter_size * (k % b_sh_wr_iters);
1152 +
1153 + // Determine "position" inside the thread-block (based on warp and
1154 + // thread-id)
1155 + int warp_id = threadIdx.x / 32;
1156 + int n_warps =
1157 + thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N
1158 +
1159 + int warp_row = warp_id / n_warps;
1160 + int warp_col = warp_id % n_warps;
1161 +
1162 + cur_k += warp_row * 16;
1163 +
1164 + int th_id = threadIdx.x % 32;
1165 + cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix
1166 +
1167 + int s_col_shift =
1168 + /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +
1169 + (th_id / 4) * act_s_col_stride;
1170 +
1171 + if (is_same_group[pipe]) {
1172 + if (k % 2 == 0) {
1173 + *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1174 + sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +
1175 + s_col_shift];
1176 + } else {
1177 + *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =
1178 + *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));
1179 + }
1180 +
1181 + for (int i = 1; i < 4; i++) {
1182 + *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1183 + *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));
1184 + }
1185 + return;
1186 + }
1187 +
1188 + int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;
1189 + int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);
1190 +
1191 + constexpr int k_frag_offsets[4] = {0, 1, 8,
1192 + 9}; // Tensor core offsets per thread
1193 +
1194 + #pragma unroll
1195 + for (int i = 0; i < 4; i++) {
1196 + int actual_k = cur_k + k_frag_offsets[i];
1197 +
1198 + int group_id = sh_g_idx_int_ptr[actual_k];
1199 + int rel_group_id = group_id - sh_first_group_id;
1200 +
1201 + *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =
1202 + sh_s[rel_group_id * s_sh_stride + s_col_shift];
1203 + }
1204 + };
1205 +
1206 + auto fetch_zp_to_registers = [&](int k, int full_pipe) {
1207 + // This code does not handle group_blocks == 0,
1208 + // which signifies act_order.
1209 + // has_zp implies AWQ, which doesn't have act_order,
1210 + static_assert(!has_zp || group_blocks != 0);
1211 +
1212 + if constexpr (has_zp && !is_zp_float) {
1213 + int pipe = full_pipe % stages;
1214 +
1215 + if constexpr (group_blocks == -1) {
1216 + // load only when starting a new slice
1217 + if (k == 0 && full_pipe == 0) {
1218 + #pragma unroll
1219 + for (int i = 0; i < num_ints_per_thread; i++) {
1220 + frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];
1221 + }
1222 + }
1223 +
1224 + } else if constexpr (group_blocks >= thread_k_blocks) {
1225 + int4* sh_zp_stage =
1226 + sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1227 + (pipe / (group_blocks / thread_k_blocks)));
1228 + for (int i = 0; i < num_ints_per_thread; i++) {
1229 + frag_qzp[k % 2][i] =
1230 + (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1231 + }
1232 + } else {
1233 + int warp_id = threadIdx.x / 32;
1234 + int n_warps = thread_n_blocks / 4;
1235 +
1236 + int warp_row = warp_id / n_warps;
1237 +
1238 + int cur_k = warp_row * 16;
1239 + cur_k += k_iter_size * (k % b_sh_wr_iters);
1240 +
1241 + int k_blocks = cur_k / 16;
1242 + int cur_group_id = 0;
1243 +
1244 + // Suppress bogus and persistent divide-by-zero warning
1245 + #pragma nv_diagnostic push
1246 + #pragma nv_diag_suppress divide_by_zero
1247 + cur_group_id = k_blocks / group_blocks;
1248 + #pragma nv_diagnostic pop
1249 +
1250 + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1251 +
1252 + sh_zp_stage += cur_group_id * zp_sh_stride;
1253 +
1254 + for (int i = 0; i < num_ints_per_thread; i++) {
1255 + frag_qzp[k % 2][i] =
1256 + (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];
1257 + }
1258 + }
1259 + }
1260 +
1261 + else if constexpr (has_zp && is_zp_float) {
1262 + int pipe = full_pipe % stages;
1263 +
1264 + if constexpr (group_blocks != -1) {
1265 + if constexpr (group_blocks >= thread_k_blocks) {
1266 + int4* sh_zp_stage =
1267 + sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *
1268 + (pipe / (group_blocks / thread_k_blocks)));
1269 + reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd];
1270 + } else {
1271 + int warp_id = threadIdx.x / 32;
1272 + int n_warps = thread_n_blocks / 4;
1273 +
1274 + int warp_row = warp_id / n_warps;
1275 +
1276 + int cur_k = warp_row * 16;
1277 + cur_k += k_iter_size * (k % b_sh_wr_iters);
1278 +
1279 + int k_blocks = cur_k / 16;
1280 + // Suppress bogus and persistent divide-by-zero warning
1281 + #pragma nv_diagnostic push
1282 + #pragma nv_diag_suppress divide_by_zero
1283 + int cur_group_id = k_blocks / group_blocks;
1284 + #pragma nv_diagnostic pop
1285 +
1286 + int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;
1287 +
1288 + reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] =
1289 + sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride];
1290 + }
1291 + }
1292 + }
1293 + };
1294 +
1295 + // Execute the actual tensor core matmul of a sub-tile.
1296 + bool is_first_matmul_in_slice = true;
1297 + auto matmul = [&](int k) {
1298 + int k2 = k % 2;
1299 + const bool is_new_zp =
1300 + ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 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 + FragB frag_zp_0;
1306 + FragB frag_zp_1;
1307 + int zp_quant_0, zp_quant_1;
1308 +
1309 + if constexpr (w_type.size_bits() == 4) {
1310 + zp_quant_0 = frag_qzp[k2][0];
1311 + zp_quant_1 = zp_quant_0 >> 8;
1312 + } else {
1313 + static_assert(w_type.size_bits() == 8);
1314 + zp_quant_0 = frag_qzp[k2][0];
1315 + zp_quant_1 = frag_qzp[k2][1];
1316 + }
1317 +
1318 + dequant<scalar_t, w_type.size_bits()>(zp_quant_0, frag_zp_0);
1319 + dequant<scalar_t, w_type.size_bits()>(zp_quant_1, frag_zp_1);
1320 +
1321 + frag_zp[0] = frag_zp_0[0];
1322 + frag_zp[1] = frag_zp_0[1];
1323 + frag_zp[2] = frag_zp_1[0];
1324 + frag_zp[3] = frag_zp_1[1];
1325 + }
1326 + }
1327 + // We have the m dimension as the inner loop in order to encourage overlapping
1328 + // dequantization and matmul operations.
1329 + #pragma unroll
1330 + for (int j = 0; j < 4; j++) {
1331 + FragB frag_b0;
1332 + FragB frag_b1;
1333 + int b_quant_0, b_quant_1;
1334 +
1335 + if constexpr (w_type.size_bits() == 4) {
1336 + b_quant_0 = frag_b_quant[k2][0][j];
1337 + b_quant_1 = b_quant_0 >> 8;
1338 + } else {
1339 + static_assert(w_type.size_bits() == 8);
1340 + int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]);
1341 + b_quant_0 = frag_b_quant_ptr[j * 2 + 0];
1342 + b_quant_1 = frag_b_quant_ptr[j * 2 + 1];
1343 + }
1344 +
1345 + dequant<scalar_t, w_type.size_bits()>(b_quant_0, frag_b0);
1346 + dequant<scalar_t, w_type.size_bits()>(b_quant_1, frag_b1);
1347 +
1348 + // Apply scale to frag_b0
1349 + if constexpr (has_act_order) {
1350 + static_assert(group_blocks != -1);
1351 + scale4<scalar_t>(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1352 + act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1353 + scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1354 + act_frag_s[k][2][j], act_frag_s[k2][3][j], 1);
1355 +
1356 + } else if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1357 + int idx = (threadIdx.x / 4) % 2;
1358 + scalar_t2 s2 = Dtype::nums2num2(
1359 + reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
1360 + reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]);
1361 + if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1362 + scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1363 + scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1364 + } else if constexpr (has_zp && !is_zp_float && group_blocks != -1) {
1365 + if (is_new_zp)
1366 + frag_zp[j] = __hmul2(frag_zp[j],
1367 + *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1368 + scale_and_sub<scalar_t>(frag_b0, frag_s[k % 2][j][0].x, frag_zp[j].x);
1369 + scale_and_sub<scalar_t>(frag_b1, frag_s[k % 2][j][0].y, frag_zp[j].y);
1370 + } else if constexpr (has_zp && is_zp_float && group_blocks != -1) {
1371 + if (is_new_zp)
1372 + frag_zpf[k2][j] = __hmul2(
1373 + frag_zpf[k2][j], *reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
1374 + scale_and_sub<scalar_t>(frag_b0, frag_s[k2][j].x, frag_zpf[k2][j].x);
1375 + scale_and_sub<scalar_t>(frag_b1, frag_s[k2][j].y, frag_zpf[k2][j].y);
1376 + } else if constexpr (group_blocks != -1) {
1377 + scale<scalar_t>(frag_b0, frag_s[k2][j], 0);
1378 + scale<scalar_t>(frag_b1, frag_s[k2][j], 1);
1379 + }
1380 +
1381 + #pragma unroll
1382 + for (int i = 0; i < thread_m_blocks; i++) {
1383 + if constexpr (m_block_size_8) {
1384 + mma_trans<scalar_t>(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]);
1385 + } else {
1386 + mma<scalar_t>(frag_a[k2][i], frag_b0, frag_c[i][j][0]);
1387 + mma<scalar_t>(frag_a[k2][i], frag_b1, frag_c[i][j][1]);
1388 + }
1389 + }
1390 + }
1391 + };
1392 +
1393 + // Since we slice across the k dimension of a tile in order to increase the
1394 + // number of warps while keeping the n dimension of a tile reasonable, we have
1395 + // multiple warps that accumulate their partial sums of the same output
1396 + // location; which we have to reduce over in the end. We do in shared memory.
1397 + auto thread_block_reduce = [&]() {
1398 + constexpr int red_off = threads / b_sh_stride_threads / 2;
1399 + if (red_off >= 1) {
1400 + int red_idx = threadIdx.x / b_sh_stride_threads;
1401 + constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;
1402 + constexpr int red_sh_delta = b_sh_stride_threads;
1403 + int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +
1404 + (threadIdx.x % b_sh_stride_threads);
1405 +
1406 + // Parallel logarithmic shared memory reduction. We make sure to avoid any
1407 + // unnecessary read or write iterations, e.g., for two warps we write only
1408 + // once by warp 1 and read only once by warp 0.
1409 +
1410 + #pragma unroll
1411 + for (int m_block = 0; m_block < thread_m_blocks; m_block++) {
1412 + #pragma unroll
1413 + for (int i = red_off; i > 0; i /= 2) {
1414 + if (i <= red_idx && red_idx < 2 * i) {
1415 + #pragma unroll
1416 + for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) {
1417 + int red_sh_wr =
1418 + red_sh_delta * j + (red_sh_rd - red_sh_stride * i);
1419 + if (i < red_off) {
1420 + float* c_rd = reinterpret_cast<float*>(
1421 + &sh_red[red_sh_delta * j + red_sh_rd]);
1422 + float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]);
1423 + #pragma unroll
1424 + for (int k = 0; k < 4; k++)
1425 + reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=
1426 + c_rd[k] + c_wr[k];
1427 + }
1428 + sh_red[red_sh_wr] =
1429 + reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];
1430 + }
1431 + }
1432 + __syncthreads();
1433 + }
1434 + if (red_idx == 0) {
1435 + #pragma unroll
1436 + for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) {
1437 + float* c_rd =
1438 + reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]);
1439 + #pragma unroll
1440 + for (int j = 0; j < 4; j++)
1441 + reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=
1442 + c_rd[j];
1443 + }
1444 + }
1445 + __syncthreads();
1446 + }
1447 + }
1448 + };
1449 +
1450 + // Since multiple threadblocks may process parts of the same column slice, we
1451 + // finally have to globally reduce over the results. As the striped
1452 + // partitioning minimizes the number of such reductions and our outputs are
1453 + // usually rather small, we perform this reduction serially in L2 cache.
1454 + auto global_reduce_fp16 = [&](bool first = false, bool last = false) {
1455 + // We are very careful here to reduce directly in the output buffer to
1456 + // maximize L2 cache utilization in this step. To do this, we write out
1457 + // results in FP16 (but still reduce with FP32 compute).
1458 + constexpr int active_threads = 32 * thread_n_blocks / 4;
1459 + bool is_th_active = threadIdx.x < active_threads;
1460 + if (!is_th_active) {
1461 + return;
1462 + }
1463 +
1464 + int c_gl_stride = prob_n / 8;
1465 + int c_gl_wr_delta_o = 8 * c_gl_stride;
1466 + int c_gl_wr_delta_i = 4 * (active_threads / 32);
1467 + int c_gl_wr;
1468 + if constexpr (m_block_size_8) {
1469 + c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) +
1470 + (threadIdx.x % 32) / 8;
1471 + c_gl_wr += (2 * thread_n_blocks) * slice_col;
1472 + } else {
1473 + c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +
1474 + 4 * (threadIdx.x / 32) + threadIdx.x % 4;
1475 + c_gl_wr += (2 * thread_n_blocks) * slice_col;
1476 + }
1477 + constexpr int c_sh_wr_delta = active_threads;
1478 + int c_sh_wr = threadIdx.x;
1479 +
1480 + if (!first) {
1481 +
1482 + #pragma unroll
1483 + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
1484 + int c_idx;
1485 + if constexpr (m_block_size_8)
1486 + c_idx = c_gl_wr + i * c_gl_stride +
1487 + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
1488 + else
1489 + c_idx =
1490 + c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
1491 + if (c_idx / c_gl_stride < block_num_valid_tokens) {
1492 + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
1493 + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
1494 + sh_red[c_sh_wr + c_sh_wr_delta * i] = C[true_idx];
1495 + }
1496 + }
1497 + }
1498 +
1499 + #pragma unroll
1500 + for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) {
1501 + if (!first) {
1502 + int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta];
1503 + #pragma unroll
1504 + for (int j = 0; j < 2 * 4; j++) {
1505 + int delta = 0;
1506 + if constexpr (m_block_size_8) {
1507 + delta = j % 2 == 1 ? -2 : 0;
1508 + }
1509 + reinterpret_cast<float*>(
1510 + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] +=
1511 + Dtype::num2float(reinterpret_cast<scalar_t*>(&c_red)[j]);
1512 + }
1513 + }
1514 + if (!last) {
1515 + int4 c;
1516 + #pragma unroll
1517 + for (int j = 0; j < 2 * 4; j++) {
1518 + int delta = 0;
1519 + if constexpr (m_block_size_8) {
1520 + delta = j % 2 == 1 ? -2 : 0;
1521 + }
1522 + reinterpret_cast<scalar_t*>(&c)[j] =
1523 + Dtype::float2num(reinterpret_cast<float*>(
1524 + &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]);
1525 + }
1526 +
1527 + int c_idx;
1528 + if constexpr (m_block_size_8)
1529 + c_idx = c_gl_wr + i * c_gl_stride +
1530 + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i;
1531 + else
1532 + c_idx =
1533 + c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);
1534 + if (c_idx / c_gl_stride < block_num_valid_tokens) {
1535 + int64_t sorted_row = sh_block_sorted_ids[c_idx / c_gl_stride];
1536 + int64_t true_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;
1537 + C[true_idx] = c;
1538 + }
1539 + }
1540 + }
1541 + };
1542 +
1543 + // Globally reduce over threadblocks that compute the same column block.
1544 + // We use a tmp C buffer to reduce in full fp32 precision.
1545 + auto global_reduce_fp32 = [&](bool first = false, bool last = false) {
1546 + constexpr int tb_m = thread_m_blocks * 16;
1547 + constexpr int tb_n = thread_n_blocks * 16;
1548 +
1549 + constexpr int c_size = tb_m * tb_n * sizeof(float) / 16;
1550 +
1551 + constexpr int active_threads = 32 * thread_n_blocks / 4;
1552 + bool is_th_active = threadIdx.x < active_threads;
1553 +
1554 + constexpr int num_floats = thread_m_blocks * 4 * 2 * 4;
1555 + constexpr int th_size = num_floats * sizeof(float) / 16;
1556 +
1557 + int c_cur_offset = locks_off * c_size;
1558 +
1559 + if (!is_th_active) {
1560 + return;
1561 + }
1562 +
1563 + if (!first) {
1564 + float* frag_c_ptr = reinterpret_cast<float*>(&frag_c);
1565 + #pragma unroll
1566 + for (int k = 0; k < th_size; k++) {
1567 + if constexpr (m_block_size_8) {
1568 + if (k % 2) continue;
1569 + } else {
1570 + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
1571 + continue;
1572 + }
1573 +
1574 + sh_red[threadIdx.x] =
1575 + C_tmp[c_cur_offset + active_threads * k + threadIdx.x];
1576 +
1577 + float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]);
1578 + #pragma unroll
1579 + for (int f = 0; f < 4; f++) {
1580 + frag_c_ptr[k * 4 + f] += sh_c_ptr[f];
1581 + }
1582 + }
1583 + }
1584 +
1585 + if (!last) {
1586 + int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c);
1587 + #pragma unroll
1588 + for (int k = 0; k < th_size; k++) {
1589 + if constexpr (m_block_size_8) {
1590 + if (k % 2) continue;
1591 + } else {
1592 + if (k / 8 * 16 + (threadIdx.x % 32) / 4 >= block_num_valid_tokens)
1593 + continue;
1594 + }
1595 +
1596 + C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k];
1597 + }
1598 + }
1599 + };
1600 +
1601 + // Write out the reduce final result in the correct layout. We only actually
1602 + // reshuffle matrix fragments in this step, the reduction above is performed
1603 + // in fragment layout.
1604 + auto write_result = [&]() {
1605 + int c_gl_stride = prob_n / 8;
1606 + constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
1607 + int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
1608 + constexpr int c_sh_rd_delta =
1609 + c_sh_stride * (threads / (2 * thread_n_blocks));
1610 +
1611 + int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1612 + (threadIdx.x % (2 * thread_n_blocks));
1613 + c_gl_wr += (2 * thread_n_blocks) * slice_col;
1614 + int c_sh_wr;
1615 + if constexpr (m_block_size_8) {
1616 + c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) +
1617 + (threadIdx.x % 32) / 4;
1618 + c_sh_wr += 64 * (threadIdx.x / 32);
1619 + } else {
1620 + c_sh_wr =
1621 + (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;
1622 + c_sh_wr += 32 * (threadIdx.x / 32);
1623 + }
1624 +
1625 + int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +
1626 + (threadIdx.x % (2 * thread_n_blocks));
1627 +
1628 + // We first reorder in shared memory to guarantee the most efficient final
1629 + // global write patterns
1630 + auto write = [&](int idx, float c0, float c1, FragS& s) {
1631 + scalar_t2 res =
1632 + Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1633 +
1634 + // For per-column quantization we finally apply the scale here (only for
1635 + // 4-bit)
1636 + if constexpr (!has_act_order && group_blocks == -1 &&
1637 + w_type.size_bits() == 4 && !has_zp) {
1638 + res = __hmul2(res, s[0]);
1639 + }
1640 +
1641 + if constexpr (m_block_size_8) {
1642 + ((scalar_t*)sh_red)[idx] = res.x;
1643 + ((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
1644 + } else {
1645 + ((scalar_t2*)sh_red)[idx] = res;
1646 + }
1647 + };
1648 +
1649 + if (threadIdx.x / 32 < thread_n_blocks / 4) {
1650 + #pragma unroll
1651 + for (int i = 0; i < thread_m_blocks; i++) {
1652 + #pragma unroll
1653 + for (int j = 0; j < 4; j++) {
1654 + if constexpr (m_block_size_8) {
1655 + int wr = c_sh_wr + 16 * j;
1656 + write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
1657 + frag_s[j / 2][2 * (j % 2) + 0]);
1658 + write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
1659 + frag_s[j / 2][2 * (j % 2) + 1]);
1660 + } else {
1661 + int wr = c_sh_wr + 8 * j;
1662 + write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
1663 + frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
1664 + write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
1665 + frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
1666 + write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
1667 + frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
1668 + write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
1669 + frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
1670 + }
1671 + }
1672 + c_sh_wr += 16 * (4 * c_sh_stride);
1673 + }
1674 + }
1675 + __syncthreads();
1676 +
1677 + #pragma unroll
1678 + for (int i = 0;
1679 + i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));
1680 + i++) {
1681 + int row = c_gl_wr / c_gl_stride;
1682 + if (row < block_num_valid_tokens) {
1683 + int64_t sorted_row = sh_block_sorted_ids[row];
1684 + int64_t true_idx = sorted_row * c_gl_stride + c_gl_wr % c_gl_stride;
1685 + scalar_t2 topk_weight_score;
1686 + if (mul_topk_weights) topk_weight_score = sh_block_topk_weights[row];
1687 + if (use_atomic_add && slice_count > 1 || mul_topk_weights) {
1688 + scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[true_idx]);
1689 + scalar_t2* sh_red_half2 =
1690 + reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);
1691 + #pragma unroll
1692 + for (int a = 0; a < 4; a++) {
1693 + scalar_t2 res = sh_red_half2[a];
1694 + if (mul_topk_weights) {
1695 + res = __hmul2(res, topk_weight_score);
1696 + }
1697 +
1698 + if (use_atomic_add && slice_count > 1) {
1699 + atomicAdd(&C_half2[a], res);
1700 + } else {
1701 + C_half2[a] = res;
1702 + };
1703 + }
1704 + } else {
1705 + C[true_idx] = sh_red[c_sh_rd];
1706 + }
1707 + c_gl_wr += c_gl_wr_delta;
1708 + c_sh_rd += c_sh_rd_delta;
1709 + }
1710 + }
1711 + __syncthreads();
1712 + };
1713 +
1714 + // Start global fetch and register load pipelines.
1715 + auto start_pipes = [&]() {
1716 +
1717 + #pragma unroll
1718 + for (int i = 0; i < stages - 1; i++) {
1719 + if (has_act_order && i == 0) {
1720 + int last_g_idx = slice_k_start + stages * tb_k * 2;
1721 + if (last_g_idx >= prob_k) {
1722 + last_g_idx = prob_k - 1;
1723 + }
1724 + fetch_act_order_scales_to_shared(true, g_idx[slice_k_start],
1725 + g_idx[last_g_idx]);
1726 + }
1727 +
1728 + if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1729 + if (i == 0) {
1730 + fetch_col_zp_to_shared();
1731 + fetch_col_scale_to_shared();
1732 + }
1733 + }
1734 + fetch_to_shared(i, i, i < slice_iters);
1735 + }
1736 +
1737 + zero_accums();
1738 + wait_for_stage();
1739 + init_same_group(0);
1740 + fetch_to_registers(0, 0);
1741 + fetch_scales_to_registers(0, 0);
1742 + fetch_zp_to_registers(0, 0);
1743 + a_gl_rd += a_gl_rd_delta_o * (stages - 1);
1744 + slice_k_start_shared_fetch += tb_k * (stages - 1);
1745 + };
1746 + if (slice_iters) {
1747 + start_pipes();
1748 + }
1749 +
1750 + // Main loop.
1751 + while (slice_iters) {
1752 + // We unroll over both the global fetch and the register load pipeline to
1753 + // ensure all shared memory accesses are static. Note that both pipelines
1754 + // have even length meaning that the next iteration will always start at
1755 + // index 0.
1756 +
1757 + #pragma unroll
1758 + for (int pipe = 0; pipe < stages;) {
1759 + #pragma unroll
1760 + for (int k = 0; k < b_sh_wr_iters; k++) {
1761 + fetch_to_registers(k + 1, pipe % stages);
1762 + fetch_scales_to_registers(k + 1, pipe);
1763 + fetch_zp_to_registers(k + 1, pipe);
1764 + if (k == b_sh_wr_iters - 2) {
1765 + fetch_to_shared((pipe + stages - 1) % stages, pipe,
1766 + slice_iters >= stages);
1767 + pipe++;
1768 + wait_for_stage();
1769 + init_same_group(pipe % stages);
1770 + }
1771 + matmul(k);
1772 + }
1773 + slice_iters--;
1774 + if (slice_iters == 0) {
1775 + break;
1776 + }
1777 + }
1778 + a_remaining_load_count_in_slice = 0;
1779 +
1780 + a_gl_rd += a_gl_rd_delta_o * stages;
1781 + slice_k_start += tb_k * stages;
1782 + slice_k_start_shared_fetch += tb_k * stages;
1783 +
1784 + if constexpr (has_act_order) {
1785 + int first_group_id = g_idx[slice_k_start];
1786 + int last_g_idx = slice_k_start + stages * tb_k * 2;
1787 + if (last_g_idx >= prob_k) {
1788 + last_g_idx = prob_k - 1;
1789 + }
1790 + int last_group_id = g_idx[last_g_idx];
1791 + if (last_group_id >= sh_first_group_id + sh_num_groups) {
1792 + fetch_act_order_scales_to_shared(false, first_group_id, last_group_id);
1793 + __syncthreads();
1794 + }
1795 + }
1796 +
1797 + // Process results and, if necessary, proceed to the next column slice.
1798 + // While this pattern may not be the most readable, other ways of writing
1799 + // the loop seemed to noticeably worse performance after compilation.
1800 + if (slice_iters == 0) {
1801 + cp_async_wait<0>();
1802 + bool last = slice_idx == slice_count - 1;
1803 + // For per-column scales, we only fetch them here in the final step before
1804 + // write-out
1805 + if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1806 + if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1807 + if (s_sh_wr_pred) {
1808 + cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
1809 + }
1810 + cp_async_fence();
1811 + }
1812 + }
1813 +
1814 + thread_block_reduce();
1815 + if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1816 + if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1817 + cp_async_wait<0>();
1818 + __syncthreads();
1819 + if (threadIdx.x / 32 < thread_n_blocks / 4) {
1820 + reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];
1821 + reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];
1822 + if constexpr (m_block_size_8) {
1823 + int idx = (threadIdx.x / 4) % 2;
1824 + scalar_t2* frag_s_half2 = reinterpret_cast<scalar_t2*>(frag_s);
1825 + #pragma unroll
1826 + for (int i = 0; i < 8; i++) {
1827 + frag_s_half2[i] = Dtype::num2num2(
1828 + reinterpret_cast<scalar_t*>(&frag_s_half2[i])[idx]);
1829 + }
1830 + }
1831 + }
1832 + }
1833 + }
1834 +
1835 + // For 8-bit channelwise, we apply the scale before the global reduction
1836 + // that converts the fp32 results to fp16 (so that we avoid possible
1837 + // overflow in fp16)
1838 + if constexpr (!has_act_order && group_blocks == -1 &&
1839 + w_type.size_bits() == 8 && !has_zp) {
1840 + if (threadIdx.x / 32 < thread_n_blocks / 4) {
1841 + #pragma unroll
1842 + for (int i = 0; i < thread_m_blocks; i++) {
1843 + #pragma unroll
1844 + for (int j = 0; j < 4; j++) {
1845 + scale_float<scalar_t>(
1846 + reinterpret_cast<float*>(&frag_c[i][j][0][0]),
1847 + frag_s[j / 2][2 * (j % 2) + 0]);
1848 + scale_float<scalar_t>(
1849 + reinterpret_cast<float*>(&frag_c[i][j][0][2]),
1850 + frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]);
1851 +
1852 + if constexpr (!m_block_size_8) {
1853 + scale_float<scalar_t>(
1854 + reinterpret_cast<float*>(&frag_c[i][j][1][0]),
1855 + frag_s[j / 2][2 * (j % 2) + 1]);
1856 + scale_float<scalar_t>(
1857 + reinterpret_cast<float*>(&frag_c[i][j][1][2]),
1858 + frag_s[j / 2][2 * (j % 2) + 1]);
1859 + }
1860 + }
1861 + }
1862 + }
1863 + }
1864 +
1865 + if (slice_count > 1 && !use_atomic_add) {
1866 + // only globally reduce if there is more than one block in a slice
1867 + barrier_acquire(&locks[locks_off], slice_idx);
1868 + if (use_fp32_reduce) {
1869 + global_reduce_fp32(slice_idx == 0, last);
1870 + } else {
1871 + global_reduce_fp16(slice_idx == 0, last);
1872 + }
1873 + barrier_release(&locks[locks_off], last);
1874 + }
1875 + if (use_atomic_add && slice_count > 1 && slice_idx != 0)
1876 + wait_negative_and_add(&locks[locks_off]);
1877 + if (last || use_atomic_add)
1878 + // only the last block in a slice actually writes the result
1879 + write_result();
1880 + if (slice_row) a_remaining_load_count_in_slice = stages;
1881 + slice_row = 0;
1882 + slice_col_par++;
1883 + slice_col++;
1884 + is_first_matmul_in_slice = true;
1885 + init_slice();
1886 + if (slice_iters) {
1887 + a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +
1888 + (threadIdx.x % a_gl_rd_delta_o);
1889 + #pragma unroll
1890 + for (int i = 0; i < b_sh_wr_iters; i++)
1891 + B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;
1892 + if (slice_col == 0) {
1893 + #pragma unroll
1894 + for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1895 + }
1896 +
1897 + // Update slice k/n for scales loading
1898 + if constexpr (has_act_order) {
1899 + slice_k_start = tb_k * slice_row;
1900 + slice_k_finish = slice_k_start + tb_k * slice_iters;
1901 + slice_k_start_shared_fetch = slice_k_start;
1902 + slice_n_offset = act_s_col_tb_stride * slice_col;
1903 +
1904 + } else {
1905 + s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
1906 + zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;
1907 + }
1908 +
1909 + start_pipes();
1910 + }
1911 + }
1912 + }
1913 +}
1914 +
1915 +} // namespace MARLIN_NAMESPACE_NAME
1916 +
1917 +#endif
csrc/moe/marlin_moe_wna16/ops.cu new
+927
@@ -0,0 +1,927 @@
1 +/*
2 + * Modified by Neural Magic
3 + * Copyright (C) Marlin.2024 Elias Frantar
4 + *
5 + * Licensed under the Apache License, Version 2.0 (the "License");
6 + * you may not use this file except in compliance with the License.
7 + * You may obtain a copy of the License at
8 + *
9 + * http://www.apache.org/licenses/LICENSE-2.0
10 + *
11 + * Unless required by applicable law or agreed to in writing, software
12 + * distributed under the License is distributed on an "AS IS" BASIS,
13 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 + * See the License for the specific language governing permissions and
15 + * limitations under the License.
16 + */
17 +
18 +/*
19 + * Adapted from https://github.com/IST-DASLab/marlin
20 + */
21 +
22 +#ifndef MARLIN_NAMESPACE_NAME
23 + #define MARLIN_NAMESPACE_NAME marlin_moe_wna16
24 +#endif
25 +
26 +#include "kernel.h"
27 +#include "core/registration.h"
28 +
29 +#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
30 + static_assert(std::is_same<scalar_t, half>::value || \
31 + std::is_same<scalar_t, nv_bfloat16>::value, \
32 + "only float16 and bfloat16 is supported");
33 +
34 +namespace MARLIN_NAMESPACE_NAME {
35 +
36 +__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
37 +
38 +using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS);
39 +
40 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
41 +
42 +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> const& c_or_none,
55 + torch::Tensor& b_q_weight, torch::Tensor& b_scales,
56 + std::optional<torch::Tensor> const& b_zeros_or_none,
57 + std::optional<torch::Tensor> const& g_idx_or_none,
58 + std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
59 + torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
60 + torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
61 + int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
62 + vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
63 + int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
64 + bool is_zp_float) {
65 + TORCH_CHECK_NOT_IMPLEMENTED(false,
66 + "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
67 + return torch::empty({1, 1});
68 +}
69 +
70 +#else
71 +
72 +// For a given "a" of size [M,K] performs a permutation of the K columns based
73 +// on the given "perm" indices.
74 +template <int moe_block_size>
75 +__global__ void permute_cols_kernel(
76 + int4 const* __restrict__ a_int4_ptr, int const* __restrict__ perm_int_ptr,
77 + int4* __restrict__ out_int4_ptr,
78 + const int32_t* __restrict__ sorted_token_ids_ptr,
79 + const int32_t* __restrict__ expert_ids_ptr,
80 + const int32_t* __restrict__ num_tokens_past_padded_ptr, int size_m,
81 + int size_k, int top_k) {
82 + int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
83 + int num_moe_blocks = div_ceil(num_tokens_past_padded, moe_block_size);
84 + int32_t block_sorted_ids[moe_block_size];
85 + int block_num_valid_tokens = 0;
86 + int64_t old_expert_id = 0;
87 + int64_t expert_id = 0;
88 + int row_stride = size_k * sizeof(half) / 16;
89 +
90 + auto read_moe_block_data = [&](int block_id) {
91 + block_num_valid_tokens = moe_block_size;
92 + int4* tmp_block_sorted_ids = reinterpret_cast<int4*>(block_sorted_ids);
93 + for (int i = 0; i < moe_block_size / 4; i++) {
94 + tmp_block_sorted_ids[i] =
95 + ((int4*)sorted_token_ids_ptr)[block_id * moe_block_size / 4 + i];
96 + }
97 + for (int i = 0; i < moe_block_size; i++) {
98 + if (block_sorted_ids[i] >= size_m * top_k) {
99 + block_num_valid_tokens = i;
100 + break;
101 + };
102 + }
103 + };
104 +
105 + auto permute_row = [&](int row) {
106 + int iters = size_k / default_threads;
107 + int rest = size_k % default_threads;
108 +
109 + int in_offset = (row / top_k) * row_stride;
110 + int out_offset = row * row_stride;
111 +
112 + half const* a_row_half =
113 + reinterpret_cast<half const*>(a_int4_ptr + in_offset);
114 + half* out_half = reinterpret_cast<half*>(out_int4_ptr + out_offset);
115 +
116 + int base_k = 0;
117 +
118 + for (int i = 0; i < iters; i++) {
119 + int cur_k = base_k + threadIdx.x;
120 + int src_pos = perm_int_ptr[cur_k];
121 +
122 + out_half[cur_k] = a_row_half[src_pos];
123 +
124 + base_k += default_threads;
125 + }
126 +
127 + if (rest) {
128 + if (threadIdx.x < rest) {
129 + int cur_k = base_k + threadIdx.x;
130 + int src_pos = perm_int_ptr[cur_k];
131 +
132 + out_half[cur_k] = a_row_half[src_pos];
133 + }
134 + }
135 + };
136 +
137 + for (int index = blockIdx.x; index < num_moe_blocks; index += gridDim.x) {
138 + old_expert_id = expert_id;
139 + int tmp_expert_id = expert_ids_ptr[index];
140 + if (tmp_expert_id == -1) continue;
141 + expert_id = tmp_expert_id;
142 + perm_int_ptr += (expert_id - old_expert_id) * size_k;
143 + read_moe_block_data(index);
144 +
145 + for (int i = 0; i < block_num_valid_tokens; i++)
146 + permute_row(block_sorted_ids[i]);
147 + }
148 +}
149 +
150 +typedef struct {
151 + int thread_k;
152 + int thread_n;
153 + int num_threads;
154 +} thread_config_t;
155 +
156 +thread_config_t small_batch_thread_configs[] = {
157 + // Ordered by priority
158 +
159 + // thread_k, thread_n, num_threads
160 + {128, 128, 256},
161 + {64, 128, 128}};
162 +
163 +thread_config_t large_batch_thread_configs[] = {
164 + // Ordered by priority
165 +
166 + // thread_k, thread_n, num_threads
167 + {64, 256, 256},
168 + {64, 128, 128}};
169 +
170 +typedef struct {
171 + int blocks_per_sm;
172 + thread_config_t tb_cfg;
173 +} exec_config_t;
174 +
175 +int get_scales_cache_size(thread_config_t const& th_config, int prob_m,
176 + int prob_n, int prob_k, int num_bits, int group_size,
177 + bool has_act_order, bool is_k_full) {
178 + bool cache_scales_chunk = has_act_order && !is_k_full;
179 +
180 + int tb_n = th_config.thread_n;
181 + int tb_k = th_config.thread_k;
182 +
183 + // Get max scale groups per thread-block
184 + int tb_groups;
185 + if (group_size == -1) {
186 + tb_groups = 1;
187 + } else if (group_size == 0) {
188 + tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size
189 + } else {
190 + tb_groups = div_ceil(tb_k, group_size);
191 + }
192 +
193 + if (cache_scales_chunk) {
194 + int load_groups =
195 + tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K
196 + load_groups = max(load_groups, 32); // We load at least 32 scale groups
197 + return load_groups * tb_n * 2;
198 +
199 + } else {
200 + int tb_scales = tb_groups * tb_n * 2;
201 +
202 + return tb_scales * pipe_stages;
203 + }
204 +}
205 +
206 +int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
207 + int prob_m, int prob_n, int prob_k, int num_bits,
208 + int group_size, bool has_act_order, bool is_k_full,
209 + int has_zp, int is_zp_float) {
210 + int pack_factor = 32 / num_bits;
211 +
212 + // Get B size
213 + int tb_k = th_config.thread_k;
214 + int tb_n = th_config.thread_n;
215 + int tb_m = thread_m_blocks * 16;
216 +
217 + // shm size for block_sorted_ids/block_topk_weights
218 + // both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
219 + int sh_block_meta_size = tb_m * 4 * 2;
220 + int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
221 + int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
222 + int sh_s_size =
223 + get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
224 + group_size, has_act_order, is_k_full);
225 + int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0;
226 + int sh_zp_size = 0;
227 + if (has_zp) {
228 + if (is_zp_float)
229 + sh_zp_size = sh_s_size;
230 + else if (num_bits == 4)
231 + sh_zp_size = sh_s_size / 4;
232 + else if (num_bits == 8)
233 + sh_zp_size = sh_s_size / 2;
234 + }
235 +
236 + int total_size = sh_a_size + sh_b_size + sh_s_size + sh_zp_size +
237 + sh_g_idx_size + sh_block_meta_size;
238 +
239 + return total_size;
240 +}
241 +
242 +bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
243 + int prob_m, int prob_n, int prob_k, int num_bits,
244 + int group_size, bool has_act_order, bool is_k_full,
245 + int has_zp, int is_zp_float, int max_shared_mem) {
246 + // Sanity
247 + if (th_config.thread_k == -1 || th_config.thread_n == -1 ||
248 + th_config.num_threads == -1) {
249 + return false;
250 + }
251 +
252 + // Verify K/N are divisible by thread K/N
253 + if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) {
254 + return false;
255 + }
256 +
257 + // Verify min for thread K/N
258 + if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) {
259 + return false;
260 + }
261 +
262 + // num_threads must be at least 128 (= 4 warps)
263 + if (th_config.num_threads < 128) {
264 + return false;
265 + }
266 +
267 + // Check that pipeline fits into cache
268 + int cache_size = get_kernel_cache_size(
269 + th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
270 + has_act_order, is_k_full, has_zp, is_zp_float);
271 + return cache_size <= max_shared_mem;
272 +}
273 +
274 + #define __GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
275 + M_BLOCK_SIZE_8, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, \
276 + NUM_THREADS, IS_ZP_FLOAT) \
277 + else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
278 + thread_n_blocks == THREAD_N_BLOCKS && \
279 + thread_k_blocks == THREAD_K_BLOCKS && \
280 + m_block_size_8 == M_BLOCK_SIZE_8 && \
281 + has_act_order == HAS_ACT_ORDER && has_zp == HAS_ZP && \
282 + group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
283 + is_zp_float == IS_ZP_FLOAT) { \
284 + kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
285 + THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
286 + pipe_stages, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS, \
287 + IS_ZP_FLOAT>; \
288 + }
289 +
290 + #define GPTQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
291 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, true, false, 0, NUM_THREADS, \
292 + false) \
293 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
294 + NUM_THREADS, false) \
295 + \
296 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, -1, \
297 + NUM_THREADS, false) \
298 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 2, \
299 + NUM_THREADS, false) \
300 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 4, \
301 + NUM_THREADS, false) \
302 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, false, 8, \
303 + NUM_THREADS, false) \
304 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
305 + NUM_THREADS, false) \
306 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
307 + NUM_THREADS, false) \
308 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
309 + NUM_THREADS, false) \
310 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
311 + NUM_THREADS, false)
312 +
313 + #define GPTQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
314 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
315 + NUM_THREADS, false) \
316 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
317 + NUM_THREADS, false) \
318 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, true, false, 0, \
319 + NUM_THREADS, false) \
320 + \
321 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
322 + NUM_THREADS, false) \
323 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
324 + NUM_THREADS, false) \
325 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
326 + NUM_THREADS, false) \
327 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
328 + NUM_THREADS, false) \
329 + \
330 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
331 + NUM_THREADS, false) \
332 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
333 + NUM_THREADS, false) \
334 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
335 + NUM_THREADS, false) \
336 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
337 + NUM_THREADS, false) \
338 + \
339 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, -1, \
340 + NUM_THREADS, false) \
341 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 2, \
342 + NUM_THREADS, false) \
343 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 4, \
344 + NUM_THREADS, false) \
345 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, false, 8, \
346 + NUM_THREADS, false)
347 +
348 + #define AWQ_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
349 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, -1, \
350 + NUM_THREADS, false) \
351 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 2, NUM_THREADS, \
352 + false) \
353 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, \
354 + false) \
355 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 8, NUM_THREADS, \
356 + false) \
357 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
358 + NUM_THREADS, false) \
359 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
360 + NUM_THREADS, false) \
361 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
362 + NUM_THREADS, false) \
363 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
364 + NUM_THREADS, false)
365 +
366 + #define AWQ_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
367 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
368 + NUM_THREADS, false) \
369 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
370 + NUM_THREADS, false) \
371 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
372 + NUM_THREADS, false) \
373 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
374 + NUM_THREADS, false) \
375 + \
376 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
377 + NUM_THREADS, false) \
378 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
379 + NUM_THREADS, false) \
380 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
381 + NUM_THREADS, false) \
382 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
383 + NUM_THREADS, false) \
384 + \
385 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, -1, \
386 + NUM_THREADS, false) \
387 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 2, \
388 + NUM_THREADS, false) \
389 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
390 + NUM_THREADS, false) \
391 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 8, \
392 + NUM_THREADS, false)
393 +
394 + // We currently have 4-bit models only with group_blocks == 4
395 + #define HQQ_GET_IF(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
396 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, false, true, 4, NUM_THREADS, \
397 + true) \
398 + __GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
399 + NUM_THREADS, true) \
400 + __GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
401 + NUM_THREADS, true) \
402 + __GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
403 + NUM_THREADS, true) \
404 + __GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, false, true, 4, \
405 + NUM_THREADS, true)
406 +
407 +template <typename scalar_t>
408 +MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
409 + int thread_m_blocks, int thread_n_blocks,
410 + int thread_k_blocks, bool m_block_size_8,
411 + bool has_act_order, bool has_zp,
412 + int group_blocks, int num_threads,
413 + bool is_zp_float) {
414 + int num_bits = q_type.size_bits();
415 + auto kernel = MarlinDefault;
416 + if (false) {
417 + }
418 + GPTQ_GET_IF_M1(vllm::kU4B8, 8, 8, 256)
419 + GPTQ_GET_IF_M1(vllm::kU4B8, 8, 4, 128)
420 +
421 + GPTQ_GET_IF_M234(vllm::kU4B8, 16, 4, 256)
422 + GPTQ_GET_IF_M234(vllm::kU4B8, 8, 4, 128)
423 +
424 + GPTQ_GET_IF_M1(vllm::kU8B128, 8, 8, 256)
425 + GPTQ_GET_IF_M1(vllm::kU8B128, 8, 4, 128)
426 +
427 + GPTQ_GET_IF_M234(vllm::kU8B128, 16, 4, 256)
428 + GPTQ_GET_IF_M234(vllm::kU8B128, 8, 4, 128)
429 +
430 + AWQ_GET_IF_M1(vllm::kU4, 8, 8, 256)
431 + AWQ_GET_IF_M1(vllm::kU4, 8, 4, 128)
432 +
433 + AWQ_GET_IF_M234(vllm::kU4, 16, 4, 256)
434 + AWQ_GET_IF_M234(vllm::kU4, 8, 4, 128)
435 +
436 + return kernel;
437 +}
438 +
439 +template <typename scalar_t>
440 +exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
441 + int prob_n, int prob_k, int thread_m_blocks,
442 + bool m_block_size_8, int num_bits,
443 + int group_size, bool has_act_order,
444 + bool is_k_full, bool has_zp,
445 + bool is_zp_float, int max_shared_mem) {
446 + exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}};
447 + thread_config_t* thread_configs = thread_m_blocks > 1
448 + ? large_batch_thread_configs
449 + : small_batch_thread_configs;
450 + int thread_configs_size =
451 + thread_m_blocks > 1
452 + ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t)
453 + : sizeof(small_batch_thread_configs) / sizeof(thread_config_t);
454 +
455 + int count = 0;
456 + constexpr int device_max_reg_size = 255 * 1024;
457 + for (int i = 0; i < thread_configs_size; i++) {
458 + thread_config_t th_config = thread_configs[i];
459 +
460 + if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k,
461 + num_bits, group_size, has_act_order, is_k_full, has_zp,
462 + is_zp_float, max_shared_mem)) {
463 + continue;
464 + }
465 +
466 + int cache_size = get_kernel_cache_size(
467 + th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits,
468 + group_size, has_act_order, is_k_full, has_zp, is_zp_float);
469 +
470 + int group_blocks = 0;
471 + if (!has_act_order) {
472 + group_blocks = group_size == -1 ? -1 : group_size / 16;
473 + }
474 +
475 + auto kernel = get_marlin_kernel<scalar_t>(
476 + q_type, thread_m_blocks, th_config.thread_n / 16,
477 + th_config.thread_k / 16, m_block_size_8, has_act_order, has_zp,
478 + group_blocks, th_config.num_threads, is_zp_float);
479 +
480 + if (kernel == MarlinDefault) continue;
481 +
482 + if (thread_m_blocks > 1) {
483 + exec_cfg = {1, th_config};
484 + break;
485 + } else {
486 + cudaFuncAttributes attr;
487 + cudaFuncGetAttributes(&attr, kernel);
488 + int reg_size = max(attr.numRegs, 1) * th_config.num_threads * 4;
489 + int allow_count = min(device_max_reg_size / reg_size,
490 + max_shared_mem / (cache_size + 1024));
491 + allow_count = max(min(allow_count, 4), 1);
492 + if (allow_count > count) {
493 + count = allow_count;
494 + exec_cfg = {count, th_config};
495 + };
496 + }
497 + }
498 +
499 + return exec_cfg;
500 +}
501 +
502 +template <typename scalar_t>
503 +void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
504 + void* zp, void* g_idx, void* perm, void* a_tmp,
505 + void* sorted_token_ids, void* expert_ids,
506 + void* num_tokens_past_padded, void* topk_weights,
507 + int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
508 + int prob_m, int prob_n, int prob_k, void* workspace,
509 + vllm::ScalarType const& q_type, bool has_act_order,
510 + bool is_k_full, bool has_zp, int num_groups, int group_size,
511 + int dev, cudaStream_t stream, int thread_k, int thread_n,
512 + int sms, bool use_atomic_add, bool use_fp32_reduce,
513 + bool is_zp_float) {
514 + int thread_m_blocks = div_ceil(moe_block_size, 16);
515 + bool m_block_size_8 = moe_block_size == 8;
516 +
517 + if (has_zp) {
518 + TORCH_CHECK(
519 + q_type == vllm::kU4 || q_type == vllm::kU8,
520 + "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
521 + } else {
522 + TORCH_CHECK(
523 + q_type == vllm::kU4B8 || q_type == vllm::kU8B128,
524 + "q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
525 + q_type.str());
526 + }
527 +
528 + TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
529 + ", ", prob_n, ", ", prob_k, "]");
530 +
531 + int group_blocks = 0;
532 + if (has_act_order) {
533 + if (is_k_full) {
534 + TORCH_CHECK(group_size != -1);
535 + group_blocks = group_size / 16;
536 + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
537 + " is not divisible by group_blocks = ", group_blocks);
538 + } else {
539 + TORCH_CHECK(group_size == 0);
540 + group_blocks = 0;
541 + }
542 + } else {
543 + if (group_size == -1) {
544 + group_blocks = -1;
545 + } else {
546 + group_blocks = group_size / 16;
547 + TORCH_CHECK(prob_k % group_blocks == 0, "prob_k = ", prob_k,
548 + " is not divisible by group_blocks = ", group_blocks);
549 + }
550 + }
551 +
552 + int num_bits = q_type.size_bits();
553 + const int4* A_ptr = (const int4*)A;
554 + const int4* B_ptr = (const int4*)B;
555 + int4* C_ptr = (int4*)C;
556 + int4* C_tmp_ptr = (int4*)C_tmp;
557 + const int4* s_ptr = (const int4*)s;
558 + const int4* zp_ptr = (const int4*)zp;
559 + const int* g_idx_ptr = (const int*)g_idx;
560 + const int* perm_ptr = (const int*)perm;
561 + int4* a_tmp_ptr = (int4*)a_tmp;
562 + const int32_t* sorted_token_ids_ptr = (const int32_t*)sorted_token_ids;
563 + const int32_t* expert_ids_ptr = (const int32_t*)expert_ids;
564 + const int32_t* num_tokens_past_padded_ptr =
565 + (const int32_t*)num_tokens_past_padded;
566 + const float* topk_weights_ptr = (const float*)topk_weights;
567 + int* locks = (int*)workspace;
568 +
569 + if (has_act_order) {
570 + // Permute A columns
571 + auto kernel = permute_cols_kernel<8>;
572 + if (moe_block_size == 8) {
573 + } else if (moe_block_size == 16)
574 + kernel = permute_cols_kernel<16>;
575 + else if (moe_block_size == 32)
576 + kernel = permute_cols_kernel<32>;
577 + else if (moe_block_size == 48)
578 + kernel = permute_cols_kernel<48>;
579 + else if (moe_block_size == 64)
580 + kernel = permute_cols_kernel<64>;
581 + else
582 + TORCH_CHECK(false, "unsupported moe_block_size ", moe_block_size);
583 +
584 + // avoid ">>>" being formatted to "> > >"
585 + // clang-format off
586 + kernel<<<sms, default_threads, 0, stream>>>(
587 + A_ptr, perm_ptr, a_tmp_ptr, sorted_token_ids_ptr, expert_ids_ptr,
588 + num_tokens_past_padded_ptr, prob_m, prob_k, top_k);
589 + // clang-format on
590 + A_ptr = a_tmp_ptr;
591 + prob_m = prob_m * top_k;
592 + top_k = 1;
593 +
594 + // If we have a full K, then we can run the non-act-order version of Marlin
595 + // (since the weight rows are reordered by increasing group ids, and by
596 + // having a full K, we have full original groups)
597 + if (is_k_full) has_act_order = false;
598 + }
599 +
600 + int max_shared_mem = 0;
601 + cudaDeviceGetAttribute(&max_shared_mem,
602 + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
603 + TORCH_CHECK(max_shared_mem > 0);
604 +
605 + // Set thread config
606 + exec_config_t exec_cfg;
607 + thread_config_t thread_tfg;
608 + if (thread_k != -1 && thread_n != -1) {
609 + thread_tfg = thread_config_t{thread_k, thread_n, default_threads};
610 + exec_cfg = exec_config_t{1, thread_tfg};
611 + TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n,
612 + " is not divisible by thread_n = ", thread_n);
613 + TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k,
614 + " is not divisible by thread_k = ", thread_k);
615 + } else {
616 + // Auto config
617 + exec_cfg = determine_exec_config<scalar_t>(
618 + q_type, prob_m, prob_n, prob_k, thread_m_blocks, m_block_size_8,
619 + num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float,
620 + max_shared_mem);
621 + thread_tfg = exec_cfg.tb_cfg;
622 + }
623 +
624 + int num_threads = thread_tfg.num_threads;
625 + thread_k = thread_tfg.thread_k;
626 + thread_n = thread_tfg.thread_n;
627 + int blocks = sms * exec_cfg.blocks_per_sm;
628 + if (exec_cfg.blocks_per_sm > 1)
629 + max_shared_mem = max_shared_mem / exec_cfg.blocks_per_sm - 1024;
630 +
631 + int thread_k_blocks = thread_k / 16;
632 + int thread_n_blocks = thread_n / 16;
633 +
634 + TORCH_CHECK(is_valid_config(thread_tfg, thread_m_blocks, prob_m, prob_n,
635 + prob_k, num_bits, group_size, has_act_order,
636 + is_k_full, has_zp, is_zp_float, max_shared_mem),
637 + "Invalid thread config: thread_m_blocks = ", thread_m_blocks,
638 + ", thread_k = ", thread_tfg.thread_k,
639 + ", thread_n = ", thread_tfg.thread_n,
640 + ", num_threads = ", thread_tfg.num_threads, " for MKN = [",
641 + prob_m, ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits,
642 + ", group_size = ", group_size,
643 + ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full,
644 + ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float,
645 + ", max_shared_mem = ", max_shared_mem);
646 +
647 + auto kernel = get_marlin_kernel<scalar_t>(
648 + q_type, thread_m_blocks, thread_n_blocks, thread_k_blocks, m_block_size_8,
649 + has_act_order, has_zp, group_blocks, num_threads, is_zp_float);
650 +
651 + if (kernel == MarlinDefault) {
652 + TORCH_CHECK(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n,
653 + ", ", prob_k, "]", ", has_act_order = ", has_act_order,
654 + ", num_groups = ", num_groups, ", group_size = ", group_size,
655 + ", thread_m_blocks = ", thread_m_blocks,
656 + ", thread_n_blocks = ", thread_n_blocks,
657 + ", thread_k_blocks = ", thread_k_blocks,
658 + ", num_bits = ", num_bits);
659 + }
660 +
661 + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
662 + max_shared_mem);
663 + // avoid ">>>" being formatted to "> > >"
664 + // clang-format off
665 + kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
666 + A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr,
667 + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
668 + topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
669 + prob_n, prob_k, locks, use_atomic_add, use_fp32_reduce);
670 + // clang-format on
671 +}
672 +
673 +} // namespace MARLIN_NAMESPACE_NAME
674 +
675 +torch::Tensor moe_wna16_marlin_gemm(
676 + torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
677 + torch::Tensor& b_q_weight, torch::Tensor& b_scales,
678 + std::optional<torch::Tensor> const& b_zeros_or_none,
679 + std::optional<torch::Tensor> const& g_idx_or_none,
680 + std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
681 + torch::Tensor& sorted_token_ids, torch::Tensor& expert_ids,
682 + torch::Tensor& num_tokens_past_padded, torch::Tensor& topk_weights,
683 + int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, bool is_ep,
684 + vllm::ScalarTypeId const& b_q_type_id, int64_t size_m, int64_t size_n,
685 + int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
686 + bool is_zp_float) {
687 + vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);
688 + int pack_factor = 32 / b_q_type.size_bits();
689 +
690 + if (moe_block_size != 8) {
691 + TORCH_CHECK(moe_block_size % 16 == 0,
692 + "unsupported moe_block_size=", moe_block_size);
693 + TORCH_CHECK(moe_block_size >= 16 && moe_block_size <= 64,
694 + "unsupported moe_block_size=", moe_block_size);
695 + }
696 +
697 + // Verify A
698 + TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0),
699 + ", size_m = ", size_m);
700 + TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1),
701 + ", size_k = ", size_k);
702 +
703 + // Verify B
704 + TORCH_CHECK(
705 + size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k,
706 + " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
707 + TORCH_CHECK((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(1),
708 + "Shape mismatch: b_q_weight.size(1) = ", b_q_weight.size(1),
709 + ", size_k = ", size_k,
710 + ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
711 + TORCH_CHECK(
712 + b_q_weight.size(2) % MARLIN_NAMESPACE_NAME::tile_size == 0,
713 + "b_q_weight.size(2) = ", b_q_weight.size(2),
714 + " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size);
715 + int actual_size_n =
716 + (b_q_weight.size(2) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor;
717 + TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n,
718 + ", actual_size_n = ", actual_size_n);
719 +
720 + // Verify device and strides
721 + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
722 + TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
723 +
724 + TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
725 + TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
726 +
727 + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
728 + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
729 +
730 + // thread_k: `k` size of a thread_tile in `weights` (can usually be left as
731 + // auto -1)
732 + int thread_k = -1;
733 + // thread_n: `n` size of a thread_tile in `weights` (can usually be left as
734 + // auto -1)
735 + int thread_n = -1;
736 + // sms: number of SMs to use for the kernel
737 + int sms = -1;
738 + cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
739 +
740 + // Alloc buffers
741 + const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
742 + auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
743 + torch::Tensor c;
744 + if (c_or_none.has_value()) {
745 + c = c_or_none.value();
746 + TORCH_CHECK(c.device().is_cuda(), "c is not on GPU");
747 + TORCH_CHECK(c.is_contiguous(), "c is not contiguous");
748 + TORCH_CHECK(c.size(0) == size_m * top_k,
749 + "Shape mismatch: c.size(0) = ", c.size(0),
750 + ", size_m * topk = ", size_m * top_k);
751 + TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1),
752 + ", size_n = ", size_n);
753 + } else {
754 + c = torch::empty({size_m * top_k, size_n}, options);
755 + }
756 +
757 + // Alloc C tmp buffer that is going to be used for the global reduce
758 + torch::Tensor c_tmp;
759 + auto options_fp32 =
760 + torch::TensorOptions().dtype(at::kFloat).device(a.device());
761 + if (use_fp32_reduce && !use_atomic_add) {
762 + // max num of threadblocks is sms * 4
763 + long max_c_tmp_size = min(
764 + (long)size_n * sorted_token_ids.size(0),
765 + (long)sms * 4 * moe_block_size * MARLIN_NAMESPACE_NAME::max_thread_n);
766 + if (moe_block_size == 8) max_c_tmp_size *= 2;
767 + c_tmp = torch::empty({max_c_tmp_size}, options_fp32);
768 + } else {
769 + c_tmp = torch::empty({0}, options_fp32);
770 + }
771 +
772 + // Detect groupsize and act_order
773 + int num_groups = -1;
774 + int group_size = -1;
775 +
776 + int rank = b_scales.sizes().size();
777 + TORCH_CHECK(rank == 3, "b_scales rank = ", rank, " is not 3");
778 + TORCH_CHECK(b_scales.size(2) == size_n, "b_scales dim 2 = ", b_scales.size(2),
779 + " is not size_n = ", size_n);
780 + num_groups = b_scales.size(1);
781 +
782 + torch::Tensor g_idx, perm, a_tmp;
783 + ;
784 + if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
785 + g_idx = g_idx_or_none.value();
786 + perm = perm_or_none.value();
787 +
788 + TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU");
789 + TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous");
790 + TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
791 + TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
792 +
793 + // Verify g_idx and perm
794 + TORCH_CHECK((g_idx.size(-1) == 0 && perm.size(-1) == 0) ||
795 + (g_idx.size(-1) == size_k && perm.size(-1) == size_k),
796 + "Unexpected g_idx.size(-1) = ", g_idx.size(-1),
797 + " and perm.size(-1) = ", perm.size(-1),
798 + ", where size_k = ", size_k);
799 + } else {
800 + g_idx = torch::empty({0}, options);
801 + perm = torch::empty({0}, options);
802 + a_tmp = torch::empty({0}, options);
803 + }
804 + bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0;
805 +
806 + if (has_act_order) {
807 + a_tmp = torch::empty({size_m * top_k, size_k}, options);
808 + if (is_k_full) {
809 + TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1");
810 + TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k,
811 + ", is not divisible by num_groups = ", num_groups);
812 + group_size = size_k / num_groups;
813 + } else {
814 + group_size = 0;
815 + }
816 +
817 + } else {
818 + a_tmp = torch::empty({0}, options);
819 + if (num_groups > 1) {
820 + TORCH_CHECK(
821 + size_k % num_groups == 0, "size_k = ", size_k,
822 + ", is not divisible by b_scales.size(1) = ", b_scales.size(1));
823 + group_size = size_k / num_groups;
824 + } else {
825 + group_size = -1;
826 + }
827 + }
828 +
829 + torch::Tensor b_zeros;
830 + if (b_zeros_or_none.has_value()) {
831 + b_zeros = b_zeros_or_none.value();
832 + TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU");
833 + TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous");
834 + } else {
835 + b_zeros = torch::empty({0}, options);
836 + }
837 + bool has_zp = b_zeros.size(-1) > 0;
838 +
839 + if (has_zp) {
840 + TORCH_CHECK(
841 + b_q_type == vllm::kU4,
842 + "b_q_type must be u4 when has_zp = True. Got = ", b_q_type.str());
843 + } else {
844 + TORCH_CHECK(
845 + b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128,
846 + "b_q_type must be uint4b8 or uint8b128 when has_zp = False. Got = ",
847 + b_q_type.str());
848 + }
849 +
850 + if (has_zp && is_zp_float) {
851 + TORCH_CHECK(a.scalar_type() == at::ScalarType::Half,
852 + "Computation type must be float16 (half) when using float zero "
853 + "points.");
854 + }
855 +
856 + // Verify b_zeros
857 + if (has_zp) {
858 + int rank = b_zeros.sizes().size();
859 + TORCH_CHECK(rank == 3, "b_zeros rank = ", rank, " is not 3");
860 + if (is_zp_float) {
861 + TORCH_CHECK(b_zeros.size(2) == size_n,
862 + "b_zeros dim 2 = ", b_zeros.size(2),
863 + " is not size_n = ", size_n);
864 + TORCH_CHECK(num_groups == b_zeros.size(1),
865 + "b_zeros dim 1 = ", b_zeros.size(1),
866 + " is not num_groups = ", num_groups);
867 + TORCH_CHECK(num_groups != -1, "num_groups must be != -1");
868 + } else {
869 + TORCH_CHECK(b_zeros.size(1) == num_groups,
870 + "b_zeros dim 1 = ", b_zeros.size(1),
871 + " is not num_groups = ", num_groups);
872 + TORCH_CHECK(b_zeros.size(2) == size_n / pack_factor,
873 + "b_zeros dim 2 = ", b_zeros.size(2),
874 + " is not size_n / pack_factor = ", size_n / pack_factor);
875 + }
876 + }
877 +
878 + // Verify workspace size
879 + TORCH_CHECK(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0,
880 + "size_n = ", size_n, ", is not divisible by min_thread_n = ",
881 + MARLIN_NAMESPACE_NAME::min_thread_n);
882 +
883 + int max_n_tiles = size_n / MARLIN_NAMESPACE_NAME::min_thread_n;
884 + int min_workspace_size = min(
885 + max_n_tiles * (int)(sorted_token_ids.size(0) / moe_block_size), sms * 4);
886 + TORCH_CHECK(workspace.numel() >= min_workspace_size,
887 + "workspace.numel = ", workspace.numel(),
888 + " is below min_workspace_size = ", min_workspace_size);
889 +
890 + int dev = a.get_device();
891 + if (a.scalar_type() == at::ScalarType::Half) {
892 + MARLIN_NAMESPACE_NAME::marlin_mm<half>(
893 + a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
894 + c_tmp.data_ptr<float>(), b_scales.data_ptr<at::Half>(),
895 + b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(),
896 + a_tmp.data_ptr<at::Half>(), sorted_token_ids.data_ptr(),
897 + expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
898 + topk_weights.data_ptr(), moe_block_size, top_k, mul_topk_weights, is_ep,
899 + size_m, size_n, size_k, workspace.data_ptr(), b_q_type, has_act_order,
900 + is_k_full, has_zp, num_groups, group_size, dev,
901 + at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
902 + use_atomic_add, use_fp32_reduce, is_zp_float);
903 + } else if (a.scalar_type() == at::ScalarType::BFloat16) {
904 + MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
905 + a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
906 + c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
907 + b_scales.data_ptr<at::BFloat16>(), b_zeros.data_ptr(), g_idx.data_ptr(),
908 + perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
909 + sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
910 + num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
911 + moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
912 + workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
913 + num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
914 + thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float);
915 + } else {
916 + TORCH_CHECK(false,
917 + "moe_wna16_marlin_gemm only supports bfloat16 and float16");
918 + }
919 +
920 + return c;
921 +}
922 +
923 +#endif
924 +
925 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
926 + m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
927 +}
csrc/moe/torch_bindings.cpp
+11 -8
@@ -43,14 +43,17 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
43 m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm);
44
45 m.def(
46 - "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
47 - "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
48 - "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, "
49 - "int b_q_type, SymInt size_m, "
50 - "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int "
51 - "topk, "
52 - "int moe_block_size, bool replicate_input, bool apply_weights)"
53 - " -> Tensor");
46 + "moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
47 + "Tensor! b_q_weight, Tensor! b_scales, Tensor? b_zeros_or_none,"
48 + "Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
49 + "Tensor sorted_token_ids,"
50 + "Tensor! expert_ids, Tensor! num_tokens_past_padded,"
51 + "Tensor! topk_weights, int moe_block_size, int top_k, "
52 + "bool mul_topk_weights, bool is_ep, int b_q_type_id,"
53 + "int size_m, int size_n, int size_k,"
54 + "bool is_full_k, bool use_atomic_add,"
55 + "bool use_fp32_reduce, bool is_zp_float) -> Tensor");
56 +
57 // conditionally compiled so impl registration is in source file
58
59 #endif
csrc/quantization/gptq_marlin/marlin.cuh
+7 -2
@@ -9,7 +9,11 @@
9 #include <cuda_runtime.h>
10 #include <iostream>
11
12 -namespace marlin {
12 +#ifndef MARLIN_NAMESPACE_NAME
13 + #define MARLIN_NAMESPACE_NAME marlin
14 +#endif
15 +
16 +namespace MARLIN_NAMESPACE_NAME {
17
18 // Marlin params
19
@@ -23,6 +27,7 @@ static constexpr int pipe_stages =
27
28 static constexpr int min_thread_n = 64;
29 static constexpr int min_thread_k = 64;
30 +static constexpr int max_thread_n = 256;
31
32 static constexpr int tile_size = 16;
33 static constexpr int max_par = 16;
@@ -84,4 +89,4 @@ __device__ inline void cp_async_wait() {
89
90 #endif
91
87 -} // namespace marlin
92 +} // namespace MARLIN_NAMESPACE_NAME
csrc/quantization/gptq_marlin/marlin_dtypes.cuh
+7 -3
@@ -5,7 +5,11 @@
5 #include <cuda_fp16.h>
6 #include <cuda_bf16.h>
7
8 -namespace marlin {
8 +#ifndef MARLIN_NAMESPACE_NAME
9 + #define MARLIN_NAMESPACE_NAME marlin
10 +#endif
11 +
12 +namespace MARLIN_NAMESPACE_NAME {
13
14 template <typename scalar_t>
15 class ScalarType {};
@@ -54,7 +58,7 @@ class ScalarType<nv_bfloat16> {
58 using FragS = Vec<nv_bfloat162, 1>;
59 using FragZP = Vec<nv_bfloat162, 4>;
60
57 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
61 +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
62 static __device__ float inline num2float(const nv_bfloat16 x) {
63 return __bfloat162float(x);
64 }
@@ -74,6 +78,6 @@ class ScalarType<nv_bfloat16> {
78 #endif
79 };
80
77 -} // namespace marlin
81 +} // namespace MARLIN_NAMESPACE_NAME
82
83 #endif
tests/kernels/test_moe.py
+137 -117
@@ -11,16 +11,14 @@ from transformers import MixtralConfig
11 from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
12
13 import vllm.model_executor.layers.fused_moe # noqa
14 -from tests.kernels.utils import (compute_max_diff, opcheck, stack_and_dev,
15 - torch_moe, torch_moe_single)
16 -from vllm import _custom_ops as ops
14 +from tests.kernels.utils import (opcheck, stack_and_dev, torch_moe,
15 + torch_moe_single)
16 from vllm.model_executor.layers.fused_moe import fused_moe
18 -from vllm.model_executor.layers.fused_moe.fused_moe import (
19 - fused_topk, moe_align_block_size)
17 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
18 from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
19 fused_moe as iterative_moe)
20 from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
23 - marlin_quantize)
21 + awq_marlin_quantize, marlin_quantize)
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 quantize_weights)
24 from vllm.model_executor.models.mixtral import MixtralMoE
@@ -287,14 +285,17 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
285 atol=mixtral_moe_tol[dtype])
286
287
290 -@pytest.mark.parametrize("m", [1, 33, 64, 222])
291 -@pytest.mark.parametrize("n", [128, 2048])
292 -@pytest.mark.parametrize("k", [128, 1024])
293 -@pytest.mark.parametrize("e", NUM_EXPERTS)
294 -@pytest.mark.parametrize("topk", TOP_KS)
288 +@pytest.mark.parametrize("m", [1, 33, 123])
289 +@pytest.mark.parametrize("n", [128, 1024])
290 +@pytest.mark.parametrize("k", [256, 2048])
291 +@pytest.mark.parametrize("e", [4, 12])
292 +@pytest.mark.parametrize("topk", [2, 3])
293 +@pytest.mark.parametrize("ep_size", [1, 4])
294 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
295 @pytest.mark.parametrize("group_size", [-1, 32, 128])
296 @pytest.mark.parametrize("act_order", [True, False])
297 @pytest.mark.parametrize("num_bits", [4, 8])
298 +@pytest.mark.parametrize("has_zp", [True, False])
299 @pytest.mark.parametrize("is_k_full", [True, False])
300 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
301 def test_fused_marlin_moe(
@@ -303,9 +304,12 @@ def test_fused_marlin_moe(
304 k: int,
305 e: int,
306 topk: int,
307 + ep_size: int,
308 + dtype: torch.dtype,
309 group_size: int,
310 act_order: bool,
311 num_bits: int,
312 + has_zp: bool,
313 is_k_full: bool,
314 ):
315 current_platform.seed_everything(7)
@@ -316,75 +320,110 @@ def test_fused_marlin_moe(
320 return
321 if group_size in (k, n):
322 return
323 + if has_zp:
324 + return
325 else:
326 if not is_k_full:
327 return
328
323 - quant_type = (scalar_types.uint4b8
324 - if num_bits == 4 else scalar_types.uint8b128)
325 - dtype = torch.float16
329 + if has_zp:
330 + # we don't build kernel for int8 with zero
331 + if num_bits == 8:
332 + return
333 + quant_type = scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
334 + else:
335 + quant_type = scalar_types.uint4b8 \
336 + if num_bits == 4 else scalar_types.uint8b128
337 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
338 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
339 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
340
341 + if ep_size > 1:
342 + local_e = e // ep_size
343 + e_ids = torch.randperm(e, device="cuda", dtype=torch.int32)[:local_e]
344 + e_map = torch.full((e, ), -1, device="cuda", dtype=torch.int32)
345 + e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
346 + w1 = w1[e_ids]
347 + w2 = w2[e_ids]
348 + else:
349 + e_map = None
350 +
351 w_ref1_l = []
352 qweight1_l = []
353 scales1_l = []
354 + zeros1_l = []
355 g_idx1_l = []
356 sort_indices1_l = []
357
358 for i in range(w1.shape[0]):
337 - test_perm = torch.randperm(k)
338 - w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = marlin_quantize(
339 - w1[i].transpose(1, 0), quant_type, group_size, act_order,
340 - test_perm)
341 - w_ref1_l.append(w_ref1)
342 - qweight1_l.append(qweight1)
343 - scales1_l.append(scales1)
344 - g_idx1_l.append(g_idx1)
345 - sort_indices1_l.append(sort_indices1)
359 + if has_zp:
360 + w_ref1, qweight1, scales1, zeros1 = awq_marlin_quantize(
361 + w1[i].transpose(1, 0), quant_type, group_size)
362 +
363 + w_ref1_l.append(w_ref1.T)
364 + qweight1_l.append(qweight1)
365 + scales1_l.append(scales1)
366 + zeros1_l.append(zeros1)
367 + else:
368 + test_perm = torch.randperm(k)
369 + quant_res = marlin_quantize(w1[i].transpose(1, 0), quant_type,
370 + group_size, act_order, test_perm)
371 + w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = quant_res
372 +
373 + w_ref1_l.append(w_ref1.T)
374 + qweight1_l.append(qweight1)
375 + scales1_l.append(scales1)
376 + g_idx1_l.append(g_idx1)
377 + sort_indices1_l.append(sort_indices1)
378
379 w_ref1 = stack_and_dev(w_ref1_l)
380 qweight1 = stack_and_dev(qweight1_l).contiguous()
381 scales1 = stack_and_dev(scales1_l)
350 - g_idx1 = stack_and_dev(g_idx1_l)
351 - sort_indices1 = stack_and_dev(sort_indices1_l)
382 + g_idx1 = stack_and_dev(g_idx1_l) if g_idx1_l else None
383 + zeros1 = stack_and_dev(zeros1_l) if zeros1_l else None
384 + sort_indices1 = stack_and_dev(sort_indices1_l) if sort_indices1_l else None
385
386 w_ref2_l = []
387 qweight2_l = []
388 scales2_l = []
389 + zeros2_l = []
390 g_idx2_l = []
391 sort_indices2_l = []
392
393 for i in range(w2.shape[0]):
360 - test_perm = torch.randperm(n)
361 - w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = marlin_quantize(
362 - w2[i].transpose(1, 0), quant_type, group_size, act_order,
363 - test_perm)
364 - w_ref2_l.append(w_ref2)
365 - qweight2_l.append(qweight2)
366 - scales2_l.append(scales2)
367 - g_idx2_l.append(g_idx2)
368 - sort_indices2_l.append(sort_indices2)
394 + if has_zp:
395 + w_ref2, qweight2, scales2, zeros2 = awq_marlin_quantize(
396 + w2[i].transpose(1, 0), quant_type, group_size)
397 +
398 + w_ref2_l.append(w_ref2.T)
399 + qweight2_l.append(qweight2)
400 + scales2_l.append(scales2)
401 + zeros2_l.append(zeros2)
402 + else:
403 + test_perm = torch.randperm(n)
404 + quant_res = marlin_quantize(w2[i].transpose(1, 0), quant_type,
405 + group_size, act_order, test_perm)
406 + w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = quant_res
407 +
408 + w_ref2_l.append(w_ref2.T)
409 + qweight2_l.append(qweight2)
410 + scales2_l.append(scales2)
411 + g_idx2_l.append(g_idx2)
412 + sort_indices2_l.append(sort_indices2)
413
414 w_ref2 = stack_and_dev(w_ref2_l)
415 qweight2 = stack_and_dev(qweight2_l).contiguous()
416 scales2 = stack_and_dev(scales2_l)
373 - g_idx2 = stack_and_dev(g_idx2_l)
374 - sort_indices2 = stack_and_dev(sort_indices2_l)
417 + g_idx2 = stack_and_dev(g_idx2_l) if g_idx2_l else None
418 + zeros2 = stack_and_dev(zeros2_l) if zeros2_l else None
419 + sort_indices2 = stack_and_dev(sort_indices2_l) if sort_indices2_l else None
420
421 score = torch.randn((m, e), device="cuda", dtype=dtype)
422
423 topk_weights, topk_ids = fused_topk(a, score, topk, False)
424
380 - triton_output = fused_moe(
381 - a,
382 - w_ref1.transpose(1, 2).contiguous(),
383 - w_ref2.transpose(1, 2).contiguous(),
384 - score,
385 - topk,
386 - renormalize=False,
387 - )
425 + torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
426 +
427 marlin_output = torch.ops.vllm.fused_marlin_moe(
428 a,
429 qweight1,
@@ -394,111 +433,91 @@ def test_fused_marlin_moe(
433 score,
434 topk_weights,
435 topk_ids,
436 + global_num_experts=e,
437 + expert_map=e_map,
438 g_idx1=g_idx1,
439 g_idx2=g_idx2,
440 sort_indices1=sort_indices1,
441 sort_indices2=sort_indices2,
442 + w1_zeros=zeros1,
443 + w2_zeros=zeros2,
444 num_bits=num_bits,
402 - is_k_full=is_k_full,
403 - )
404 -
405 - assert compute_max_diff(marlin_output, triton_output) < 4e-2
406 -
407 - if ops.supports_moe_ops:
408 - token_expert_indicies = torch.empty(m,
409 - topk,
410 - dtype=torch.int32,
411 - device=a.device)
412 -
413 - opcheck(torch.ops._moe_C.topk_softmax, (
414 - topk_weights,
415 - topk_ids,
416 - token_expert_indicies,
417 - score.float(),
418 - ))
419 -
420 - block_size_m = 4
421 -
422 - sorted_token_ids, _, _ = moe_align_block_size(topk_ids, block_size_m,
423 - e)
445 + is_k_full=is_k_full)
446
425 - max_workspace_size = ((m + 255) // 256) * (max(2 * n, k) // 64) * 16
426 - workspace = torch.zeros(max_workspace_size,
427 - dtype=torch.int,
428 - device="cuda",
429 - requires_grad=False)
430 -
431 - zp = torch.empty((0, 0),
432 - dtype=dtype,
433 - device="cuda",
434 - requires_grad=False)
435 - opcheck(torch.ops._moe_C.marlin_gemm_moe,
436 - (a, qweight1, sorted_token_ids, topk_weights, topk_ids,
437 - scales1, zp, g_idx1, sort_indices1, workspace, quant_type.id,
438 - m, 2 * n, k, True, e, topk, block_size_m, True, False))
447 + torch.testing.assert_close(marlin_output, torch_output, atol=2e-2, rtol=0)
448
449
450 @pytest.mark.skip("This test is here for the sake of debugging, "
451 "don't run it in automated tests.")
443 -@pytest.mark.parametrize("m", [64, 512, 222, 33, 1])
444 -@pytest.mark.parametrize("n", [128, 2048, 256, 1024])
445 -@pytest.mark.parametrize("k", [128, 1024, 512])
446 -@pytest.mark.parametrize("e", [8, 64])
447 -@pytest.mark.parametrize("topk", [2, 6])
448 -@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
452 +@pytest.mark.parametrize("m", [1, 33, 123])
453 +@pytest.mark.parametrize("n", [128, 1024])
454 +@pytest.mark.parametrize("k", [256, 2048])
455 +@pytest.mark.parametrize("e", [4, 12])
456 +@pytest.mark.parametrize("topk", [2, 3])
457 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
458 +@pytest.mark.parametrize("group_size", [-1, 32, 128])
459 @pytest.mark.parametrize("act_order", [True, False])
460 @pytest.mark.parametrize("num_bits", [4, 8])
461 +@pytest.mark.parametrize("has_zp", [True, False])
462 @pytest.mark.parametrize("is_k_full", [True, False])
452 -@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
453 -def test_single_marlin_moe_multiply(
454 - m: int,
455 - n: int,
456 - k: int,
457 - e: int,
458 - topk: int,
459 - group_size: int,
460 - act_order: bool,
461 - num_bits: int,
462 - is_k_full: bool,
463 -):
464 -
463 +def test_single_marlin_moe_multiply(m: int, n: int, k: int, e: int, topk: int,
464 + dtype: torch.dtype, group_size: int,
465 + act_order: bool, num_bits: int,
466 + has_zp: bool, is_k_full: bool):
467 # Filter act_order
468 if act_order:
469 if group_size == -1:
470 return
469 - if group_size == k:
471 + if group_size in (k, n):
472 + return
473 + if has_zp:
474 return
475 else:
476 if not is_k_full:
477 return
478
475 - quant_type = (scalar_types.uint4b8
476 - if num_bits == 4 else scalar_types.uint8b128)
477 - dtype = torch.float16
479 + if has_zp:
480 + quant_type = scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
481 + else:
482 + quant_type = scalar_types.uint4b8 \
483 + if num_bits == 4 else scalar_types.uint8b128
484 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
485 w = torch.randn((e, n, k), device="cuda", dtype=dtype) / 10
486
487 w_ref_l = []
482 - qweights_l = []
488 + qweight_l = []
489 scales_l = []
490 + zeros_l = []
491 g_idx_l = []
492 sort_indices_l = []
493
494 for i in range(w.shape[0]):
488 - test_perm = torch.randperm(k)
489 - w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
490 - w[i].transpose(1, 0), quant_type, group_size, act_order, test_perm)
491 - w_ref_l.append(w_ref)
492 - qweights_l.append(qweight)
493 - scales_l.append(scales)
494 - g_idx_l.append(g_idx)
495 - sort_indices_l.append(sort_indices)
495 + if has_zp:
496 + w_ref, qweight, scales, zeros = awq_marlin_quantize(
497 + w[i].transpose(1, 0), quant_type, group_size)
498 +
499 + w_ref_l.append(w_ref.T)
500 + qweight_l.append(qweight)
501 + scales_l.append(scales)
502 + zeros_l.append(zeros)
503 + else:
504 + test_perm = torch.randperm(k)
505 + w_ref, qweight, scales, g_idx, sort_indices, _ = marlin_quantize(
506 + w[i].transpose(1, 0), quant_type, group_size, act_order,
507 + test_perm)
508 +
509 + w_ref_l.append(w_ref.T)
510 + qweight_l.append(qweight)
511 + scales_l.append(scales)
512 + g_idx_l.append(g_idx)
513 + sort_indices_l.append(sort_indices)
514
515 w_ref = stack_and_dev(w_ref_l)
498 - qweight = stack_and_dev(qweights_l).contiguous()
516 + qweight = stack_and_dev(qweight_l).contiguous()
517 scales = stack_and_dev(scales_l)
500 - g_idx = stack_and_dev(g_idx_l)
501 - sort_indices = stack_and_dev(sort_indices_l)
518 + g_idx = stack_and_dev(g_idx_l) if g_idx_l else None
519 + zeros = stack_and_dev(zeros_l) if zeros_l else None
520 + sort_indices = stack_and_dev(sort_indices_l) if sort_indices_l else None
521
522 score = torch.randn((m, e), device="cuda", dtype=dtype)
523 marlin_output = torch.ops.vllm.single_marlin_moe(
@@ -510,13 +529,14 @@ def test_single_marlin_moe_multiply(
529 renormalize=False,
530 g_idx=g_idx,
531 sort_indices=sort_indices,
532 + w_zeros=zeros,
533 num_bits=num_bits,
534 is_k_full=is_k_full,
535 )
536
517 - torch_output = torch_moe_single(a, w_ref.transpose(1, 2), score, topk)
537 + torch_output = torch_moe_single(a, w_ref, score, topk)
538
519 - assert compute_max_diff(marlin_output, torch_output) < 1e-2
539 + torch.testing.assert_close(marlin_output, torch_output, atol=2e-2, rtol=0)
540
541
542 def test_moe_align_block_size_opcheck():
vllm/_custom_ops.py
+46
@@ -1245,6 +1245,29 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1245 token_expert_indicies, gating_output)
1246
1247
1248 +def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1249 + b_qweight: torch.Tensor, b_scales: torch.Tensor,
1250 + b_qzeros: Optional[torch.Tensor],
1251 + g_idx: Optional[torch.Tensor],
1252 + perm: Optional[torch.Tensor],
1253 + workspace: torch.Tensor,
1254 + sorted_token_ids: torch.Tensor,
1255 + expert_ids: torch.Tensor,
1256 + num_tokens_past_padded: torch.Tensor,
1257 + topk_weights: torch.Tensor, moe_block_size: int,
1258 + top_k: int, mul_topk_weights: bool, is_ep: bool,
1259 + b_q_type: ScalarType, size_m: int, size_n: int,
1260 + size_k: int, is_k_full: bool, use_atomic_add: bool,
1261 + use_fp32_reduce: bool,
1262 + is_zp_float: bool) -> torch.Tensor:
1263 + return torch.ops._moe_C.moe_wna16_marlin_gemm(
1264 + input, output, b_qweight, b_scales, b_qzeros, g_idx, perm, workspace,
1265 + sorted_token_ids, expert_ids, num_tokens_past_padded, topk_weights,
1266 + moe_block_size, top_k, mul_topk_weights, is_ep, b_q_type.id, size_m,
1267 + size_n, size_k, is_k_full, use_atomic_add, use_fp32_reduce,
1268 + is_zp_float)
1269 +
1270 +
1271 if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
1272
1273 @register_fake("_moe_C::marlin_gemm_moe")
@@ -1263,6 +1286,29 @@ if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
1286 dtype=a.dtype,
1287 device=a.device)
1288
1289 + @register_fake("_moe_C::moe_wna16_marlin_gemm")
1290 + def moe_wna16_marlin_gemm_fake(input: torch.Tensor,
1291 + output: Optional[torch.Tensor],
1292 + b_qweight: torch.Tensor,
1293 + b_scales: torch.Tensor,
1294 + b_qzeros: Optional[torch.Tensor],
1295 + g_idx: Optional[torch.Tensor],
1296 + perm: Optional[torch.Tensor],
1297 + workspace: torch.Tensor,
1298 + sorted_token_ids: torch.Tensor,
1299 + expert_ids: torch.Tensor,
1300 + num_tokens_past_padded: torch.Tensor,
1301 + topk_weights: torch.Tensor,
1302 + moe_block_size: int, top_k: int,
1303 + mul_topk_weights: bool, is_ep: bool,
1304 + b_q_type: ScalarType, size_m: int,
1305 + size_n: int, size_k: int, is_k_full: bool,
1306 + use_atomic_add: bool, use_fp32_reduce: bool,
1307 + is_zp_float: bool) -> torch.Tensor:
1308 + return torch.empty((size_m * top_k, size_n),
1309 + dtype=input.dtype,
1310 + device=input.device)
1311 +
1312
1313 def reshape_and_cache(
1314 key: torch.Tensor,
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+173 -160
@@ -5,17 +5,16 @@ from typing import Optional
5
6 import torch
7
8 +import vllm._custom_ops as ops
9 from vllm.model_executor.layers.fused_moe.fused_moe import (
10 fused_topk, moe_align_block_size, try_get_optimal_moe_config)
10 -from vllm.platforms import current_platform
11 from vllm.scalar_type import scalar_types
12 from vllm.utils import direct_register_custom_op
13
14
15 def get_scalar_type(num_bits: int, has_zp: bool):
16 if has_zp:
17 - assert num_bits == 4
18 - return scalar_types.uint4
17 + return scalar_types.uint4 if num_bits == 4 else scalar_types.uint8
18 else:
19 return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128
20
@@ -27,9 +26,12 @@ def single_marlin_moe(
26 gating_output: torch.Tensor,
27 topk: int,
28 renormalize: bool,
29 + global_num_experts: int = -1,
30 + expert_map: Optional[torch.Tensor] = None,
31 g_idx: Optional[torch.Tensor] = None,
32 sort_indices: Optional[torch.Tensor] = None,
33 w_zeros: Optional[torch.Tensor] = None,
34 + workspace: Optional[torch.Tensor] = None,
35 num_bits: int = 8,
36 is_k_full: bool = True,
37 ) -> torch.Tensor:
@@ -62,7 +64,7 @@ def single_marlin_moe(
64 assert gating_output.shape[1] == w.shape[0], "Number of experts mismatch"
65 assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
66 assert w.is_contiguous(), "Expert weights must be contiguous"
65 - assert hidden_states.dtype == torch.float16
67 + assert hidden_states.dtype in [torch.float16, torch.bfloat16]
68 assert num_bits in [4, 8]
69
70 M, K = hidden_states.shape
@@ -83,39 +85,54 @@ def single_marlin_moe(
85
86 block_size_m = config['BLOCK_SIZE_M']
87
86 - sorted_token_ids, _, _ = moe_align_block_size(topk_ids, block_size_m, E)
87 -
88 - max_workspace_size = (N // 64) * 16
89 - workspace = torch.zeros(max_workspace_size,
90 - dtype=torch.int,
91 - device=hidden_states.device,
92 - requires_grad=False)
93 -
94 - has_zero_point = w_zeros is not None
95 - if w_zeros is None:
96 - w_zeros = torch.empty((0, 0),
97 - dtype=hidden_states.dtype,
98 - device=hidden_states.device,
99 - requires_grad=False)
100 -
101 - if g_idx is None:
102 - g_idx = torch.empty((0, 0),
103 - dtype=torch.int32,
104 - device=hidden_states.device,
105 - requires_grad=False)
106 -
107 - if sort_indices is None:
108 - sort_indices = torch.empty((0),
109 - dtype=torch.int32,
110 - device=hidden_states.device,
111 - requires_grad=False)
112 -
113 - scalar_type = get_scalar_type(num_bits, has_zero_point)
88 + if global_num_experts == -1:
89 + global_num_experts = E
90 + sorted_token_ids, expert_ids, num_tokens_post_padded = \
91 + moe_align_block_size(topk_ids, block_size_m, E, expert_map)
92 +
93 + if workspace is None:
94 + max_workspace_size = (max(2 * N, K) // 64) * \
95 + (sorted_token_ids.size(0) // block_size_m)
96 + device = hidden_states.device
97 + sms = torch.cuda.get_device_properties(device).multi_processor_count
98 + max_workspace_size = min(max_workspace_size, sms)
99 + workspace = torch.zeros(max_workspace_size,
100 + dtype=torch.int,
101 + device=device,
102 + requires_grad=False)
103 +
104 + scalar_type = get_scalar_type(num_bits, w_zeros is not None)
105 + intermediate_cache = torch.empty(
106 + (M * topk_ids.shape[1], N),
107 + device=hidden_states.device,
108 + dtype=hidden_states.dtype,
109 + )
110
115 - intermediate_cache = torch.ops._moe_C.marlin_gemm_moe(
116 - hidden_states, w, sorted_token_ids, topk_weights, topk_ids, scales,
117 - w_zeros, g_idx, sort_indices, workspace, scalar_type.id, M, N, K,
118 - is_k_full, E, topk, block_size_m, True, False)
111 + ops.moe_wna16_marlin_gemm(hidden_states,
112 + intermediate_cache,
113 + w,
114 + scales,
115 + w_zeros,
116 + g_idx,
117 + sort_indices,
118 + workspace,
119 + sorted_token_ids,
120 + expert_ids,
121 + num_tokens_post_padded,
122 + topk_weights,
123 + moe_block_size=block_size_m,
124 + top_k=topk,
125 + mul_topk_weights=False,
126 + is_ep=expert_map is not None,
127 + b_q_type=scalar_type,
128 + size_m=M,
129 + size_n=N,
130 + size_k=K,
131 + is_k_full=is_k_full,
132 + use_atomic_add=False,
133 + use_fp32_reduce=True,
134 + is_zp_float=False)
135 + intermediate_cache = intermediate_cache.view(-1, topk, N)
136
137 return torch.sum(intermediate_cache.view(*intermediate_cache.shape), dim=1)
138
@@ -127,9 +144,12 @@ def single_marlin_moe_fake(
144 gating_output: torch.Tensor,
145 topk: int,
146 renormalize: bool,
147 + global_num_experts: int = -1,
148 + expert_map: Optional[torch.Tensor] = None,
149 g_idx: Optional[torch.Tensor] = None,
150 sort_indices: Optional[torch.Tensor] = None,
151 w_zeros: Optional[torch.Tensor] = None,
152 + workspace: Optional[torch.Tensor] = None,
153 num_bits: int = 8,
154 is_k_full: bool = True,
155 ) -> torch.Tensor:
@@ -144,24 +164,26 @@ direct_register_custom_op(
164 )
165
166
147 -def fused_marlin_moe(
148 - hidden_states: torch.Tensor,
149 - w1: torch.Tensor,
150 - w2: torch.Tensor,
151 - w1_scale: torch.Tensor,
152 - w2_scale: torch.Tensor,
153 - gating_output: torch.Tensor,
154 - topk_weights: torch.Tensor,
155 - topk_ids: torch.Tensor,
156 - g_idx1: Optional[torch.Tensor] = None,
157 - g_idx2: Optional[torch.Tensor] = None,
158 - sort_indices1: Optional[torch.Tensor] = None,
159 - sort_indices2: Optional[torch.Tensor] = None,
160 - w1_zeros: Optional[torch.Tensor] = None,
161 - w2_zeros: Optional[torch.Tensor] = None,
162 - num_bits: int = 8,
163 - is_k_full: bool = True,
164 -) -> torch.Tensor:
167 +def fused_marlin_moe(hidden_states: torch.Tensor,
168 + w1: torch.Tensor,
169 + w2: torch.Tensor,
170 + w1_scale: torch.Tensor,
171 + w2_scale: torch.Tensor,
172 + gating_output: torch.Tensor,
173 + topk_weights: torch.Tensor,
174 + topk_ids: torch.Tensor,
175 + global_num_experts: int = -1,
176 + expert_map: Optional[torch.Tensor] = None,
177 + g_idx1: Optional[torch.Tensor] = None,
178 + g_idx2: Optional[torch.Tensor] = None,
179 + sort_indices1: Optional[torch.Tensor] = None,
180 + sort_indices2: Optional[torch.Tensor] = None,
181 + w1_zeros: Optional[torch.Tensor] = None,
182 + w2_zeros: Optional[torch.Tensor] = None,
183 + workspace: Optional[torch.Tensor] = None,
184 + num_bits: int = 8,
185 + is_k_full: bool = True,
186 + inplace: bool = False) -> torch.Tensor:
187 """
188 This function computes a Mixture of Experts (MoE) layer using two sets of
189 weights, w1 and w2, and top-k gating mechanism.
@@ -196,27 +218,12 @@ def fused_marlin_moe(
218 1] == w1.shape[1] * 16, "Hidden size mismatch w1"
219 assert hidden_states.shape[1] == w2.shape[2] // (
220 num_bits // 2), "Hidden size mismatch w2"
199 - assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
221 assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
222 assert w1.is_contiguous(), "Expert weights1 must be contiguous"
223 assert w2.is_contiguous(), "Expert weights2 must be contiguous"
203 - assert hidden_states.dtype == torch.float16
224 + assert hidden_states.dtype in [torch.float16, torch.bfloat16]
225 assert num_bits in [4, 8]
226
206 - has_no_act_order = (g_idx1 is None and g_idx2 is None
207 - and sort_indices1 is None and sort_indices2 is None)
208 - has_all_act_order = (g_idx1 is not None and g_idx2 is not None
209 - and sort_indices1 is not None
210 - and sort_indices2 is not None)
211 - assert has_no_act_order or has_all_act_order, (
212 - "g_idx and sorted_indices "
213 - "must be all not None or must be all None")
214 -
215 - has_no_zp = w1_zeros is None and w2_zeros is None
216 - has_all_zp = w1_zeros is not None and w2_zeros is not None
217 - assert has_no_zp or has_all_zp, ("zero points must be both not None or "
218 - "must be both None")
219 -
227 M, K = hidden_states.shape
228 E = w1.shape[0]
229 N = w2.shape[1] * 16
@@ -234,122 +241,128 @@ def fused_marlin_moe(
241
242 block_size_m = config["BLOCK_SIZE_M"]
243
237 - sorted_token_ids, _, _ = moe_align_block_size(topk_ids, block_size_m, E)
238 -
239 - max_workspace_size = (max(2 * N, K) // 64) * 16
240 - workspace = torch.zeros(max_workspace_size,
241 - dtype=torch.int,
242 - device=current_platform.device_type,
243 - requires_grad=False)
244 -
245 - if has_no_zp:
246 - w1_zeros = torch.empty((0, 0),
247 - dtype=hidden_states.dtype,
248 - device=hidden_states.device,
249 - requires_grad=False)
250 - w2_zeros = torch.empty((0, 0),
251 - dtype=hidden_states.dtype,
252 - device=hidden_states.device,
253 - requires_grad=False)
254 -
255 - if has_no_act_order:
256 - g_idx1 = torch.empty((0, 0),
257 - dtype=torch.int32,
258 - device=hidden_states.device,
259 - requires_grad=False)
260 - g_idx2 = torch.empty((0, 0),
261 - dtype=torch.int32,
262 - device=hidden_states.device,
263 - requires_grad=False)
264 - sort_indices1 = torch.empty((0),
265 - dtype=torch.int32,
266 - device=hidden_states.device,
267 - requires_grad=False)
268 - sort_indices2 = torch.empty((0, 0),
269 - dtype=torch.int32,
270 - device=hidden_states.device,
271 - requires_grad=False)
272 -
273 - scalar_type1 = get_scalar_type(num_bits, has_all_zp)
274 - scalar_type2 = get_scalar_type(num_bits, has_all_zp)
244 + if global_num_experts == -1:
245 + global_num_experts = E
246 + sorted_token_ids, expert_ids, num_tokens_post_padded = \
247 + moe_align_block_size(topk_ids, block_size_m, global_num_experts,
248 + expert_map)
249 +
250 + if workspace is None:
251 + max_workspace_size = (max(2 * N, K) // 64) * \
252 + (sorted_token_ids.size(0) // block_size_m)
253 + device = hidden_states.device
254 + sms = torch.cuda.get_device_properties(device).multi_processor_count
255 + max_workspace_size = min(max_workspace_size, sms * 4)
256 + workspace = torch.zeros(max_workspace_size,
257 + dtype=torch.int,
258 + device=device,
259 + requires_grad=False)
260 +
261 + scalar_type1 = get_scalar_type(num_bits, w1_zeros is not None)
262 + scalar_type2 = get_scalar_type(num_bits, w2_zeros is not None)
263
264 intermediate_cache2 = torch.empty(
265 (M * topk_ids.shape[1], N),
266 device=hidden_states.device,
267 dtype=hidden_states.dtype,
268 )
269 + intermediate_cache13 = torch.empty(
270 + (M * topk_ids.shape[1] * max(2 * N, K), ),
271 + device=hidden_states.device,
272 + dtype=hidden_states.dtype,
273 + )
274 + intermediate_cache1 = intermediate_cache13[:M * topk_ids.shape[1] * 2 * N]
275 + intermediate_cache1 = intermediate_cache1.view(-1, 2 * N)
276 + intermediate_cache3 = intermediate_cache13[:M * topk_ids.shape[1] * K]
277 + intermediate_cache3 = intermediate_cache3.view(-1, K)
278 +
279 + use_atomic_add = hidden_states.dtype == torch.half or \
280 + torch.cuda.get_device_capability(hidden_states.device)[0] >= 9
281
282 - intermediate_cache1 = torch.ops._moe_C.marlin_gemm_moe(
282 + intermediate_cache1 = ops.moe_wna16_marlin_gemm(
283 hidden_states,
284 + intermediate_cache1,
285 w1,
285 - sorted_token_ids,
286 - topk_weights,
287 - topk_ids,
286 w1_scale,
287 w1_zeros,
288 g_idx1,
289 sort_indices1,
290 workspace,
293 - scalar_type1.id,
294 - M,
295 - 2 * N,
296 - K,
297 - is_k_full,
298 - E,
299 - topk,
300 - block_size_m,
301 - True,
302 - False,
303 - )
291 + sorted_token_ids,
292 + expert_ids,
293 + num_tokens_post_padded,
294 + topk_weights,
295 + moe_block_size=block_size_m,
296 + top_k=topk,
297 + mul_topk_weights=False,
298 + is_ep=expert_map is not None,
299 + b_q_type=scalar_type1,
300 + size_m=M,
301 + size_n=2 * N,
302 + size_k=K,
303 + is_k_full=is_k_full,
304 + use_atomic_add=use_atomic_add,
305 + use_fp32_reduce=True,
306 + is_zp_float=False)
307
308 torch.ops._C.silu_and_mul(intermediate_cache2,
309 intermediate_cache1.view(-1, 2 * N))
310
308 - intermediate_cache3 = torch.ops._moe_C.marlin_gemm_moe(
311 + if expert_map is not None:
312 + intermediate_cache3.zero_()
313 +
314 + intermediate_cache3 = ops.moe_wna16_marlin_gemm(
315 intermediate_cache2,
316 + intermediate_cache3,
317 w2,
311 - sorted_token_ids,
312 - topk_weights,
313 - topk_ids,
318 w2_scale,
319 w2_zeros,
320 g_idx2,
321 sort_indices2,
322 workspace,
319 - scalar_type2.id,
320 - M,
321 - K,
322 - N,
323 - is_k_full,
324 - E,
325 - topk,
326 - block_size_m,
327 - False,
328 - True,
329 - )
330 -
323 + sorted_token_ids,
324 + expert_ids,
325 + num_tokens_post_padded,
326 + topk_weights,
327 + moe_block_size=block_size_m,
328 + top_k=1,
329 + mul_topk_weights=True,
330 + is_ep=expert_map is not None,
331 + b_q_type=scalar_type2,
332 + size_m=M * topk,
333 + size_n=K,
334 + size_k=N,
335 + is_k_full=is_k_full,
336 + use_atomic_add=use_atomic_add,
337 + use_fp32_reduce=True,
338 + is_zp_float=False).view(-1, topk, K)
339 +
340 + output = hidden_states if inplace else torch.empty_like(hidden_states)
341 return torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
332 - dim=1)
333 -
334 -
335 -def fused_marlin_moe_fake(
336 - hidden_states: torch.Tensor,
337 - w1: torch.Tensor,
338 - w2: torch.Tensor,
339 - w1_scale: torch.Tensor,
340 - w2_scale: torch.Tensor,
341 - gating_output: torch.Tensor,
342 - topk_weights: torch.Tensor,
343 - topk_ids: torch.Tensor,
344 - g_idx1: Optional[torch.Tensor] = None,
345 - g_idx2: Optional[torch.Tensor] = None,
346 - sort_indices1: Optional[torch.Tensor] = None,
347 - sort_indices2: Optional[torch.Tensor] = None,
348 - w1_zeros: Optional[torch.Tensor] = None,
349 - w2_zeros: Optional[torch.Tensor] = None,
350 - num_bits: int = 8,
351 - is_k_full: bool = True,
352 -) -> torch.Tensor:
342 + dim=1,
343 + out=output)
344 +
345 +
346 +def fused_marlin_moe_fake(hidden_states: torch.Tensor,
347 + w1: torch.Tensor,
348 + w2: torch.Tensor,
349 + w1_scale: torch.Tensor,
350 + w2_scale: torch.Tensor,
351 + gating_output: torch.Tensor,
352 + topk_weights: torch.Tensor,
353 + topk_ids: torch.Tensor,
354 + global_num_experts: int = -1,
355 + expert_map: Optional[torch.Tensor] = None,
356 + g_idx1: Optional[torch.Tensor] = None,
357 + g_idx2: Optional[torch.Tensor] = None,
358 + sort_indices1: Optional[torch.Tensor] = None,
359 + sort_indices2: Optional[torch.Tensor] = None,
360 + w1_zeros: Optional[torch.Tensor] = None,
361 + w2_zeros: Optional[torch.Tensor] = None,
362 + workspace: Optional[torch.Tensor] = None,
363 + num_bits: int = 8,
364 + is_k_full: bool = True,
365 + inplace: bool = False) -> torch.Tensor:
366 return torch.empty_like(hidden_states)
367
368
vllm/model_executor/layers/fused_moe/fused_moe.py
+12 -8
@@ -773,6 +773,18 @@ def get_default_config(
773 config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
774 else:
775 config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
776 + elif is_marlin:
777 + for block_size_m in [8, 16, 32, 48, 64]:
778 + if M * topk / E / block_size_m < 0.9:
779 + break
780 + return {"BLOCK_SIZE_M": block_size_m}
781 + elif M <= E:
782 + config = {
783 + "BLOCK_SIZE_M": 16,
784 + "BLOCK_SIZE_N": 32,
785 + "BLOCK_SIZE_K": 64,
786 + "GROUP_SIZE_M": 1,
787 + }
788 else:
789 config = {
790 "BLOCK_SIZE_M": 64,
@@ -780,14 +792,6 @@ def get_default_config(
792 "BLOCK_SIZE_K": 32,
793 "GROUP_SIZE_M": 8,
794 }
783 - # A heuristic: fused marlin works faster with this config for small M
784 - if M <= E or (is_marlin and M <= 32):
785 - config = {
786 - "BLOCK_SIZE_M": 16,
787 - "BLOCK_SIZE_N": 32,
788 - "BLOCK_SIZE_K": 64,
789 - "GROUP_SIZE_M": 1,
790 - }
795 return config
796
797
vllm/model_executor/layers/fused_moe/layer.py
+1
@@ -472,6 +472,7 @@ class FusedMoE(torch.nn.Module):
472 self.global_num_experts = num_experts
473
474 assert intermediate_size % self.tp_size == 0
475 + self.hidden_size = hidden_size
476 self.intermediate_size_per_partition = intermediate_size // self.tp_size
477 self.reduce_results = reduce_results
478 self.renormalize = renormalize
vllm/model_executor/layers/quantization/awq_marlin.py
+22 -13
@@ -17,14 +17,13 @@ from vllm.model_executor.layers.quantization.awq import (AWQConfig,
17 is_layer_skipped_awq)
18 from vllm.model_executor.layers.quantization.base_config import (
19 QuantizationConfig, QuantizeMethodBase)
20 -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
20 from vllm.model_executor.layers.quantization.utils import replace_parameter
21 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
22 apply_awq_marlin_linear, awq_to_marlin_zero_points, check_marlin_supported,
24 - check_marlin_supports_layer, marlin_make_empty_g_idx,
25 - marlin_make_workspace, marlin_moe_permute_scales, marlin_permute_scales,
26 - moe_awq_to_marlin_zero_points, verify_marlin_supported,
27 - verify_marlin_supports_shape)
23 + check_marlin_supports_layer, check_moe_marlin_supports_layer,
24 + marlin_make_empty_g_idx, marlin_make_workspace, marlin_moe_permute_scales,
25 + marlin_permute_scales, moe_awq_to_marlin_zero_points,
26 + verify_marlin_supported, verify_marlin_supports_shape)
27 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
28 from vllm.model_executor.parameter import (GroupQuantScaleParameter,
29 PackedvLLMParameter)
@@ -136,12 +135,15 @@ class AWQMarlinConfig(QuantizationConfig):
135 self.full_config).get_quant_method(layer, prefix)
136 return AWQMarlinLinearMethod(self)
137 elif isinstance(layer, FusedMoE):
139 - if layer.local_num_experts > 32:
140 - # For MoEs with many experts the moe_wna16 kernel is faster
138 + from vllm.model_executor.layers.quantization.moe_wna16 import (
139 + MoeWNA16Config)
140 + if not check_moe_marlin_supports_layer(layer, self.group_size):
141 + logger.warning_one(
142 + f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
143 + "Falling back to Moe WNA16 kernels.")
144 return MoeWNA16Config.from_config(
145 self.full_config).get_quant_method(layer, prefix)
143 - else:
144 - return AWQMoEMethod(self)
146 + return AWQMoEMethod(self)
147 return None
148
149 @classmethod
@@ -391,6 +393,13 @@ class AWQMoEMethod(FusedMoEMethodBase):
393 layer.register_parameter("w2_qzeros", w2_qzeros)
394 set_weight_attrs(w2_qzeros, extra_weight_attrs)
395
396 + device = layer.w13_qweight.device
397 + sms = torch.cuda.get_device_properties(device).multi_processor_count
398 + layer.workspace = torch.zeros((sms * 4, ),
399 + dtype=torch.int,
400 + device=device,
401 + requires_grad=False)
402 +
403 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
404 num_experts = layer.w13_qweight.shape[0]
405 device = layer.w13_qweight.device
@@ -473,10 +482,7 @@ class AWQMoEMethod(FusedMoEMethodBase):
482 activation: str = "silu",
483 ) -> torch.Tensor:
484 assert activation == "silu", "Only SiLU activation is supported."
476 - if expert_map is not None:
477 - raise NotImplementedError(
478 - "Expert Parallelism is not supported for "
479 - "fused Marlin MoE method.")
485 +
486 if apply_router_weight_on_input:
487 raise NotImplementedError(
488 "Apply router weight on input is not supported for"
@@ -503,7 +509,10 @@ class AWQMoEMethod(FusedMoEMethodBase):
509 router_logits,
510 topk_weights,
511 topk_ids,
512 + global_num_experts=global_num_experts,
513 + expert_map=expert_map,
514 w1_zeros=layer.w13_qzeros,
515 w2_zeros=layer.w2_qzeros,
516 + workspace=layer.workspace,
517 num_bits=self.quant_config.weight_bits,
518 )
vllm/model_executor/layers/quantization/gptq_marlin.py
+23 -14
@@ -15,13 +15,13 @@ from vllm.model_executor.layers.quantization.base_config import (
15 QuantizationConfig, QuantizeMethodBase)
16 from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
17 MPLinearLayerConfig, choose_mp_linear_kernel)
18 -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
18 from vllm.model_executor.layers.quantization.utils import replace_parameter
19 from vllm.model_executor.layers.quantization.utils.gptq_utils import (
20 get_linear_quant_method)
21 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
23 - check_marlin_supported, marlin_moe_permute_scales,
24 - marlin_repeat_scales_on_all_ranks, verify_marlin_supported)
22 + check_marlin_supported, check_moe_marlin_supports_layer,
23 + marlin_moe_permute_scales, marlin_repeat_scales_on_all_ranks,
24 + verify_marlin_supported)
25 from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
26 GroupQuantScaleParameter,
27 PackedColumnParameter,
@@ -153,12 +153,15 @@ class GPTQMarlinConfig(QuantizationConfig):
153 def get_quant_method(self, layer: torch.nn.Module,
154 prefix: str) -> Optional["QuantizeMethodBase"]:
155 if isinstance(layer, FusedMoE):
156 - if layer.local_num_experts > 32:
157 - # For MoEs with many experts the moe_wna16 kernel is faster
156 + from vllm.model_executor.layers.quantization.moe_wna16 import (
157 + MoeWNA16Config)
158 + if not check_moe_marlin_supports_layer(layer, self.group_size):
159 + logger.warning_one(
160 + f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
161 + "Falling back to Moe WNA16 kernels.")
162 return MoeWNA16Config.from_config(
163 self.full_config).get_quant_method(layer, prefix)
160 - else:
161 - return GPTQMarlinMoEMethod(self)
164 + return GPTQMarlinMoEMethod(self)
165 return get_linear_quant_method(self, layer, prefix,
166 GPTQMarlinLinearMethod)
167
@@ -408,7 +411,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
411 torch.empty(num_experts,
412 scales_size13,
413 2 * intermediate_size_per_partition,
411 - dtype=torch.half),
414 + dtype=params_dtype),
415 requires_grad=False,
416 )
417 layer.register_parameter("w13_scales", w13_scales)
@@ -418,7 +421,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
421 torch.empty(num_experts,
422 scales_size2,
423 hidden_size,
421 - dtype=torch.half),
424 + dtype=params_dtype),
425 requires_grad=False,
426 )
427 layer.register_parameter("w2_scales", w2_scales)
@@ -493,6 +496,13 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
496 w2_g_idx_sort_indices)
497 set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
498
499 + device = layer.w13_qweight.device
500 + sms = torch.cuda.get_device_properties(device).multi_processor_count
501 + layer.workspace = torch.zeros((sms * 4, ),
502 + dtype=torch.int,
503 + device=device,
504 + requires_grad=False)
505 +
506 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
507
508 # Process act_order
@@ -601,10 +611,6 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
611 "Apply router weight on input is not supported for"
612 "fused Marlin MoE method.")
613
604 - # The input must currently be float16
605 - orig_dtype = x.dtype
606 - x = x.half()
607 -
614 topk_weights, topk_ids = FusedMoE.select_experts(
615 hidden_states=x,
616 router_logits=router_logits,
@@ -626,9 +632,12 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
632 router_logits,
633 topk_weights,
634 topk_ids,
635 + global_num_experts=global_num_experts,
636 + expert_map=expert_map,
637 g_idx1=layer.w13_g_idx,
638 g_idx2=layer.w2_g_idx,
639 sort_indices1=layer.w13_g_idx_sort_indices,
640 sort_indices2=layer.w2_g_idx_sort_indices,
641 num_bits=self.quant_config.quant_type.size_bits,
634 - is_k_full=self.is_k_full).to(orig_dtype)
642 + workspace=layer.workspace,
643 + is_k_full=self.is_k_full)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+13
@@ -151,6 +151,19 @@ def check_marlin_supports_layer(layer: LinearBase, group_size: int) \
151 group_size=group_size)[0]
152
153
154 +def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) \
155 + -> bool:
156 + hidden_size = layer.hidden_size
157 + intermediate_size_per_partition = layer.intermediate_size_per_partition
158 +
159 + # gate-up: (n, k) = (intermediate_size_per_partition * 2, hidden_size)
160 + # down: (n, k) = (hidden_size, intermediate_size_per_partition)
161 + # moe marlin requires n % 128 == 0 and k % 64 == 0
162 + return hidden_size % 128 == 0 and \
163 + intermediate_size_per_partition % max(64, group_size) == 0 and \
164 + group_size in [-1, 32, 64, 128]
165 +
166 +
167 def marlin_make_workspace(output_size_per_partition: int,
168 device: torch.device) -> torch.Tensor:
169 max_workspace_size = (output_size_per_partition //