Repository: NVIDIA/TensorRT-LLM
Pull request #4047: feat: Add heuristic for GroupRMSNorm kernel selection.
Author: login: SimengLiu-nv
Created: 2025-05-04T23:58:55Z | Merged: 2025-05-13T00:52:53Z

# [feat] Add heuristic for GroupRMSNorm kernel selection.
## Description
This PR is a follow up of [feat: Add group_rms_norm kernel to normalize multiple inputs in a single operator.](https://github.com/NVIDIA/TensorRT-LLM/pull/3438)

This PR implements a logistic regression model to dynamically select between:
- GroupRMSNormBase: Allocates warps proportional to sum of dimensions (better SM occupancy in most cases)
- GroupRMSNormLargeBatch: Allocates warps proportional to max dimension (better block scheduling in large batch scenarios)

Selection heuristic considers batch size, allocated warps, and scheduling efficiency on the current GPU architecture. Models for Compute Capability 9.x and 10.x are trained base on nsys kernel runtime data. The default kernel selection is the base kernel.

The python operator group_rms_norm will use the heuristic by default. User can pick to use the base or large batch kernels as well.

## Test Coverage
Accuracy tests for group_rms_norm with all kernels.

## Benchmark
Here is the kernel runtime comparison using flashinfer RMS kernel vs the GroupRMSNorm ones on B200.
Assume that the dimension of input_a > input_b. The baseline is the flashinfer kernel with input_a. I benchmarked and find out that running another kernel in a different stream to process input_b does not slow down the kernel processing input_a.

For the following table. GroupRMSLargeBatch performs better when batch size >= 256 with dims = [5120, 1024] (Llama4 config). 

:blush: The heuristic picks the right kernels for all batch sizes.

Batch Size | FlashInfer(ns) | GroupRMSBase(ns) | GroupRMSLargeBatch(ns) | Speedup_Base | Speedup_LargeBatch | Heuristic
-- | -- | -- | -- | -- | -- | --
1 | 2089.9 | 1843.6 | 1917.7 | 1.1336 | 1.0898 | Base
8 | 2253.1 | 1966.9 | 2068.9 | 1.1455 | 1.0890 | Base
16 | 2265.3 | 2091.4 | 2201.3 | 1.0832 | 1.0291 | Base
32 | 2387.1 | 2153.4 | 2245.2 | 1.1085 | 1.0632 | Base
64 | 2421.3 | 2169.6 | 2270.1 | 1.1160 | 1.0666 | Base
128 | 2538.7 | 2261.7 | 2339.2 | 1.1225 | 1.0853 | Base
256 | 3011.5 | 2773.1 | 2670.6 | 1.0860 | 1.1276 | LargeBatch
512 | 4619.7 | 4365.7 | 4242.4 | 1.0582 | 1.0889 | LargeBatch
1024 | 7145.9 | 7233.0 | 6453.1 | 0.9880 | 1.1074 | LargeBatch
2048 | 12246.2 | 12543.2 | 10656.6 | 0.9763 | 1.1492 | LargeBatch
4096 | 25732.8 | 28818.1 | 22852.0 | 0.8929 | 1.1261 | LargeBatch
8192 | 51040.1 | 59599.5 | 46570.5 | 0.8564 | 1.0960 | LargeBatch
16384 | 98505.7 | 116169.4 | 90321.5 | 0.8479 | 1.0906 | LargeBatch

For the following table. GroupRMSLargeBatch performs better when batch size > 4096 with dims = [1536, 512] (DS3 config). The heuristic picks the large batch kernel for batch size > 2048.  

:sweat: Wrong kernel is picked for batch size 2048 but the runtime difference is only 1.5%.
Batch Size | FlashInfer(ns) | GroupRMSBase(ns) | GroupRMSLargeBatch(ns) | Speedup_Base | Speedup_LargeBatch | Heuristic
-- | -- | -- | -- | -- | -- | --
1 | 1719.7 | 1473.7 | 1658.6 | 1.1669 | 1.0368 | Base
8 | 1819.6 | 1561.5 | 1839.2 | 1.1653 | 0.9893 | Base
16 | 1874.4 | 1625.5 | 1894.7 | 1.1531 | 0.9893 | Base
32 | 2069.5 | 1812.7 | 2097.9 | 1.1417 | 0.9865 | Base
64 | 2098.7 | 1853.1 | 2111.0 | 1.1325 | 0.9942 | Base
128 | 2113.3 | 1875.9 | 2144.3 | 1.1266 | 0.9855 | Base
256 | 2183.2 | 1956.6 | 2195.6 | 1.1158 | 0.9944 | Base
512 | 2419.8 | 2187.0 | 2407.4 | 1.1064 | 1.0052 | Base
1024 | 3114.4 | 2881.8 | 3037.5 | 1.0807 | 1.0253 | Base
2048 | 4699.5 | 4404.2 | 4471.4 | 1.0670 | 1.0510 | LargeBatch
4096 | 7809.9 | 7096.5 | 6954.2 | 1.1005 | 1.1230 | LargeBatch
8192 | 13565.4 | 12894.2 | 12151.6 | 1.0521 | 1.1163 | LargeBatch
16384 | 28830.2 | 28066.2 | 26844.1 | 1.0272 | 1.0740 | LargeBatch




Diff:
diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
index 1b86abaeefe2..b13c8e100fdf 100644
--- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
+++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
@@ -25,9 +25,66 @@
 
 namespace tensorrt_llm::kernels::group_rms_norm
 {
+// Helper function to calculate the number of warps to launch for GroupRMSNormBase
+template <typename DType, int n>
+uint32_t calculateNumWarpsBase(GroupRMSParams<n> const& params)
+{
+    constexpr uint32_t kPackedSize = sizeof(float4) / sizeof(DType);
+    uint32_t input_chunk_per_warp = 32 * kPackedSize;
+
+    // Calculate rounded input dimensions and total input length
+    int rounded_input_dims[n];
+    for (uint32_t i = 0; i < params.num_inputs; i++)
+    {
+        // Make rounded_input_dims[i] a multiple of 32 * kPackedSize
+        rounded_input_dims[i]
+            = (params.input_last_dims[i] + input_chunk_per_warp - 1) / input_chunk_per_warp * input_chunk_per_warp;
+    }
+
+    // Calculate total warps needed
+    uint32_t total_input_length = std::accumulate(rounded_input_dims, rounded_input_dims + params.num_inputs, 0);
+    uint32_t total_warps_needed = total_input_length / input_chunk_per_warp;
+    return std::min<uint32_t>(32, total_warps_needed);
+}
+
+template <typename DType, int n>
+struct LargeBatchWarpsInfo
+{
+    uint32_t num_warps_to_launch;   // Total warps to launch
+    uint32_t num_warps_to_launch_0; // Warps for first input
+    uint32_t num_warps_to_launch_1; // Warps for second input
+    uint32_t rounds_0;              // Rounds for first input
+    uint32_t rounds_1;              // Rounds for second input
+};
+
+// Helper function to calculate the number of warps to launch for GroupRMSNormKernelLargeBatch
+template <typename DType, int n>
+LargeBatchWarpsInfo<DType, n> calculateNumWarpsLargeBatch(GroupRMSParams<n> const& params)
+{
+    constexpr uint32_t kPackedSize = sizeof(float4) / sizeof(DType);
+    uint32_t input_chunk_per_warp = 32 * kPackedSize;
+
+    // Calculate warps needed for each input
+    uint32_t warps_needed_0 = (params.input_last_dims[0] + input_chunk_per_warp - 1) / input_chunk_per_warp;
+    uint32_t warps_needed_1 = (params.input_last_dims[1] + input_chunk_per_warp - 1) / input_chunk_per_warp;
+
+    LargeBatchWarpsInfo<DType, n> info;
+    info.num_warps_to_launch_0 = std::min((uint32_t) 32, warps_needed_0);
+    info.num_warps_to_launch_1 = std::min((uint32_t) 32, warps_needed_1);
+
+    // Use the maximum of the two for the final warps to launch
+    info.num_warps_to_launch = std::max(info.num_warps_to_launch_0, info.num_warps_to_launch_1);
+
+    // Calculate rounds needed for each input
+    info.rounds_0 = (warps_needed_0 + info.num_warps_to_launch_0 - 1) / info.num_warps_to_launch_0;
+    info.rounds_1 = (warps_needed_1 + info.num_warps_to_launch_1 - 1) / info.num_warps_to_launch_1;
+
+    return info;
+}
+
 // Allocate more warps to deal with the second input
 template <typename DType, typename PackedType, int n, bool EnableWeights, bool MultiRounds>
-__global__ void GroupRMSNormKernel(GroupRMSParams<n> params, int rounds)
+__global__ void GroupRMSNormBaseKernel(GroupRMSParams<n> params, int rounds)
 {
     const uint32_t batch_idx = blockIdx.x; // Maps to batch size
     constexpr uint32_t warp_size = 32;
@@ -510,7 +567,7 @@ __global__ void GroupRMSNormKernelLargeBatch(
 }
 
 template <typename DType, int n, bool EnableWeights>
-void GroupRMSNormKernel(GroupRMSParams<n> params)
+void GroupRMSNormBaseKernel(GroupRMSParams<n>& params)
 {
     // Kernel assertions
     constexpr uint32_t kPackedSize = sizeof(float4) / sizeof(DType);
@@ -535,7 +592,7 @@ void GroupRMSNormKernel(GroupRMSParams<n> params)
     // Calculate total warps to launch and rounds needed
     uint32_t total_input_length = std::accumulate(rounded_input_dims, rounded_input_dims + params.num_inputs, 0);
     uint32_t total_warps_needed = total_input_length / input_chunk_per_warp;
-    uint32_t num_warps_to_launch = std::min<uint32_t>(32, total_warps_needed);
+    uint32_t num_warps_to_launch = calculateNumWarpsBase<DType, n>(params);
     uint32_t rounds = (total_warps_needed + num_warps_to_launch - 1) / num_warps_to_launch; // ceil_div
 
     // Calculate warp_prefix_sum
@@ -571,26 +628,26 @@ void GroupRMSNormKernel(GroupRMSParams<n> params)
     if (rounds > 1)
     {
         TLLM_CUDA_CHECK(
-            cudaLaunchKernelEx(&cfg, GroupRMSNormKernel<DType, float4, n, EnableWeights, true>, params, rounds));
+            cudaLaunchKernelEx(&cfg, GroupRMSNormBaseKernel<DType, float4, n, EnableWeights, true>, params, rounds));
     }
     else
     {
         TLLM_CUDA_CHECK(
-            cudaLaunchKernelEx(&cfg, GroupRMSNormKernel<DType, float4, n, EnableWeights, false>, params, rounds));
+            cudaLaunchKernelEx(&cfg, GroupRMSNormBaseKernel<DType, float4, n, EnableWeights, false>, params, rounds));
     }
 }
 
 template <int n>
-void GroupRMSNormKernelLauncher(GroupRMSParams<n> params)
+void GroupRMSNormBaseKernelLauncher(GroupRMSParams<n>& params)
 {
 #define GROUP_RMS_NORM_DISPATCH(DTYPE)                                                                                 \
     if (params.enable_weights)                                                                                         \
     {                                                                                                                  \
-        return GroupRMSNormKernel<DTYPE, n, true>(params);                                                             \
+        return GroupRMSNormBaseKernel<DTYPE, n, true>(params);                                                         \
     }                                                                                                                  \
     else                                                                                                               \
     {                                                                                                                  \
-        return GroupRMSNormKernel<DTYPE, n, false>(params);                                                            \
+        return GroupRMSNormBaseKernel<DTYPE, n, false>(params);                                                        \
     }
 
     switch (params.dtype)
@@ -604,13 +661,13 @@ void GroupRMSNormKernelLauncher(GroupRMSParams<n> params)
 #undef GROUP_RMS_NORM_DISPATCH
 }
 
-#define INSTANTIATE_GROUP_RMS_NORM(n) template void GroupRMSNormKernelLauncher<n>(GroupRMSParams<n> params);
+#define INSTANTIATE_GROUP_RMS_NORM_BASE(n) template void GroupRMSNormBaseKernelLauncher<n>(GroupRMSParams<n> & params);
 
-INSTANTIATE_GROUP_RMS_NORM(1)
-INSTANTIATE_GROUP_RMS_NORM(2)
+INSTANTIATE_GROUP_RMS_NORM_BASE(1)
+INSTANTIATE_GROUP_RMS_NORM_BASE(2)
 
 template <typename DType, int n, bool EnableWeights>
-void GroupRMSNormKernelLargeBatch(GroupRMSParams<n> params)
+void GroupRMSNormKernelLargeBatch(GroupRMSParams<n>& params)
 {
     // Kernel assertions
     constexpr uint32_t kPackedSize = sizeof(float4) / sizeof(DType);
@@ -626,18 +683,11 @@ void GroupRMSNormKernelLargeBatch(GroupRMSParams<n> params)
             i, params.input_last_dims[i], kPackedSize);
     }
 
