Revise the code according to comment

Author: ChristinaZ

csrc/trtllm_fused_moe_routing_deepseek.cu

@@ -189,19 +189,19 @@ __global__ void routingMainKernel(KernelParams params) {
       __syncthreads();
       if (warpIdx == 0) {
         int constexpr NumInterTopKPerThread = (NumInterTopK * NumExpertWarps - 1) / WarpSize + 1;
-        float intermidiateScore[NumInterTopKPerThread];
-        int32_t intermidiateExpert[NumInterTopKPerThread];
+        float intermediateScore[NumInterTopKPerThread];
+        int32_t intermediateExpert[NumInterTopKPerThread];
         for (int i = laneIdx; i < NumInterTopKPerThread * WarpSize; i += WarpSize) {
           int ii = i / WarpSize;
           if (i < NumInterTopK) {
-            intermidiateScore[ii] = smemInterTopScores[i];
-            intermidiateExpert[ii] = smemInterTopExperts[i];
+            intermediateScore[ii] = smemInterTopScores[i];
+            intermediateExpert[ii] = smemInterTopExperts[i];
           } else {
-            intermidiateScore[ii] = invalidScoreFloat;
-            intermidiateExpert[ii] = KernelParams::MaxNumExperts - 1;
+            intermediateScore[ii] = invalidScoreFloat;
+            intermediateExpert[ii] = KernelParams::MaxNumExperts - 1;
           }
         }
-        topk::reduceTopK(warp, topScores, topExperts, intermidiateScore, intermidiateExpert,
+        topk::reduceTopK(warp, topScores, topExperts, intermediateScore, intermediateExpert,
                          /* minValue */ invalidScoreFloat, params.mTopK);
       }
     } else {

include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh

@@ -419,12 +419,13 @@ __global__ void __launch_bounds__(KernelParams::MaxNumExperts)
     PackedScoreIdx<OutputT> scoreIdx;
     int idx;
     if (params.mPtrTopKIds != nullptr) {
+      // If params.mPtrTopKIds != nullptr, we don't need to store the weights
       idx = params.mPtrTopKIds[expandedIdx];
     } else {
-      // If params.mPtrTopKIds != nullptr, we don't need to store the weights
+      scoreIdx = params.mPtrTopKPacked[expandedIdx];
+      idx = scoreIdx.idx;
       if (params.mPtrTopKWeights != nullptr) {
-        scoreIdx = params.mPtrTopKPacked[expandedIdx];
-        idx = scoreIdx.idx;
+        // For now, params.mPtrTopKWeights shouldn't be nullptr.
         params.mPtrTopKWeights[expandedIdx] = static_cast<OutputT>(scoreIdx.score);
       }
     }

include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh

@@ -53,8 +53,7 @@ struct TopKRedType {
   static __host__ __device__ inline TypeCmp makeCmpVal(TypeExpW val, int32_t idx = 0) {
     auto valueBits = cub::Traits<TypeExpW>::TwiddleIn(
         reinterpret_cast<typename cub::Traits<TypeExpW>::UnsignedBits&>(val));
-    TypeCmp compactTmp;
-    memcpy(&compactTmp, &valueBits, sizeof(valueBits));
+    TypeCmp compactTmp = valueBits;
     compactTmp = (compactTmp << moveBits) | (0xFFFF & (maxIdx - idx));
     // Use 65535 minus idx to give higher priority to elements with smaller indices.
     return compactTmp;
@@ -203,6 +202,9 @@ __forceinline__ __device__ void reduceTopK(cg::thread_block_tile<WarpSize> const
   static_assert(K < WarpSize, "Top K must have K < WarpSize");
   static_assert(N > 0, "Top K must have N > 0");
   static_assert(N <= 16, "Only support candidates number less than or equal to 16*32=512");
+  static_assert(
+      N <= 4 || N % 4 == 0,
+      "Only support candidates number is a multiple of 4*32=128 or less than or equal to 4");
   using RedType = TopKRedType<Type>;
 
   if constexpr (N <= 4) {

include/flashinfer/trtllm/fused_moe/runner.h

@@ -265,7 +265,7 @@ struct MoERunnerArgs {
   // Hidden dimension output of MoE block. It is not padded.
   // If not provided it is the same as hidden_size.
   std::optional<int32_t> hidden_size_output;
-  // TODO: only compiled routing kernel supports top_k = 8
+  // Now support top_k<=10
   int32_t top_k{0};
   int32_t n_group{0};
   // TODO: only compiled routing kernel supports topk_group = 4