diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu
index 064b76c9cd42..ea4ff67ef3e4 100644
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -20,6 +20,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "../cuda_compat.h"
+#include <cuda/std/functional>
 
 #ifndef USE_ROCM
     #include <cub/util_type.cuh>
@@ -62,7 +63,7 @@ __launch_bounds__(TPB) __global__
 
     const int thread_row_offset = blockIdx.x * num_cols;
 
-    cub::Sum sum;
+    cuda::std::plus<float> sum;
     float threadData(-FLT_MAX);
 
     // Don't touch finished rows.