From 226b452a2059166a00d4d3f620288fced4720c55 Mon Sep 17 00:00:00 2001
From: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Date: Tue, 22 Jul 2025 11:22:10 -0400
Subject: [PATCH] Revert "[Refactor] Fix Compile Warning #1444-D (#21208)"
 (#21384)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
---
 csrc/moe/topk_softmax_kernels.cu | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu
index ea4ff67ef3e40..064b76c9cd427 100644
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -20,7 +20,6 @@
 #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>
@@ -63,7 +62,7 @@ __launch_bounds__(TPB) __global__
 
     const int thread_row_offset = blockIdx.x * num_cols;
 
-    cuda::std::plus<float> sum;
+    cub::Sum sum;
     float threadData(-FLT_MAX);
 
     // Don't touch finished rows.