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, 19:45 UTC
[Kernel] moe wna16 cuda kernel (#13321)
Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Jinzhen Lin committed
Mar 11, 2025 at 08:12 UTC
90e88ab75632745c137647bf710d63997529fb89
7 files changed
+698
-1
CMakeLists.txt
+8
@@ -558,6 +558,7 @@ target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
558
set(VLLM_MOE_EXT_SRC
559
"csrc/moe/torch_bindings.cpp"
560
"csrc/moe/moe_align_sum_kernels.cu"
561
+ "csrc/moe/moe_wna16.cu"
562
"csrc/moe/topk_softmax_kernels.cu")
563
564
set_gencode_flags_for_srcs(
@@ -565,6 +566,13 @@ set_gencode_flags_for_srcs(
566
CUDA_ARCHS "${CUDA_ARCHS}")
567
568
if(VLLM_GPU_LANG STREQUAL "CUDA")
569
+ set(VLLM_MOE_WNA16_SRC
570
+ "csrc/moe/moe_wna16.cu")
571
+
572
+ set_gencode_flags_for_srcs(
573
+ SRCS "${VLLM_MOE_WNA16_SRC}"
574
+ CUDA_ARCHS "${CUDA_ARCHS}")
575
+
576
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}")
577
if (MARLIN_MOE_ARCHS)
578
set(MARLIN_MOE_SRC
csrc/moe/moe_ops.h
+10
@@ -18,3 +18,13 @@ void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
18
torch::Tensor sorted_token_ids,
19
torch::Tensor experts_ids,
20
torch::Tensor num_tokens_post_pad);
21
+
22
+torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
23
+ torch::Tensor b_qweight, torch::Tensor b_scales,
24
+ std::optional<torch::Tensor> b_qzeros,
25
+ std::optional<torch::Tensor> topk_weights,
26
+ torch::Tensor sorted_token_ids,
27
+ torch::Tensor expert_ids,
28
+ torch::Tensor num_tokens_post_pad, int64_t top_k,
29
+ int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
30
+ int64_t BLOCK_SIZE_K, int64_t bit);
csrc/moe/moe_wna16.cu
new
+346
@@ -0,0 +1,346 @@
1
+
2
+#include <torch/all.h>
3
+#include <c10/cuda/CUDAGuard.h>
4
+#include <ATen/cuda/CUDAContext.h>
5
+#include <cuda_runtime.h>
6
+
7
+#include <cuda_fp16.h>
8
+#include <cuda_bf16.h>
9
+#include "moe_wna16_utils.h"
10
+
11
+#define DIVIDE(x, size) (((x) + (size) - 1) / (size))
12
+
13
+template <typename scalar_t, int bit, int GROUPS>
14
+__global__ void moe_wna16_gemm_kernel(
15
+ const scalar_t* __restrict__ input, scalar_t* __restrict__ output,
16
+
17
+ const uint32_t* __restrict__ qweight, const scalar_t* __restrict__ scales,
18
+ const uint32_t* __restrict__ qzeros,
19
+
20
+ const float* __restrict__ topk_weights,
21
+ const int32_t* __restrict__ sorted_token_ids,
22
+ const int32_t* __restrict__ expert_ids,
23
+ const int32_t* __restrict__ num_tokens_post_pad,
24
+
25
+ uint16_t num_experts, uint16_t group_size, uint16_t top_k, uint32_t size_m,
26
+ uint32_t size_n, uint32_t size_k, uint16_t BLOCK_SIZE_M,
27
+ uint16_t BLOCK_SIZE_N, uint16_t BLOCK_SIZE_K, bool has_zp,
28
+ bool mul_topk_weight) {
29
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
30
+ if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
31
+ return;
32
+ } else {
33
+#endif
34
+
35
+ using Dtype = ScalarType<scalar_t>;
36
+ using scalar_t2 = typename ScalarType<scalar_t>::scalar_t2;
37
+
38
+ if (blockIdx.x * BLOCK_SIZE_M >= num_tokens_post_pad[0]) return;
39
+
40
+ const int32_t offset_n = blockIdx.y * BLOCK_SIZE_N + threadIdx.x;
41
+ const int32_t offset_k = blockIdx.z * BLOCK_SIZE_K;
42
+
43
+ const int32_t expert_id = expert_ids[blockIdx.x];
44
+
45
+ int32_t num_valid_tokens = 0;
46
+ extern __shared__ uint16_t block_input_tmp[];
47
+ scalar_t* block_input = reinterpret_cast<scalar_t*>(block_input_tmp);
48
+ scalar_t2* block_input_half2 = reinterpret_cast<scalar_t2*>(block_input);
49
+
50
+ // load BLOCK_SIZE_M * BLOCK_SIZE_K into shared memory
51
+ for (int m = 0; m < BLOCK_SIZE_M; m++) {
52
+ const int32_t offset_m = blockIdx.x * BLOCK_SIZE_M + m;
53
+ const int32_t token_index = sorted_token_ids[offset_m];
54
+ if (token_index / top_k >= size_m) break;
55
+
56
+ num_valid_tokens = m + 1;
57
+ if (blockIdx.z == 0 && offset_n < size_n)
58
+ output[token_index * size_n + offset_n] = Dtype::int2num(0);
59
+
60
+ if (expert_id != -1) {
61
+ int k_per_thread = DIVIDE(BLOCK_SIZE_K, BLOCK_SIZE_N);
62
+ for (int i = 0; i < k_per_thread; i++) {
63
+ int k = BLOCK_SIZE_N * i + threadIdx.x;
64
+ if (k >= BLOCK_SIZE_K) break;
65
+ if (offset_k + k >= size_k) break;
66
+
67
+ // load input to shared memory
68
+ // use a special layout to fit the layout of dequanted-weight
69
+ int origin_k;
70
+ if constexpr (bit == 4) {
71
+ // [0, 4, 1, 5, 2, 6, 3, 7]
72
+ int8_t order = (threadIdx.x % 2) * 4 + ((threadIdx.x % 8) / 2);
73
+ origin_k = BLOCK_SIZE_N * i + threadIdx.x / 8 * 8 + order;
74
+ } else {
75
+ // [0, 2, 1, 3]
76
+ int8_t order = (threadIdx.x % 2) * 2 + ((threadIdx.x % 4) / 2);
77
+ origin_k = BLOCK_SIZE_N * i + threadIdx.x / 4 * 4 + order;
78
+ }
79
+
80
+ origin_k += token_index / top_k * size_k + blockIdx.z * BLOCK_SIZE_K;
81
+ block_input[m * BLOCK_SIZE_K + k] = input[origin_k];
82
+ }
83
+ }
84
+ }
85
+
86
+ if (expert_id == -1) return;
87
+ __syncthreads();
88
+ if (threadIdx.x >= BLOCK_SIZE_N || offset_n >= size_n) return;
89
+
90
+ float res[64]; // assume BLOCK_SIZE_M <= 64
91
+ scalar_t2 res2;
92
+ scalar_t2 scale_f2;
93
+ scalar_t2 qzero_f2;
94
+
95
+ // note that (size_n * size_k * expert_id) may greater than 2 ** 31
96
+ constexpr int8_t pack_factor = 32 / bit;
97
+ const uint64_t expert_offset = ((uint64_t)size_n) * size_k * expert_id;
98
+ const uint32_t* expert_qweight = qweight + expert_offset / pack_factor;
99
+ const scalar_t* expert_scales = scales + expert_offset / group_size;
100
+ const uint32_t* expert_qzeros =
101
+ qzeros + expert_offset / group_size / pack_factor;
102
+
103
+ // load 4*int32 one time: 4 int32 = 128 bit = 1 float4
104
+ // weight would be loaded in loop
105
+ uint32_t expert_qweight_tmp[4];
106
+ float4* expert_qweight_tmp_float4 =
107
+ reinterpret_cast<float4*>(expert_qweight_tmp);
108
+
109
+ // load all required scales one time
110
+ scalar_t expert_scales_groups[GROUPS];
111
+ int scales_offset_tmp =
112
+ (offset_n * size_k + offset_k) / group_size / GROUPS;
113
+ if constexpr (GROUPS == 1) {
114
+ *expert_scales_groups = expert_scales[scales_offset_tmp];
115
+ } else if constexpr (GROUPS == 2) {
116
+ float* expert_scales_groups_tmp =
117
+ reinterpret_cast<float*>(expert_scales_groups);
118
+ *expert_scales_groups_tmp =
119
+ reinterpret_cast<const float*>(expert_scales)[scales_offset_tmp];
120
+ } else if constexpr (GROUPS == 4) {
121
+ float2* expert_scales_groups_tmp =
122
+ reinterpret_cast<float2*>(expert_scales_groups);
123
+ *expert_scales_groups_tmp =
124
+ reinterpret_cast<const float2*>(expert_scales)[scales_offset_tmp];
125
+ } else if constexpr (GROUPS == 8) {
126
+ float4* expert_scales_groups_tmp =
127
+ reinterpret_cast<float4*>(expert_scales_groups);
128
+ *expert_scales_groups_tmp =
129
+ reinterpret_cast<const float4*>(expert_scales)[scales_offset_tmp];
130
+ }
131
+
132
+ // load all required qzeros one time
133
+ uint8_t expert_qzeros_groups[GROUPS];
134
+ if (!has_zp) {
135
+ if constexpr (bit == 4) {
136
+ qzero_f2 = Dtype::num2num2(Dtype::int2num(8));
137
+ } else {
138
+ qzero_f2 = Dtype::num2num2(Dtype::int2num(128));
139
+ }
140
+ } else {
141
+ int qzeros_offset_tmp =
142
+ (offset_n / (8 / bit)) * (size_k / group_size / GROUPS) +
143
+ offset_k / group_size / GROUPS;
144
+ if constexpr (GROUPS == 1) {
145
+ uint8_t* expert_qzeros_groups_tmp =
146
+ reinterpret_cast<uint8_t*>(expert_qzeros_groups);
147
+ *expert_qzeros_groups_tmp =
148
+ reinterpret_cast<const uint8_t*>(expert_qzeros)[qzeros_offset_tmp];
149
+ } else if constexpr (GROUPS == 2) {
150
+ uint16_t* expert_qzeros_groups_tmp =
151
+ reinterpret_cast<uint16_t*>(expert_qzeros_groups);
152
+ *expert_qzeros_groups_tmp =
153
+ reinterpret_cast<const uint16_t*>(expert_qzeros)[qzeros_offset_tmp];
154
+ } else if constexpr (GROUPS == 4) {
155
+ uint32_t* expert_qzeros_groups_tmp =
156
+ reinterpret_cast<uint32_t*>(expert_qzeros_groups);
157
+ *expert_qzeros_groups_tmp =
158
+ reinterpret_cast<const uint32_t*>(expert_qzeros)[qzeros_offset_tmp];
159
+ } else if constexpr (GROUPS == 8) {
160
+ uint64_t* expert_qzeros_groups_tmp =
161
+ reinterpret_cast<uint64_t*>(expert_qzeros_groups);
162
+ *expert_qzeros_groups_tmp =
163
+ reinterpret_cast<const uint64_t*>(expert_qzeros)[qzeros_offset_tmp];
164
+ }
165
+ }
166
+
167
+ for (int tmp_k = 0; tmp_k < BLOCK_SIZE_K / pack_factor; tmp_k++) {
168
+ int k = offset_k + tmp_k * pack_factor;
169
+ if (k >= size_k) break;
170
+ const int32_t weight_offset = offset_n * size_k + k;
171
+
172
+ if (tmp_k % 4 == 0) {
173
+ *expert_qweight_tmp_float4 = reinterpret_cast<const float4*>(
174
+ expert_qweight)[weight_offset / pack_factor / 4];
175
+ }
176
+
177
+ if (tmp_k % (group_size / pack_factor) == 0) {
178
+ scalar_t scale_f =
179
+ expert_scales_groups[tmp_k / (group_size / pack_factor)];
180
+ scale_f2 = Dtype::num2num2(scale_f);
181
+
182
+ if (has_zp) {
183
+ uint8_t qzero =
184
+ expert_qzeros_groups[tmp_k / (group_size / pack_factor)];
185
+ if constexpr (bit == 4) {
186
+ qzero = (qzero >> ((threadIdx.x % 2) * 4)) & 0xF;
187
+ }
188
+ qzero_f2 = Dtype::num2num2(Dtype::int2num(qzero));
189
+ }
190
+ }
191
+
192
+ scalar_t2 weight_half2[16 / bit];
193
+ dequant<scalar_t2, bit>(expert_qweight_tmp[tmp_k % 4], weight_half2);
194
+
195
+ for (int m = 0; m < num_valid_tokens; m++) {
196
+ res2 = {};
197
+
198
+#pragma unroll
199
+ for (int i = 0; i < 16 / bit; i++) {
200
+ int32_t offset_input = m * BLOCK_SIZE_K / 2 + tmp_k * (16 / bit) + i;
201
+ res2 = __hfma2(__hmul2(__hsub2(weight_half2[i], qzero_f2), scale_f2),
202
+ block_input_half2[offset_input], res2);
203
+ }
204
+
205
+ if (tmp_k == 0) {
206
+ res[m] = Dtype::num2float(res2.x) + Dtype::num2float(res2.y);
207
+ } else {
208
+ res[m] += Dtype::num2float(res2.x) + Dtype::num2float(res2.y);
209
+ }
210
+ }
211
+ }
212
+
213
+ for (int m = 0; m < num_valid_tokens; ++m) {
214
+ const int32_t token_index =
215
+ sorted_token_ids[blockIdx.x * BLOCK_SIZE_M + m];
216
+ if (mul_topk_weight) {
217
+ res[m] *= topk_weights[token_index];
218
+ }
219
+ atomicAdd(&output[token_index * size_n + offset_n],
220
+ Dtype::float2num(res[m]));
221
+ }
222
+
223
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
224
+ }
225
+#endif
226
+}
227
+
228
+template <typename scalar_t>
229
+void run_moe_wna16_gemm(const scalar_t* input, scalar_t* output,
230
+ const uint32_t* b_qweight, const scalar_t* b_scales,
231
+ const uint32_t* b_qzeros, const float* topk_weights,
232
+ const int32_t* sorted_token_ids,
233
+ const int32_t* expert_ids,
234
+ const int32_t* num_tokens_post_pad, int num_experts,
235
+ int group_size, int num_token_blocks, int top_k,
236
+ int size_m, int size_n, int size_k, int BLOCK_SIZE_M,
237
+ int BLOCK_SIZE_N, int BLOCK_SIZE_K, int bit,
238
+ bool has_zp, bool mul_topk_weight) {
239
+ dim3 blockDim, gridDim;
240
+ blockDim.x = BLOCK_SIZE_N;
241
+ blockDim.y = 1;
242
+ blockDim.z = 1;
243
+ gridDim.x = num_token_blocks;
244
+ gridDim.y = DIVIDE(size_n, BLOCK_SIZE_N);
245
+ gridDim.z = DIVIDE(size_k, BLOCK_SIZE_K);
246
+
247
+ auto kernel = moe_wna16_gemm_kernel<scalar_t, 4, 1>;
248
+ if (bit == 4) {
249
+ if (BLOCK_SIZE_K / group_size == 2) {
250
+ kernel = moe_wna16_gemm_kernel<scalar_t, 4, 2>;
251
+ } else if (BLOCK_SIZE_K / group_size == 4) {
252
+ kernel = moe_wna16_gemm_kernel<scalar_t, 4, 4>;
253
+ } else if (BLOCK_SIZE_K / group_size == 8) {
254
+ kernel = moe_wna16_gemm_kernel<scalar_t, 4, 8>;
255
+ }
256
+ } else {
257
+ if (BLOCK_SIZE_K / group_size == 1) {
258
+ kernel = moe_wna16_gemm_kernel<scalar_t, 8, 1>;
259
+ } else if (BLOCK_SIZE_K / group_size == 2) {
260
+ kernel = moe_wna16_gemm_kernel<scalar_t, 8, 2>;
261
+ } else if (BLOCK_SIZE_K / group_size == 4) {
262
+ kernel = moe_wna16_gemm_kernel<scalar_t, 8, 4>;
263
+ } else if (BLOCK_SIZE_K / group_size == 8) {
264
+ kernel = moe_wna16_gemm_kernel<scalar_t, 8, 8>;
265
+ }
266
+ }
267
+
268
+ const int shared_mem_size = BLOCK_SIZE_M * BLOCK_SIZE_K * 2;
269
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
270
+ kernel<<<gridDim, blockDim, shared_mem_size, stream>>>(
271
+ input, output, b_qweight, b_scales, b_qzeros, topk_weights,
272
+ sorted_token_ids, expert_ids, num_tokens_post_pad, num_experts,
273
+ group_size, top_k, size_m, size_n, size_k, BLOCK_SIZE_M, BLOCK_SIZE_N,
274
+ BLOCK_SIZE_K, has_zp, mul_topk_weight);
275
+}
276
+
277
+torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
278
+ torch::Tensor b_qweight, torch::Tensor b_scales,
279
+ std::optional<torch::Tensor> b_qzeros,
280
+ std::optional<torch::Tensor> topk_weights,
281
+ torch::Tensor sorted_token_ids,
282
+ torch::Tensor expert_ids,
283
+ torch::Tensor num_tokens_post_pad, int64_t top_k,
284
+ int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
285
+ int64_t BLOCK_SIZE_K, int64_t bit) {
286
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
287
+ auto options =
288
+ torch::TensorOptions().dtype(input.dtype()).device(input.device());
289
+
290
+ const int num_experts = b_qweight.size(0);
291
+ const int size_m = input.size(0);
292
+ const int size_n = b_qweight.size(1);
293
+ const int size_k = input.size(1);
294
+ const int group_size = size_k / b_scales.size(2);
295
+
296
+ int64_t EM = sorted_token_ids.size(0);
297
+ if (size_m <= BLOCK_SIZE_M) {
298
+ EM = min(EM, size_m * BLOCK_SIZE_M * top_k);
299
+ }
300
+ const int num_token_blocks = (EM + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M;
301
+
302
+ const uint32_t* b_qzeros_ptr;
303
+ if (b_qzeros.has_value())
304
+ b_qzeros_ptr = (const uint32_t*)b_qzeros.value().data_ptr<uint8_t>();
305
+ const float* topk_weights_ptr;
306
+ if (topk_weights.has_value())
307
+ topk_weights_ptr = (const float*)topk_weights.value().data_ptr();
308
+
309
+ int groups_per_block_row = BLOCK_SIZE_K / group_size;
310
+ TORCH_CHECK(bit == 4 || bit == 8, "bit must be 4 or 8");
311
+ TORCH_CHECK(size_k % BLOCK_SIZE_K == 0,
312
+ "size_k must divisible by BLOCK_SIZE_K");
313
+ TORCH_CHECK(BLOCK_SIZE_K % group_size == 0,
314
+ "BLOCK_SIZE_K must divisible by group_size");
315
+ TORCH_CHECK(BLOCK_SIZE_M <= 64, "BLOCK_SIZE_M must less or equal to 64");
316
+ TORCH_CHECK(groups_per_block_row == 1 || groups_per_block_row == 2 ||
317
+ groups_per_block_row == 4 || groups_per_block_row == 8,
318
+ "BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]");
319
+
320
+ if (input.scalar_type() == at::ScalarType::Half) {
321
+ run_moe_wna16_gemm<half>(
322
+ (const half*)input.data_ptr<at::Half>(),
323
+ (half*)output.data_ptr<at::Half>(),
324
+ (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
325
+ (const half*)b_scales.data_ptr<at::Half>(), b_qzeros_ptr,
326
+ topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
327
+ expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
328
+ num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
329
+ size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
330
+ b_qzeros.has_value(), topk_weights.has_value());
331
+ } else if (input.scalar_type() == at::ScalarType::BFloat16) {
332
+ run_moe_wna16_gemm<nv_bfloat16>(
333
+ (const nv_bfloat16*)input.data_ptr<at::BFloat16>(),
334
+ (nv_bfloat16*)output.data_ptr<at::BFloat16>(),
335
+ (const uint32_t*)b_qweight.data_ptr<uint8_t>(),
336
+ (const nv_bfloat16*)b_scales.data_ptr<at::BFloat16>(), b_qzeros_ptr,
337
+ topk_weights_ptr, sorted_token_ids.data_ptr<int32_t>(),
338
+ expert_ids.data_ptr<int32_t>(), num_tokens_post_pad.data_ptr<int32_t>(),
339
+ num_experts, group_size, num_token_blocks, top_k, size_m, size_n,
340
+ size_k, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, bit,
341
+ b_qzeros.has_value(), topk_weights.has_value());
342
+ } else {
343
+ TORCH_CHECK(false, "moe_wna16_gemm only supports bfloat16 and float16");
344
+ }
345
+ return output;
346
+}
csrc/moe/moe_wna16_utils.h
new
+200
@@ -0,0 +1,200 @@
1
+
2
+#include <cuda_fp16.h>
3
+#include <cuda_bf16.h>
4
+
5
+template <typename scalar_t>
6
+class ScalarType {};
7
+
8
+template <>
9
+class ScalarType<half> {
10
+ public:
11
+ using scalar_t = half;
12
+ using scalar_t2 = half2;
13
+
14
+ static __device__ float inline num2float(const half x) {
15
+ return __half2float(x);
16
+ }
17
+
18
+ static __device__ half2 inline num2num2(const half x) {
19
+ return __half2half2(x);
20
+ }
21
+
22
+ static __device__ half2 inline nums2num2(const half x1, const half x2) {
23
+ return __halves2half2(x1, x2);
24
+ }
25
+
26
+ static __host__ __device__ half inline float2num(const float x) {
27
+ return __float2half(x);
28
+ }
29
+
30
+ static __host__ __device__ half inline int2num(const float x) {
31
+ return __int2half_rn(x);
32
+ }
33
+
34
+ static __host__ __device__ float2 inline num22float2(const half2 x) {
35
+ return __half22float2(x);
36
+ }
37
+
38
+ static __host__ __device__ half2 inline float22num2(const float2 x) {
39
+ return __float22half2_rn(x);
40
+ }
41
+};
42
+
43
+template <>
44
+class ScalarType<nv_bfloat16> {
45
+ public:
46
+ using scalar_t = nv_bfloat16;
47
+ using scalar_t2 = nv_bfloat162;
48
+
49
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
50
+ static __device__ float inline num2float(const nv_bfloat16 x) {
51
+ return __bfloat162float(x);
52
+ }
53
+
54
+ static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
55
+ return __bfloat162bfloat162(x);
56
+ }
57
+
58
+ static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1,
59
+ const nv_bfloat16 x2) {
60
+ return __halves2bfloat162(x1, x2);
61
+ }
62
+
63
+ static __host__ __device__ nv_bfloat16 inline float2num(const float x) {
64
+ return __float2bfloat16(x);
65
+ }
66
+
67
+ static __host__ __device__ nv_bfloat16 inline int2num(const float x) {
68
+ return __int2bfloat16_rn(x);
69
+ }
70
+
71
+ static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) {
72
+ return __bfloat1622float2(x);
73
+ }
74
+
75
+ static __host__ __device__ nv_bfloat162 inline float22num2(const float2 x) {
76
+ return __float22bfloat162_rn(x);
77
+ }
78
+#endif
79
+};
80
+
81
+template <int lut>
82
+__device__ inline int lop3(int a, int b, int c) {
83
+ int res;
84
+ asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
85
+ : "=r"(res)
86
+ : "r"(a), "r"(b), "r"(c), "n"(lut));
87
+ return res;
88
+}
89
+
90
+template <int start_byte, int mask>
91
+__device__ inline uint32_t prmt(uint32_t a) {
92
+ uint32_t res;
93
+ asm volatile("prmt.b32 %0, %1, %2, %3;\n"
94
+ : "=r"(res)
95
+ : "r"(a), "n"(start_byte), "n"(mask));
96
+ return res;
97
+}
98
+
99
+template <typename scalar_t2, int bit>
100
+__device__ inline void dequant(int q, scalar_t2* res) {}
101
+
102
+template <>
103
+__device__ inline void dequant<half2, 4>(int q, half2* res) {
104
+ const int LO = 0x000f000f;
105
+ const int HI = 0x00f000f0;
106
+ const int EX = 0x64006400;
107
+ const int SUB = 0x64006400;
108
+ const int MUL = 0x2c002c00;
109
+ const int ADD = 0xd400d400;
110
+
111
+ int lo0 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, LO, EX);
112
+ int hi0 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, HI, EX);
113
+ q >>= 8;
114
+ int lo1 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, LO, EX);
115
+ int hi1 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, HI, EX);
116
+
117
+ res[0] = __hsub2(*reinterpret_cast<half2*>(&lo0),
118
+ *reinterpret_cast<const half2*>(&SUB));
119
+ res[1] = __hfma2(*reinterpret_cast<half2*>(&hi0),
120
+ *reinterpret_cast<const half2*>(&MUL),
121
+ *reinterpret_cast<const half2*>(&ADD));
122
+ res[2] = __hsub2(*reinterpret_cast<half2*>(&lo1),
123
+ *reinterpret_cast<const half2*>(&SUB));
124
+ res[3] = __hfma2(*reinterpret_cast<half2*>(&hi1),
125
+ *reinterpret_cast<const half2*>(&MUL),
126
+ *reinterpret_cast<const half2*>(&ADD));
127
+}
128
+
129
+template <>
130
+__device__ inline void dequant<half2, 8>(int q, half2* res) {
131
+ static constexpr uint32_t mask_for_elt_01 = 0x5250;
132
+ static constexpr uint32_t mask_for_elt_23 = 0x5351;
133
+ static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
134
+
135
+ uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
136
+ uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
137
+
138
+ static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
139
+
140
+ res[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
141
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
142
+ res[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
143
+ *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
144
+}
145
+
146
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
147
+template <>
148
+__device__ inline void dequant<nv_bfloat162, 4>(int q, nv_bfloat162* res) {
149
+ static constexpr uint32_t MASK = 0x000f000f;
150
+ static constexpr uint32_t EX = 0x43004300;
151
+
152
+ int lo0 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
153
+ q >>= 4;
154
+ int hi0 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
155
+ q >>= 4;
156
+ int lo1 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
157
+ q >>= 4;
158
+ int hi1 = lop3 < (0xf0 & 0xcc) | 0xaa > (q, MASK, EX);
159
+
160
+ static constexpr uint32_t MUL = 0x3F803F80;
161
+ static constexpr uint32_t ADD = 0xC300C300;
162
+
163
+ res[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo0),
164
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
165
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
166
+ res[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi0),
167
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
168
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
169
+ res[2] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo1),
170
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
171
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
172
+ res[3] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi1),
173
+ *reinterpret_cast<const nv_bfloat162*>(&MUL),
174
+ *reinterpret_cast<const nv_bfloat162*>(&ADD));
175
+}
176
+
177
+template <>
178
+__device__ inline void dequant<nv_bfloat162, 8>(int q, nv_bfloat162* res) {
179
+ float fp32_intermediates[4];
180
+ uint32_t* fp32_intermediates_casted =
181
+ reinterpret_cast<uint32_t*>(fp32_intermediates);
182
+
183
+ static constexpr uint32_t fp32_base = 0x4B000000;
184
+ fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650);
185
+ fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652);
186
+ fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651);
187
+ fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653);
188
+
189
+ fp32_intermediates[0] -= 8388608.f;
190
+ fp32_intermediates[1] -= 8388608.f;
191
+ fp32_intermediates[2] -= 8388608.f;
192
+ fp32_intermediates[3] -= 8388608.f;
193
+
194
+ uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(res);
195
+ bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0],
196
+ fp32_intermediates_casted[1], 0x7632);
197
+ bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2],
198
+ fp32_intermediates_casted[3], 0x7632);
199
+}
200
+#endif
csrc/moe/torch_bindings.cpp
+10
@@ -31,6 +31,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
31
" Tensor! num_tokens_post_pad) -> ()");
32
m.impl("sgl_moe_align_block_size", torch::kCUDA, &sgl_moe_align_block_size);
33
34
+ m.def(
35
+ "moe_wna16_gemm(Tensor input, Tensor! output, Tensor b_qweight, "
36
+ "Tensor b_scales, Tensor? b_qzeros, "
37
+ "Tensor? topk_weights, Tensor sorted_token_ids, "
38
+ "Tensor expert_ids, Tensor num_tokens_post_pad, "
39
+ "int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, "
40
+ "int bit) -> Tensor");
41
+
42
+ m.impl("moe_wna16_gemm", torch::kCUDA, &moe_wna16_gemm);
43
+
44
#ifndef USE_ROCM
45
m.def(
46
"marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
vllm/_custom_ops.py
+15
@@ -1098,6 +1098,21 @@ def sgl_moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
1098
experts_ids, num_tokens_post_pad)
1099
1100
1101
+def moe_wna16_gemm(input: torch.Tensor, output: torch.Tensor,
1102
+ b_qweight: torch.Tensor, b_scales: torch.Tensor,
1103
+ b_qzeros: Optional[torch.Tensor],
1104
+ topk_weights: Optional[torch.Tensor],
1105
+ sorted_token_ids: torch.Tensor, experts_ids: torch.Tensor,
1106
+ num_tokens_post_pad: torch.Tensor, top_k: int,
1107
+ BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int,
1108
+ bit: int) -> torch.Tensor:
1109
+ torch.ops._moe_C.moe_wna16_gemm(input, output, b_qweight, b_scales,
1110
+ b_qzeros, topk_weights, sorted_token_ids,
1111
+ experts_ids, num_tokens_post_pad, top_k,
1112
+ BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
1113
+ bit)
1114
+
1115
+
1116
def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1117
token_expert_indicies: torch.Tensor,
1118
gating_output: float) -> None:
vllm/model_executor/layers/fused_moe/fused_moe.py
+109
-1
@@ -719,6 +719,33 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
719
assert B_scale is not None and B_scale.ndim == 3
720
assert B_zp is None or B_zp.ndim == 3
721
722
+ use_moe_wna16_cuda = should_moe_wna16_use_cuda(
723
+ num_valid_tokens=topk_ids.numel(),
724
+ group_size=block_shape[1],
725
+ num_experts=B.shape[0],
726
+ bit=4 if use_int4_w4a16 else 8)
727
+ config = config.copy()
728
+ config.update(
729
+ get_moe_wna16_block_config(config=config,
730
+ use_moe_wna16_cuda=use_moe_wna16_cuda,
731
+ num_valid_tokens=topk_ids.numel(),
732
+ size_k=A.shape[1],
733
+ size_n=B.shape[1],
734
+ num_experts=B.shape[1],
735
+ group_size=block_shape[1],
736
+ real_top_k=topk_ids.shape[1],
737
+ block_size_m=config["BLOCK_SIZE_M"]))
738
+
739
+ if use_moe_wna16_cuda:
740
+ bit = 4 if use_int4_w4a16 else 8
741
+ ops.moe_wna16_gemm(A, C, B, B_scale, B_zp,
742
+ topk_weights if mul_routed_weight else None,
743
+ sorted_token_ids, expert_ids,
744
+ num_tokens_post_padded, top_k,
745
+ config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"],
746
+ config["BLOCK_SIZE_K"], bit)
747
+ return
748
+
749
fused_moe_kernel_gptq_awq[grid](
750
A,
751
B,
@@ -852,6 +879,70 @@ def get_moe_configs(
879
return None
880
881
882
+def get_moe_wna16_block_config(config: Dict[str,
883
+ int], use_moe_wna16_cuda: bool,
884
+ num_valid_tokens: int, size_k: int, size_n: int,
885
+ num_experts: int, group_size: int,
886
+ real_top_k: int, block_size_m: int):
887
+ if "BLOCK_SIZE_N" in config and "BLOCK_SIZE_K" in config:
888
+ # optimal block config is set
889
+ return {}
890
+ if not use_moe_wna16_cuda:
891
+ # triton moe wna16 kernel
892
+ if num_valid_tokens // real_top_k == 1:
893
+ # if bs=1, use a smaller BLOCK_SIZE_N
894
+ return {"BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}
895
+ else:
896
+ return {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}
897
+ else:
898
+ # cuda moe wna16 kernel
899
+ # set default block_size 128, and increase them when num_blocks
900
+ # is too large.
901
+ block_size_n = 128
902
+ block_size_k = 128
903
+ if block_size_k <= group_size:
904
+ block_size_k = group_size
905
+
906
+ num_n_blocks = size_k // block_size_k
907
+ num_k_blocks = size_n // block_size_k
908
+ num_m_blocks = (num_valid_tokens + block_size_m - 1) / block_size_m + \
909
+ num_experts
910
+ if num_valid_tokens // real_top_k <= block_size_m:
911
+ num_m_blocks = min(num_m_blocks, num_valid_tokens)
912
+ num_blocks = num_m_blocks * num_n_blocks * num_k_blocks
913
+
914
+ if size_k % 256 == 0 and num_blocks >= 256 and \
915
+ block_size_k < 256:
916
+ block_size_k = 256
917
+ num_blocks = num_blocks // (256 // block_size_k)
918
+
919
+ if num_m_blocks <= 16 and size_k % (block_size_k * 2) == 0 and \
920
+ size_k % (block_size_k * 2) == 0 and block_size_k <= 512 and \
921
+ num_blocks >= 512:
922
+ block_size_k = block_size_k * 2
923
+ num_blocks = num_blocks // 2
924
+
925
+ if num_blocks > 1024:
926
+ block_size_n = 256
927
+ num_n_blocks = num_n_blocks // 2
928
+ num_blocks = num_blocks // 2
929
+
930
+ if size_n <= 1024 and num_blocks >= 1024:
931
+ # The kernel performance got much better with BLOCK_SIZE_N=1024
932
+ # when num_blocks is large, event when N is small.
933
+ # Not sure why, maybe it force the CUDA SM process only one block
934
+ # at the same time.
935
+ block_size_n = 1024
936
+
937
+ return {"BLOCK_SIZE_N": block_size_n, "BLOCK_SIZE_K": block_size_k}
938
+
939
+
940
+def should_moe_wna16_use_cuda(num_valid_tokens: int, group_size: int,
941
+ num_experts: int, bit: int):
942
+ return bit == 4 and group_size in [32, 64, 128] and \
943
+ num_valid_tokens / num_experts <= 6
944
+
945
+
946
def get_default_config(
947
M: int,
948
E: int,
@@ -873,6 +964,21 @@ def get_default_config(
964
"num_warps": 4,
965
"num_stages": 3,
966
}
967
+ elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
968
+ # moe wna16 kernels
969
+ # only set BLOCK_SIZE_M
970
+ # BLOCK_SIZE_N and BLOCK_SIZE_K would be set later
971
+ bit = 4 if dtype == "int4_w4a16" else 8
972
+ use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk,
973
+ block_shape[1], E, bit)
974
+ if use_moe_wna16_cuda:
975
+ config = {"BLOCK_SIZE_M": min(16, M)}
976
+ elif M <= 20:
977
+ config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1}
978
+ elif M <= 40:
979
+ config = {"BLOCK_SIZE_M": 32, "GROUP_SIZE_M": 1}
980
+ else:
981
+ config = {"BLOCK_SIZE_M": 64, "GROUP_SIZE_M": 1}
982
else:
983
config = {
984
"BLOCK_SIZE_M": 64,
@@ -907,6 +1013,8 @@ def try_get_optimal_moe_config(
1013
else:
1014
# First try to load optimal config from the file
1015
E, _, N = w2_shape
1016
+ if dtype == "int4_w4a16":
1017
+ N = N * 2
1018
block_n = block_shape[0] if block_shape else 0
1019
block_k = block_shape[1] if block_shape else 0
1020
configs = get_moe_configs(E, N, dtype, block_n, block_k)
@@ -1027,7 +1135,7 @@ def get_config_dtype_str(dtype: torch.dtype,
1135
elif use_int8_w8a16:
1136
return "int8_w8a16"
1137
elif use_int4_w4a16:
1030
- return "int4_w8a16"
1138
+ return "int4_w4a16"
1139
elif dtype == torch.float:
1140
# avoiding cases where kernel fails when float32 MoE
1141
# use fp16/bfloat16 configs