[TRTLLM-15316][feat] Rubin trtllmgen batchedGemm MoE

.gitattributes

@@ -15,6 +15,7 @@ docs/source/blogs/media/tech_blog10_baseline_performance_detail.png filter=lfs d
 docs/source/blogs/media/tech_blog10_full_strategy_performance.png filter=lfs diff=lfs merge=lfs -text
 docs/source/blogs/media/tech_blog10_context_wait_performance.png  filter=lfs diff=lfs merge=lfs -text
 cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo_cubin.cpp filter=lfs diff=lfs merge=lfs -text
+cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo.h filter=lfs diff=lfs merge=lfs -text
 cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/cubin/xqa_kernel_cubin.cpp filter=lfs diff=lfs merge=lfs -text
 docs/source/blogs/media/tech_blog26_deepseek_v4_hybrid_attention.png filter=lfs diff=lfs merge=lfs -text
 docs/source/blogs/media/tech_blog26_deepseek_v4_mhc_moe.png filter=lfs diff=lfs merge=lfs -text

cpp/tensorrt_llm/common/envUtils.cpp

@@ -20,6 +20,7 @@
 #include "tensorrt_llm/common/cudaUtils.h"
 #include "tensorrt_llm/common/logger.h"
 #include "tensorrt_llm/common/stringUtils.h"
+#include <atomic>
 #include <cstddef>
 #include <cstdlib>
 #include <mutex>
@@ -274,6 +275,20 @@ bool getEnvEnableTrtllmgenMoeRoutingRenormPDL()
     return enabled;
 }
 
+static std::atomic<bool> gFineGrainedSyncDisabledOverride{false};
+
+void setFineGrainedSyncDisabledOverride(bool disabled)
+{
+    gFineGrainedSyncDisabledOverride.store(disabled, std::memory_order_relaxed);
+}
+
+bool getEnvUseFineGrainedSync()
+{
+    // Deliberately uncached: tests flip the env var between cases within one process.
+    return !gFineGrainedSyncDisabledOverride.load(std::memory_order_relaxed)
+        && getBoolEnv("TLLM_USE_FINE_GRAINED_SYNC");
+}
+
 bool getEnvUseUCXKvCache()
 {
     static bool const useUCXKVCache = getBoolEnv("TRTLLM_USE_UCX_KVCACHE");

cpp/tensorrt_llm/common/envUtils.h

@@ -65,6 +65,12 @@ bool getEnvEnableCascadeMmha();
 // Set TRTLLM_ENABLE_TRTLLMGEN_MOE_ROUTING_RENORM_PDL=1 to re-enable.
 bool getEnvEnableTrtllmgenMoeRoutingRenormPDL();
 
+// Set TLLM_USE_FINE_GRAINED_SYNC=1 to select fine-grained sync MoE kernel variants on SM107.
+bool getEnvUseFineGrainedSync();
+
+// Forces getEnvUseFineGrainedSync() to false while set; used by the autotuner during profiling.
+void setFineGrainedSyncDisabledOverride(bool disabled);
+
 template <typename KernelFn, typename... Args>
 inline void launchWithPdlWhenEnabled(char const* name, KernelFn kernelFn, dim3 grid, dim3 block, size_t dynamicShmSize,
     cudaStream_t stream, Args&&... args)

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/CMakeLists.txt

@@ -19,7 +19,7 @@ file(GLOB_RECURSE SRC_CPP *.cpp)
 file(GLOB_RECURSE SRC_CU *.cu)
 
 add_library(trtllm_gen_batched_gemm OBJECT)
-filter_source_cuda_architectures(SOURCE_LIST SRC_CPP ARCHS 100 103 100f)
+filter_source_cuda_architectures(SOURCE_LIST SRC_CPP ARCHS 100 103 107 100f)
 
 target_sources(trtllm_gen_batched_gemm PRIVATE ${SRC_CPP} ${SRC_CU})
 target_compile_definitions(trtllm_gen_batched_gemm
@@ -36,4 +36,4 @@ tllm_add_cubin_archive_sources(
   trtllm_gen_batched_gemm
   "${CMAKE_CURRENT_SOURCE_DIR}/trtllmGen_bmm_export/cubins"
   NAMESPACE batchedGemm tensorrt_llm ${TRTLLM_ABI_NAMESPACE} kernels
-  ARCHS 100 103 100f)
+  ARCHS 100 103 107 100f)

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

@@ -47,8 +47,7 @@ constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
 {
     // Blackwell family. Family-compatible (sm_100f) cubins load on any SM in
     // [100, 110); arch-specific ones are locked to the SM they were built for.
-    // The batchedGemm drop ships Sm100f, Sm100a and Sm103a only, so family
-    // members other than 100/103 are served by Sm100f alone.
+    // The batchedGemm drop ships Sm100f, Sm100a, Sm103a and Sm107a.
     // Range kept in sync with tensorrt_llm::common::isSM100Family (cudaUtils.h),
     // which is not constexpr and so cannot be reused here.
     if (gpuSM >= 100 && gpuSM < 110)
@@ -65,6 +64,10 @@ constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
         {
             return kernelSM == SmVersion::Sm103a;
         }
+        if (gpuSM == 107)
+        {
+            return kernelSM == SmVersion::Sm107a;
+        }
         return false;
     }
     else if (gpuSM == 90)
@@ -249,31 +252,9 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
             continue;
         }
 
