fix: fast redux detection in trtllm gen routing kernel (NVIDIA#5941)

Signed-off-by: Yuan Tong <[email protected]>

Author: tongyuantongyu

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernel.cuh

@@ -28,8 +28,6 @@
 
 #include <type_traits>
 
-#include "tensorrt_llm/kernels/archCondition.h"
-
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 namespace moe::dev
@@ -53,10 +51,6 @@ static constexpr int NumWarpsHist = NumThreadsHist / WarpSize;
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
-static constexpr bool TLLM_GEN_HAS_FAST_REDUX = tensorrt_llm::kernels::arch::is_major_v<10>;
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
 static __device__ inline float sigmoid_accurate(float x)
 {
     return 0.5f * tanhf(0.5f * x) + 0.5f;

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingKernelTopK.cuh

@@ -20,6 +20,8 @@
 #include <cooperative_groups/reduce.h>
 #include <cub/cub.cuh>
 
+#include "tensorrt_llm/kernels/archCondition.h"
+
 namespace moe::dev::routing
 {
 
@@ -86,12 +88,8 @@ struct TopKRedType
 
     __device__ inline TypeCmp reduce(cg::thread_block_tile<WarpSize> const& warp)
     {
-#if defined(TLLM_GEN_HAS_FAST_REDUX)
-        static constexpr bool UseCg = false;
-#else
-        static constexpr bool UseCg = true;
-#endif
-        if constexpr (UseCg || sizeof(TypeCmp) == 8)
+        static constexpr bool hasFastRedux = tensorrt_llm::kernels::arch::is_major_v<10>;
+        if constexpr (!hasFastRedux || sizeof(TypeCmp) == 8)
         {
             return cg::reduce(warp, compVal, cg::greater<TypeCmp>{});
         }