[None][feat] fuse shared to sparse experts in TRT-LLM Gen MoE (#11143)

Signed-off-by: Nikita Korobov <14355239+nekorobov@users.noreply.github.com>

Author: nekorobov

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/RoutingDeepSeek.cu

@@ -269,17 +269,31 @@ __global__ void routingMainKernel(KernelParams params)
             auto finalScore = OutputT{scoreNorm * params.mRouteScale / redNorm};
 
             // write expert idx out already
-            auto idxTopK = blockIdx.x * params.mTopK + laneIdx;
+            auto idxTopK = blockIdx.x * params.mTotalExpertsPerToken + laneIdx;
+            auto idxShared = blockIdx.x * params.mTotalExpertsPerToken + params.mTopK + laneIdx;
             if (laneIdx < params.mTopK && params.mPtrTopKPacked != nullptr)
             {
                 PackedScoreIdx<OutputT> packedScore{static_cast<OutputT>(finalScore), static_cast<int16_t>(expertIdx)};
                 params.mPtrTopKPacked[idxTopK] = packedScore;
             }
 
+            if (laneIdx < params.mNumFusedSharedExperts && params.mPtrTopKPacked != nullptr)
+            {
+                PackedScoreIdx<OutputT> packedScore{
+                    static_cast<OutputT>(1.0F), static_cast<int16_t>(params.mNumExperts + laneIdx)};
+                params.mPtrTopKPacked[idxShared] = packedScore;
+            }
+
             if (laneIdx < params.mTopK && params.mPtrTopKWeights != nullptr && params.mPtrTopKIds == nullptr)
             {
                 params.mPtrTopKWeights[idxTopK] = finalScore;
             }
+
+            // Write score of 1.0 for shared expert if enabled
+            if (laneIdx < params.mNumFusedSharedExperts && params.mPtrTopKWeights != nullptr)
+            {
+                params.mPtrTopKWeights[idxShared] = static_cast<OutputT>(1.0F);
+            }
         }
     }
 }
@@ -629,9 +643,15 @@ void run(Data& data, void* stream)
             "If permuted index is required, `mPtrTopKPacked` or `mPtrTopKIds` is also required");
     TLLM_CHECK_WITH_INFO(!data.mUseRoutingSoftmax, "Routing with softmax not implemented yet");
     int const numBlocks = data.mNumTokens;
-    int const numThreadsHist = getMaxNumExperts(data.mNumExperts);
+    int const numExperts = data.mNumExperts + data.mNumFusedSharedExperts;
+    int const topK = data.mTopK + data.mNumFusedSharedExperts;
+    int const numThreadsHist = getMaxNumExperts(numExperts);
+    int const maxNumTopExperts = getMaxNumExperts(numExperts);
+
+    // Number of threads in the cluster.
+    int numThreadsPerCluster = numThreadsHist * NumBlocksPerCluster;
 
