@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8a3cd90af5
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 14, 05:30 UTC

[Kernel] Add fused grouped_topk kernel for MoE (#23274)

Signed-off-by: Xin Yang <xyangx@amazon.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Xin Yang committed Aug 25, 2025 at 11:47 UTC 8a3cd90af534c39425ebfdfd295eea0a4582d541
8 files changed +909 -2
CMakeLists.txt
+3 -1
@@ -817,7 +817,9 @@ set(VLLM_MOE_EXT_SRC
817 "csrc/moe/topk_softmax_kernels.cu")
818
819 if(VLLM_GPU_LANG STREQUAL "CUDA")
820 - list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu")
820 + list(APPEND VLLM_MOE_EXT_SRC
821 + "csrc/moe/moe_wna16.cu"
822 + "csrc/moe/grouped_topk_kernels.cu")
823 endif()
824
825 if(VLLM_GPU_LANG STREQUAL "CUDA")
csrc/moe/grouped_topk_kernels.cu new
+757
@@ -0,0 +1,757 @@
1 +/*
2 + * Adapted from
3 + * https://github.com/NVIDIA/TensorRT-LLM/blob/v0.21.0/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu
4 + * Copyright (c) 2025, The vLLM team.
5 + * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
6 + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
7 + *
8 + * Licensed under the Apache License, Version 2.0 (the "License");
9 + * you may not use this file except in compliance with the License.
10 + * You may obtain a copy of the License at
11 + *
12 + * http://www.apache.org/licenses/LICENSE-2.0
13 + *
14 + * Unless required by applicable law or agreed to in writing, software
15 + * distributed under the License is distributed on an "AS IS" BASIS,
16 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 + * See the License for the specific language governing permissions and
18 + * limitations under the License.
19 + */
20 +#include <c10/cuda/CUDAStream.h>
21 +#include <torch/all.h>
22 +#include <cuda_fp16.h>
23 +#include <cuda_bf16.h>
24 +#include <cooperative_groups.h>
25 +#include <cooperative_groups/reduce.h>
26 +namespace cg = cooperative_groups;
27 +
28 +namespace vllm {
29 +namespace moe {
30 +
31 +constexpr unsigned FULL_WARP_MASK = 0xffffffff;
32 +constexpr int32_t WARP_SIZE = 32;
33 +constexpr int32_t BLOCK_SIZE = 512;
34 +constexpr int32_t NUM_WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;
35 +
36 +namespace warp_topk {
37 +
38 +template <int size, typename T>
39 +__host__ __device__ constexpr T round_up_to_multiple_of(T len) {
40 + if (len == 0) {
41 + return 0;
42 + }
43 + return ((len - 1) / size + 1) * size;
44 +}
45 +
46 +template <typename T>
47 +constexpr __host__ __device__ bool isPowerOf2(T v) {
48 + return (v && !(v & (v - 1)));
49 +}
50 +
51 +template <bool greater, typename T>
52 +__forceinline__ __device__ bool is_better_than(T val, T baseline) {
53 + return (val > baseline && greater) || (val < baseline && !greater);
54 +}
55 +
56 +template <bool greater, typename T, typename idxT>
57 +__forceinline__ __device__ bool is_better_than(T val, T baseline, idxT index,
58 + idxT baseline_index) {
59 + bool res = (val > baseline && greater) || (val < baseline && !greater);
60 + if (val == baseline) {
61 + res = (index < baseline_index && greater) ||
62 + (index < baseline_index && !greater);
63 + }
64 + return res;
65 +}
66 +
67 +template <typename T, typename idxT>
68 +int calc_smem_size_for_block_wide(int num_of_warp, int64_t k) {
69 + int64_t cache_topk = (sizeof(T) + sizeof(idxT)) * num_of_warp * k;
70 + int64_t n = std::max<int>(num_of_warp / 2 * k, num_of_warp * WARP_SIZE);
71 + return max(cache_topk,
72 + round_up_to_multiple_of<256>(n * sizeof(T)) + n * sizeof(idxT));
73 +}
74 +
75 +template <int size, bool ascending, bool reverse, typename T, typename idxT,
76 + bool is_stable>
77 +struct BitonicMerge {
78 + // input should be a bitonic sequence, and sort it to be a monotonic sequence
79 + __device__ static void merge(T* __restrict__ val_arr,
80 + idxT* __restrict__ idx_arr) {
81 + static_assert(isPowerOf2(size));
82 + static_assert(size >= 2 * WARP_SIZE);
83 + constexpr int arr_len = size / WARP_SIZE;
84 +
85 + constexpr int stride = arr_len / 2;
86 + for (int i = 0; i < stride; ++i) {
87 + int const other_i = i + stride;
88 + T& val = val_arr[i];
89 + T& other_val = val_arr[other_i];
90 + bool is_better;
91 + if constexpr (is_stable) {
92 + is_better = is_better_than<ascending>(val, other_val, idx_arr[i],
93 + idx_arr[other_i]);
94 + } else {
95 + is_better = is_better_than<ascending>(val, other_val);
96 + }
97 +
98 + if (is_better) {
99 + T tmp = val;
100 + val = other_val;
101 + other_val = tmp;
102 +
103 + idxT tmp2 = idx_arr[i];
104 + idx_arr[i] = idx_arr[other_i];
105 + idx_arr[other_i] = tmp2;
106 + }
107 + }
108 +
109 + BitonicMerge<size / 2, ascending, reverse, T, idxT, is_stable>::merge(
110 + val_arr, idx_arr);
111 + BitonicMerge<size / 2, ascending, reverse, T, idxT, is_stable>::merge(
112 + val_arr + arr_len / 2, idx_arr + arr_len / 2);
113 + }
114 +};
115 +
116 +template <int size, bool ascending, typename T, typename idxT, bool is_stable>
117 +struct BitonicSort {
118 + __device__ static void sort(T* __restrict__ val_arr,
119 + idxT* __restrict__ idx_arr) {
120 + static_assert(isPowerOf2(size));
121 + static_assert(size >= 2 * WARP_SIZE);
122 + constexpr int arr_len = size / WARP_SIZE;
123 +
124 + BitonicSort<size / 2, true, T, idxT, is_stable>::sort(val_arr, idx_arr);
125 + BitonicSort<size / 2, false, T, idxT, is_stable>::sort(
126 + val_arr + arr_len / 2, idx_arr + arr_len / 2);
127 + BitonicMerge<size, ascending, ascending, T, idxT, is_stable>::merge(
128 + val_arr, idx_arr);
129 + }
130 +};
131 +
132 +template <bool ascending, typename T, typename idxT, bool is_stable>
133 +struct BitonicSort<32, ascending, T, idxT, is_stable> {
134 + __device__ static void sort(T* __restrict__ val_arr,
135 + idxT* __restrict__ idx_arr) {
136 + int const lane = threadIdx.x % WARP_SIZE;
137 +
138 + // ascending doesn't matter before merging since all we need is a bitonic
139 + // sequence
140 + for (int stage = 0; stage < 4; ++stage) {
141 + for (int stride = (1 << stage); stride > 0; stride /= 2) {
142 + bool reverse = (lane >> stage) & 2;
143 + bool is_second = lane & stride;
144 +
145 + T other = __shfl_xor_sync(FULL_WARP_MASK, *val_arr, stride);
146 + idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, *idx_arr, stride);
147 +
148 + bool is_better;
149 + if constexpr (is_stable) {
150 + if constexpr (ascending) {
151 + is_better = ((*val_arr > other) ||
152 + ((*val_arr == other) && (*idx_arr < other_idx))) !=
153 + (reverse != is_second);
154 + } else {
155 + is_better = ((*val_arr > other) ||
156 + ((*val_arr == other) && (*idx_arr > other_idx))) !=
157 + (reverse != is_second);
158 + }
159 + } else {
160 + is_better = (*val_arr != other &&
161 + (*val_arr > other) != (reverse != is_second));
162 + }
163 + if (is_better) {
164 + *val_arr = other;
165 + *idx_arr = other_idx;
166 + }
167 + }
168 + }
169 +
170 + BitonicMerge<32, ascending, ascending, T, idxT, is_stable>::merge(val_arr,
171 + idx_arr);
172 + }
173 +};
174 +
175 +template <bool ascending, bool reverse, typename T, typename idxT,
176 + bool is_stable>
177 +struct BitonicMerge<32, ascending, reverse, T, idxT, is_stable> {
178 + __device__ static void merge(T* __restrict__ val_arr,
179 + idxT* __restrict__ idx_arr) {
180 + int const lane = threadIdx.x % WARP_SIZE;
181 + for (int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
182 + bool is_second = lane & stride;
183 + T& val = *val_arr;
184 + T other = __shfl_xor_sync(FULL_WARP_MASK, val, stride);
185 + idxT& idx = *idx_arr;
186 + idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, idx, stride);
187 +
188 + bool is_better;
189 + if constexpr (is_stable) {
190 + if constexpr (ascending) {
191 + is_better = ((*val_arr > other) ||
192 + ((*val_arr == other) && (*idx_arr < other_idx))) ==
193 + (reverse != is_second); // for min
194 + } else {
195 + is_better = ((*val_arr > other) ||
196 + ((*val_arr == other) && (*idx_arr > other_idx))) ==
197 + (reverse != is_second); // for max
198 + }
199 + } else {
200 + is_better =
201 + (val != other && ((val > other) == (ascending != is_second)));
202 + }
203 +
204 + if (is_better) {
205 + val = other;
206 + idx = other_idx;
207 + }
208 + }
209 + }
210 +};
211 +
212 +template <int capacity, bool greater, typename T, typename idxT, bool is_stable>
213 +class WarpSort {
214 + public:
215 + __device__ WarpSort(idxT k, T dummy)
216 + : lane_(threadIdx.x % WARP_SIZE), k_(k), dummy_(dummy) {
217 + static_assert(capacity >= WARP_SIZE && isPowerOf2(capacity));
218 +
219 + for (int i = 0; i < max_arr_len_; ++i) {
220 + val_arr_[i] = dummy_;
221 + idx_arr_[i] = 0;
222 + }
223 + }
224 +
225 + // load and merge k sorted values
226 + __device__ void load_sorted(T const* __restrict__ in,
227 + idxT const* __restrict__ in_idx, idxT start) {
228 + idxT idx = start + WARP_SIZE - 1 - lane_;
229 + for (int i = max_arr_len_ - 1; i >= 0; --i, idx += WARP_SIZE) {
230 + if (idx < start + k_) {
231 + T t = in[idx];
232 + bool is_better;
233 + if constexpr (is_stable) {
234 + is_better =
235 + is_better_than<greater>(t, val_arr_[i], in_idx[idx], idx_arr_[i]);
236 + } else {
237 + is_better = is_better_than<greater>(t, val_arr_[i]);
238 + }
239 + if (is_better) {
240 + val_arr_[i] = t;
241 + idx_arr_[i] = in_idx[idx];
242 + }
243 + }
244 + }
245 +
246 + BitonicMerge<capacity, greater, !greater, T, idxT, is_stable>::merge(
247 + val_arr_, idx_arr_);
248 + }
249 +
250 + __device__ void dump(T* __restrict__ out, idxT* __restrict__ out_idx) const {
251 + for (int i = 0; i < max_arr_len_; ++i) {
252 + idxT out_i = i * WARP_SIZE + lane_;
253 + if (out_i < k_) {
254 + out[out_i] = val_arr_[i];
255 + out_idx[out_i] = idx_arr_[i];
256 + }
257 + }
258 + }
259 +
260 + __device__ void dumpIdx(idxT* __restrict__ out_idx) const {
261 + for (int i = 0; i < max_arr_len_; ++i) {
262 + idxT out_i = i * WARP_SIZE + lane_;
263 + if (out_i < k_) {
264 + out_idx[out_i] = idx_arr_[i];
265 + }
266 + }
267 + }
268 +
269 + protected:
270 + static constexpr int max_arr_len_ = capacity / WARP_SIZE;
271 +
272 + T val_arr_[max_arr_len_];
273 + idxT idx_arr_[max_arr_len_];
274 +
275 + int const lane_;
276 + idxT const k_;
277 + T const dummy_;
278 +
279 +}; // end class WarpSort
280 +
281 +template <int capacity, bool greater, typename T, typename idxT, bool is_stable>
282 +class WarpSelect : public WarpSort<capacity, greater, T, idxT, is_stable> {
283 + public:
284 + __device__ WarpSelect(idxT k, T dummy)
285 + : WarpSort<capacity, greater, T, idxT, is_stable>(k, dummy),
286 + k_th_(dummy),
287 + k_th_lane_((k - 1) % WARP_SIZE) {
288 + extern __shared__ char smem_buf[]; // extern __shared__ T smem_buf[];
289 +
290 + int const num_of_warp = blockDim.x / WARP_SIZE;
291 + int const warp_id = threadIdx.x / WARP_SIZE;
292 + val_smem_ = reinterpret_cast<T*>(smem_buf);
293 + val_smem_ += warp_id * WARP_SIZE;
294 + idx_smem_ = reinterpret_cast<idxT*>(
295 + smem_buf +
296 + round_up_to_multiple_of<256>(num_of_warp * sizeof(T) * WARP_SIZE));
297 + idx_smem_ += warp_id * WARP_SIZE;
298 + }
299 +
300 + __device__ void add(T const* in, idxT start, idxT end) {
301 + idxT const end_for_fullwarp =
302 + round_up_to_multiple_of<WARP_SIZE>(end - start) + start;
303 + for (idxT i = start + lane_; i < end_for_fullwarp; i += WARP_SIZE) {
304 + T val = (i < end) ? in[i] : dummy_;
305 + add(val, i);
306 + }
307 + }
308 +
309 + __device__ void add(T val, idxT idx) {
310 + bool do_add;
311 + if constexpr (is_stable) {
312 + do_add = is_better_than<greater>(val, k_th_, idx, k_th_idx_);
313 + } else {
314 + do_add = is_better_than<greater>(val, k_th_);
315 + }
316 +
317 + uint32_t mask = __ballot_sync(FULL_WARP_MASK, do_add);
318 + if (mask == 0) {
319 + return;
320 + }
321 +
322 + int pos = smem_buf_len_ + __popc(mask & ((0x1u << lane_) - 1));
323 + if (do_add && pos < WARP_SIZE) {
324 + val_smem_[pos] = val;
325 + idx_smem_[pos] = idx;
326 + do_add = false;
327 + }
328 + smem_buf_len_ += __popc(mask);
329 + if (smem_buf_len_ >= WARP_SIZE) {
330 + __syncwarp();
331 + merge_buf_(val_smem_[lane_], idx_smem_[lane_]);
332 + smem_buf_len_ -= WARP_SIZE;
333 + }
334 + if (do_add) {
335 + pos -= WARP_SIZE;
336 + val_smem_[pos] = val;
337 + idx_smem_[pos] = idx;
338 + }
339 + __syncwarp();
340 + }
341 +
342 + __device__ void done() {
343 + if (smem_buf_len_) {
344 + T val = (lane_ < smem_buf_len_) ? val_smem_[lane_] : dummy_;
345 + idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0;
346 + merge_buf_(val, idx);
347 + }
348 +
349 + // after done(), smem is used for merging results among warps
350 + __syncthreads();
351 + }
352 +
353 + private:
354 + __device__ void set_k_th_() {
355 + k_th_ = __shfl_sync(FULL_WARP_MASK, val_arr_[max_arr_len_ - 1], k_th_lane_);
356 + if constexpr (is_stable) {
357 + k_th_idx_ =
358 + __shfl_sync(FULL_WARP_MASK, idx_arr_[max_arr_len_ - 1], k_th_lane_);
359 + }
360 + }
361 +
362 + __device__ void merge_buf_(T val, idxT idx) {
363 + BitonicSort<WARP_SIZE, greater, T, idxT, is_stable>::sort(&val, &idx);
364 +
365 + T& old = val_arr_[max_arr_len_ - 1];
366 +
367 + bool is_better;
368 + if constexpr (is_stable) {
369 + is_better =
370 + is_better_than<greater>(val, old, idx, idx_arr_[max_arr_len_ - 1]);
371 + } else {
372 + is_better = is_better_than<greater>(val, old);
373 + }
374 +
375 + if (is_better) {
376 + old = val;
377 + idx_arr_[max_arr_len_ - 1] = idx;
378 + }
379 +
380 + BitonicMerge<capacity, greater, !greater, T, idxT, is_stable>::merge(
381 + val_arr_, idx_arr_);
382 +
383 + set_k_th_();
384 + }
385 +
386 + using WarpSort<capacity, greater, T, idxT, is_stable>::max_arr_len_;
387 + using WarpSort<capacity, greater, T, idxT, is_stable>::val_arr_;
388 + using WarpSort<capacity, greater, T, idxT, is_stable>::idx_arr_;
389 + using WarpSort<capacity, greater, T, idxT, is_stable>::lane_;
390 + using WarpSort<capacity, greater, T, idxT, is_stable>::k_;
391 + using WarpSort<capacity, greater, T, idxT, is_stable>::dummy_;
392 +
393 + T* val_smem_;
394 + idxT* idx_smem_;
395 + int smem_buf_len_ = 0;
396 +
397 + T k_th_;
398 + idxT k_th_idx_;
399 + int const k_th_lane_;
400 +}; // end class WarpSelect
401 +} // namespace warp_topk
402 +
403 +template <typename T_OUT, typename T_IN>
404 +__device__ inline T_OUT cuda_cast(T_IN val) {
405 + return val;
406 +}
407 +
408 +template <>
409 +__device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val) {
410 + return __bfloat162float(val);
411 +}
412 +
413 +template <typename T>
414 +__device__ void topk_with_k2(T* output, T const* input,
415 + cg::thread_block_tile<32> const& tile,
416 + int32_t const lane_id,
417 + int const num_experts_per_group) {
418 + // Get the top2 per thread
419 + T largest = -INFINITY;
420 + T second_largest = -INFINITY;
421 +
422 + if (num_experts_per_group > WARP_SIZE) {
423 + for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
424 + T value = input[i];
425 + if (value > largest) {
426 + second_largest = largest;
427 + largest = value;
428 + } else if (value > second_largest) {
429 + second_largest = value;
430 + }
431 + }
432 + } else {
433 + for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
434 + largest = input[i];
435 + }
436 + }
437 +
438 + __syncwarp(); // Ensure all threads have valid data before reduction
439 + // Get the top2 warpwise
440 + T max1 = cg::reduce(tile, largest, cg::greater<T>());
441 +
442 + T max2 = max1;
443 + bool equal_to_max1 = (max1 == largest);
444 +
445 + int count_max1 = __popc(__ballot_sync(FULL_WARP_MASK, equal_to_max1));
446 +
447 + if (count_max1 == 1) {
448 + largest = (largest == max1) ? second_largest : largest;
449 + max2 = cg::reduce(tile, largest, cg::greater<T>());
450 + }
451 +
452 + if (lane_id == 0) {
453 + *output = max1 + max2;
454 + }
455 +}
456 +
457 +template <typename T>
458 +__global__ void topk_with_k2_kernel(T* output, T* input,
459 + int64_t const num_tokens,
460 + int64_t const num_cases,
461 + int64_t const n_group,
462 + int64_t const num_experts_per_group) {
463 + int32_t warp_id = threadIdx.x / WARP_SIZE;
464 + int32_t lane_id = threadIdx.x % WARP_SIZE;
465 +
466 + int32_t case_id = blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;
467 + if (case_id < num_cases) {
468 + input += case_id * num_experts_per_group;
469 + output += case_id;
470 +
471 + cg::thread_block block = cg::this_thread_block();
472 + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);
473 +
474 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
475 + asm volatile("griddepcontrol.wait;");
476 +#endif
477 + topk_with_k2(output, input, tile, lane_id, num_experts_per_group);
478 + }
479 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
480 + asm volatile("griddepcontrol.launch_dependents;");
481 +#endif
482 +}
483 +
484 +template <typename T, typename IdxT>
485 +__global__ void group_idx_and_topk_idx_kernel(
486 + T* scores, T const* group_scores, T* topk_values, IdxT* topk_indices,
487 + T* scores_with_bias, int64_t const num_tokens, int64_t const n_group,
488 + int64_t const topk_group, int64_t const topk, int64_t const num_experts,
489 + int64_t const num_experts_per_group, bool renormalize,
490 + double routed_scaling_factor) {
491 + int32_t warp_id = threadIdx.x / WARP_SIZE;
492 + int32_t lane_id = threadIdx.x % WARP_SIZE;
493 + int32_t case_id =
494 + blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id; // one per token
495 + scores_with_bias += case_id * num_experts;
496 + scores += case_id * num_experts;
497 + group_scores += case_id * n_group;
498 + topk_values += case_id * topk;
499 + topk_indices += case_id * topk;
500 +
501 + int32_t align_num_experts_per_group =
502 + warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);
503 +
504 + cg::thread_block block = cg::this_thread_block();
505 + cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);
506 +
507 + extern __shared__ char smem_buf[]; // NOTE: reuse the shared memory here to
508 + // store the target topk idx
509 + int32_t* s_topk_idx = reinterpret_cast<int32_t*>(smem_buf);
510 + T* s_topk_value =
511 + reinterpret_cast<T*>(s_topk_idx + NUM_WARPS_PER_BLOCK * topk) +
512 + warp_id * topk;
513 + s_topk_idx += warp_id * topk;
514 +
515 + T value = cuda::std::numeric_limits<T>::min();
516 + T topk_group_value = cuda::std::numeric_limits<T>::min();
517 + int32_t num_equalto_topkth_group;
518 +
519 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
520 + asm volatile("griddepcontrol.wait;"); // I think all prolog can be put before
521 + // acqbulk because it's ptr arithmetic
522 +#endif
523 +
524 + if (case_id < num_tokens) {
525 + // calculate group_idx
526 + int32_t target_num_min = WARP_SIZE - n_group + topk_group;
527 + if (lane_id < n_group &&
528 + (isfinite(cuda_cast<float, T>(
529 + group_scores[lane_id])))) // The check is necessary to avoid
530 + // abnormal input
531 + {
532 + value = group_scores[lane_id];
533 + }
534 +
535 + int count_equal_to_top_value = WARP_SIZE - n_group;
536 + int pre_count_equal_to_top_value = 0;
537 + // Use loop to find the largset top_group
538 + while (count_equal_to_top_value < target_num_min) {
539 + __syncwarp(); // Ensure all threads have valid data before reduction
540 + topk_group_value = cg::reduce(tile, value, cg::greater<T>());
541 + if (value == topk_group_value) {
542 + value = cuda::std::numeric_limits<T>::min();
543 + }
544 + pre_count_equal_to_top_value = count_equal_to_top_value;
545 + count_equal_to_top_value = __popc(__ballot_sync(
546 + FULL_WARP_MASK, (value == cuda::std::numeric_limits<T>::min())));
547 + }
548 + num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
549 + }
550 + __syncthreads();
551 +
552 + warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
553 + /* is_stable */ true>
554 + queue((int32_t)topk, -INFINITY);
555 +
556 + int count_equalto_topkth_group = 0;
557 + bool if_proceed_next_topk =
558 + (topk_group_value != cuda::std::numeric_limits<T>::min());
559 + if (case_id < num_tokens && if_proceed_next_topk) {
560 + for (int i_group = 0; i_group < n_group; i_group++) {
561 + if ((group_scores[i_group] > topk_group_value) ||
562 + ((group_scores[i_group] == topk_group_value) &&
563 + (count_equalto_topkth_group < num_equalto_topkth_group))) {
564 + int32_t offset = i_group * num_experts_per_group;
565 + for (int32_t i = lane_id; i < align_num_experts_per_group;
566 + i += WARP_SIZE) {
567 + T candidates =
568 + (i < num_experts_per_group) && isfinite(cuda_cast<float, T>(
569 + scores_with_bias[offset + i]))
570 + ? scores_with_bias[offset + i]
571 + : cuda::std::numeric_limits<T>::min();
572 + queue.add(candidates, offset + i);
573 + }
574 + if (group_scores[i_group] == topk_group_value) {
575 + count_equalto_topkth_group++;
576 + }
577 + }
578 + }
579 + queue.done();
580 + __syncwarp();
581 + // Get the topk_idx
582 + queue.dumpIdx(s_topk_idx);
583 + __syncwarp();
584 + }
585 +
586 + // Load the valid score value
587 + // Calculate the summation
588 + float topk_sum = 1e-20;
589 + if (case_id < num_tokens && if_proceed_next_topk) {
590 + for (int i = lane_id;
591 + i < warp_topk::round_up_to_multiple_of<WARP_SIZE>(topk);
592 + i += WARP_SIZE) {
593 + T value =
594 + i < topk
595 + ? scores[s_topk_idx[i]]
596 + : cuda_cast<T, float>(0.0f); // Load the valid value of expert
597 + if (i < topk) {
598 + s_topk_value[i] = value;
599 + }
600 + topk_sum += reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
601 + }
602 + }
603 +
604 + __syncthreads();
605 +
606 + if (case_id < num_tokens) {
607 + if (if_proceed_next_topk) {
608 + for (int i = lane_id; i < topk; i += WARP_SIZE) {
609 + float value;
610 + if (renormalize) {
611 + value = cuda_cast<float, T>(s_topk_value[i]) / topk_sum *
612 + routed_scaling_factor;
613 + } else {
614 + value = cuda_cast<float, T>(s_topk_value[i]) * routed_scaling_factor;
615 + }
616 + topk_indices[i] = s_topk_idx[i];
617 + topk_values[i] = cuda_cast<T, float>(value);
618 + }
619 + } else {
620 + for (int i = lane_id; i < topk; i += WARP_SIZE) {
621 + topk_indices[i] = i;
622 + topk_values[i] = cuda_cast<T, float>(1.0f / topk);
623 + }
624 + }
625 + // Note: when if_proceed_next_topk==false, choose the first 8 experts as the
626 + // default result.
627 + }
628 +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
629 + asm volatile("griddepcontrol.launch_dependents;");
630 +#endif
631 +}
632 +
633 +template <typename T, typename IdxT>
634 +void invokeNoAuxTc(T* scores, T* group_scores, T* topk_values,
635 + IdxT* topk_indices, T* scores_with_bias,
636 + int64_t const num_tokens, int64_t const num_experts,
637 + int64_t const n_group, int64_t const topk_group,
638 + int64_t const topk, bool const renormalize,
639 + double const routed_scaling_factor, bool enable_pdl = false,
640 + cudaStream_t const stream = 0) {
641 + int64_t num_cases = num_tokens * n_group;
642 + int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
643 + auto* kernel_instance1 = &topk_with_k2_kernel<T>;
644 + cudaLaunchConfig_t config;
645 + config.gridDim = topk_with_k2_num_blocks;
646 + config.blockDim = BLOCK_SIZE;
647 + config.dynamicSmemBytes = 0;
648 + config.stream = stream;
649 + cudaLaunchAttribute attrs[1];
650 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
651 + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
652 + config.numAttrs = 1;
653 + config.attrs = attrs;
654 + cudaLaunchKernelEx(&config, kernel_instance1, group_scores, scores_with_bias,
655 + num_tokens, num_cases, n_group, num_experts / n_group);
656 +
657 + int64_t topk_with_k_group_num_blocks =
658 + (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;
659 + size_t dynamic_smem_in_bytes =
660 + warp_topk::calc_smem_size_for_block_wide<T, int32_t>(NUM_WARPS_PER_BLOCK,
661 + topk);
662 + auto* kernel_instance2 = &group_idx_and_topk_idx_kernel<T, IdxT>;
663 + config.gridDim = topk_with_k_group_num_blocks;
664 + config.blockDim = BLOCK_SIZE;
665 + config.dynamicSmemBytes = dynamic_smem_in_bytes;
666 + config.stream = stream;
667 + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
668 + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
669 + config.numAttrs = 1;
670 + config.attrs = attrs;
671 + cudaLaunchKernelEx(&config, kernel_instance2, scores, group_scores,
672 + topk_values, topk_indices, scores_with_bias, num_tokens,
673 + n_group, topk_group, topk, num_experts,
674 + num_experts / n_group, renormalize, routed_scaling_factor);
675 +}
676 +
677 +#define INSTANTIATE_NOAUX_TC(T, IdxT) \
678 + template void invokeNoAuxTc<T, IdxT>( \
679 + T * scores, T * group_scores, T * topk_values, IdxT * topk_indices, \
680 + T * scores_with_bias, int64_t const num_tokens, \
681 + int64_t const num_experts, int64_t const n_group, \
682 + int64_t const topk_group, int64_t const topk, bool const renormalize, \
683 + double const routed_scaling_factor, bool enable_pdl, \
684 + cudaStream_t const stream);
685 +
686 +INSTANTIATE_NOAUX_TC(float, int32_t);
687 +INSTANTIATE_NOAUX_TC(half, int32_t);
688 +INSTANTIATE_NOAUX_TC(__nv_bfloat16, int32_t);
689 +} // end namespace moe
690 +} // namespace vllm
691 +
692 +std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
693 + torch::Tensor const& scores, torch::Tensor const& scores_with_bias,
694 + int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize,
695 + double routed_scaling_factor) {
696 + auto data_type = scores_with_bias.scalar_type();
697 + auto input_size = scores_with_bias.sizes();
698 + int64_t num_tokens = input_size[0];
699 + int64_t num_experts = input_size[1];
700 + TORCH_CHECK(input_size.size() == 2, "scores_with_bias must be a 2D Tensor");
701 + TORCH_CHECK(num_experts % n_group == 0,
702 + "num_experts should be divisible by n_group");
703 + TORCH_CHECK(n_group <= 32,
704 + "n_group should be smaller than or equal to 32 for now");
705 + TORCH_CHECK(topk <= 32, "topk should be smaller than or equal to 32 for now");
706 +
707 + torch::Tensor group_scores = torch::empty(
708 + {num_tokens, n_group}, torch::dtype(data_type).device(torch::kCUDA));
709 + torch::Tensor topk_values = torch::empty(
710 + {num_tokens, topk}, torch::dtype(data_type).device(torch::kCUDA));
711 + torch::Tensor topk_indices = torch::empty(
712 + {num_tokens, topk}, torch::dtype(torch::kInt32).device(torch::kCUDA));
713 +
714 + auto stream = c10::cuda::getCurrentCUDAStream(scores_with_bias.get_device());
715 +
716 + switch (data_type) {
717 + case torch::kFloat16:
718 + // Handle Float16
719 + vllm::moe::invokeNoAuxTc<half, int32_t>(
720 + reinterpret_cast<half*>(scores.mutable_data_ptr()),
721 + reinterpret_cast<half*>(group_scores.mutable_data_ptr()),
722 + reinterpret_cast<half*>(topk_values.mutable_data_ptr()),
723 + reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
724 + reinterpret_cast<half*>(scores_with_bias.data_ptr()), num_tokens,
725 + num_experts, n_group, topk_group, topk, renormalize,
726 + routed_scaling_factor, false, stream);
727 + break;
728 + case torch::kFloat32:
729 + // Handle Float32
730 + vllm::moe::invokeNoAuxTc<float, int32_t>(
731 + reinterpret_cast<float*>(scores.mutable_data_ptr()),
732 + reinterpret_cast<float*>(group_scores.mutable_data_ptr()),
733 + reinterpret_cast<float*>(topk_values.mutable_data_ptr()),
734 + reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
735 + reinterpret_cast<float*>(scores_with_bias.data_ptr()), num_tokens,
736 + num_experts, n_group, topk_group, topk, renormalize,
737 + routed_scaling_factor, false, stream);
738 + break;
739 + case torch::kBFloat16:
740 + // Handle BFloat16
741 + vllm::moe::invokeNoAuxTc<__nv_bfloat16, int32_t>(
742 + reinterpret_cast<__nv_bfloat16*>(scores.mutable_data_ptr()),
743 + reinterpret_cast<__nv_bfloat16*>(group_scores.mutable_data_ptr()),
744 + reinterpret_cast<__nv_bfloat16*>(topk_values.mutable_data_ptr()),
745 + reinterpret_cast<int32_t*>(topk_indices.mutable_data_ptr()),
746 + reinterpret_cast<__nv_bfloat16*>(scores_with_bias.data_ptr()),
747 + num_tokens, num_experts, n_group, topk_group, topk, renormalize,
748 + routed_scaling_factor, false, stream);
749 + break;
750 + default:
751 + // Handle other data types
752 + throw std::invalid_argument(
753 + "Invalid dtype, only supports float16, float32, and bfloat16");
754 + break;
755 + }
756 + return {topk_values, topk_indices};
757 +}
csrc/moe/moe_ops.h
+5
@@ -22,6 +22,11 @@ torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
22 torch::Tensor num_tokens_post_pad, int64_t top_k,
23 int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
24 int64_t BLOCK_SIZE_K, int64_t bit);
25 +
26 +std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
27 + torch::Tensor const& scores, torch::Tensor const& scores_with_bias,
28 + int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize,
29 + double routed_scaling_factor);
30 #endif
31
32 bool moe_permute_unpermute_supported();
csrc/moe/torch_bindings.cpp
+6
@@ -78,6 +78,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
78 "output_tensor) -> ()");
79 m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
80
81 + // Apply grouped topk routing to select experts.
82 + m.def(
83 + "grouped_topk(Tensor scores, Tensor scores_with_bias, int n_group, int "
84 + "topk_group, int topk, bool renormalize, float "
85 + "routed_scaling_factor) -> (Tensor, Tensor)");
86 + m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
87 #endif
88 }
89
tests/kernels/moe/test_grouped_topk.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the MoE grouped topk kernel
4 +
5 +Run `pytest tests/kernels/moe/test_grouped_topk.py`.
6 +"""
7 +import pytest
8 +import torch
9 +
10 +from vllm.model_executor.layers.fused_moe.fused_moe import (fused_grouped_topk,
11 + grouped_topk)
12 +from vllm.platforms import current_platform
13 +
14 +
15 +@pytest.mark.skipif(not current_platform.is_cuda(),
16 + reason="This test is skipped on non-CUDA platform.")
17 +@pytest.mark.parametrize("n_token", [1, 33, 64])
18 +@pytest.mark.parametrize("n_hidden", [1024, 2048])
19 +@pytest.mark.parametrize("n_expert", [16])
20 +@pytest.mark.parametrize("topk", [2])
21 +@pytest.mark.parametrize("renormalize", [True, False])
22 +@pytest.mark.parametrize("num_expert_group", [8])
23 +@pytest.mark.parametrize("topk_group", [2])
24 +@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
25 +@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5])
26 +@pytest.mark.parametrize("dtype",
27 + [torch.float16, torch.bfloat16, torch.float32])
28 +def test_grouped_topk(monkeypatch: pytest.MonkeyPatch, n_token: int,
29 + n_hidden: int, n_expert: int, topk: int,
30 + renormalize: bool, num_expert_group: int,
31 + topk_group: int, scoring_func: str,
32 + routed_scaling_factor: float, dtype: torch.dtype):
33 + current_platform.seed_everything(0)
34 + hidden_states = torch.randn((n_token, n_hidden),
35 + dtype=dtype,
36 + device="cuda")
37 + gating_output = torch.randn((n_token, n_expert),
38 + dtype=dtype,
39 + device="cuda")
40 + e_score_correction_bias = torch.randn((n_expert, ),
41 + dtype=torch.float32,
42 + device="cuda")
43 +
44 + with monkeypatch.context() as m:
45 + m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
46 + baseline_topk_weights, baseline_topk_ids = grouped_topk(
47 + hidden_states=hidden_states,
48 + gating_output=gating_output,
49 + topk=topk,
50 + renormalize=renormalize,
51 + num_expert_group=num_expert_group,
52 + topk_group=topk_group,
53 + scoring_func=scoring_func,
54 + routed_scaling_factor=routed_scaling_factor,
55 + e_score_correction_bias=e_score_correction_bias)
56 +
57 + test_topk_weights, test_topk_ids = fused_grouped_topk(
58 + hidden_states=hidden_states,
59 + gating_output=gating_output,
60 + topk=topk,
61 + renormalize=renormalize,
62 + num_expert_group=num_expert_group,
63 + topk_group=topk_group,
64 + scoring_func=scoring_func,
65 + routed_scaling_factor=routed_scaling_factor,
66 + e_score_correction_bias=e_score_correction_bias)
67 +
68 + if renormalize:
69 + torch.testing.assert_close(baseline_topk_weights,
70 + test_topk_weights,
71 + atol=2e-2,
72 + rtol=0)
73 + torch.testing.assert_close(baseline_topk_ids,
74 + test_topk_ids,
75 + atol=0,
76 + rtol=0)
vllm/_custom_ops.py
+11
@@ -1502,6 +1502,17 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1502 gating_output)
1503
1504
1505 +def grouped_topk(scores: torch.Tensor, scores_with_bias: torch.Tensor,
1506 + num_expert_group: int, topk_group: int, topk: int,
1507 + renormalize: bool, routed_scaling_factor: float):
1508 + if not current_platform.is_cuda():
1509 + raise NotImplementedError("The fused grouped_topk kernel is only "
1510 + "available on CUDA platforms")
1511 + return torch.ops._moe_C.grouped_topk(scores, scores_with_bias,
1512 + num_expert_group, topk_group, topk,
1513 + renormalize, routed_scaling_factor)
1514 +
1515 +
1516 def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1517 b_qweight: torch.Tensor,
1518 b_bias: Optional[torch.Tensor],
vllm/envs.py
+6
@@ -131,6 +131,7 @@ if TYPE_CHECKING:
131 VLLM_USE_DEEP_GEMM: bool = False
132 VLLM_USE_DEEP_GEMM_E8M0: bool = True
133 VLLM_SKIP_DEEP_GEMM_WARMUP: bool = False
134 + VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
135 VLLM_USE_FLASHINFER_MOE_FP8: bool = False
136 VLLM_USE_FLASHINFER_MOE_FP4: bool = False
137 VLLM_FLASHINFER_MOE_BACKEND: str = "throughput"
@@ -963,6 +964,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
964 "VLLM_SKIP_DEEP_GEMM_WARMUP":
965 lambda: bool(int(os.getenv("VLLM_SKIP_DEEP_GEMM_WARMUP", "0"))),
966
967 + # Whether to use fused grouped_topk used for MoE expert selection.
968 + "VLLM_USE_FUSED_MOE_GROUPED_TOPK":
969 + lambda: bool(int(os.getenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "1"))),
970 +
971 # Allow use of FlashInfer MoE kernels for fused moe ops.
972 "VLLM_USE_FLASHINFER_MOE_FP8":
973 lambda: bool(int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP8", "0"))),
@@ -1229,6 +1234,7 @@ def compute_hash() -> str:
1234 "VLLM_DISABLED_KERNELS",
1235 "VLLM_USE_DEEP_GEMM",
1236 "VLLM_USE_TRTLLM_FP4_GEMM",
1237 + "VLLM_USE_FUSED_MOE_GROUPED_TOPK",
1238 "VLLM_USE_FLASHINFER_MOE_FP8",
1239 "VLLM_USE_FLASHINFER_MOE_FP4",
1240 "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8",
vllm/model_executor/layers/fused_moe/fused_moe.py
+45 -1
@@ -949,8 +949,23 @@ def grouped_topk(
949 num_expert_group: int = 0,
950 topk_group: int = 0,
951 scoring_func: str = "softmax",
952 - e_score_correction_bias: Optional[torch.Tensor] = None
952 + routed_scaling_factor: float = 1.0,
953 + e_score_correction_bias: Optional[torch.Tensor] = None,
954 ) -> tuple[torch.Tensor, torch.Tensor]:
955 + if envs.VLLM_USE_FUSED_MOE_GROUPED_TOPK and \
956 + current_platform.is_cuda() and \
957 + num_expert_group <= 32 and topk <= 32 and \
958 + e_score_correction_bias is not None:
959 + return fused_grouped_topk(
960 + hidden_states=hidden_states,
961 + gating_output=gating_output,
962 + topk=topk,
963 + renormalize=renormalize,
964 + e_score_correction_bias=e_score_correction_bias,
965 + num_expert_group=num_expert_group,
966 + topk_group=topk_group,
967 + scoring_func=scoring_func,
968 + routed_scaling_factor=routed_scaling_factor)
969
970 assert hidden_states.size(0) == gating_output.size(0), (
971 "Number of tokens mismatch")
@@ -996,9 +1011,38 @@ def grouped_topk(
1011 if renormalize:
1012 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
1013
1014 + topk_weights = topk_weights * routed_scaling_factor
1015 return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
1016
1017
1018 +def fused_grouped_topk(
1019 + hidden_states: torch.Tensor,
1020 + gating_output: torch.Tensor,
1021 + topk: int,
1022 + renormalize: bool,
1023 + e_score_correction_bias: torch.Tensor,
1024 + num_expert_group: int = 0,
1025 + topk_group: int = 0,
1026 + scoring_func: str = "softmax",
1027 + routed_scaling_factor: float = 1.0,
1028 +) -> tuple[torch.Tensor, torch.Tensor]:
1029 + assert hidden_states.size(0) == gating_output.size(0), (
1030 + "Number of tokens mismatch")
1031 +
1032 + if scoring_func == "softmax":
1033 + scores = torch.softmax(gating_output, dim=-1)
1034 + elif scoring_func == "sigmoid":
1035 + scores = gating_output.sigmoid()
1036 + else:
1037 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
1038 +
1039 + scores_with_bias = scores + e_score_correction_bias.unsqueeze(0)
1040 + topk_values, topk_indices = ops.grouped_topk(
1041 + scores, scores_with_bias.to(scores.dtype), num_expert_group,
1042 + topk_group, topk, renormalize, routed_scaling_factor)
1043 + return topk_values.to(torch.float32), topk_indices.to(torch.int32)
1044 +
1045 +
1046 def get_config_dtype_str(
1047 dtype: torch.dtype,
1048 use_int4_w4a16: Optional[bool] = False,