-    // Calculate warps needed for each input
-    uint32_t input_chunk_per_warp = 32 * kPackedSize;
-    uint32_t warps_needed_0 = (params.input_last_dims[0] + input_chunk_per_warp - 1) / input_chunk_per_warp;
-    uint32_t warps_needed_1 = (params.input_last_dims[1] + input_chunk_per_warp - 1) / input_chunk_per_warp;
-    uint32_t num_warps_to_launch_0 = std::min((uint32_t) 32, warps_needed_0);
-    uint32_t num_warps_to_launch_1 = std::min((uint32_t) 32, warps_needed_1);
-
-    // Calculate rounds needed for each input based on the number of warps we'll launch
-    uint32_t rounds_0 = (warps_needed_0 + num_warps_to_launch_0 - 1) / num_warps_to_launch_0;
-    uint32_t rounds_1 = (warps_needed_1 + num_warps_to_launch_1 - 1) / num_warps_to_launch_1;
-
-    uint32_t num_warps_to_launch = std::max(num_warps_to_launch_0, num_warps_to_launch_1);
+    // Calculate warps information
+    auto warpInfo = calculateNumWarpsLargeBatch<DType, n>(params);
+    uint32_t num_warps_to_launch = warpInfo.num_warps_to_launch;
+    uint32_t rounds_0 = warpInfo.rounds_0;
+    uint32_t rounds_1 = warpInfo.rounds_1;
 
     dim3 grid_dim(params.batch_size);
     dim3 block_dim(32, num_warps_to_launch);