-        auto sm = configs[i].mSm;
-        if (sm != SmVersion::Sm100f)
-        {
-            int smVersion = tensorrt_llm::common::getSMVersion();
-            if (smVersion == 100)
-            {
-                if (!acceptIf(sm == SmVersion::Sm100a,
-                        fmtstr("SM version 100 requires Sm100a (kernel has: %d)", static_cast<int>(sm))))
-                {
-                    continue;
-                }
-            }
-            else if (smVersion == 103)
-            {
-                if (!acceptIf(sm == SmVersion::Sm103a,
-                        fmtstr("SM version 103 requires Sm103a (kernel has: %d)", static_cast<int>(sm))))
-                {
-                    continue;
-                }
-            }
-        }
-
         if (options.mUseDeepSeekFp8)
         {
-            if (!acceptIf(options.mUseShuffledMatrix == false, "useShuffledMatrixA should be false for DeepSeek Fp8"))
+            if (!acceptIf(options.mUseShuffledMatrix == false, "useShuffledMatrix should be false for DeepSeek Fp8"))
             {
                 continue;
             }
@@ -302,6 +283,31 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunne
             continue;
         }
 
+        if (!acceptIf(options.mFineGrainedConsumerA == mOptions.fineGrainedConsumerA,
+                fmtstr("fineGrainedConsumerA mismatch (kernel: %d, expected: %d)", options.mFineGrainedConsumerA,
+                    mOptions.fineGrainedConsumerA)))
+        {
+            continue;
+        }
+        if (!acceptIf(options.mFineGrainedConsumerB == mOptions.fineGrainedConsumerB,
+                fmtstr("fineGrainedConsumerB mismatch (kernel: %d, expected: %d)", options.mFineGrainedConsumerB,
+                    mOptions.fineGrainedConsumerB)))
+        {
+            continue;
+        }
+        if (!acceptIf(options.mFineGrainedForceValid == mOptions.fineGrainedForceValid,
+                fmtstr("fineGrainedForceValid mismatch (kernel: %d, expected: %d)", options.mFineGrainedForceValid,
+                    mOptions.fineGrainedForceValid)))
+        {
+            continue;
+        }
+        if (!acceptIf(options.mFineGrainedProducer == mOptions.fineGrainedProducer,
+                fmtstr("fineGrainedProducer mismatch (kernel: %d, expected: %d)", options.mFineGrainedProducer,
+                    mOptions.fineGrainedProducer)))
+        {
+            continue;
+        }
+
         // Kernel passed all filters
         mPassingConfigIndices.push_back(i);
     }
@@ -459,7 +465,7 @@ void TrtllmGenBatchedGemmRunner::run(int32_t m, int32_t n, int32_t k, int32_t va
     bmm.runInitBeforeWorldSync(config, gemmData, static_cast<void*>(stream));
 
     auto const err = bmm.run(config, workspace, gemmData, static_cast<void*>(stream), multiProcessorCount,
-        tensorrt_llm::common::getEnvEnablePDL(), nullptr, globalTrtllmGenBatchedGemmModuleCache);
+        tensorrt_llm::common::getEnvEnablePDL(), nullptr, std::ref(globalTrtllmGenBatchedGemmModuleCache));
 
     CUresult cuErr = static_cast<CUresult>(err);
     char const* cuErrStr = nullptr;

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