-    bool const useSingleCluster = data.mNumTokens <= 1024;
+    bool const useSingleCluster = data.mNumTokens <= 1024 && data.mNumTokens * topK <= numThreadsPerCluster;
     if (!useSingleCluster)
     {
         // Reset the global histograms (not used in single-cluster code path).
@@ -658,15 +678,15 @@ void run(Data& data, void* stream)
     int const numBlocksCoop = smCount - 8;
 
     // Maximum number of tokens supported by the kernel using a cooperative launch.
-    int const maxTokensCoop = (numBlocksCoop * numThreadsHist * 64) / data.mTopK;
+    int const maxTokensCoop = (numBlocksCoop * numThreadsHist * 64) / topK;
     if (data.mPtrTopKIds == nullptr)
     {
         TLLM_CHECK_WITH_INFO(data.mNumExperts >= MaxSupportedTopExperts,
             "Routing kernel expects %d to be at most #experts %d", MaxSupportedTopExperts, data.mNumExperts);
         TLLM_CHECK_WITH_INFO(data.mNumExperts <= MaxSupportedExpertCount,
             "Routing kernel expects #experts %d  <= #threads %d", data.mNumExperts, MaxSupportedExpertCount);
-        TLLM_CHECK_WITH_INFO(data.mTopK <= MaxSupportedTopExperts, "Routing kernel expects topK experts <= %d, got %d",
-            MaxSupportedTopExperts, data.mTopK);
+        TLLM_CHECK_WITH_INFO(topK <= MaxSupportedTopExperts, "Routing kernel expects topK experts <= %d, got %d",
+            MaxSupportedTopExperts, topK);
 
         // Routing needs to be executed - validate routing kernel constraints
         if (data.mNumExpertGroups > 1)
@@ -690,6 +710,16 @@ void run(Data& data, void* stream)
                 data.mNumExpertGroups);
             TLLM_CHECK_WITH_INFO(data.mNumExperts % 4 == 0, "Routing kernel expects #experts %d to be a multiple of 4.",
                 data.mNumExperts);
+
+            TLLM_CHECK_WITH_INFO(data.mNumFusedSharedExperts <= WarpSize,
+                "Number of fused shared experts (%d must be less than warp size.", WarpSize);
+
+            if (data.mNumFusedSharedExperts > 0)
+            {
+                // Disabling due to lack of testing
+                // TLLM_CHECK_WITH_INFO(
+                //     data.mPtrTopKPacked == nullptr, "Shared expert fusion is not compatible with packed scores");
+            }
         }
 
         int const numThreadsMain = max(data.mNumExpertGroups * WarpSize, getMaxNumExperts(data.mNumExperts));
@@ -707,6 +737,14 @@ void run(Data& data, void* stream)
             stream, data.mNumExpertGroups > 1, /*forceFloatInput=*/false);
     }
 
+    if (data.mNumFusedSharedExperts > 0)
+    {
+        data.mNumExperts += data.mNumFusedSharedExperts;
+        data.mTopK += data.mNumFusedSharedExperts;
+        data.mNumLocalExperts += data.mNumFusedSharedExperts;
+        // data.mLocalExpertsStartIdx += data.mNumFusedSharedExperts;
+    }
+
     if (data.mPtrPermutedIdxSize != nullptr)
     {
         if (useSingleCluster)
@@ -725,7 +763,7 @@ void run(Data& data, void* stream)
         }
         else
         {
-            const int32_t expandedIdxSize = data.mNumTokens * data.mTopK;
+            const int32_t expandedIdxSize = data.mNumTokens * topK;
             const int32_t histogramEltsPerBlock = 8 * numThreadsHist;
             const int32_t offsetEltsPerBlock = NumEltsPerOffsetTilePerThread * numThreadsHist;
 

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

@@ -107,6 +107,12 @@ struct DataBase
     int32_t mLocalExpertsStartIdx;
     int32_t mLocalExpertsStrideLog2;
     int32_t mNumLocalExperts;
+
+    /// For fused shared expert
+    int32_t mNumFusedSharedExperts;
+    int32_t mSharedExpertTokenOffset;
+    int32_t mSharedExpertNumTokens;
+    int32_t mTotalExpertsPerToken;
 };
 
 template <typename InputT_, typename OutputT_, int MaxNumExperts_, bool isPow2_, bool UsePdl_>
@@ -141,6 +147,11 @@ struct KernelParamsBase
     int32_t mLocalExpertsStrideLog2 = 0;
     int32_t mNumLocalExperts = 0;
 
+    int32_t mNumFusedSharedExperts;
+    int32_t mSharedExpertTokenOffset;
+    int32_t mSharedExpertNumTokens;
+    int32_t mTotalExpertsPerToken;
+
     // Public initialization function - make it a template to accept different Data types
     template <typename DataType>
     void setBaseParams(DataType const& data)
@@ -165,6 +176,11 @@ struct KernelParamsBase
         mLocalExpertsStartIdx = data.mLocalExpertsStartIdx;
         mLocalExpertsStrideLog2 = data.mLocalExpertsStrideLog2;
         mNumLocalExperts = data.mNumLocalExperts;
