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, 22:35 UTC
[Refactor] Fix Compile Warning #1444-D (#21462)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 1, 2025 at 09:10 UTC
a59cd9d9f7fd89e19beeffb7e7f89437d413eafb
1 file changed
+4
-2
csrc/moe/topk_softmax_kernels.cu
+4
-2
@@ -24,9 +24,12 @@
24
#ifndef USE_ROCM
25
#include <cub/util_type.cuh>
26
#include <cub/cub.cuh>
27
+ #include <cuda/std/functional>
28
+ using AddOp = cuda::std::plus<float>;
29
#else
30
#include <hipcub/util_type.hpp>
31
#include <hipcub/hipcub.hpp>
32
+ using AddOp = cub::Sum;
33
#endif
34
35
#define MAX(a, b) ((a) > (b) ? (a) : (b))
@@ -62,7 +65,6 @@ __launch_bounds__(TPB) __global__
65
66
const int thread_row_offset = blockIdx.x * num_cols;
67
65
- cub::Sum sum;
68
float threadData(-FLT_MAX);
69
70
// Don't touch finished rows.
@@ -92,7 +94,7 @@ __launch_bounds__(TPB) __global__
94
threadData += exp((static_cast<float>(input[idx]) - float_max));
95
}
96
95
- const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
97
+ const auto Z = BlockReduce(tmpStorage).Reduce(threadData, AddOp());
98
99
if (threadIdx.x == 0)
100
{