bugfix: partially fix tests/test_trtllm_gen_fused_moe.py unit test failure (#1724)

<!-- .github/pull_request_template.md -->

## 📌 Description

To fix the `Routing kernel expects top K 6 to be <= #warps 2` error
happens with `tests/test_trtllm_gen_fused_moe.py` on GB200 and GB300
Moreover, revert number of experts template in [PR
1696](#1696)

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: nv-yunzheq

csrc/trtllm_fused_moe_routing_deepseek.cu

@@ -22,10 +22,13 @@ namespace routingDeepSeek {
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+static constexpr int NumThreads = 384;
+static constexpr int NumWarps = NumThreads / WarpSize;
 static constexpr int NumTopGroupScores = 2;
 static constexpr int MaxNumTopExperts = 8;
+static constexpr int MaxNumTopGroupsDefault = 16;
 
-__host__ __device__ int constexpr getMaxNumTopGroups(const bool useGroups, const int numExperts) {
+__host__ __device__ int getMaxNumTopGroups(const bool useGroups, const int numExperts) {
   if (useGroups || numExperts <= 256) {
     return 4;
   } else {
@@ -38,10 +41,8 @@ __global__ void routingMainKernel(KernelParams params) {
   // declare types
   using OutputT = typename KernelParams::OutputT;
   using InputT = typename KernelParams::InputT;
-  static constexpr int NumThreads = KernelParams::NumExperts;  // DeepSeek uses 1 thread per expert
   static constexpr int NumWarps = NumThreads / WarpSize;
-  constexpr int MaxNumTopGroups =
-      getMaxNumTopGroups(KernelParams::UseGroups, KernelParams::NumExperts);
+  int MaxNumTopGroups = getMaxNumTopGroups(KernelParams::UseGroups, params.mNumExperts);
 
   // declare shared memory structure
   // number of experts is bounded by number of threads
@@ -71,19 +72,19 @@ __global__ void routingMainKernel(KernelParams params) {
 
   // load bias already; each warp represents one expert group
   auto threadExpert = threadIdx.x;
-  bool expertSelected = threadExpert < KernelParams::NumExperts;
+  bool expertSelected = threadExpert < params.mNumExperts;
   if constexpr (KernelParams::UseGroups) {
     threadExpert = warpIdx * params.mNumExpertsPerGroup + laneIdx;
     expertSelected = laneIdx < params.mNumExpertsPerGroup;
   }
-  auto scoreIdx = int64_t{blockIdx.x} * int64_t{KernelParams::NumExperts} + threadExpert;
+  auto scoreIdx = int64_t{blockIdx.x} * int64_t{params.mNumExperts} + threadExpert;
   auto biasVal = expertSelected ? params.mPtrRoutingBias[threadExpert] : invalidScore;
 
   // initialize the mPtrExpertCounts
   if (params.mPtrExpertCounts) {
     int32_t globalThreadIdx = blockIdx.x * NumThreads + threadIdx.x;
     int32_t globalThreadStride = gridDim.x * NumThreads;
-    int32_t expertCountsNum = 2 * KernelParams::NumExperts;
+    int32_t expertCountsNum = 2 * params.mNumExperts;
     initArr(globalThreadIdx, expertCountsNum, globalThreadStride, params.mPtrExpertCounts, 0);
   }
 
@@ -118,10 +119,10 @@ __global__ void routingMainKernel(KernelParams params) {
   // registers for top group score reduction
   float topExpGroupScores[NumTopGroupScores];
   [[maybe_unused]] int32_t topExpGroupIdx[NumTopGroupScores];
-  float topGroups[MaxNumTopGroups];  // bound of params.mNumLimitedGroups
-  int32_t topGroupIdx[MaxNumTopGroups];
-  float expertScoreGroup[MaxNumTopGroups];
-  int32_t expertIdxGroup[MaxNumTopGroups];
+  float topGroups[MaxNumTopGroupsDefault];  // bound of params.mNumLimitedGroups
+  int32_t topGroupIdx[MaxNumTopGroupsDefault];
+  float expertScoreGroup[MaxNumTopGroupsDefault];
+  int32_t expertIdxGroup[MaxNumTopGroupsDefault];
   float topScores[MaxNumTopExperts];  // bound of params.mTopK
   int32_t topExperts[MaxNumTopExperts];
 
@@ -168,15 +169,13 @@ __global__ void routingMainKernel(KernelParams params) {
     } else {
       // without groups, each thread just takes `MaxNumTopGroups` experts
 
-#pragma unroll
       for (int ii = 0; ii < MaxNumTopGroups; ++ii) {
         auto expertIdx = ii * WarpSize + laneIdx;
         expertIdxGroup[ii] = expertIdx;
         expertScoreGroup[ii] =
-            expertIdx < KernelParams::NumExperts ? smemScoreBias[expertIdx] : invalidScoreFloat;
+            expertIdx < params.mNumExperts ? smemScoreBias[expertIdx] : invalidScoreFloat;
       }
     }
-
     topk::reduceTopK(warp, topScores, topExperts, expertScoreGroup, expertIdxGroup,
                      /* minValue */ invalidScoreFloat, params.mTopK);
 
@@ -214,10 +213,8 @@ __global__ void routingMainKernel(KernelParams params) {
 template <typename KernelParams>
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
 __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1)
-    __launch_bounds__(KernelParams::NumExperts) routingIndicesClusterKernel(KernelParams params) {
+    routingIndicesClusterKernel(KernelParams params) {
   using OutputT = typename KernelParams::OutputT;
-  static constexpr int NumThreads = KernelParams::NumExperts;  // DeepSeek uses 1 thread per expert
-  static constexpr int NumWarps = NumThreads / WarpSize;
 
   int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0);
   int32_t const clusterBlockRank = blockIdx.x;
@@ -241,9 +238,7 @@ __global__ void routingIndicesClusterKernel(KernelParams params) {
 
 template <typename KernelParams>
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-__global__ void __launch_bounds__(KernelParams::NumExperts)
-    routingIndicesCoopKernel(KernelParams params) {
-  static constexpr int NumThreads = KernelParams::NumExperts;  // DeepSeek uses 1 thread per expert
+__global__ void routingIndicesCoopKernel(KernelParams params) {
   static constexpr int NumWarps = NumThreads / WarpSize;
   // number of experts is bounded by number of threads
   __shared__ int32_t __attribute((aligned(128))) smemExpertCount[NumThreads];
@@ -336,16 +331,15 @@ __global__ void __launch_bounds__(KernelParams::NumExperts)
   int32_t const localExpertCount = smemExpertCount[threadIdx.x];
 
   int32_t blockExpertOffset = 0;
-  if (threadIdx.x < KernelParams::NumExperts) {
+  if (threadIdx.x < params.mNumExperts) {
     blockExpertOffset = atomicAdd(&params.mPtrExpertCounts[threadIdx.x], localExpertCount);
   }
 
   // Sync to wait for completion of the histogram reduction.
   grid.sync();
 
   // Get total count for this expert.
-  int32_t count =
-      (threadIdx.x < KernelParams::NumExperts) ? params.mPtrExpertCounts[threadIdx.x] : 0;
+  int32_t count = (threadIdx.x < params.mNumExperts) ? params.mPtrExpertCounts[threadIdx.x] : 0;
 
   // Note: the scan is redundant in all CTAs, but doing it in only 1 CTA would be worse for latency.
 
@@ -422,16 +416,13 @@ __global__ void routingIndicesCoopKernel(KernelParams params) {
 #endif
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int NumExperts>
 void runImpl(Data& data, void* stream) {
-  static constexpr int NumThreads = NumExperts;  // DeepSeek: 1 thread per expert
   static constexpr int NumWarps = NumThreads / WarpSize;
-  const int MaxNumTopGroups = getMaxNumTopGroups(data.mNumExpertGroups > 1, NumExperts);
+  int MaxNumTopGroups = getMaxNumTopGroups(data.mNumExpertGroups > 1, data.mNumExperts);
 
   // Validate that the template parameter matches the data
-  TORCH_CHECK(data.mNumExperts == NumExperts, "DeepSeek routing kernel expects exactly ",
-              NumExperts, " experts, got ", data.mNumExperts);
+  // TORCH_CHECK(data.mNumExperts == NumExperts, "DeepSeek routing kernel expects exactly ",
+  //             NumExperts, " experts, got ", data.mNumExperts);
   TORCH_CHECK(data.mPtrExpertIdx != nullptr || data.mPtrPermutedIdxSize != nullptr ||
                   data.mPtrExpertWeights != nullptr,
               "Routing kernel requires at least one output parameter");
@@ -503,22 +494,21 @@ void runImpl(Data& data, void* stream) {
   LAUNCH_ROUTING_WITH_EXTRA_FLAG(data,
                                  /*coopLaunch=*/false, routingMainKernel, numBlocks, NumThreads,
                                  /*smemSize=*/0,  // No dynamic smem
-                                 stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true,
-                                 NumExperts);
+                                 stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true);
 
   if (data.mPtrPermutedIdxSize != nullptr) {
     if (useSingleCluster) {
-      LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-          data,
-          /*coopLaunch=*/false, routingIndicesClusterKernel, NumBlocksPerCluster, NumThreads,
-          /*smemSize=*/0,  // No dynamic smem
-          stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true, NumExperts);
+      LAUNCH_ROUTING_WITH_EXTRA_FLAG(data,
+                                     /*coopLaunch=*/false, routingIndicesClusterKernel,
+                                     NumBlocksPerCluster, NumThreads,
+                                     /*smemSize=*/0,  // No dynamic smem
+                                     stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true);
     } else if (data.mNumTokens <= maxTokensCoop) {
-      LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-          data,
-          /*coopLaunch=*/true, routingIndicesCoopKernel, numBlocksCoop, NumThreads,
-          /*smemSize=*/0,  // No dynamic smem
-          stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true, NumExperts);
+      LAUNCH_ROUTING_WITH_EXTRA_FLAG(data,
+                                     /*coopLaunch=*/true, routingIndicesCoopKernel, numBlocksCoop,
+                                     NumThreads,
+                                     /*smemSize=*/0,  // No dynamic smem
+                                     stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true);
     } else {
       const int32_t expandedIdxSize = data.mNumTokens * data.mTopK;
 
@@ -533,31 +523,21 @@ void runImpl(Data& data, void* stream) {
       int const numBlocksOffsets =
           std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks);
 
-      LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-          data,
-          /*coopLaunch=*/false, routingIndicesHistogramKernel, numBlocksHistogram, NumThreads,
-          /*smemSize=*/0,  // No dynamic smem
-          stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true, NumExperts);
-      LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-          data,
-          /*coopLaunch=*/false, routingIndicesOffsetsKernel, numBlocksOffsets, NumThreads,
-          /*smemSize=*/0,  // No dynamic smem
-          stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true, NumExperts);
+      LAUNCH_ROUTING_WITH_EXTRA_FLAG(data,
+                                     /*coopLaunch=*/false, routingIndicesHistogramKernel,
+                                     numBlocksHistogram, NumThreads,
+                                     /*smemSize=*/0,  // No dynamic smem
+                                     stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true);
+      LAUNCH_ROUTING_WITH_EXTRA_FLAG(data,
+                                     /*coopLaunch=*/false, routingIndicesOffsetsKernel,
+                                     numBlocksOffsets, NumThreads,
+                                     /*smemSize=*/0,  // No dynamic smem
+                                     stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/true);
     }
   }
 }
 
-void run(Data& data, void* stream) {
-  if (data.mNumExperts == 72) {
-    runImpl<72>(data, stream);
-  } else if (data.mNumExperts == 256) {
-    runImpl<256>(data, stream);
-  } else if (data.mNumExperts == 384) {
-    runImpl<384>(data, stream);
-  } else {
-    TORCH_CHECK(false, "Unsupported number of experts: ", data.mNumExperts);
-  }
-}
+void run(Data& data, void* stream) { runImpl(data, stream); }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 

csrc/trtllm_fused_moe_routing_llama4.cu

@@ -110,13 +110,13 @@ __global__ void __launch_bounds__(WarpSize) routingIndicesWarpKernel(KernelParam
     // for each token, we load the scores then use `reduceTopK` for this.
     // each thread works on 4 experts, so a local reduction is done before
     for (int tokenIdx = 0; tokenIdx < params.mNumTokens; ++tokenIdx) {
-      auto scoreOffset = tokenIdx * KernelParams::NumExperts;
+      auto scoreOffset = tokenIdx * params.mNumExperts;
       int32_t warpMaxExpertIdx[MaxNumTopExperts];
       InputT warpMaxScore[MaxNumTopExperts];
 
       // Use routingTopKExperts function instead of inline logic
       routingTopKExperts<InputT, ExpertsPerThread>(warp, warpMaxScore, warpMaxExpertIdx,
-                                                   threadIdx.x, KernelParams::NumExperts,
+                                                   threadIdx.x, params.mNumExperts,
                                                    params.mPtrScores + scoreOffset);
 
       if (cute::elect_one_sync()) {
@@ -285,7 +285,7 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu
 
   // TODO(mjoux): expand to more tokens (possibly)
   auto warpTokenIdx = clusterBlockRank * NumWarps + warpIdx;
-  auto scoreOffset = warpTokenIdx * KernelParams::NumExperts;
+  auto scoreOffset = warpTokenIdx * params.mNumExperts;
   bool validToken = warpTokenIdx < params.mNumTokens;
   InputT minScore = InputT{-INFINITY};
 
@@ -306,7 +306,7 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu
 
     if (validToken) {
       routingTopKExperts<InputT, MaxNumExperts / WarpSize>(warp, warpMaxScore, warpMaxExpertIdx,
-                                                           laneIdx, KernelParams::NumExperts,
+                                                           laneIdx, params.mNumExperts,
                                                            params.mPtrScores + scoreOffset);
       if (cute::elect_one_sync()) {
         auto finalScore = OutputT{sigmoid_accurate(float{warpMaxScore[0]})};
@@ -351,7 +351,7 @@ __global__ void __launch_bounds__(NumThreadsHist)
   auto warp = cg::tiled_partition<WarpSize>(block);
 
   // initialize the mPtrExpertCounts
-  int32_t expertCountsNum = 2 * KernelParams::NumExperts;
+  int32_t expertCountsNum = 2 * params.mNumExperts;
   int32_t globalThreadIdx = blockIdx.x * NumThreads + threadIdx.x;
   int32_t globalThreadStride = gridDim.x * NumThreads;
   initArr(globalThreadIdx, expertCountsNum, globalThreadStride, params.mPtrExpertCounts, 0);
@@ -367,12 +367,12 @@ __global__ void __launch_bounds__(NumThreadsHist)
   // in this case, each warp represents a token, and we use a grid-stride loop
   // over all warps/tokens
   for (int tokenIdx = globalWarpIdx; tokenIdx < params.mNumTokens; tokenIdx += globalWarpStride) {
-    auto scoreOffset = tokenIdx * KernelParams::NumExperts;
+    auto scoreOffset = tokenIdx * params.mNumExperts;
     int32_t warpMaxExpertIdx[MaxNumTopExperts];
     InputT warpMaxScore[MaxNumTopExperts];
 
     routingTopKExperts<InputT, MaxNumExperts / WarpSize>(warp, warpMaxScore, warpMaxExpertIdx,
-                                                         laneIdx, KernelParams::NumExperts,
+                                                         laneIdx, params.mNumExperts,
                                                          params.mPtrScores + scoreOffset);
 
     if (cute::elect_one_sync()) {
@@ -384,25 +384,19 @@ __global__ void __launch_bounds__(NumThreadsHist)
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int NumExperts>
 void runImpl(Data const& data, void* stream) {
-  // Validate that the runtime value matches the template parameter
-  TORCH_CHECK(data.mNumExperts == NumExperts, "Llama4 routing kernel expects #experts ",
-              data.mNumExperts, " to match template parameter ", NumExperts);
-
   TORCH_CHECK(data.mPtrExpertIdx != nullptr || data.mPtrScores != nullptr,
               "Routing kernel requires at least one input parameter");
   TORCH_CHECK(data.mPtrPermutedIdxSize != nullptr && data.mPtrCtaIdxXyToBatchIdx != nullptr &&
                   data.mPtrCtaIdxXyToMnLimit != nullptr && data.mPtrNumNonExitingCtas != nullptr,
               "Llama4 routing kernel expects permuted idx and grouped Gemm launch config buffers");
   TORCH_CHECK(data.mTopK <= MaxNumTopExperts,
               "Routing kernel expects topK experts <= ", MaxNumTopExperts, ", got ", data.mTopK);
-  TORCH_CHECK(NumExperts <= MaxNumExperts, "Routing kernel expects #experts ", NumExperts,
-              " to be at most max #experts ", MaxNumExperts);
+  TORCH_CHECK(data.mNumExperts <= MaxNumExperts, "Routing kernel expects #experts ",
+              data.mNumExperts, " to be at most max #experts ", MaxNumExperts);
   static_assert(MaxNumExperts <= NumThreads, "#experts must be bounded by #threads");
   static_assert(MaxNumExperts <= NumThreadsHist, "#experts must be bounded by #threads");
-  TORCH_CHECK(NumExperts % 4 == 0, "Routing kernel expects #experts ", NumExperts,
+  TORCH_CHECK(data.mNumExperts % 4 == 0, "Routing kernel expects #experts ", data.mNumExperts,
               " to be a multiple of 4.");
   TORCH_CHECK(data.mPaddingLog2 < 8, "Routing kernel expects padding log2 < 8, got ",
               data.mPaddingLog2);
@@ -424,13 +418,13 @@ void runImpl(Data const& data, void* stream) {
     LAUNCH_ROUTING(data,
                    /*coopLaunch=*/false, routingIndicesWarpKernel, 1, WarpSize,
                    /*smemSize=*/0,  // No dynamic smem
-                   stream, NumExperts);
+                   stream);
   } else if (useSingleCluster) {
     LAUNCH_ROUTING(data,
                    /*coopLaunch=*/false, routingIndicesClusterKernel, NumBlocksPerCluster,
                    NumThreads,
                    /*smemSize=*/0,  // No dynamic smem
-                   stream, NumExperts);
+                   stream);
   } else {
     const uint32_t expandedIdxSize = data.mNumTokens * data.mTopK;
 
@@ -450,23 +444,23 @@ void runImpl(Data const& data, void* stream) {
                      /*coopLaunch=*/false, routingIndicesHistogramScoresKernel, maxNumBlocks,
                      NumThreadsHist,
                      /*smemSize=*/0,  // No dynamic smem
-                     stream, NumExperts);
+                     stream);
     } else {
       // Reset the global histograms.
       CHECK_CUDA_ERROR(cudaMemsetAsync(data.mPtrExpertCounts, 0,
-                                       static_cast<size_t>(2 * NumExperts) * sizeof(int32_t),
+                                       static_cast<size_t>(2 * data.mNumExperts) * sizeof(int32_t),
                                        (cudaStream_t)stream));
     }
     LAUNCH_ROUTING(data,
                    /*coopLaunch=*/false, routingIndicesHistogramKernel, numBlocksHistogram,
                    NumThreadsHist,
                    /*smemSize=*/0,  // No dynamic smem
-                   stream, NumExperts);
+                   stream);
     LAUNCH_ROUTING(data,
                    /*coopLaunch=*/false, routingIndicesOffsetsKernel, numBlocksOffsets,
                    NumThreadsHist,
                    /*smemSize=*/0,  // No dynamic smem
-                   stream, NumExperts);
+                   stream);
   }
 }
 
@@ -481,18 +475,7 @@ void run(Data const& data, void* stream) {
   TORCH_CHECK(data.mPaddingLog2 < 8, "Routing kernel expects padding log2 < 8, got ",
               data.mPaddingLog2);
 
-  // Dispatch to the appropriate template instantiation based on the number of experts
-  switch (data.mNumExperts) {
-    case 16:
-      runImpl<16>(data, stream);
-      break;
-    case 128:
-      runImpl<128>(data, stream);
-      break;
-    default:
-      TORCH_CHECK(false, "Unsupported number of experts: ", data.mNumExperts,
-                  ". Supported values are: 16, 128");
-  }
+  runImpl(data, stream);
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////