@@ -103,6 +103,10 @@ struct TrtllmGenBatchedGemmRunnerOptions
     bool transposeMmaOutput{false};
     int32_t tileSize{8};
     int32_t epilogueTileM{128};
+    bool fineGrainedConsumerA{false};
+    bool fineGrainedConsumerB{false};
+    bool fineGrainedForceValid{false};
+    bool fineGrainedProducer{false};
 };
 
 class TrtllmGenBatchedGemmRunner

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllm/gen/CommonUtils.h

@@ -0,0 +1,63 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+namespace batchedGemm
+{
+
+namespace trtllm
+{
+namespace gen
+{
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// TMA OOB optimization constants.
+//
+// CUDA Programming Guide states that "globalDim must be non-zero and less than or equal to 2^32".
+// In practice, the kernel acts funny with TMA shape of 2^32 so we use 2^31.
+constexpr unsigned long TmaDimMax = 1UL << 31;
+// Chosen so that LargeN * XLargeN * sizeof(dtype) >= 2^64 which causes overflow and effectively
+// becomes 0. As sizeof(dtype) can be as small as 0.5B, we choose LargeN = 2^30 and XLargeN = 2^35
+// so overflow can happen.
+constexpr unsigned long LargeN = 1UL << 30;
+// Used in TMA stride. Should be less than 2^40.
+constexpr unsigned long XLargeN = 1UL << 35;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+inline T ceilDiv(T m, T n)
+{
+    return (m + n - T(1)) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+inline T roundUp(T m, T n)
+{
+    return ceilDiv(m, n) * n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace gen
+} // namespace trtllm
+
+} // namespace batchedGemm

cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllm/gen/CudaArchDecl.h

@@ -0,0 +1,136 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cassert>
+#include <cstring>
+#include <string>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Be careful when modifying this file as it is included by the generated kernels. For example, do
+// not add TLLM_CHECK_* constructs in this file. Thanks!
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace batchedGemm
+{
+
+namespace trtllm
+{
+namespace gen
+{
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class CudaArch
+{
+    // Hopper
+    Sm90a = 0,
+    // Blackwell
+    Sm100a,
+    // Blackwell-family
+    Sm100f,
+    // Blackwell Ultra
+    Sm103a,
+    // SM107
+    Sm107a,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline bool isArchHopper(CudaArch cudaArch)
+{
+    return cudaArch == CudaArch::Sm90a;
+}
+
+inline bool isArchRubin(CudaArch cudaArch)
+{
+    // Note: when compiling with a Blackwell target compatible with Rubin such as 100f,
+    // no Rubin-specific features shall be used.
+    if (cudaArch == CudaArch::Sm107a)
+    {
+        return true;
+    }
+    return false;
+}
+
+inline bool isArchBlackwell(CudaArch cudaArch)
+{
+    if (cudaArch == CudaArch::Sm107a)
+    {
+        return true;
+    }
+    return cudaArch == CudaArch::Sm100a || cudaArch == CudaArch::Sm100f || cudaArch == CudaArch::Sm103a;
+}
+
+inline bool isArchBlackwellUltra(CudaArch cudaArch)
+{
+    return cudaArch == CudaArch::Sm103a;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline std::string cudaArchToString(CudaArch cudaArch, bool isFull = true)
+{
+    switch (cudaArch)
+    {
+    case CudaArch::Sm90a: return isFull ? "90a" : "90";
+    case CudaArch::Sm100a: return isFull ? "100a" : "100";
+    case CudaArch::Sm100f: return isFull ? "100f" : "100";
+    case CudaArch::Sm103a: return isFull ? "103a" : "103";
+    case CudaArch::Sm107a: return isFull ? "107a" : "107";
+    default: assert(false); return "";
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline CudaArch stringToCudaArch(std::string const& str)
+{
+    if (str == "90a")
+    {
+        return CudaArch::Sm90a;
+    }
+    else if (str == "100a")
+    {
+        return CudaArch::Sm100a;
+    }
+    else if (str == "100f")
+    {
+        return CudaArch::Sm100f;
+    }
+    else if (str == "103a")
+    {
+        return CudaArch::Sm103a;
+    }
+    else if (str == "107a")
+    {
+        return CudaArch::Sm107a;
+    }
+    else
+    {
+        assert(false);
+        return CudaArch::Sm100a;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace gen
+} // namespace trtllm
+} // namespace batchedGemm