@@ -660,30 +710,30 @@ void GroupRMSNormKernelLargeBatch(GroupRMSParams<n> params)
     {
         TLLM_CUDA_CHECK(
             cudaLaunchKernelEx(&cfg, GroupRMSNormKernelLargeBatch<DType, float4, n, EnableWeights, true, true>, params,
-                rounds_0, rounds_1, num_warps_to_launch_0, num_warps_to_launch_1));
+                rounds_0, rounds_1, warpInfo.num_warps_to_launch_0, warpInfo.num_warps_to_launch_1));
     }
     else if (MultiRounds_0 && !MultiRounds_1)
     {
         TLLM_CUDA_CHECK(
             cudaLaunchKernelEx(&cfg, GroupRMSNormKernelLargeBatch<DType, float4, n, EnableWeights, true, false>, params,
-                rounds_0, rounds_1, num_warps_to_launch_0, num_warps_to_launch_1));
+                rounds_0, rounds_1, warpInfo.num_warps_to_launch_0, warpInfo.num_warps_to_launch_1));
     }
     else if (!MultiRounds_0 && MultiRounds_1)
     {
         TLLM_CUDA_CHECK(
             cudaLaunchKernelEx(&cfg, GroupRMSNormKernelLargeBatch<DType, float4, n, EnableWeights, false, true>, params,
-                rounds_0, rounds_1, num_warps_to_launch_0, num_warps_to_launch_1));
+                rounds_0, rounds_1, warpInfo.num_warps_to_launch_0, warpInfo.num_warps_to_launch_1));
     }
     else
     {
         TLLM_CUDA_CHECK(
             cudaLaunchKernelEx(&cfg, GroupRMSNormKernelLargeBatch<DType, float4, n, EnableWeights, false, false>,
-                params, rounds_0, rounds_1, num_warps_to_launch_0, num_warps_to_launch_1));
+                params, rounds_0, rounds_1, warpInfo.num_warps_to_launch_0, warpInfo.num_warps_to_launch_1));
     }
 }
 
 template <int n>
-void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
+void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n>& params)
 {
 #define GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(DTYPE)                                                                     \
     if (params.enable_weights)                                                                                         \
@@ -707,8 +757,123 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 }
 
 #define INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(n)                                                                      \
-    template void GroupRMSNormKernelLargeBatchLauncher<n>(GroupRMSParams<n> params);
+    template void GroupRMSNormKernelLargeBatchLauncher<n>(GroupRMSParams<n> & params);
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    bool found_match = false;
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            found_match = true;
+            return p > 0.5f;
+        }
+    }
+    if (!found_match)
+    {
+        TLLM_LOG_INFO(
+            "GroupRMSNorm: Failed to find heuristic for GPU compute capability %d. Falling back to the base kernel.",
+            sm_major);
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n>& params)
+{
+    if (params.num_inputs == 1)
+    {
+        GroupRMSNormBaseKernelLauncher<n>(params);
+    }
+    else if (params.num_inputs == 2)
+    {
+        int num_warps_per_sm = 64;
+        uint32_t base_warps;
+        uint32_t large_batch_warps;
+
+        // Choose the appropriate DType
+        switch (params.dtype)
+        {
+        case nvinfer1::DataType::kHALF:
+            base_warps = calculateNumWarpsBase<half, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<half, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kBF16:
+            base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kFLOAT:
+            base_warps = calculateNumWarpsBase<float, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<float, n>(params).num_warps_to_launch;
+            break;
+        default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); return;
+        }
+
+        int concurrent_block_per_sm_base = std::floor(num_warps_per_sm / base_warps);
+        int concurrent_block_per_sm_large_batch = std::floor(num_warps_per_sm / large_batch_warps);
+
+        /*
+         * Kernel Selection Logic:
+         * We use trained Logistic Regression models to determine which kernel variant to use based on performance
+         * characteristics:
+         *
+         * - base_warps: Proportional to the sum of last dimensions of inputs
+         * - large_batch_warps: Proportional to the max of last dimensions of inputs
+         *
+         * Trade-offs:
+         * - With equal concurrent blocks per SM, base_warps achieves better compute efficiency
+         * - However, large_batch_warps allows more concurrent blocks to be scheduled:
+         *   - concurrent_block_per_sm_base: Maximum blocks of base kernel schedulable per SM
+         *   - concurrent_block_per_sm_large_batch: Maximum blocks of large batch kernel schedulable per SM
+         *
+         * The large batch kernel is preferred when the scheduling efficiency advantage outweighs
+         * the compute efficiency advantage of the base kernel, particularly at larger batch sizes.
+         */
+        if (concurrent_block_per_sm_large_batch > concurrent_block_per_sm_base)
+        {
+            float scheduling_efficiency_ratio
+                = float(concurrent_block_per_sm_large_batch) / float(concurrent_block_per_sm_base);
+            if (prefer_base_kernel(params.batch_size, base_warps, scheduling_efficiency_ratio))
+            {
+                GroupRMSNormBaseKernelLauncher<n>(params);
+            }
+            else
+            {
+                GroupRMSNormKernelLargeBatchLauncher<n>(params);
+            }
+        }
+        else
+        {
+            GroupRMSNormBaseKernelLauncher<n>(params);
+        }
+    }
+    else
+    {
+        // Unsupported number of inputs
+        TLLM_CHECK_WITH_INFO(false, "Unsupported number of inputs for GroupRMSNorm");
+    }
+}
+
+#define INSTANTIATE_GROUP_RMS_NORM_WITH_HEURISTIC(n)                                                                   \
+    template void GroupRMSNormKernelLauncherWithHeuristic<n>(GroupRMSParams<n> & params);
+
+INSTANTIATE_GROUP_RMS_NORM_WITH_HEURISTIC(1)
+INSTANTIATE_GROUP_RMS_NORM_WITH_HEURISTIC(2)
+
 } // namespace tensorrt_llm::kernels::group_rms_norm
diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h
index cda52fa7d27b..c121705f6dcd 100644
--- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h
+++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h
@@ -17,6 +17,7 @@
 #include <NvInferRuntime.h>
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
+#include <map>
 
 #include "tensorrt_llm/common/assert.h"
 #include "tensorrt_llm/common/cudaUtils.h"
@@ -44,10 +45,32 @@ struct GroupRMSParams
     cudaStream_t stream;
 };
 
+// Logistic regression model for predicting when the large batch kernel is faster
+struct Model
+{
+    float batch_size;
+    // Number of warps to launch for the base kernel
+    float base_warps;
+    // Ratio of the concurrent_block_per_sm for the large batch kernel vs the base kernel
+    float scheduling_efficiency_ratio;
+    // Intercept of the logistic regression model
+    float intercept;
+};
+
+// Trained parameters for the logistic regression model
+// For each major version of the Compute Capability
+inline std::map<int, Model> gpu_models = {
+    {10, {-0.004011f, -0.180179f, -0.396733f, 6.714080f}},
+    {9, {-0.006522f, -0.178540f, -0.558174f, 8.210834f}},
+};
+
+template <int n>
+void GroupRMSNormBaseKernelLauncher(GroupRMSParams<n>& params);
+
 template <int n>
