[TRTLLM-14701][feat] Update trtllm-gen batchedGemm kernel drop for Kimi K3 MoE

.pre-commit-config.yaml

@@ -1561,7 +1561,7 @@ static-analysis-files: &static_analysis_files |
     )$
 
 # Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx)
-exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)'
+exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$)'
 
 default_install_hook_types: [pre-commit, commit-msg]
 repos:

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <optional>
 #include <set>
 #include <unistd.h>
 #include <vector>
@@ -61,6 +62,22 @@ constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
     return false;
 }
 
+// Maps the runner-local gated ActType to the generated batchedGemm::gemmGatedAct::ActType.
+// The two enums evolve independently (the generator inserted SiTuGlu before None, shifting
+// values), so a value-preserving cast is not safe. Relu2/Silu are element-wise activations
+// (EltwiseActType) with no fused gated kernel counterpart and must never reach this mapping.
+static batchedGemm::gemmGatedAct::ActType toGemmGatedActType(ActType actType)
+{
+    switch (actType)
+    {
+    case ActType::SwiGlu: return batchedGemm::gemmGatedAct::ActType::SwiGlu;
+    case ActType::SiTu: return batchedGemm::gemmGatedAct::ActType::SiTuGlu;
+    case ActType::Relu2:
+    case ActType::Silu: break;
+    }
+    TLLM_THROW("ActType %d has no fused gated kernel mapping", static_cast<int>(actType));
+}
+
 static inline bool skipQuirks(BatchedGemmConfig const& config)
 {
     // Skip kernels that are known to hang/crash. Keep a record here for future reference.
@@ -117,6 +134,14 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
         rejectReason.resize(bmm.getNumBatchedGemmConfigs());
     }
 
+    // Only fused gated kernels compare the gated ActType; resolve the mapping once so a
+    // misconfigured runner (element-wise ActType + fusedAct) fails loudly at construction.
+    std::optional<batchedGemm::gemmGatedAct::ActType> gatedActType;
+    if (mOptions.fusedAct)
+    {
+        gatedActType = toGemmGatedActType(mOptions.actType);
+    }
+
     int gpuSM = tensorrt_llm::common::getSMVersion();
     for (size_t i = 0; i < bmm.getNumBatchedGemmConfigs(); ++i)
     {
@@ -149,6 +174,13 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
             continue;
         }
 
