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
Use macro guard CUDA functions for back compatibility in grouped_topk_kernel.cu (#25346)
Signed-off-by: Ming Yang <minos.future@gmail.com> Signed-off-by: Rahul Tuli <rtuli@redhat.com> Co-authored-by: Rahul Tuli <rtuli@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com> Co-authored-by: Ye (Charlotte) Qi <yeq@meta.com>
Ming Yang committed
Sep 23, 2025 at 09:45 UTC
527821d191e05e0414695fa8aea24ecaeb3d609c
1 file changed
+14
-6
csrc/moe/grouped_topk_kernels.cu
+14
-6
@@ -418,6 +418,15 @@ __device__ inline T neg_inf() {
418
return cuda_cast<T, float>(-cuda::std::numeric_limits<float>::infinity());
419
}
420
421
+template <typename T>
422
+__device__ inline bool is_finite(const T val) {
423
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
424
+ return cuda::std::isfinite(val);
425
+#else
426
+ return isfinite(cuda_cast<float, T>(val));
427
+#endif
428
+}
429
+
430
template <typename T>
431
__device__ void topk_with_k2(T* output, T const* input,
432
cg::thread_block_tile<32> const& tile,
@@ -533,7 +542,7 @@ __global__ void group_idx_and_topk_idx_kernel(
542
// calculate group_idx
543
int32_t target_num_min = WARP_SIZE - n_group + topk_group;
544
// The check is necessary to avoid abnormal input
536
- if (lane_id < n_group && cuda::std::isfinite(group_scores[lane_id])) {
545
+ if (lane_id < n_group && is_finite(group_scores[lane_id])) {
546
value = group_scores[lane_id];
547
}
548
@@ -568,11 +577,10 @@ __global__ void group_idx_and_topk_idx_kernel(
577
int32_t offset = i_group * num_experts_per_group;
578
for (int32_t i = lane_id; i < align_num_experts_per_group;
579
i += WARP_SIZE) {
571
- T candidates =
572
- (i < num_experts_per_group) &&
573
- cuda::std::isfinite(scores_with_bias[offset + i])
574
- ? scores_with_bias[offset + i]
575
- : neg_inf<T>();
580
+ T candidates = (i < num_experts_per_group) &&
581
+ is_finite(scores_with_bias[offset + i])
582
+ ? scores_with_bias[offset + i]
583
+ : neg_inf<T>();
584
queue.add(candidates, offset + i);
585
}
586
if (group_scores[i_group] == topk_group_value) {