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][Bugfix] Fix grouped topk cu (#24146)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com>
Qiming Zhang committed
Sep 3, 2025 at 21:37 UTC
e919d6f549f4da22fa60ea394f00aaf93ef23aa0
1 file changed
+7
-6
csrc/moe/grouped_topk_kernels.cu
+7
-6
@@ -28,6 +28,7 @@ namespace cg = cooperative_groups;
28
namespace vllm {
29
namespace moe {
30
31
+constexpr float kNegInfinity = INFINITY * -1;
32
constexpr unsigned FULL_WARP_MASK = 0xffffffff;
33
constexpr int32_t WARP_SIZE = 32;
34
constexpr int32_t BLOCK_SIZE = 512;
@@ -512,8 +513,8 @@ __global__ void group_idx_and_topk_idx_kernel(
513
warp_id * topk;
514
s_topk_idx += warp_id * topk;
515
515
- T value = cuda::std::numeric_limits<T>::min();
516
- T topk_group_value = cuda::std::numeric_limits<T>::min();
516
+ T value = kNegInfinity;
517
+ T topk_group_value = kNegInfinity;
518
int32_t num_equalto_topkth_group;
519
520
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -539,11 +540,11 @@ __global__ void group_idx_and_topk_idx_kernel(
540
__syncwarp(); // Ensure all threads have valid data before reduction
541
topk_group_value = cg::reduce(tile, value, cg::greater<T>());
542
if (value == topk_group_value) {
542
- value = cuda::std::numeric_limits<T>::min();
543
+ value = kNegInfinity;
544
}
545
pre_count_equal_to_top_value = count_equal_to_top_value;
546
count_equal_to_top_value = __popc(__ballot_sync(
546
- FULL_WARP_MASK, (value == cuda::std::numeric_limits<T>::min())));
547
+ FULL_WARP_MASK, (value == cuda_cast<T, float>(kNegInfinity))));
548
}
549
num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
550
}
@@ -555,7 +556,7 @@ __global__ void group_idx_and_topk_idx_kernel(
556
557
int count_equalto_topkth_group = 0;
558
bool if_proceed_next_topk =
558
- (topk_group_value != cuda::std::numeric_limits<T>::min());
559
+ (topk_group_value != cuda_cast<T, float>(kNegInfinity));
560
if (case_id < num_tokens && if_proceed_next_topk) {
561
for (int i_group = 0; i_group < n_group; i_group++) {
562
if ((group_scores[i_group] > topk_group_value) ||
@@ -568,7 +569,7 @@ __global__ void group_idx_and_topk_idx_kernel(
569
(i < num_experts_per_group) && isfinite(cuda_cast<float, T>(
570
scores_with_bias[offset + i]))
571
? scores_with_bias[offset + i]
571
- : cuda::std::numeric_limits<T>::min();
572
+ : cuda_cast<T, float>(kNegInfinity);
573
queue.add(candidates, offset + i);
574
}
575
if (group_scores[i_group] == topk_group_value) {