flashinfer-python 0.4.0rc2__tar.gz → 0.4.0rc3__tar.gz

{flashinfer_python-0.4.0rc2/flashinfer_python.egg-info → flashinfer_python-0.4.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.4.0rc2
+Version: 0.4.0rc3
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0
@@ -153,7 +153,7 @@ FlashInfer also provides C++ API and TVM bindings, please refer to [documentatio
 
 ## GPU Support
 
-FlashInfer currently provides support for NVIDIA SM architectures 80 and higher and beta support for 103, 110, 120, and 121.
+FlashInfer currently provides support for NVIDIA SM architectures 75 and higher and beta support for 103, 110, 120, and 121.
 
 ## Adoption
 

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/README.md

@@ -125,7 +125,7 @@ FlashInfer also provides C++ API and TVM bindings, please refer to [documentatio
 
 ## GPU Support
 
-FlashInfer currently provides support for NVIDIA SM architectures 80 and higher and beta support for 103, 110, 120, and 121.
+FlashInfer currently provides support for NVIDIA SM architectures 75 and higher and beta support for 103, 110, 120, and 121.
 
 ## Adoption
 

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/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,7 +331,7 @@ __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);
   }
 
@@ -344,8 +339,7 @@ __global__ void __launch_bounds__(KernelParams::NumExperts)
   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); }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/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,13 +384,7 @@ __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 &&
@@ -398,11 +392,11 @@ void runImpl(Data const& data, void* stream) {
               "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);
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/csrc/trtllm_fused_moe_routing_renormalize.cu

@@ -187,7 +187,7 @@ __global__ void __launch_bounds__(NumThreadsHist)
 #endif  // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
 
   // 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);
@@ -206,12 +206,12 @@ __global__ void __launch_bounds__(NumThreadsHist)
   BaseType warpTopKScore[MaxNumTopExperts];
   int32_t warpTopKExpertIdx[MaxNumTopExperts];
   for (int tokenIdx = globalWarpIdx; tokenIdx < params.mNumTokens; tokenIdx += globalWarpStride) {
-    auto scoreOffset = tokenIdx * KernelParams::NumExperts;
+    auto scoreOffset = tokenIdx * params.mNumExperts;
 
     routingTopKExperts<BaseType, InputT, VecSize, KernelParams::DoSoftmaxBeforeTopK>(
         warp, allScores, allExpertIdx, warpTopKScore, warpTopKExpertIdx, laneIdx,
-        KernelParams::NumExperts, params.mTopK, params.mPtrScores + scoreOffset,
-        params.mNormTopkProb, params.mApplySoftmaxAfterTopK);
+        params.mNumExperts, params.mTopK, params.mPtrScores + scoreOffset, params.mNormTopkProb,
+        params.mApplySoftmaxAfterTopK);
 
     if (laneIdx < params.mTopK) {
       PackedScoreIdx<OutputT> packedScore{static_cast<OutputT>(warpTopKScore[laneIdx]),
@@ -223,12 +223,7 @@ __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, "Renormalize 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 &&
@@ -236,11 +231,11 @@ void runImpl(Data const& data, void* stream) {
               "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);
@@ -257,10 +252,10 @@ void runImpl(Data const& data, void* stream) {
   }
 
   if (useSingleCluster) {
-    LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-        data, false, routingIndicesClusterKernel, NumBlocksPerCluster, NumThreads,
-        /*smemSize=*/0,  // No dynamic smem
-        stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false, NumExperts);
+    LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesClusterKernel, NumBlocksPerCluster,
+                                   NumThreads,
+                                   /*smemSize=*/0,  // No dynamic smem
+                                   stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false);
   } else {
     uint32_t const expandedIdxSize = data.mNumTokens * data.mTopK;
 
@@ -276,24 +271,24 @@ void runImpl(Data const& data, void* stream) {
         std::min((expandedIdxSize + offsetEltsPerBlock - 1) / offsetEltsPerBlock, maxNumBlocks);
 
     if (data.mPtrScores != nullptr) {
-      LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-          data, false, routingIndicesHistogramScoresKernel, maxNumBlocks, NumThreadsHist,
-          /*smemSize=*/0,  // No dynamic smem
-          stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false, NumExperts);
+      LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesHistogramScoresKernel, maxNumBlocks,
+                                     NumThreadsHist,
+                                     /*smemSize=*/0,  // No dynamic smem
+                                     stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false);
     } 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_WITH_EXTRA_FLAG(
-        data, false, routingIndicesHistogramKernel, numBlocksHistogram, NumThreadsHist,
-        /*smemSize=*/0,  // No dynamic smem
-        stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false, NumExperts);
-    LAUNCH_ROUTING_WITH_EXTRA_FLAG(
-        data, false, routingIndicesOffsetsKernel, numBlocksOffsets, NumThreadsHist,
-        /*smemSize=*/0,  // No dynamic smem
-        stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false, NumExperts);
+    LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesHistogramKernel, numBlocksHistogram,
+                                   NumThreadsHist,
+                                   /*smemSize=*/0,  // No dynamic smem
+                                   stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false);
+    LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, false, routingIndicesOffsetsKernel, numBlocksOffsets,
+                                   NumThreadsHist,
+                                   /*smemSize=*/0,  // No dynamic smem
+                                   stream, data.mDoSoftmaxBeforeTopK, /*forceFloatInput=*/false);
   }
 }
 
