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] Better inf handling for grouped topk cu (#24886)
Signed-off-by: lumina37 <starry.qvq@gmail.com>
Lumina committed
Sep 18, 2025 at 13:53 UTC
81b16a2bc922e837267db7216a274c4d89a2cc0c
1 file changed
+24
-20
csrc/moe/grouped_topk_kernels.cu
+24
-20
@@ -21,6 +21,7 @@
21
#include <torch/all.h>
22
#include <cuda_fp16.h>
23
#include <cuda_bf16.h>
24
+#include <cuda/std/limits>
25
#include <cooperative_groups.h>
26
#include <cooperative_groups/reduce.h>
27
namespace cg = cooperative_groups;
@@ -28,7 +29,6 @@ namespace cg = cooperative_groups;
29
namespace vllm {
30
namespace moe {
31
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;
@@ -411,14 +411,21 @@ __device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val) {
411
return __bfloat162float(val);
412
}
413
414
+template <typename T>
415
+__device__ inline T neg_inf() {
416
+ // cuda::std::numeric_limits<T>::infinity() returns `0` for [T=bf16 or fp16]
417
+ // so we need to cast from fp32
418
+ return cuda_cast<T, float>(-cuda::std::numeric_limits<float>::infinity());
419
+}
420
+
421
template <typename T>
422
__device__ void topk_with_k2(T* output, T const* input,
423
cg::thread_block_tile<32> const& tile,
424
int32_t const lane_id,
425
int const num_experts_per_group) {
426
// Get the top2 per thread
420
- T largest = -INFINITY;
421
- T second_largest = -INFINITY;
427
+ T largest = neg_inf<T>();
428
+ T second_largest = neg_inf<T>();
429
430
if (num_experts_per_group > WARP_SIZE) {
431
for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
@@ -513,8 +520,8 @@ __global__ void group_idx_and_topk_idx_kernel(
520
warp_id * topk;
521
s_topk_idx += warp_id * topk;
522
516
- T value = kNegInfinity;
517
- T topk_group_value = kNegInfinity;
523
+ T value = neg_inf<T>();
524
+ T topk_group_value = neg_inf<T>();
525
int32_t num_equalto_topkth_group;
526
527
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
@@ -525,11 +532,8 @@ __global__ void group_idx_and_topk_idx_kernel(
532
if (case_id < num_tokens) {
533
// calculate group_idx
534
int32_t target_num_min = WARP_SIZE - n_group + topk_group;
528
- if (lane_id < n_group &&
529
- (isfinite(cuda_cast<float, T>(
530
- group_scores[lane_id])))) // The check is necessary to avoid
531
- // abnormal input
532
- {
535
+ // The check is necessary to avoid abnormal input
536
+ if (lane_id < n_group && cuda::std::isfinite(group_scores[lane_id])) {
537
value = group_scores[lane_id];
538
}
539
@@ -540,11 +544,11 @@ __global__ void group_idx_and_topk_idx_kernel(
544
__syncwarp(); // Ensure all threads have valid data before reduction
545
topk_group_value = cg::reduce(tile, value, cg::greater<T>());
546
if (value == topk_group_value) {
543
- value = kNegInfinity;
547
+ value = neg_inf<T>();
548
}
549
pre_count_equal_to_top_value = count_equal_to_top_value;
546
- count_equal_to_top_value = __popc(__ballot_sync(
547
- FULL_WARP_MASK, (value == cuda_cast<T, float>(kNegInfinity))));
550
+ count_equal_to_top_value =
551
+ __popc(__ballot_sync(FULL_WARP_MASK, (value == neg_inf<T>())));
552
}
553
num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
554
}
@@ -552,11 +556,10 @@ __global__ void group_idx_and_topk_idx_kernel(
556
557
warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t,
558
/* is_stable */ true>
555
- queue((int32_t)topk, -INFINITY);
559
+ queue((int32_t)topk, neg_inf<T>());
560
561
int count_equalto_topkth_group = 0;
558
- bool if_proceed_next_topk =
559
- (topk_group_value != cuda_cast<T, float>(kNegInfinity));
562
+ bool if_proceed_next_topk = topk_group_value != neg_inf<T>();
563
if (case_id < num_tokens && if_proceed_next_topk) {
564
for (int i_group = 0; i_group < n_group; i_group++) {
565
if ((group_scores[i_group] > topk_group_value) ||
@@ -566,10 +569,10 @@ __global__ void group_idx_and_topk_idx_kernel(
569
for (int32_t i = lane_id; i < align_num_experts_per_group;
570
i += WARP_SIZE) {
571
T candidates =
569
- (i < num_experts_per_group) && isfinite(cuda_cast<float, T>(
570
- scores_with_bias[offset + i]))
572
+ (i < num_experts_per_group) &&
573
+ cuda::std::isfinite(scores_with_bias[offset + i])
574
? scores_with_bias[offset + i]
572
- : cuda_cast<T, float>(kNegInfinity);
575
+ : neg_inf<T>();
576
queue.add(candidates, offset + i);
577
}
578
if (group_scores[i_group] == topk_group_value) {
@@ -598,7 +601,8 @@ __global__ void group_idx_and_topk_idx_kernel(
601
if (i < topk) {
602
s_topk_value[i] = value;
603
}
601
- topk_sum += reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
604
+ topk_sum +=
605
+ cg::reduce(tile, cuda_cast<float, T>(value), cg::plus<float>());
606
}
607
}
608