[TRTLLM-15316][feat] Rubin sm107 trtllm-gen gemms

.pre-commit-config.yaml

@@ -1543,7 +1543,7 @@ legacy-files: &legacy_files |
 # list; the hook's own `files:` pattern only gates *when* the hook triggers.
 
 # 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/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$)'
+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$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)'
 
 default_install_hook_types: [pre-commit, commit-msg]
 repos:

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

@@ -19,7 +19,7 @@ file(GLOB_RECURSE SRC_CPP *.cpp)
 file(GLOB_RECURSE SRC_CU *.cu)
 
 add_library(trtllm_gen_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_gemm PRIVATE ${SRC_CPP} ${SRC_CU})
 target_compile_definitions(trtllm_gen_gemm PUBLIC TLLM_GEN_EXPORT_INTERFACE
@@ -32,5 +32,5 @@ set_property(TARGET trtllm_gen_gemm PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
 include(tllm_cubin_archive)
 tllm_add_cubin_archive_sources(
   trtllm_gen_gemm "${CMAKE_CURRENT_SOURCE_DIR}/trtllmGen_gemm_export/cubins"
-  NAMESPACE gemm tensorrt_llm kernels
-  ARCHS 100 103 100f)
+  NAMESPACE gemm tensorrt_llm ${TRTLLM_ABI_NAMESPACE} kernels
+  ARCHS 100 103 107 100f)

cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/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.
@@ -40,7 +40,12 @@ static GemmInterface::ModuleCache globalTrtllmGenGemmModuleCache;
 
 constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
 {
-    if (gpuSM == 103)
+    if (gpuSM == 107)
+    {
+        // SM107 can run SM100f family kernels or SM107a-specific kernels
+        return kernelSM == SmVersion::Sm107a || kernelSM == SmVersion::Sm100f;
+    }
+    else if (gpuSM == 103)
     {
         return kernelSM == SmVersion::Sm103a || kernelSM == SmVersion::Sm100f;
     }
@@ -52,6 +57,11 @@ constexpr bool isSMCompatible(int gpuSM, SmVersion kernelSM)
     {
         return kernelSM == SmVersion::Sm90a;
     }
+    // Redirect SM100 family (major version 10) to family kernels
+    else if (tensorrt_llm::common::isSM100Family(gpuSM))
+    {
+        return kernelSM == SmVersion::Sm100f;
+    }
     return true;
 }
 
@@ -88,6 +98,9 @@ size_t TrtllmGenGemmRunner::getWorkspaceSizeInBytes(int32_t m, int32_t n, int32_
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = mOptions.transposeMmaOutput ? n : m;
+    gemmData.mProblemDimensions.mValidN = mOptions.transposeMmaOutput ? m : n;
+    gemmData.mProblemDimensions.mValidK = k;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 
@@ -118,6 +131,9 @@ void TrtllmGenGemmRunner::run(int32_t m, int32_t n, int32_t k, void const* a, fl
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = mOptions.transposeMmaOutput ? n : m;
+    gemmData.mProblemDimensions.mValidN = mOptions.transposeMmaOutput ? m : n;
+    gemmData.mProblemDimensions.mValidK = k;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 
@@ -161,12 +177,17 @@ void TrtllmGenGemmRunner::selectGemmConfig(int32_t m, int32_t n, int32_t k)
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = mOptions.transposeMmaOutput ? n : m;
+    gemmData.mProblemDimensions.mValidN = mOptions.transposeMmaOutput ? m : n;
+    gemmData.mProblemDimensions.mValidK = k;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 
+    int const gpuSM = tensorrt_llm::common::getSMVersion();
+
     std::vector<int32_t> sortedIndices = mPassingConfigIndices;
     std::sort(sortedIndices.begin(), sortedIndices.end(),
-        [&configs, &gemmData](int32_t idx0, int32_t idx1)
+        [&configs, &gemmData, gpuSM](int32_t idx0, int32_t idx1)
         {
             auto const& optionsA = configs[idx0].mOptions;
             auto const& optionsB = configs[idx1].mOptions;
@@ -214,7 +235,15 @@ void TrtllmGenGemmRunner::selectGemmConfig(int32_t m, int32_t n, int32_t k)
                 return optionsA.mNumSlicesForSplitK > optionsB.mNumSlicesForSplitK;
             }
 
-            return true;
+            // On SM107, prefer Sm107a variants (2x mmaK) over Sm100f fallbacks; the generated
+            // table lists Sm100f first, so they would otherwise never be selected.
+            if (gpuSM == 107 && configs[idx0].mSm != configs[idx1].mSm)
+            {
+                return configs[idx0].mSm == SmVersion::Sm107a;
+            }
+
+            // Elements are equivalent - must return false for strict weak ordering
+            return false;
         });
 
     for (auto const& configIndex : sortedIndices)

cpp/tensorrt_llm/kernels/trtllmGenKernels/gemm/trtllmGen_gemm_export/Enums.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
+ * 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");
@@ -93,13 +93,90 @@ enum class BiasType : uint32_t
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+enum class FusedBiasShuffleMode : uint32_t
+{
+    // BiasType::Mn layout is fully prepared by host preprocessing.
+    None = 0,
+    // The host skips only the bias shuffle and the kernel applies it while loading bias.
+    Shuffle,
+    // The host skips both the gated-act bias reorder and the bias shuffle; the kernel applies both
+    // while loading bias.
+    ReorderAndShuffle,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Type of the element-wise activation to apply after the Gemm
+enum class EltwiseActType
+{
+    None = 0,
+    // Gelu is defined as the following operation:
+    // act = x0 * phi(x0)
+    // where x0 is the output of the Gemm
+    // phi is the CDF of standard normal distribution approximated by
+    // phi(x) = 0.5 * (1 + tanh(0.7978845608028654 * (x + 0.044715 * x * x * x)))
+    Gelu,
+    // Relu2 (also known as squared Relu) is defined as the following operation:
+    // act = relu(x0) ^ 2
+    // where x0 is the output of the Gemm.
+    Relu2,
+    // Silu is defined as the following operation:
+    // act = x0 * sigmoid(x0)
+    // where x0 is the output of the Gemm.
+    Silu,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
 enum class TileScheduler
 {
     // Static scheduler (Non-persistent).
     Static = 0,
-    // Dynamic persistent scheduler. This is either based on an atomically incremented global work id
-    // prior to SM100 archs, or the HW supported work id scheduler based on UGETNEXTWORKID for SM100+.
+    // Dynamic persistent scheduler for SM100+.
     Persistent,
+    // Static persistent scheduler. Launches a fixed grid size based on the number of SMs and uses
+    // the underlying PersistentTileSchedulerSm90 for static work distribution. Each CTA iterates
+    // through tiles and exits the loop by setting is_valid_tile to false when work is exhausted.
+    StaticPersistent,
+    // Dynamic persistent scheduler for SM90+ using atomicAdd on a global counter.
+    // Uses DynamicPersistentPipelinedTileSchedulerSm90 which enables work-stealing among CTAs
+    // by atomically fetching work tile indices from a global counter.
+    PersistentSm90,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Which task to fuse the persistent scheduler with.
+enum class SchedHostTask
+{
+    // The persistent scheduler has its own dedicated task.
+    Self = 0,
+    // The persistent scheduler is fused with the task loading A.
+    LoadA,
+    // The persistent scheduler is fused with the task loading B.
+    LoadB,
+    // The persistent scheduler is fused with the task loading scaling factors for A.
+    LoadSfA,
+    // The persistent scheduler is fused with the task loading scaling factors for B.
+    LoadSfB,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class CtaSwizzleType : uint32_t
+{
+    // Rasterize CTAs along the M dimension.
+    RasterizeAlongM = 0,
+    // Rasterize CTAs along the N dimension.
+    RasterizeAlongN,
+    // Swizzle CTAs in zig-zag pattern along M dimension, Zig-zag width is 2.
+    ZigZagAlongM2,
+    // Swizzle CTAs in zig-zag pattern along N dimension, Zig-zag width is 2.
+    ZigZagAlongN2,
+    // Swizzle CTAs in zig-zag pattern along M dimension, Zig-zag width is 4.
+    ZigZagAlongM4,
+    // Swizzle CTAs in zig-zag pattern along N dimension, Zig-zag width is 4.
+    ZigZagAlongN4,
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -136,6 +213,56 @@ BIAS_TYPE_FUNCTION(Mn)
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+#define FUSED_BIAS_SHUFFLE_MODE_FUNCTION(Mode)                                                                         \
+    inline bool isFusedBiasShuffleMode##Mode(FusedBiasShuffleMode mode)                                                \
+    {                                                                                                                  \
+        return (mode == FusedBiasShuffleMode::Mode);                                                                   \
+    }
+
+FUSED_BIAS_SHUFFLE_MODE_FUNCTION(None)
+FUSED_BIAS_SHUFFLE_MODE_FUNCTION(Shuffle)
+FUSED_BIAS_SHUFFLE_MODE_FUNCTION(ReorderAndShuffle)
+
+#undef FUSED_BIAS_SHUFFLE_MODE_FUNCTION
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline bool usesFusedBiasShuffle(FusedBiasShuffleMode mode)
+{
+    return !isFusedBiasShuffleModeNone(mode);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline bool usesFusedBiasReorder(FusedBiasShuffleMode mode)
+{
+    return isFusedBiasShuffleModeReorderAndShuffle(mode);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Helper function to check if a scheduler is persistent.
+inline bool isPersistentScheduler(TileScheduler scheduler)
+{
+    return scheduler == TileScheduler::Persistent || scheduler == TileScheduler::StaticPersistent
+        || scheduler == TileScheduler::PersistentSm90;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Helper function to check if CTA rasterization order is compatible with clean early exit of the
+// kernel. Clean early exit requires CTA indices to increase monotonically along the batch
+// dimension, so when a CTA exits the kernel early, it exits with all valid tiles already done.
+// Zigzag or batch-major patterns are NOT compatible because they may cause valid tiles to be
+// skipped when exiting early.
+inline bool supportsCleanEarlyExit(CtaSwizzleType swizzleType, bool batchM, TileScheduler /* scheduler */)
+{
+    return (
+        batchM ? (swizzleType == CtaSwizzleType::RasterizeAlongN) : (swizzleType == CtaSwizzleType::RasterizeAlongM));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
 } // namespace gemm
 
 } // namespace gemm