@@ -308,48 +303,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 8:
-      runImpl<8>(data, stream);
-      break;
-    case 16:
-      runImpl<16>(data, stream);
-      break;
-    case 24:
-      runImpl<24>(data, stream);
-      break;
-    case 32:
-      runImpl<32>(data, stream);
-      break;
-    case 48:
-      runImpl<48>(data, stream);
-      break;
-    case 64:
-      runImpl<64>(data, stream);
-      break;
-    case 72:
-      runImpl<72>(data, stream);
-      break;
-    case 96:
-      runImpl<96>(data, stream);
-      break;
-    case 128:
-      runImpl<128>(data, stream);
-      break;
-    case 192:
-      runImpl<192>(data, stream);
-      break;
-    case 256:
-      runImpl<256>(data, stream);
-      break;
-    case 384:
-      runImpl<384>(data, stream);
-      break;
-    default:
-      TORCH_CHECK(false, "Unsupported number of experts: ", data.mNumExperts,
-                  ". Supported values are: 8, 16, 24, 32, 48, 64, 72, 96, 128, 192, 256, 384");
-  }
+  runImpl(data, stream);
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/flashinfer/__init__.py

@@ -50,6 +50,9 @@ from .decode import (
 from .decode import (
     CUDAGraphBatchDecodeWithPagedKVCacheWrapper as CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
 )
+from .decode import (
+    fast_decode_plan as fast_decode_plan,
+)
 from .decode import cudnn_batch_decode_with_kv_cache as cudnn_batch_decode_with_kv_cache
 from .decode import single_decode_with_kv_cache as single_decode_with_kv_cache
 from .fp4_quantization import (
@@ -143,6 +146,7 @@ try:
 
     flashinfer_cubin_version = flashinfer_cubin.__version__
 
+    print("__version__", __version__)
     if __version__ != flashinfer_cubin_version:
         raise RuntimeError(
             f"flashinfer-cubin version ({flashinfer_cubin_version}) does not match "

flashinfer_python-0.4.0rc3/flashinfer/_build_meta.py

@@ -0,0 +1 @@
+__version__ = '0.4.0rc3'

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/flashinfer/aot.py

@@ -26,7 +26,6 @@ import shutil
 from itertools import product
 from pathlib import Path
 from typing import List, Tuple, Iterator
-import importlib
 
 import torch
 import torch.version
@@ -676,11 +675,7 @@ def main():
     # Start with default configuration
     project_root = Path(__file__).resolve().parents[1]
     config = get_default_config()
-    out_dir = (
-        project_root / "aot-ops"
-        if importlib.util.find_spec("flashinfer") is None
-        else jit_env.FLASHINFER_WORKSPACE_DIR / "cached_ops"
-    )
+    out_dir = project_root / "aot-ops"
     build_dir = project_root / "build" / "aot"
 
     # Override with command line arguments
@@ -723,17 +718,13 @@ def main():
     sm_capabilities = detect_sm_capabilities()
 
     # Update data dir
-    # if flashinfer installed, do not change these paths, otherwise change to the local paths
-
-    if importlib.util.find_spec("flashinfer") is None:
-        # NOTE(Zihao): this script can run before or after flashinfer is installed, if flashinfer is not installed, use the local paths
-        jit_env.FLASHINFER_CSRC_DIR = project_root / "csrc"
-        jit_env.FLASHINFER_INCLUDE_DIR = project_root / "include"
-        jit_env.CUTLASS_INCLUDE_DIRS = [
-            project_root / "3rdparty" / "cutlass" / "include",
-            project_root / "3rdparty" / "cutlass" / "tools" / "util" / "include",
-        ]
-        jit_env.SPDLOG_INCLUDE_DIR = project_root / "3rdparty" / "spdlog" / "include"
+    jit_env.FLASHINFER_CSRC_DIR = project_root / "csrc"
+    jit_env.FLASHINFER_INCLUDE_DIR = project_root / "include"
+    jit_env.CUTLASS_INCLUDE_DIRS = [
+        project_root / "3rdparty" / "cutlass" / "include",
+        project_root / "3rdparty" / "cutlass" / "tools" / "util" / "include",
+    ]
+    jit_env.SPDLOG_INCLUDE_DIR = project_root / "3rdparty" / "spdlog" / "include"
 
     # Update workdir
     jit_env.FLASHINFER_WORKSPACE_DIR = build_dir

{flashinfer_python-0.4.0rc2 → flashinfer_python-0.4.0rc3}/flashinfer/artifacts.py

@@ -18,7 +18,7 @@ import os
 import re
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
-
+from urllib.parse import urljoin
 import requests  # type: ignore[import-untyped]
 import shutil
 
@@ -94,6 +94,7 @@ class MetaInfoHash:
 
 
 def get_cubin_file_list():
+    base = FLASHINFER_CUBINS_REPOSITORY.rstrip("/")
     cubin_files = [
         (ArtifactPath.TRTLLM_GEN_FMHA + "include/flashInferMetaInfo", ".h"),
         (ArtifactPath.TRTLLM_GEN_GEMM + "include/flashinferMetaInfo", ".h"),
@@ -108,7 +109,7 @@ def get_cubin_file_list():
         cubin_files += [
             (kernel + name, extension)
             for name, extension in get_available_cubin_files(
-                FLASHINFER_CUBINS_REPOSITORY + "/" + kernel
+                urljoin(base + "/", kernel)
             )
         ]
     return cubin_files