-void GroupRMSNormKernelLauncher(GroupRMSParams<n> params);
+void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n>& params);
 
 template <int n>
-void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params);
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n>& params);
 
 } // namespace tensorrt_llm::kernels::group_rms_norm
diff --git a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp
index 05c7ec16635f..4cdffe636392 100644
--- a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp
+++ b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp
@@ -31,8 +31,8 @@
 namespace torch_ext
 {
 
-void groupRMSNorm(torch::TensorList const& inputs, torch::TensorList const& outputs, torch::TensorList const& weights,
-    double eps, double weight_bias)
+void groupRMSNormBase(torch::TensorList const& inputs, torch::TensorList const& outputs,
+    torch::TensorList const& weights, double eps, double weight_bias)
 {
     TORCH_CHECK(!inputs.empty(), "Input tensor list cannot be empty.");
     TORCH_CHECK(!outputs.empty(), "Output tensor list cannot be empty.");
@@ -103,7 +103,7 @@ void groupRMSNorm(torch::TensorList const& inputs, torch::TensorList const& outp
         case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break;                               \
         default: TORCH_CHECK(false, "Unsupported data type");                                                          \
         }                                                                                                              \
-        tensorrt_llm::kernels::group_rms_norm::GroupRMSNormKernelLauncher<n>(params);                                  \
+        tensorrt_llm::kernels::group_rms_norm::GroupRMSNormBaseKernelLauncher<n>(params);                              \
         break;                                                                                                         \
     }
 
@@ -188,11 +188,102 @@ void groupRMSNormLargeBatch(torch::TensorList const& inputs, torch::TensorList c
     tensorrt_llm::kernels::group_rms_norm::GroupRMSNormKernelLargeBatchLauncher<2>(params);
 }
 
+void groupRMSNormHeuristic(torch::TensorList const& inputs, torch::TensorList const& outputs,
+    torch::TensorList const& weights, double eps, double weight_bias)
+{
+    TORCH_CHECK(!inputs.empty(), "Input tensor list cannot be empty.");
+    TORCH_CHECK(!outputs.empty(), "Output tensor list cannot be empty.");
+    uint32_t const num_inputs = inputs.size();
+    TORCH_CHECK(num_inputs <= 2, "Heuristic kernel only supports up to 2 input tensors.");
+    auto const first_input = inputs[0];
+    TORCH_CHECK(first_input.dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+    TORCH_CHECK(first_input.sizes()[0] > 0, "Batch size must be greater than 0.");
+    TORCH_CHECK(first_input.is_cuda(), "Inputs must be CUDA tensors.");
+    uint32_t const batch_size = first_input.sizes()[0];
+    auto const dtype = first_input.scalar_type();
+    auto const device = first_input.device();
+
+    if (!weights.empty())
+    {
+        TORCH_CHECK(weights.size() == num_inputs, "Weights list size must match inputs list size.");
+    }
+
+    for (size_t i = 0; i < num_inputs; ++i)
+    {
+        TORCH_CHECK(inputs[i].sizes()[0] == batch_size, "Inputs must have the same batch size.");
+        TORCH_CHECK(inputs[i].dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+        TORCH_CHECK(inputs[i].device() == device, "Inputs must be on the same device.");
+        TORCH_CHECK(inputs[i].scalar_type() == dtype, "Inputs must be of the same type.");
+        TORCH_CHECK(outputs[i].dim() == 2, "Outputs must be 2D tensors [batch_size, hidden_dim].");
+        TORCH_CHECK(outputs[i].device() == device, "Outputs must be on the same device.");
+        TORCH_CHECK(outputs[i].scalar_type() == dtype, "Outputs must be of the same type.");
+        TORCH_CHECK(outputs[i].sizes()[0] == batch_size, "Outputs must have the same batch size.");
+        TORCH_CHECK(
+            outputs[i].sizes()[1] == inputs[i].sizes()[1], "Outputs and inputs must have the same last dimension.");
+        TORCH_CHECK(inputs[i].strides()[0] == outputs[i].strides()[0], "Inputs and outputs must have the same stride.");
+        TORCH_CHECK(inputs[i].strides()[1] == 1, "Inputs must be contiguous along the last dimension.");
+        TORCH_CHECK(outputs[i].strides()[1] == 1, "Outputs must be contiguous along the last dimension.");
+        if (!weights.empty())
+        {
+            TORCH_CHECK(
+                inputs[i].sizes()[1] == weights[i].sizes()[0], "Inputs and weights must have the same last dimension.");
+        }
+    }
+
+    // Dispatch based on number of inputs - templates require compile-time constants
+#define DISPATCH_HEURISTIC_INPUTS(n)                                                                                   \
+    {                                                                                                                  \
+        tensorrt_llm::kernels::group_rms_norm::GroupRMSParams<n> params;                                               \
+        for (size_t i = 0; i < n; ++i)                                                                                 \
+        {                                                                                                              \
+            params.inputs[i] = reinterpret_cast<float4*>(inputs[i].data_ptr());                                        \
+            params.input_last_dims[i] = inputs[i].sizes()[1];                                                          \
+            params.input_strides[i] = inputs[i].strides()[0];                                                          \
+            params.output_strides[i] = outputs[i].strides()[0];                                                        \
+            params.outputs[i] = reinterpret_cast<float4*>(outputs[i].mutable_data_ptr());                              \
+            if (!weights.empty())                                                                                      \
+            {                                                                                                          \
+                params.weights[i] = reinterpret_cast<float4 const*>(weights[i].data_ptr());                            \
+            }                                                                                                          \
+        }                                                                                                              \
+                                                                                                                       \
+        /* Set remaining params */                                                                                     \
+        params.batch_size = batch_size;                                                                                \
+        params.num_inputs = n;                                                                                         \
+        params.eps = static_cast<float>(eps);                                                                          \
+        params.weight_bias = static_cast<float>(weight_bias);                                                          \
+        params.enable_weights = !weights.empty();                                                                      \
+        params.stream = at::cuda::getCurrentCUDAStream(inputs[0].get_device());                                        \
+                                                                                                                       \
+        /* Handle dtype conversion */                                                                                  \
+        switch (dtype)                                                                                                 \
+        {                                                                                                              \
+        case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break;                                 \
+        case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break;                             \
+        case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break;                               \
+        default: TORCH_CHECK(false, "Unsupported data type");                                                          \
+        }                                                                                                              \
+                                                                                                                       \
+        /* Use the heuristic launcher that will decide between regular and large batch kernels */                      \
+        tensorrt_llm::kernels::group_rms_norm::GroupRMSNormKernelLauncherWithHeuristic<n>(params);                     \
+        break;                                                                                                         \
+    }
+
+    switch (num_inputs)
+    {
+    case 1: DISPATCH_HEURISTIC_INPUTS(1)
+    case 2: DISPATCH_HEURISTIC_INPUTS(2)
+    default: TORCH_CHECK(false, "Unsupported number of inputs (max 2)");
+    }
+#undef DISPATCH_HEURISTIC_INPUTS
+}
+
 } // namespace torch_ext
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
-    m.impl("group_rms_norm", &torch_ext::groupRMSNorm);
-    // TODO: Add heuristics for when to use the large batch kernel and merge into one torch Op.
+    m.impl("group_rms_norm_base", &torch_ext::groupRMSNormBase);
     m.impl("group_rms_norm_large_batch", &torch_ext::groupRMSNormLargeBatch);
+    // Use groupRMSNormHeuristic which automatically selects between regular and large batch kernels
+    m.impl("group_rms_norm_heuristic", &torch_ext::groupRMSNormHeuristic);
 }
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index 8f548bb5f545..a3121a385597 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -218,9 +218,9 @@ def _(ep_size: int):
     def _(max_sm_count: int):
         pass
 
-    @torch.library.custom_op("trtllm::group_rms_norm",
+    @torch.library.custom_op("trtllm::group_rms_norm_base",
                              mutates_args=("outputs", ))
-    def group_rms_norm(
+    def group_rms_norm_base(
         inputs: List[torch.Tensor],
         outputs: List[torch.Tensor],
         weights: List[torch.Tensor],
@@ -229,7 +229,7 @@ def group_rms_norm(
     ) -> None:
         pass
 
-    @group_rms_norm.register_fake
+    @group_rms_norm_base.register_fake
     def _(
         inputs: List[torch.Tensor],
         outputs: List[torch.Tensor],
@@ -259,3 +259,25 @@ def _(
         weight_bias: float,
     ) -> List[torch.Tensor]:
         return outputs
+
+    # Use groupRMSNormHeuristic which automatically selects between regular and large batch kernels
+    @torch.library.custom_op("trtllm::group_rms_norm_heuristic",
+                             mutates_args=("outputs", ))
+    def group_rms_norm_heuristic(
+        inputs: List[torch.Tensor],
+        outputs: List[torch.Tensor],
+        weights: List[torch.Tensor],
+        eps: float,
+        weight_bias: float,
+    ) -> None:
+        pass
+
+    @group_rms_norm_heuristic.register_fake
+    def _(
+        inputs: List[torch.Tensor],
+        outputs: List[torch.Tensor],
+        weights: List[torch.Tensor],
+        eps: float,
+        weight_bias: float,
+    ) -> List[torch.Tensor]:
+        return outputs
diff --git a/tensorrt_llm/_torch/modules/rms_norm.py b/tensorrt_llm/_torch/modules/rms_norm.py
index d1280e279825..9dc6c3910686 100644
--- a/tensorrt_llm/_torch/modules/rms_norm.py
+++ b/tensorrt_llm/_torch/modules/rms_norm.py
@@ -1,3 +1,4 @@
+import enum
 from typing import Optional, Tuple, Union
 
 import torch
@@ -58,34 +59,61 @@ def forward(
             return hidden_states, residual
 
 
+class GroupRMSNormKernelSelection(enum.Enum):
+    heuristic = 0
+    base = 1
+    large_batch = 2
+
+
 def group_rms_norm(
         inputs: list[torch.Tensor],
         weights: Optional[list[torch.Tensor]] = [],
         eps: Optional[float] = 1e-5,
         weight_bias: Optional[float] = 0.0,
+        kernel: GroupRMSNormKernelSelection = GroupRMSNormKernelSelection.
+    heuristic,
         outputs: Optional[list[torch.Tensor]] = None) -> list[torch.Tensor]:
-    '''
-    Group RMS Normalization for multiple inputs.
-    '''
-    out = outputs
-    if out is None:
-        out = [torch.empty_like(input) for input in inputs]
-    torch.ops.trtllm.group_rms_norm(inputs, out, weights, eps, weight_bias)
-    return out
+    '''Group RMS Normalization optimized for up to 2 inputs.
 
+    This function applies RMS normalization to multiple inputs simultaneously,
+    achieving better performance than normalizing each tensor separately with multi-stream.
 
-def group_rms_norm_large_batch(
-        inputs: list[torch.Tensor],
-        weights: Optional[list[torch.Tensor]] = [],
-        eps: Optional[float] = 1e-5,
-        weight_bias: Optional[float] = 0.0,
-        outputs: Optional[list[torch.Tensor]] = None) -> list[torch.Tensor]:
-    '''
-    Group RMS Normalization for 2 inputs with large batch size.
+    Args:
+        inputs: List of input tensors to normalize
+        weights: Optional list of weight tensors corresponding to each input
+        eps: Small constant added to variance for numerical stability
+        weight_bias: Optional bias added to weights during normalization
+        kernel: Kernel selection strategy:
+            - heuristic: Automatically selects optimal kernel based on inputs and hardware
+            - base: Uses base kernel (optimal for most cases)
+            - large_batch: Uses large batch kernel (may be better for large batches)
+        outputs: Optional pre-allocated output tensors (created if None)
+
+    Returns:
+        List of normalized tensors with the same shapes as inputs
+
+    Technical Details:
+        Available kernel implementations:
+        - Base kernel: Allocates warps proportional to the sum of last dimensions,
+          providing better SM occupancy for most workloads.
+        - Large batch kernel: Allocates warps proportional to the maximum last dimension,
+          which can be more efficient for large batch sizes with 2 inputs.
+
+        The heuristic mode uses a logistic regression model trained on benchmark data
+        to dynamically select the optimal kernel based on batch size, input dimensions,
+        and GPU architecture. This selection is optimized for compute capabilities 9.x and 10.x.
     '''
     out = outputs
     if out is None:
         out = [torch.empty_like(input) for input in inputs]
-    torch.ops.trtllm.group_rms_norm_large_batch(inputs, out, weights, eps,
-                                                weight_bias)
+    match kernel:
+        case GroupRMSNormKernelSelection.heuristic:
+            torch.ops.trtllm.group_rms_norm_heuristic(inputs, out, weights, eps,
+                                                      weight_bias)
+        case GroupRMSNormKernelSelection.base:
+            torch.ops.trtllm.group_rms_norm_base(inputs, out, weights, eps,
+                                                 weight_bias)
+        case GroupRMSNormKernelSelection.large_batch:
+            torch.ops.trtllm.group_rms_norm_large_batch(inputs, out, weights,
+                                                        eps, weight_bias)
     return out
diff --git a/tests/unittest/_torch/test_group_rmn_norm.py b/tests/unittest/_torch/test_group_rmn_norm.py
index b7bda5abbeb2..c44978aa1e6e 100644
--- a/tests/unittest/_torch/test_group_rmn_norm.py
+++ b/tests/unittest/_torch/test_group_rmn_norm.py
@@ -16,30 +16,21 @@
 import pytest
 import torch
 
-from tensorrt_llm._torch.modules.rms_norm import (RMSNorm, group_rms_norm,
-                                                  group_rms_norm_large_batch)
+from tensorrt_llm._torch.modules.rms_norm import RMSNorm, group_rms_norm
 
 
-@torch.inference_mode()
-@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda x: f"batch:{x}")
-@pytest.mark.parametrize("hidden_dims",
-                         [[256], [8448], [256, 512], [8448, 1024]],
-                         ids=lambda x: f"dims:{'-'.join(str(d) for d in x)}")
-@pytest.mark.parametrize("eps", [1e-6, 1e-5], ids=lambda x: f"eps:{x}")
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16],
-                         ids=["fp16", "bf16"])
-@pytest.mark.parametrize("enable_weights", [True, False],
-                         ids=lambda x: f"enable_weights:{x}")
-def test_group_rms_norm(batch_size, hidden_dims, eps, dtype, enable_weights):
-    """Compare group_rms_norm with RMSNorm."""
+def _prepare_rms_test_data(batch_size, hidden_dims, eps, dtype, enable_weights):
+    """Common setup for RMSNorm tests."""
     assert torch.cuda.is_available(), "This test requires CUDA"
     device = "cuda"
 
+    # Create input tensors
     inputs = [
         torch.randn((batch_size, dim), dtype=dtype, device=device)
         for dim in hidden_dims
     ]
-    weights = []
+
+    # Create weights
     if enable_weights:
         weights = [
             torch.randn((dim), dtype=dtype, device=device)
@@ -49,9 +40,8 @@ def test_group_rms_norm(batch_size, hidden_dims, eps, dtype, enable_weights):
         weights = [
             torch.ones((dim), dtype=dtype, device=device) for dim in hidden_dims
         ]
-    group_outputs = [torch.empty_like(input) for input in inputs]
 
-    # RMSNorm for reference
+    # Generate reference outputs
     ref_outputs = []
     if enable_weights:
         for i, dim in enumerate(hidden_dims):
@@ -62,39 +52,111 @@ def test_group_rms_norm(batch_size, hidden_dims, eps, dtype, enable_weights):
             norm = RMSNorm(hidden_size=dim, eps=eps, dtype=dtype, device=device)
             ref_outputs.append(norm(inputs[i]))
 
-    # Test torch.ops.trtllm.group_rms_norm
-    torch.ops.trtllm.group_rms_norm(inputs,
-                                    group_outputs,
-                                    weights if enable_weights else [],
-                                    eps=eps,
-                                    weight_bias=0.0)
+    return inputs, weights, ref_outputs
+
+
+def _verify_outputs(test_outputs, ref_outputs):
+    """Verify that test outputs match reference outputs."""
+    assert len(test_outputs) == len(ref_outputs), \
+        f"Expected {len(ref_outputs)} outputs, got {len(test_outputs)}"
+
+    # Verify each output matches reference
+    for i, (test_out, ref_out) in enumerate(zip(test_outputs, ref_outputs)):
+        torch.testing.assert_close(test_out, ref_out, rtol=1e-2, atol=1e-2)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda x: f"batch:{x}")
+@pytest.mark.parametrize("hidden_dims",
+                         [[256], [8448], [256, 512], [8448, 1024]],
+                         ids=lambda x: f"dims:{'-'.join(str(d) for d in x)}")
+@pytest.mark.parametrize("eps", [1e-6, 1e-5], ids=lambda x: f"eps:{x}")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16],
+                         ids=["fp16", "bf16"])
+@pytest.mark.parametrize("enable_weights", [True, False],
+                         ids=lambda x: f"enable_weights:{x}")
+def test_group_rms_norm_heuristic(batch_size, hidden_dims, eps, dtype,
+                                  enable_weights):
+    """Compare group_rms_norm with RMSNorm."""
+    inputs, weights, ref_outputs = _prepare_rms_test_data(
+        batch_size, hidden_dims, eps, dtype, enable_weights)
 
     # Test tensorrt_llm._torch.modules.rms_norm.group_rms_norm
     if enable_weights:
-        group_outputs_op = group_rms_norm(inputs, weights=weights, eps=eps)
+        group_outputs_heuristic = group_rms_norm(inputs,
+                                                 weights=weights,
+                                                 eps=eps)
     else:
-        group_outputs_op = group_rms_norm(inputs, eps=eps)
+        group_outputs_heuristic = group_rms_norm(inputs, eps=eps)
 
-    assert len(group_outputs) == len(ref_outputs), \
-        f"Expected {len(ref_outputs)} outputs, got {len(group_outputs)}"
-    assert len(group_outputs_op) == len(ref_outputs), \
-        f"Expected {len(ref_outputs)} outputs, got {len(group_outputs_op)}"
+    _verify_outputs(group_outputs_heuristic, ref_outputs)
 
-    # Verify each output matches reference
-    for i, (group_out, ref_out) in enumerate(zip(group_outputs, ref_outputs)):
-        torch.testing.assert_close(group_out, ref_out, rtol=1e-2, atol=1e-2)
-        torch.testing.assert_close(group_outputs_op[i],
-                                   ref_out,
-                                   rtol=1e-2,
-                                   atol=1e-2)
-    if len(hidden_dims) == 2:
-        large_batch_group_outputs = []
-        if enable_weights:
-            large_batch_group_outputs = group_rms_norm_large_batch(
-                inputs, weights=weights, eps=eps)
-        else:
-            large_batch_group_outputs = group_rms_norm_large_batch(inputs,
-                                                                   eps=eps)
-        for i, (group_out, ref_out) in enumerate(
-                zip(large_batch_group_outputs, ref_outputs)):
-            torch.testing.assert_close(group_out, ref_out, rtol=1e-2, atol=1e-2)
+
+@torch.inference_mode()
+@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda x: f"batch:{x}")
+@pytest.mark.parametrize("hidden_dims",
+                         [[256], [8448], [256, 512], [8448, 1024]],
+                         ids=lambda x: f"dims:{'-'.join(str(d) for d in x)}")
+@pytest.mark.parametrize("eps", [1e-6, 1e-5], ids=lambda x: f"eps:{x}")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16],
+                         ids=["fp16", "bf16"])
+@pytest.mark.parametrize("enable_weights", [True, False],
+                         ids=lambda x: f"enable_weights:{x}")
+def test_group_rms_norm_base(batch_size, hidden_dims, eps, dtype,
+                             enable_weights):
+    """Compare group_rms_norm_base with RMSNorm."""
+    inputs, weights, ref_outputs = _prepare_rms_test_data(
+        batch_size, hidden_dims, eps, dtype, enable_weights)
+
+    # Create output tensors
+    group_outputs_base = [torch.empty_like(input) for input in inputs]
+
+    # Test base implementation
+    if enable_weights:
+        torch.ops.trtllm.group_rms_norm_base(inputs,
+                                             group_outputs_base,
+                                             weights,
+                                             eps=eps,
+                                             weight_bias=0.0)
+    else:
+        torch.ops.trtllm.group_rms_norm_base(inputs,
+                                             group_outputs_base, [],
+                                             eps=eps,
+                                             weight_bias=0.0)
+
+    _verify_outputs(group_outputs_base, ref_outputs)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("batch_size", [1, 4], ids=lambda x: f"batch:{x}")
+@pytest.mark.parametrize("hidden_dims", [[256, 512], [8448, 1024]],
+                         ids=lambda x: f"dims:{'-'.join(str(d) for d in x)}")
+@pytest.mark.parametrize("eps", [1e-6, 1e-5], ids=lambda x: f"eps:{x}")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16],
+                         ids=["fp16", "bf16"])
+@pytest.mark.parametrize("enable_weights", [True, False],
+                         ids=lambda x: f"enable_weights:{x}")
+def test_group_rms_norm_large_batch(batch_size, hidden_dims, eps, dtype,
+                                    enable_weights):
+    """Compare group_rms_norm_large_batch with RMSNorm."""
+    inputs, weights, ref_outputs = _prepare_rms_test_data(
+        batch_size, hidden_dims, eps, dtype, enable_weights)
+
+    # Create output tensors
+    group_outputs_large_batch = [torch.empty_like(input) for input in inputs]
+
+    # Test large batch implementation
+    if enable_weights:
+        torch.ops.trtllm.group_rms_norm_large_batch(inputs,
+                                                    group_outputs_large_batch,
+                                                    weights,
+                                                    eps=eps,
+                                                    weight_bias=0.0)
+    else:
+        torch.ops.trtllm.group_rms_norm_large_batch(inputs,
+                                                    group_outputs_large_batch,
+                                                    [],
+                                                    eps=eps,
+                                                    weight_bias=0.0)
+
+    _verify_outputs(group_outputs_large_batch, ref_outputs)


Reviews:
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-07T04:49:06Z
user: login: hlu1
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T00:10:30Z
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:38:41Z
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:39:01Z
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:39:29Z
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:39:58Z
user: login: SimengLiu-nv
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:40:50Z
user: login: hlu1
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T19:21:22Z

Review comments:
diff_hunk: @@ -44,10 +45,32 @@ struct GroupRMSParams
     cudaStream_t stream;
 };
 
+// Logistic regression model for predicting when the large batch kernel is faster
+struct Model
+{
+    float batch_size;
+    // Number of warps to launch for the base kernel
+    float base_warps;
+    // Ratio of the block_limit_warps for the large batch kernel vs the base kernel
+    float scheduling_efficiency_ratio;
+    // Intercept of the logistic regression model
+    float intercept;
+};
+
+// Trained parameters for the logistic regression model
+// For each major version of the Compute Capability
+inline std::map<int, Model> gpu_models = {
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h
user: login: SimengLiu-nv
body: Trained 2 sets of parameters for sm9.x and sm10.x using all combinations of dims. The accuracy is 90%. Since at the switch point the performance diff is not significant, the accuracy could be acceptable. In addition, user can set the [GroupRMSNormKernelSelection](https://github.com/NVIDIA/TensorRT-LLM/pull/4047/files#diff-21b7c92739eae6897b6d105c1dfd9d84d0bbc513d8b8529301cc9916aab0c62eR62) to bypass the heuristic.
created_at: 2025-05-07T04:49:06Z
author_association: COLLABORATOR
diff_hunk: @@ -188,11 +188,102 @@ void groupRMSNormLargeBatch(torch::TensorList const& inputs, torch::TensorList c
     tensorrt_llm::kernels::group_rms_norm::GroupRMSNormKernelLargeBatchLauncher<2>(params);
 }
 
+void groupRMSNormHeuristic(torch::TensorList const& inputs, torch::TensorList const& outputs,
+    torch::TensorList const& weights, double eps, double weight_bias)
+{
+    TORCH_CHECK(!inputs.empty(), "Input tensor list cannot be empty.");
+    TORCH_CHECK(!outputs.empty(), "Output tensor list cannot be empty.");
+    uint32_t const num_inputs = inputs.size();
+    TORCH_CHECK(num_inputs <= 2, "Heuristic kernel only supports up to 2 input tensors.");
+    auto const first_input = inputs[0];
+    TORCH_CHECK(first_input.dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+    TORCH_CHECK(first_input.sizes()[0] > 0, "Batch size must be greater than 0.");
+    TORCH_CHECK(first_input.is_cuda(), "Inputs must be CUDA tensors.");
+    uint32_t const batch_size = first_input.sizes()[0];
+    auto const dtype = first_input.scalar_type();
+    auto const device = first_input.device();
+
+    if (!weights.empty())
+    {
+        TORCH_CHECK(weights.size() == num_inputs, "Weights list size must match inputs list size.");
+    }
+
+    for (size_t i = 0; i < num_inputs; ++i)
+    {
+        TORCH_CHECK(inputs[i].sizes()[0] == batch_size, "Inputs must have the same batch size.");
+        TORCH_CHECK(inputs[i].dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+        TORCH_CHECK(inputs[i].device() == device, "Inputs must be on the same device.");
+        TORCH_CHECK(inputs[i].scalar_type() == dtype, "Inputs must be of the same type.");
path: cpp/tensorrt_llm/thop/groupRmsNormOp.cpp
user: login: hlu1
body: Move these two and line 218-220 to another loop and start looping from i=1. 
created_at: 2025-05-07T23:03:33Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: hlu1
body: ```suggestion
void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n>& params)
```
created_at: 2025-05-07T23:09:52Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
+{
+    if (params.num_inputs == 1)
+    {
+        GroupRMSNormBaseKernelLauncher<n>(params);
+    }
+    else if (params.num_inputs == 2)
+    {
+        int num_warps_per_sm = 64;
+        uint32_t base_warps;
+        uint32_t large_batch_warps;
+
+        // Choose the appropriate DType
+        switch (params.dtype)
+        {
+        case nvinfer1::DataType::kHALF:
+            base_warps = calculateNumWarpsBase<half, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<half, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kBF16:
+            base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kFLOAT:
+            base_warps = calculateNumWarpsBase<float, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<float, n>(params).num_warps_to_launch;
+            break;
+        default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); return;
+        }
+
+        int block_limit_warps_base = std::floor(num_warps_per_sm / base_warps);
+        int block_limit_warps_large_batch = std::floor(num_warps_per_sm / large_batch_warps);
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: hlu1
body: `block_limit_warps` is a bit confusing. Can you rename it?
created_at: 2025-05-08T00:03:47Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
+{
+    if (params.num_inputs == 1)
+    {
+        GroupRMSNormBaseKernelLauncher<n>(params);
+    }
+    else if (params.num_inputs == 2)
+    {
+        int num_warps_per_sm = 64;
+        uint32_t base_warps;
+        uint32_t large_batch_warps;
+
+        // Choose the appropriate DType
+        switch (params.dtype)
+        {
+        case nvinfer1::DataType::kHALF:
+            base_warps = calculateNumWarpsBase<half, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<half, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kBF16:
+            base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kFLOAT:
+            base_warps = calculateNumWarpsBase<float, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<float, n>(params).num_warps_to_launch;
+            break;
+        default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); return;
+        }
+
+        int block_limit_warps_base = std::floor(num_warps_per_sm / base_warps);
+        int block_limit_warps_large_batch = std::floor(num_warps_per_sm / large_batch_warps);
+
+        // If the large batch kernel is more efficient for block scheduling
+        // Use trained Logistic Regression models to determine kernel selection
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: hlu1
body: This comment isn't very easy to understand. Please rephrase.
created_at: 2025-05-08T00:07:13Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: hlu1
body: This is also the fallback path for SM80 and 120. Can you add an info level log here for sm outside the range of `gpu_models`?
created_at: 2025-05-08T00:09:06Z
author_association: COLLABORATOR
diff_hunk: @@ -188,11 +188,102 @@ void groupRMSNormLargeBatch(torch::TensorList const& inputs, torch::TensorList c
     tensorrt_llm::kernels::group_rms_norm::GroupRMSNormKernelLargeBatchLauncher<2>(params);
 }
 
+void groupRMSNormHeuristic(torch::TensorList const& inputs, torch::TensorList const& outputs,
+    torch::TensorList const& weights, double eps, double weight_bias)
+{
+    TORCH_CHECK(!inputs.empty(), "Input tensor list cannot be empty.");
+    TORCH_CHECK(!outputs.empty(), "Output tensor list cannot be empty.");
+    uint32_t const num_inputs = inputs.size();
+    TORCH_CHECK(num_inputs <= 2, "Heuristic kernel only supports up to 2 input tensors.");
+    auto const first_input = inputs[0];
+    TORCH_CHECK(first_input.dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+    TORCH_CHECK(first_input.sizes()[0] > 0, "Batch size must be greater than 0.");
+    TORCH_CHECK(first_input.is_cuda(), "Inputs must be CUDA tensors.");
+    uint32_t const batch_size = first_input.sizes()[0];
+    auto const dtype = first_input.scalar_type();
+    auto const device = first_input.device();
+
+    if (!weights.empty())
+    {
+        TORCH_CHECK(weights.size() == num_inputs, "Weights list size must match inputs list size.");
+    }
+
+    for (size_t i = 0; i < num_inputs; ++i)
+    {
+        TORCH_CHECK(inputs[i].sizes()[0] == batch_size, "Inputs must have the same batch size.");
+        TORCH_CHECK(inputs[i].dim() == 2, "Inputs must be 2D tensors [batch_size, hidden_dim].");
+        TORCH_CHECK(inputs[i].device() == device, "Inputs must be on the same device.");
+        TORCH_CHECK(inputs[i].scalar_type() == dtype, "Inputs must be of the same type.");
path: cpp/tensorrt_llm/thop/groupRmsNormOp.cpp
user: login: SimengLiu-nv
body: Outputs and weights need to loop with i=0. May I leave it unchanged?
created_at: 2025-05-12T18:38:40Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: SimengLiu-nv
body: Thanks for pointing out!
created_at: 2025-05-12T18:39:01Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
+{
+    if (params.num_inputs == 1)
+    {
+        GroupRMSNormBaseKernelLauncher<n>(params);
+    }
+    else if (params.num_inputs == 2)
+    {
+        int num_warps_per_sm = 64;
+        uint32_t base_warps;
+        uint32_t large_batch_warps;
+
+        // Choose the appropriate DType
+        switch (params.dtype)
+        {
+        case nvinfer1::DataType::kHALF:
+            base_warps = calculateNumWarpsBase<half, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<half, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kBF16:
+            base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kFLOAT:
+            base_warps = calculateNumWarpsBase<float, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<float, n>(params).num_warps_to_launch;
+            break;
+        default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); return;
+        }
+
+        int block_limit_warps_base = std::floor(num_warps_per_sm / base_warps);
+        int block_limit_warps_large_batch = std::floor(num_warps_per_sm / large_batch_warps);
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: SimengLiu-nv
body: Renamed to `concurrent_block_per_sm_xxx`.
created_at: 2025-05-12T18:39:29Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
+}
+
+template <int n>
+void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams<n> params)
+{
+    if (params.num_inputs == 1)
+    {
+        GroupRMSNormBaseKernelLauncher<n>(params);
+    }
+    else if (params.num_inputs == 2)
+    {
+        int num_warps_per_sm = 64;
+        uint32_t base_warps;
+        uint32_t large_batch_warps;
+
+        // Choose the appropriate DType
+        switch (params.dtype)
+        {
+        case nvinfer1::DataType::kHALF:
+            base_warps = calculateNumWarpsBase<half, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<half, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kBF16:
+            base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch;
+            break;
+        case nvinfer1::DataType::kFLOAT:
+            base_warps = calculateNumWarpsBase<float, n>(params);
+            large_batch_warps = calculateNumWarpsLargeBatch<float, n>(params).num_warps_to_launch;
+            break;
+        default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); return;
+        }
+
+        int block_limit_warps_base = std::floor(num_warps_per_sm / base_warps);
+        int block_limit_warps_large_batch = std::floor(num_warps_per_sm / large_batch_warps);
+
+        // If the large batch kernel is more efficient for block scheduling
+        // Use trained Logistic Regression models to determine kernel selection
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: SimengLiu-nv
body: Replaced with more detailed block comments.
created_at: 2025-05-12T18:39:58Z
author_association: COLLABORATOR
diff_hunk: @@ -654,4 +704,95 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams<n> params)
 
 INSTANTIATE_GROUP_RMS_NORM_LARGE_BATCH(2)
 
+int getComputeCapabilityMajor()
+{
+    int device;
+    TLLM_CUDA_CHECK(cudaGetDevice(&device));
+    cudaDeviceProp prop;
+    TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+    return prop.major;
+}
+
+bool prefer_base_kernel(int batch, int base_warps, float scheduling_efficiency_ratio)
+{
+    int sm_major = getComputeCapabilityMajor();
+    for (auto const& [known_model, model] : gpu_models)
+    {
+        if (sm_major == known_model)
+        {
+            float p = model.batch_size * batch + model.base_warps * base_warps
+                + model.scheduling_efficiency_ratio * scheduling_efficiency_ratio + model.intercept;
+            p = 1.0f / (1.0f + std::exp(-p));
+            return p > 0.5f;
+        }
+    }
+    return true;
path: cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu
user: login: SimengLiu-nv
body: Added TLLM_LOG_INFO if no matching heuristics.
created_at: 2025-05-12T18:40:50Z
author_association: COLLABORATOR

PR comments:
user: login: SimengLiu-nv
created_at: 2025-05-05T00:05:49Z
body: /bot run --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-05-05T00:11:33Z
body: [PR_Github #4041](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/4041/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-05-05T03:20:09Z
body: [PR_Github #4041](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/4041/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #2870](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/2870/) completed with status: 'SUCCESS'

author_association: COLLABORATOR
user: login: Kefeng-Duan
created_at: 2025-05-07T00:49:24Z
body: @hlu1  @SimengLiu-nv  
do you think should we make such heuristic for hopper?
author_association: COLLABORATOR
user: login: SimengLiu-nv
created_at: 2025-05-12T18:41:15Z
body: /bot run --disable-fail-fast
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-05-12T18:47:52Z
body: [PR_Github #4896](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/4896/) [ run ] triggered by Bot
author_association: COLLABORATOR
user: login: tensorrt-cicd
created_at: 2025-05-13T00:52:51Z
body: [PR_Github #4896](https://nv/trt-llm-cicd/job/helpers/job/PR_Github/4896/) [ run ] completed with state `SUCCESS`
[/LLM/main/L0_MergeRequest_PR pipeline #3548](https://nv/trt-llm-cicd/job/main/job/L0_MergeRequest_PR/3548/) completed with status: 'SUCCESS'

author_association: COLLABORATOR