+
+        mNumFusedSharedExperts = data.mNumFusedSharedExperts;
+        mSharedExpertTokenOffset = data.mSharedExpertTokenOffset;
+        mSharedExpertNumTokens = data.mSharedExpertNumTokens;
+        mTotalExpertsPerToken = data.mTotalExpertsPerToken;
     }
 };
 

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu

@@ -61,8 +61,8 @@ Runner::Runner(int32_t tileTokensDim)
 }
 
 void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK,
-    int32_t nGroup, int32_t topkGroup, int32_t localExpertOffset, int32_t localNumExperts, float routedScalingFactor,
-    int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize,
+    int32_t numFusedSharedExpert, int32_t nGroup, int32_t topkGroup, int32_t localExpertOffset, int32_t localNumExperts,
+    float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram, int32_t* permutedIdxSize,
     int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx,
     void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx,
     int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, btg::Dtype dtypeElt, bool useRoutingScalesOnInput,
@@ -76,6 +76,8 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3
         routingData.mDtypeExpW = btg::Dtype::Bfloat16;
         routingData.mUsePdl = true;
 
+        int32_t const totalExpertsPerToken = topK + numFusedSharedExpert;
+
         // output:
         routingData.mPtrTopKPacked = routingExpertIndexes;
         routingData.mPtrExpertCounts = expertCountHistogram;
@@ -96,16 +98,35 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3
         routingData.mPtrTopKIds = expertIds;
         routingData.mNumTokens = numTokens;
         routingData.mNumExperts = numExperts;
+        routingData.mNumFusedSharedExperts = numFusedSharedExpert;
         routingData.mNumExpertGroups = nGroup;
         routingData.mNumLimitedGroups = topkGroup;
         routingData.mTopK = topK;
+        routingData.mTotalExpertsPerToken = totalExpertsPerToken;
         routingData.mPaddingLog2 = computeLog2(mTileTokensDim);
         routingData.mTileTokensDim = mTileTokensDim;
         routingData.mLocalExpertsStartIdx = localExpertOffset;
         routingData.mLocalExpertsStrideLog2 = 0;
         routingData.mNumLocalExperts = localNumExperts;
         routingData.mRouteScale = routedScalingFactor;
         routingData.mUseRoutingSoftmax = false;
+
+        // TODO Should these be passed directly instead? This does assume a constant number of experts per device
+        int32_t const numDevices = numExperts / localNumExperts;
+        int32_t const deviceIndex = localExpertOffset / localNumExperts;
+        int32_t const baseTokensPerDevice = numTokens / numDevices;
+        int32_t const remainingTokens = numTokens % numDevices;
+
+        if (deviceIndex < remainingTokens)
+        {
+            routingData.mSharedExpertTokenOffset = (baseTokensPerDevice + 1) * deviceIndex;
+            routingData.mSharedExpertNumTokens = baseTokensPerDevice + 1;
+        }
+        else
+        {
+            routingData.mSharedExpertTokenOffset = remainingTokens + deviceIndex * baseTokensPerDevice;
+            routingData.mSharedExpertNumTokens = baseTokensPerDevice;
+        }
         moe::dev::routing::routingDeepSeek::run(routingData, stream);
     }
     else if (routingMethodType == RoutingMethodType::Llama4)
@@ -115,6 +136,8 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3
         {
             TLLM_LOG_WARNING("For Llama routing method, nGroup/topkGroup is ignored, got %d/%d.", nGroup, topkGroup);
         }
+        TLLM_CHECK_WITH_INFO(numFusedSharedExpert == 0, "Llama routing method does not support fusing shared expert");
+
         moe::dev::routing::routingLlama4::Data routingData;
         routingData.mDtypeExpW = btg::Dtype::Bfloat16;
         routingData.mUsePdl = true;