+        // The host runner does not wire the multicast completion barrier pointers; reject any
+        // future metadata that enables C multicast instead of launching with null barriers.
+        if (!acceptIf(!options.mUseCMultiCast, "mUseCMultiCast is not supported by the host runner"))
+        {
+            continue;
+        }
+
         if (!acceptIf(dtypeMatch,
                 fmtstr("dtypeAB mismatch (kernel: %s/%s, expected canonical: %s/%s)",
                     tg::dtypeToString(options.mDtypeA).c_str(), tg::dtypeToString(options.mDtypeB).c_str(),
@@ -235,9 +267,9 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
 
         if (options.mFusedAct)
         {
-            if (!acceptIf(options.mActType == static_cast<batchedGemm::gemmGatedAct::ActType>(mOptions.actType),
+            if (!acceptIf(options.mActType == *gatedActType,
                     fmtstr("actType mismatch (kernel: %d, expected: %d)", static_cast<int>(options.mActType),
-                        static_cast<int>(mOptions.actType))))
+                        static_cast<int>(*gatedActType))))
             {
                 continue;
             }
@@ -305,7 +337,7 @@ size_t TrtllmGenBatchedGemmRunner::getWorkspaceSizeInBytes(int32_t m, int32_t n,
     std::vector<int32_t> const& batchedTokens, int32_t numTokens, int32_t numBatches, int32_t maxNumCtasInBatchDim,
     int32_t configIndex) const
 {
-    BatchedGemmData gemmData;
+    BatchedGemmData gemmData{};
 
     auto bmm = BatchedGemmInterface();
 
@@ -330,7 +362,7 @@ void TrtllmGenBatchedGemmRunner::run(int32_t m, int32_t n, int32_t k, int32_t va
 {
     auto bmm = BatchedGemmInterface();
 
-    BatchedGemmData gemmData;
+    BatchedGemmData gemmData{};
 
     auto const configs = bmm.getBatchedGemmConfigs();
 
@@ -489,7 +521,7 @@ std::vector<int64_t> TrtllmGenBatchedGemmRunner::getValidConfigIndices(int32_t m
 
     int32_t multiProcessorCount = tensorrt_llm::common::getMultiProcessorCount();
 
-    BatchedGemmData gemmData;
+    BatchedGemmData gemmData{};
 
     // Sanitize optional valid dimensions
     validM = validM <= 0 ? m : validM;
@@ -549,7 +581,7 @@ std::vector<int64_t> TrtllmGenBatchedGemmRunner::getValidConfigIndices(int32_t m
         // prefer persistent tile scheduler.
         if (optionsA.mTileScheduler != optionsB.mTileScheduler)
         {
-            BatchedGemmData gemmData;
+            BatchedGemmData gemmData{};
             setProblemDimensions(gemmData, optionsA.mTransposeMmaOutput, m, n, k, batchedTokens, numTokens, numBatches,
                 maxNumCtasInBatchDim, validM, validN, validK);
             auto options = bmm.getOptionsFromConfigAndData(configs[idx0], gemmData);
@@ -608,7 +640,7 @@ bool TrtllmGenBatchedGemmRunner::isValidConfigIndex(int32_t configIndex, int32_t
     auto const bmm = BatchedGemmInterface();
     auto const configs = bmm.getBatchedGemmConfigs();
 
-    BatchedGemmData gemmData;
+    BatchedGemmData gemmData{};
 
     // Sanitize optional valid dimensions
     validM = validM <= 0 ? m : validM;

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2020-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,8 +28,11 @@ TRTLLM_NAMESPACE_BEGIN
 namespace kernels
 {
 
-// Keep this in sync with the ActType in
-// cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/GemmGatedActOptions.h
+// Backend-local activation type. The numeric values here are part of the Python/thop
+// contract (see ActType_TrtllmGen in tensorrt_llm/_torch/utils.py) and are intentionally
+// decoupled from the generated batchedGemm::gemmGatedAct::ActType, whose values shift
+// between generator exports. KernelRunner.cpp maps gated entries explicitly; never cast
+// between the two enums by value.
 enum class ActType
 {
     // For ActType == SwiGlu, ideally we would like to have something like
@@ -43,9 +46,24 @@ enum class ActType
     // GatedSilu is a special case of SwiGlu where the alpha is 1.0 and the beta is 0.0.
     SwiGlu,
     Relu2,
-    Silu
+    Silu,
+    // SiTu gated activation (Kimi K3). Gate on x1, matching the SwiGlu convention:
+    //    left     = beta  * tanh(x0 / beta)
+    //    right    = alpha * tanh(x1 / alpha) * sigmoid(x1)
+    //    gatedAct = left * right
+    // alpha/beta come from the per-expert mPtrGatedActAlpha/mPtrGatedActBeta runtime
+    // parameters and must be > 0.
+    SiTu
 };
 
+// Gated activations combine gate and up projections inside the FC1 kernel; the FC1
+// logical output width is 2 * intermediateSize. Element-wise activations (Relu2/Silu)
+// run on a single projection.
+constexpr bool isGatedActType(ActType actType)
+{
+    return actType == ActType::SwiGlu || actType == ActType::SiTu;
+}
+
 // Type of the element-wise activation to apply after the Gemm
 enum class EltwiseActType
 {

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/BatchedGemmInterface.h

@@ -163,8 +163,7 @@ struct BatchedGemmData
         // This is used for either:
         //   * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is
         //   Dtype::Float32
-        //   * When the routing scales are applied to the input activations (only when output is not
-        //   transposed). The dtype is Dtype::Bfloat16
+        //   * One-sided per-token scaling. The dtype is Dtype::Bfloat16 by default.
         //
         // if (batchM (A is activations)):
         //     Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)]
@@ -243,8 +242,7 @@ struct BatchedGemmData
         // This is used for either:
         //   * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is
         //   Dtype::Float32
-        //   * When the routing scales are applied to the input activations (only when output is
-        //   transposed). The dtype is Dtype::Bfloat16
+        //   * One-sided per-token scaling. The dtype is Dtype::Bfloat16 by default.
         //
         // if (batchM (B is weights)):
         //     Logical shape is [B, divUpMul(N, tileN)]
@@ -478,6 +476,35 @@ struct BatchedGemmData
         // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)]
         int32_t const* mPtrCtaIdxXyToMnLimit;
 
+        // Map from permuted token index to expanded token index.
+        // Used for MoE finalize direct-store path when mUseCMultiCast is enabled.
+        // The dtype is int32_t. Set to nullptr when not using CMulticast.
+        int32_t const* mPtrPermutedIdxToExpandedIdx{nullptr};
+
+        //////////////////////////////////////////////////////////////////////////////////////////////////
+        //
+        // Multicast barrier parameters
+        //
+        //////////////////////////////////////////////////////////////////////////////////////////////////
+
+        // When mUseCMultiCast is set, we require a barrier in global memory on each
+        // GPU to track completion of the multicast stores.
+        //
+        // The value of the barrier must be set to zero before the kernel is executed on any GPU.
+        // The barrier is self-resetting. It does not have to be reset by another kernel.
+        // When mUseCMultiCast is false, a nullptr may be passed.
+        uint32_t* mPtrMulticastCompletionBarUc;
+        uint32_t* mPtrMulticastCompletionBarMc;
+
+        //////////////////////////////////////////////////////////////////////////////////////////////////
+        //
+        // MoE finalize parameters for direct register-to-gmem store
+        //
+        //////////////////////////////////////////////////////////////////////////////////////////////////
+
+        // Pointer to expert weights for MoE finalize.
+        // The dtype depends on the use case.
+        void const* mPtrExpertWeightsPtr{nullptr};
         // Global counter for SW-emulated dynamic tile scheduling. When dynamic scheduling is enabled,
         // must be initialized to gridDim.x * gridDim.y (gridDim.z is reserved for K-dim in CGA splitK)
         // before each kernel launch.
@@ -558,8 +585,9 @@ class BatchedGemmInterface
 
     //////////////////////////////////////////////////////////////////////////////////////////////////
 
-    BatchedGemmInterface(bool const exportsCubin = false, int32_t const numRotations = 1)
-        : mExportsCubin(exportsCubin)
+    BatchedGemmInterface(int32_t rankId = 0, bool const exportsCubin = false, int32_t const numRotations = 1)
+        : mRankId(rankId)
+        , mExportsCubin(exportsCubin)
         , mNumRotations(numRotations)
     {
     }
@@ -655,7 +683,12 @@ class BatchedGemmInterface
             batchedGemmData.mInputBuffers.mPtrScaleGate, batchedGemmData.mInputBuffers.mPtrClampLimit,
             batchedGemmData.mInputBuffers.mPtrGatedActAlpha, batchedGemmData.mInputBuffers.mPtrGatedActBeta,
             batchedGemmData.mInputBuffers.mPtrRouteMap, dPtrRowMax, dPtrRowMaxBars,
-            batchedGemmData.mInputBuffers.mPtrNumNonExitingCtas, batchedGemmData.mInputBuffers.mPtrTotalNumPaddedTokens,
+            batchedGemmData.mProblemDimensions.mRank, batchedGemmData.mProblemDimensions.mWorldSize,
+            batchedGemmData.mInputBuffers.mPtrMulticastCompletionBarUc,
+            batchedGemmData.mInputBuffers.mPtrMulticastCompletionBarMc,
+            batchedGemmData.mInputBuffers.mPtrPermutedIdxToExpandedIdx,
+            batchedGemmData.mInputBuffers.mPtrExpertWeightsPtr, batchedGemmData.mInputBuffers.mPtrNumNonExitingCtas,
+            batchedGemmData.mInputBuffers.mPtrTotalNumPaddedTokens,
             batchedGemmData.mInputBuffers.mPtrCtaIdxXyToBatchIdx, batchedGemmData.mInputBuffers.mPtrCtaIdxXyToMnLimit,
             numCtaBatch, batchedGemmData.mInputBuffers.mPtrDynamicTileCounter);
 
@@ -925,8 +958,14 @@ class BatchedGemmInterface
         auto options = getOptionsFromConfigAndData(config, data);
 
         // Check options without modifications.
+        //
+        // NOTE(TRT-LLM local fix over producer output a2ad0544-dirty): the generator
+        // inserted the `worldSize` parameter into checkAndUpdateBatchedGemmOptions
+        // without updating this call, so the literal `false` bound positionally to
+        // `worldSize` (=0) and `updateOptions` silently took its default `true`.
+        // Pass both arguments explicitly. Report upstream before the next re-export.
         return checkAndUpdateBatchedGemmOptions(options, config.mSm,
-            /* updateOptions */ false);
+            /* worldSize */ 1, /* updateOptions */ false);
     }
 
     //////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1013,6 +1052,8 @@ class BatchedGemmInterface
     //////////////////////////////////////////////////////////////////////////////////////////////////
 
 private:
+    // The rank id of the current device in the multi-gpu space.
+    int32_t mRankId;
     // Whether to export the cubin file.
     bool mExportsCubin;
     // The number of rotations.