[None][feat] Optimize trtllmgen moe routing (#15656)

Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>

Author: jiahanc

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingCustom.cu

@@ -709,23 +709,78 @@ void launchDynBlockKernel(Data const& data, uint32_t numThreadsHist, void* strea
 //
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <typename KernelParams>
-#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
-__global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(NumThreads)
-    routingIndicesClusterKernel(KernelParams params)
+static constexpr int ClusterBlockDim256 = NumExperts256Experts;
+static constexpr int ClusterBlockDim512 = NumExperts512Experts;
+static constexpr int ClusterBlockDim1024 = NumThreads;
+static constexpr int MaxNumTokensClusterScores256 = NumBlocksPerCluster * (ClusterBlockDim256 / WarpSize);
+static constexpr int MaxNumTokensClusterScores512 = NumBlocksPerCluster * (ClusterBlockDim512 / WarpSize);
+
+template <typename TierT, typename TierListT>
+struct PrependTier;
+
+template <typename TierT, typename... Tiers>
+struct PrependTier<TierT, TierList<Tiers...>>
+{
+    using type = TierList<TierT, Tiers...>;
+};
+
+template <int ClusterBlockDim, typename TierListT>
+struct FilterClusterTiers;
+
+template <int ClusterBlockDim>
+struct FilterClusterTiers<ClusterBlockDim, TierList<>>
+{
+    using type = TierList<>;
+};
+
+template <int ClusterBlockDim, typename First, typename... Rest>
+struct FilterClusterTiers<ClusterBlockDim, TierList<First, Rest...>>
+{
+    using Tail = typename FilterClusterTiers<ClusterBlockDim, TierList<Rest...>>::type;
+    static constexpr bool IsValid = First::kExperts <= ClusterBlockDim || First::kExperts % ClusterBlockDim == 0;
+    using type = std::conditional_t<IsValid, typename PrependTier<First, Tail>::type, Tail>;
+};
+
+template <int ClusterBlockDim, typename PreProc, typename PostProc>
+struct ClusterPolicyTraits
+{
+    using Pairs = typename FilterClusterTiers<ClusterBlockDim, typename PolicyTraits<PreProc, PostProc>::Pairs>::type;
+};
+
+template <>
+struct ClusterPolicyTraits<ClusterBlockDim1024, NoOpPreprocess, SoftmaxPostprocess>
+{
+    using Pairs = TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, Tier<256, 16>, Tier<512, 8>,
+        Tier<512, 16>, Tier<512, 22>, Tier<512, 32>, Tier<576, 8>, Tier<768, 32>, Tier<1024, 32>, Tier<2048, 32>>;
+};
+
+template <>
+struct ClusterPolicyTraits<ClusterBlockDim512, NoOpPreprocess, SoftmaxPostprocess>
+{
+    using Pairs = TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, Tier<256, 16>, Tier<512, 8>,
+        Tier<512, 16>, Tier<512, 22>, Tier<512, 32>, Tier<1024, 32>, Tier<1536, 32>, Tier<2048, 32>>;
+};
+
+template <>
+struct ClusterPolicyTraits<ClusterBlockDim256, NoOpPreprocess, SoftmaxPostprocess>
+{
+    using Pairs = TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, Tier<256, 16>, Tier<512, 8>,
+        Tier<512, 16>, Tier<512, 22>, Tier<512, 32>, Tier<768, 32>, Tier<1024, 32>, Tier<1536, 32>, Tier<2048, 32>>;
+};
+
+template <typename KernelParams, typename BaseType, int ClusterBlockDim, int ClusterNumWarps>
+__device__ __forceinline__ void routingIndicesClusterKernelBody(
+    KernelParams params, PackedScoreIdx<BaseType>* smemPackedScoreIdx)
 {
     using OutputT = typename KernelParams::OutputT;
     using InputT = typename KernelParams::InputT;
-    using BaseType = typename KernelParams::ExpertSelectPolicy::template BaseType<InputT>;
     using TypePacked = PackedScoreIdx<BaseType>;
     static constexpr int VecSize = KernelParams::MaxNumExperts / WarpSize;
 
-    __shared__ TypePacked __attribute((aligned(128))) smemPackedScoreIdx[NumWarps * KernelParams::MaxNumTopExperts];
-
     uint32_t const clusterBlockRank = blockIdx.x;
     int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0);
     int32_t const laneIdx = cutlass::arch::LaneId();
-    auto warpTokenIdx = clusterBlockRank * NumWarps + warpIdx;
+    auto warpTokenIdx = clusterBlockRank * ClusterNumWarps + warpIdx;
     auto scoreOffset = warpTokenIdx * params.mNumExperts;
     bool validToken = warpTokenIdx < params.mNumTokens;
     auto block = cg::this_thread_block();
@@ -758,24 +813,136 @@ __global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(Nu
 
     if (params.mPtrScores != nullptr)
     {
-        routingPermutation<KernelParams, BaseType, NumThreads, NumWarps, KernelParams::MaxNumTopExperts,
+        routingPermutation<KernelParams, BaseType, ClusterBlockDim, ClusterNumWarps, KernelParams::MaxNumTopExperts,
             /*LoadExpertIdxFromGlobal=*/false>(params, smemPackedScoreIdx, warpIdx, clusterBlockRank);
     }
     else
     {
-        routingPermutation<KernelParams, BaseType, NumThreads, NumWarps, KernelParams::MaxNumTopExperts,
+        routingPermutation<KernelParams, BaseType, ClusterBlockDim, ClusterNumWarps, KernelParams::MaxNumTopExperts,
             /*LoadExpertIdxFromGlobal=*/true>(params, smemPackedScoreIdx, warpIdx, clusterBlockRank);
     }
 }
+
+template <typename KernelParams, int ClusterBlockDim = NumThreads>
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+__global__ void __cluster_dims__(NumBlocksPerCluster, 1, 1) __launch_bounds__(ClusterBlockDim)
+    routingIndicesClusterKernel(KernelParams params)
+{
+    using InputT = typename KernelParams::InputT;
+    using BaseType = typename KernelParams::ExpertSelectPolicy::template BaseType<InputT>;
+    using TypePacked = PackedScoreIdx<BaseType>;
+    static constexpr int NumWarpsBlock = ClusterBlockDim / WarpSize;
+    static_assert(ClusterBlockDim % WarpSize == 0);
+    static_assert(ClusterBlockDim <= NumThreads);
+    __shared__ TypePacked __attribute((aligned(128)))
+    smemPackedScoreIdx[NumWarpsBlock * KernelParams::MaxNumTopExperts];
+    routingIndicesClusterKernelBody<KernelParams, BaseType, ClusterBlockDim, NumWarpsBlock>(params, smemPackedScoreIdx);
+}
 #else
-__global__ void __launch_bounds__(NumThreads) routingIndicesClusterKernel(KernelParams /* params */)
+__global__ void __launch_bounds__(ClusterBlockDim) routingIndicesClusterKernel(KernelParams /* params */)
 {
     assert(false && "routingIndicesClusterKernel is only supported on SM90+ architectures");
 }
 #endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
 
+template <typename KernelParams, int ClusterBlockDim>
+void launchClusterKernelInstance(Data const& data, void* stream)
+{
+    static_assert(ClusterBlockDim % WarpSize == 0);
+    static_assert(ClusterBlockDim <= NumThreads);
+
+    cudaLaunchConfig_t config{};
+    config.gridDim = NumBlocksPerCluster;
+    config.blockDim = ClusterBlockDim;
+    config.dynamicSmemBytes = 0;
+    config.stream = (cudaStream_t) stream;
+
+    cudaLaunchAttribute attributes[2] = {};
+    attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+    attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl);
+    attributes[1].id = cudaLaunchAttributeCooperative;
+    attributes[1].val.cooperative = 0;
+    config.attrs = attributes;
+    config.numAttrs = 2;
+
+    auto params = KernelParams::setKernelParams(data);
+    auto kernelTyped = routingIndicesClusterKernel<KernelParams, ClusterBlockDim>;
+    TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernelTyped, params));
+}
+
+template <int ClusterBlockDim, typename PreProc, typename PostProc, int MaxNumExperts, int MaxNumTopExperts>
+void launchClusterKernelForTier(Data const& data, void* stream)
+{
+    using ExpertSelect = TopKExpertSelect<PreProc, PostProc>;
+    if (data.mDtypeOutput == tg::Dtype::Fp32)
+    {
+        using ParamsT = KernelParams<float, float, MaxNumExperts, MaxNumTopExperts, ExpertSelect>;
+        launchClusterKernelInstance<ParamsT, ClusterBlockDim>(data, stream);
+    }
+    else if (data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mDtypeInput == tg::Dtype::Fp32)
+    {
+        using ParamsT = KernelParams<float, __nv_bfloat16, MaxNumExperts, MaxNumTopExperts, ExpertSelect>;
+        launchClusterKernelInstance<ParamsT, ClusterBlockDim>(data, stream);
+    }
+    else if (data.mDtypeOutput == tg::Dtype::Bfloat16 && data.mDtypeInput == tg::Dtype::Bfloat16)
+    {
+        using ParamsT = KernelParams<__nv_bfloat16, __nv_bfloat16, MaxNumExperts, MaxNumTopExperts, ExpertSelect>;
+        launchClusterKernelInstance<ParamsT, ClusterBlockDim>(data, stream);
+    }
+    else
+    {
+        TLLM_LOG_ERROR("Unsupported dtype combination: dtypeOutput=%d, dtypeInput=%d",
+            static_cast<int>(data.mDtypeOutput), static_cast<int>(data.mDtypeInput));
+    }
+}
+
+template <int ClusterBlockDim, typename PreProc, typename PostProc>
+void launchClusterKernelForPolicy(Data const& data, void* stream)
+{
+    using Pairs = typename ClusterPolicyTraits<ClusterBlockDim, PreProc, PostProc>::Pairs;
+    bool dispatched = dispatchTierPairs(static_cast<Pairs*>(nullptr), data,
+        [&](auto eTag, auto kTag)
+        {
+            launchClusterKernelForTier<ClusterBlockDim, PreProc, PostProc, decltype(eTag)::value,
+                decltype(kTag)::value>(data, stream);
+        });
+    if (!dispatched)
+    {
+        TLLM_LOG_ERROR("No tier covers numExperts=%d topK=%d", data.mNumExperts, data.mTopK);
+    }
+}
+
+template <int ClusterBlockDim>
+void launchClusterKernelForBlockDim(Data const& data, void* stream)
+{
+    dispatchRoutingPolicy(data,
+        [&](auto preProc, auto postProc)
+        { launchClusterKernelForPolicy<ClusterBlockDim, decltype(preProc), decltype(postProc)>(data, stream); });
+}
+
 void launchClusterKernel(Data const& data, void* stream)
 {
+    // Each warp owns one token, so the reduced-thread cluster variants have lower token capacity.
+    // Use them only where the requested token count fits; otherwise keep the original 1024-thread launch.
+    if (data.mNumTokens <= MaxNumTokensClusterScores256)
+    {
+        launchClusterKernelForBlockDim<ClusterBlockDim256>(data, stream);
+        return;
+    }
+    if (data.mNumTokens <= MaxNumTokensClusterScores512)
+    {
+        launchClusterKernelForBlockDim<ClusterBlockDim512>(data, stream);
+        return;
+    }
+
+    bool const useNoOpSoftmaxScores = data.mPtrScores != nullptr && data.mPreprocessType == RoutingPreprocessType::None
+        && data.mPostprocessType == RoutingPostprocessType::Softmax;
+    if (useNoOpSoftmaxScores)
+    {
+        launchClusterKernelForPolicy<ClusterBlockDim1024, NoOpPreprocess, SoftmaxPostprocess>(data, stream);
+        return;
+    }
+
     LAUNCH_ROUTING_CUSTOM(data, false, routingIndicesClusterKernel, NumBlocksPerCluster, NumThreads,
         /*smemSize=*/0, // No dynamic smem
         stream);
@@ -788,15 +955,74 @@ void launchClusterKernel(Data const& data, void* stream)
 //
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+template <int MaxNumExperts, int MaxNumTopExperts>
+struct HistogramScoresLaunchConfig : DefaultRoutingLaunchConfig<MaxNumExperts, MaxNumTopExperts>
+{
+    static constexpr int DefaultBlockDim = DefaultRoutingLaunchConfig<MaxNumExperts, MaxNumTopExperts>::BlockDim;
+    static constexpr int HistogramScoresBlockDim = NumExperts256Experts;
+
+    // This kernel uses one warp per token and keeps per-warp arrays sized by both the expert tier and
+    // the topK tier. The 256-expert tier already launches with 256 threads; for larger tiers, fewer
+    // warps per CTA gives each thread more register headroom while preserving total warp-level
+    // parallelism by scaling the grid cap below.
+    static constexpr bool UseHistogramScoresBlockDim = DefaultBlockDim > HistogramScoresBlockDim;
+    static constexpr int BlockDim = UseHistogramScoresBlockDim ? HistogramScoresBlockDim : DefaultBlockDim;
+
+    static_assert(BlockDim % WarpSize == 0);
+    static_assert(BlockDim <= NumThreads);
+
+    static int blockDim(Data const& /*data*/, int /*numThreads*/)
+    {
+        return BlockDim;
+    }
+
+    static int gridDim(Data const& data, int numBlocks, int /*blockDim*/)
+    {
+        if constexpr (UseHistogramScoresBlockDim)
+        {
+            static constexpr int NumWarpsBlock = BlockDim / WarpSize;
+            static constexpr int MaxBlockScale = (DefaultBlockDim + BlockDim - 1) / BlockDim;
+            int const tokenBlocks = (static_cast<int>(data.mNumTokens) + NumWarpsBlock - 1) / NumWarpsBlock;
+            int const scaledMaxBlocks = numBlocks * MaxBlockScale;
+            int const selectedBlocks = tokenBlocks < scaledMaxBlocks ? tokenBlocks : scaledMaxBlocks;
+            return selectedBlocks > 0 ? selectedBlocks : 1;
+        }
+        else
+        {
+            return DefaultRoutingLaunchConfig<MaxNumExperts, MaxNumTopExperts>::gridDim(
+                data, numBlocks, DefaultBlockDim);
+        }
+    }
+};
+
+template <typename ExpertSelect, int MaxNumExperts, int MaxNumTopExperts>
+struct HistogramScoresKernelConfig : HistogramScoresLaunchConfig<MaxNumExperts, MaxNumTopExperts>
+{
+};
+
+template <typename PreProc, typename PostProc>
+struct HistogramScoresPolicyTraits : PolicyTraits<PreProc, PostProc>
+{
+};
+
+template <>
+struct HistogramScoresPolicyTraits<NoOpPreprocess, SoftmaxPostprocess>
+{
+    using Pairs
+        = TierList<Tier<128, 4>, Tier<128, 8>, Tier<160, 8>, Tier<256, 8>, Tier<256, 16>, Tier<512, 8>, Tier<512, 16>,
+            Tier<512, 22>, Tier<512, 32>, Tier<576, 8>, Tier<768, 32>, Tier<1024, 32>, Tier<1536, 32>, Tier<2048, 32>>;
+};
+
 template <typename KernelParams>
-__global__ void __launch_bounds__(KernelParams::MaxNumExperts <= 1024 ? KernelParams::MaxNumExperts : 1024)
+__global__ void __launch_bounds__(HistogramScoresKernelConfig<typename KernelParams::ExpertSelectPolicy,
+    KernelParams::MaxNumExperts, KernelParams::MaxNumTopExperts>::BlockDim)
     routingIndicesHistogramScoresKernel(KernelParams params)
 {
     using OutputT = typename KernelParams::OutputT;
     using InputT = typename KernelParams::InputT;
     using BaseType = typename KernelParams::ExpertSelectPolicy::template BaseType<InputT>;
-    // Cap actual thread count at 1024 when MaxNumExperts > 1024.
-    static constexpr int NumThreadsBlock = KernelParams::MaxNumExperts <= 1024 ? KernelParams::MaxNumExperts : 1024;
+    static constexpr int NumThreadsBlock = HistogramScoresKernelConfig<typename KernelParams::ExpertSelectPolicy,
+        KernelParams::MaxNumExperts, KernelParams::MaxNumTopExperts>::BlockDim;
 
     // VecSize stays based on MaxNumExperts — each warp still processes all experts for one token.
     static constexpr int VecSize = KernelParams::MaxNumExperts / WarpSize;
@@ -855,9 +1081,30 @@ __global__ void __launch_bounds__(KernelParams::MaxNumExperts <= 1024 ? KernelPa
 
 static void launchHistogramScoresKernel(Data const& data, uint32_t maxNumBlocks, uint32_t numThreadsHist, void* stream)
 {
-    LAUNCH_ROUTING_CUSTOM(data, false, routingIndicesHistogramScoresKernel, maxNumBlocks, numThreadsHist,
-        /*smemSize=*/0, // No dynamic smem
-        stream);
+    dispatchRoutingPolicy(data,
+        [&](auto preProc, auto postProc)
+        {
+            using PreProc = decltype(preProc);
+            using PostProc = decltype(postProc);
+            using Pairs = typename HistogramScoresPolicyTraits<PreProc, PostProc>::Pairs;
+            bool dispatched = dispatchTierPairs(static_cast<Pairs*>(nullptr), data,
+                [&](auto eTag, auto kTag)
+                {
+                    using ExpertSelect = TopKExpertSelect<PreProc, PostProc>;
+                    using LaunchConfig
+                        = HistogramScoresKernelConfig<ExpertSelect, decltype(eTag)::value, decltype(kTag)::value>;
+                    int const effectiveThreads = LaunchConfig::blockDim(data, static_cast<int>(numThreadsHist));
+                    int const effectiveBlocks
+                        = LaunchConfig::gridDim(data, static_cast<int>(maxNumBlocks), effectiveThreads);
+                    LAUNCH_ROUTING_WITH_POLICIES(data, false, routingIndicesHistogramScoresKernel, effectiveBlocks,
+                        effectiveThreads,
+                        /*smemSize=*/0, stream, PreProc, PostProc, decltype(eTag)::value, decltype(kTag)::value);
+                });
+            if (!dispatched)
+            {
+                TLLM_LOG_ERROR("No tier covers numExperts=%d topK=%d", data.mNumExperts, data.mTopK);
+            }
+        });
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////