@@ -159,6 +182,9 @@ void Runner::run(void* routingLogits, void* routingBias, int32_t numTokens, int3
     else if (routingMethodType == RoutingMethodType::Renormalize /* default */
         || routingMethodType == RoutingMethodType::RenormalizeNaive /* Softmax -> TopK */)
     {
+        TLLM_CHECK_WITH_INFO(
+            numFusedSharedExpert == 0, "Renormalize routing method does not support fusing shared expert");
+
         moe::dev::routing::routingRenormalize::Data routingData;
 
         //
@@ -434,6 +460,9 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace
     moe::dev::convertsf::Data& convertSfData, moe::dev::activation::Data& activationData,
     moe::dev::finalize::Data& finalizeData)
 {
+    int32_t const totalNumExperts = args.num_experts + args.num_fused_shared_experts;
+    int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts;
+
     // Setup sf conversion data if needed
     convertSfData.inSfPtr = args.hidden_states_scale;
     convertSfData.outSfPtr = workspace.hidden_states_scale_linear;
@@ -452,7 +481,7 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace
     activationData.inDqSfsPtr = workspace.gemm1_output_scale;
     activationData.outDqSfsPtr = workspace.activation_output_scale;
     activationData.innerDim = args.intermediate_size * 2;
-    activationData.topK = args.top_k;
+    activationData.topK = totalExpertsPerToken; // TODO Rename topK in activation data struct
     activationData.numTokens = args.num_tokens;
     activationData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx;
 
@@ -479,8 +508,8 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace
         }
         finalizeData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx;
         finalizeData.numTokens = args.num_tokens;
-        finalizeData.numExperts = args.num_experts;
-        finalizeData.topK = args.top_k;
+        finalizeData.numExperts = totalNumExperts; // TODO Is this used?
+        finalizeData.topK = totalExpertsPerToken;  // TODO Rename topK in finalize data struct
         // We want to fuse unpadding into the finalize kernel, so we need to use the output hidden size.
         finalizeData.hiddenDim = args.valid_hidden_size.value_or(args.hidden_size);
         finalizeData.hiddenDimPadded = args.output_hidden_size.value_or(args.hidden_size);
@@ -490,12 +519,15 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace
 
 std::tuple<int32_t, int32_t> Runner::getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const
 {
+    int32_t const totalLocalExperts = args.local_num_experts + args.num_fused_shared_experts;
+    int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts;
+
     auto const& config = mPassingConfigs[configIndex];
 
-    auto workspace_size_fc1 = static_cast<int32_t>(mPermuteGemm1.getWorkspaceSizeInBytes(args.top_k, args.hidden_size,
-        args.intermediate_size, args.local_num_experts, args.num_tokens, config.gemm1Config));
-    auto workspace_size_fc2 = static_cast<int32_t>(mGemm2.getWorkspaceSizeInBytes(args.top_k, args.hidden_size,
-        args.intermediate_size, args.local_num_experts, args.num_tokens, config.gemm2Config));
+    auto workspace_size_fc1 = static_cast<int32_t>(mPermuteGemm1.getWorkspaceSizeInBytes(totalExpertsPerToken,
+        args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens, config.gemm1Config));
+    auto workspace_size_fc2 = static_cast<int32_t>(mGemm2.getWorkspaceSizeInBytes(totalExpertsPerToken,
+        args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens, config.gemm2Config));
     return std::make_tuple(workspace_size_fc1, workspace_size_fc2);
 }
 
@@ -530,7 +562,6 @@ std::vector<int64_t> Runner::getValidConfigIndices(int32_t topK, int32_t hiddenS
 int64_t Runner::getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize, int32_t intermediateSize,
     int32_t numLocalExperts, int32_t numTokens, int32_t validHiddenSize, int32_t validIntermediateSize) const
 {
-
     int32_t indexGemm1 = mPermuteGemm1.getDefaultValidConfigIndex(
         topK, hiddenSize, intermediateSize, numLocalExperts, numTokens, validHiddenSize, validIntermediateSize);
     int32_t indexGemm2 = mGemm2.getDefaultValidConfigIndex(
@@ -553,14 +584,17 @@ void Runner::run(
     sync_check_cuda_error(stream);
     setOpsData(args, workspace, convertSfData, activationData, finalizeData);
 
+    int32_t const totalLocalExperts = args.local_num_experts + args.num_fused_shared_experts;
+    int32_t const totalExpertsPerToken = args.top_k + args.num_fused_shared_experts;
+
     void* hidden_states_scale_linear{args.hidden_states_scale};
 
     auto const& config = mPassingConfigs[configIndex];
 
     mPermuteGemm1.run(args.hidden_states, hidden_states_scale_linear, args.gemm1_weights, args.gemm1_weights_scale,
         workspace.expert_weights, args.output1_scales_scalar, args.output1_scales_gate_scalar, args.gemm1_bias,
         args.gemm1_alpha, args.gemm1_beta, args.gemm1_clamp_limit, workspace.gemm1_output, workspace.gemm1_output_scale,
-        args.top_k, args.hidden_size, args.intermediate_size, args.local_num_experts, args.num_tokens,
+        totalExpertsPerToken, args.hidden_size, args.intermediate_size, totalLocalExperts, args.num_tokens,
         workspace.permuted_idx_to_token_idx, workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens,
         workspace.cta_idx_xy_to_batch_idx, workspace.cta_idx_xy_to_mn_limit, workspace.bmm1_workspace,
         args.mUseRoutingScalesOnInput, device, stream, config.gemm1Config,
@@ -581,11 +615,11 @@ void Runner::run(
 
     // Run gemm2
     mGemm2.run(gemm2_input, gemm2_input_scale, args.gemm2_weights, args.gemm2_weights_scale, args.output2_scales_scalar,
-        args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, args.top_k,
-        args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, args.local_num_experts,
-        args.num_tokens, workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens,
-        workspace.cta_idx_xy_to_batch_idx, workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream,
-        config.gemm2Config, args.valid_hidden_size.value_or(args.hidden_size),
+        args.gemm2_bias, workspace.gemm2_output, workspace.gemm2_output_scale, totalExpertsPerToken,
+        args.output_hidden_size.value_or(args.hidden_size), args.intermediate_size, totalLocalExperts, args.num_tokens,
+        workspace.num_non_exiting_ctas, workspace.total_num_padded_tokens, workspace.cta_idx_xy_to_batch_idx,
+        workspace.cta_idx_xy_to_mn_limit, workspace.bmm2_workspace, device, stream, config.gemm2Config,
+        args.valid_hidden_size.value_or(args.hidden_size),
         args.valid_intermediate_size.value_or(args.intermediate_size));
 
     // Run finalize

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h

@@ -147,13 +147,13 @@ class Runner
     explicit Runner(int32_t tileTokensDim);
 
     void run(void* routingLogits, void* routingBias, int32_t numTokens, int32_t numExperts, int32_t topK,
-        int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset, int32_t localNumExperts,
-        float routedScalingFactor, int32_t* routingExpertIndexes, int32_t* expertCountHistogram,
-        int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx,
-        int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds, int32_t* numTokensPerExpert,
-        int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas,
-        batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput, bool useDeepSeekFp8,
-        RoutingMethodType routingMethodType, cudaStream_t stream);
+        int32_t numFusedSharedExpert, int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset,
+        int32_t localNumExperts, float routedScalingFactor, int32_t* routingExpertIndexes,
+        int32_t* expertCountHistogram, int32_t* permutedIdxSize, int32_t* expandedIdxToPermutedIdx,
+        int32_t* permutedIdxToExpandedIdx, int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* expertIds,
+        int32_t* numTokensPerExpert, int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit,
+        int32_t* numNonExitingCtas, batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput,
+        bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream);
 
 private:
     int32_t mTileTokensDim;
@@ -268,6 +268,7 @@ struct MoERunnerArgs
 
     int32_t num_tokens{0};
     int32_t num_experts{0};
+    int32_t num_fused_shared_experts{0};
     // Hidden dimension input of MoE block. It might be padded.
     int32_t hidden_size{0};
     // Hidden dimension output of MoE block. It might be padded.