[None][feat] Add test-only grouped MLA generation selection probe

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h

@@ -95,6 +95,21 @@ constexpr bool isSMCompatible(int gpuSM, int kernelSM)
     return gpuSM == kernelSM;
 }
 
+#if defined(TLLM_FMHA_TEST_HOOKS)
+// Test-only result of probeKernelSelectionForTesting: the kernel the autotuner would launch.
+struct TllmGenFmhaSelectedKernel
+{
+    // Resolved cubin function name (empty when mFound is false).
+    std::string mFuncName;
+    bool mFound = false;
+    // True when the autotuner chose the NVRTC path instead of a precompiled cubin.
+    bool mUsedNvrtc = false;
+    // Grouping flags from the matched kernelMeta.
+    bool mGroupsHeadsQ = false;
+    bool mGroupsTokensHeadsQ = false;
+};
+#endif // TLLM_FMHA_TEST_HOOKS
+
 class TllmGenFmhaKernel
 {
 
@@ -291,6 +306,58 @@ class TllmGenFmhaKernel
         return std::make_pair(true, info);
     }
 
+#if defined(TLLM_FMHA_TEST_HOOKS)
+    // Test-only: report which kernelMeta the autotuner + hash lookup resolve to for these params,
+    // without launching the kernel.
+    TllmGenFmhaSelectedKernel probeKernelSelectionForTesting(RunnerParams const& params) const
+    {
+        TllmGenFmhaSelectedKernel result;
+        if (params.mHeadDimQk % 8 != 0 || params.mHeadDimV % 8 != 0)
+        {
+            return result;
+        }
+        if (params.mMaxSeqLenQ == 0 || params.mBatchSize == 0
+            || (!isContextKernel(params.mKernelType) && params.mMaxSeqLenKv == 0))
+        {
+            return result;
+        }
+        int32_t ctaDim = 512;
+        FmhaOptions options;
+        FmhaOptionsFromArgs optionsFromArgs;
+        parseOptionsFromRunnerParams(params, options);
+        options.mCudaArch = intToCudaArch(mSM);
+
+        FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount);
+        std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel();
+
+        checkFmhaOptions(options, optionsFromArgs);
+        updateFmhaOptions(options, optionsFromArgs);
+
+        computeNumCtas(options, params.mMultiProcessorCount);
+
+        if (shouldUseNvrtc(options))
+        {
+            result.mUsedNvrtc = true;
+            return result;
+        }
+
+        algoFilterForCubinPath(options);
+        auto [hashId, info] = hashFromFmhaOptions(options);
+
+        auto const findIter = mFunctions.find(hashId);
+        if (findIter == mFunctions.end())
+        {
+            return result;
+        }
+        auto const& kernelMeta = mKernelMeta[findIter->second.mMetaInfoIndex];
+        result.mFound = true;
+        result.mFuncName = kernelMeta.mFuncName != nullptr ? std::string(kernelMeta.mFuncName) : std::string{};
+        result.mGroupsHeadsQ = kernelMeta.mGroupsHeadsQ;
+        result.mGroupsTokensHeadsQ = kernelMeta.mGroupsTokensHeadsQ;
+        return result;
+    }
+#endif // TLLM_FMHA_TEST_HOOKS
+
     void algoFilterForCubinPath(FmhaOptions& options) const
     {
         if (!isContextKernel(options.mFmhaKernelType) && options.mMaskType == TrtllmGenAttentionMaskType::Dense
@@ -999,6 +1066,9 @@ class TllmGenFmhaKernel
 
         options.mEnablesAutoTuner = true;
         options.mIsMlaGen = isMlaGenKernel(params);
+        // Let the autotuner pick the grouped-token Q64 MLA generation kernel when its
+        // capability predicate matches (causal spec decode, supported head ratios/dtypes).
+        options.mSelectsGroupedMla = options.mIsMlaGen;
         options.mDtypeQ = dataTypeToDtype(mDtypeQ);
         options.mDtypeKv = dataTypeToDtype(mDtypeK);
         options.mDtypeK = dataTypeToDtype(mDtypeK);

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunner.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, 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.
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cuda_runtime.h>
+#include <string>
 
 #include "fmhaKernels.h"
 #include "fmhaRunnerParams.h"
@@ -50,6 +51,14 @@ class TllmGenFmhaRunner
     // Run the fmha kernel.
     void run(TllmGenFmhaRunnerParams const&);
 
+#if defined(TLLM_FMHA_TEST_HOOKS)
+    // Test-only: probe which cubin the autotuner would select for these params, without launching.
+    TllmGenFmhaSelectedKernel probeKernelSelectionForTesting(TllmGenFmhaRunnerParams const& runnerParams) const
+    {
+        return mKernel->probeKernelSelectionForTesting(runnerParams);
+    }
+#endif // TLLM_FMHA_TEST_HOOKS
+
 private:
     // The input/output datatype.
     Data_type mDtypeQ, mDtypeK, mDtypeV, mDtypeOut;

cpp/tests/unit_tests/kernels/CMakeLists.txt

@@ -113,3 +113,6 @@ endif()
 add_gtest(eaglePackDataTest eaglePackDataTest.cpp)
 add_gtest(sparseKvCacheTest sparseKvCacheTest.cu)
 add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp)
+add_gtest(kimiMlaGroupedSelectionTest kimiMlaGroupedSelectionTest.cpp)
+# Enables the test-only kernel-selection probe in fmhaRunner.h / fmhaKernels.h.
+target_compile_definitions(kimiMlaGroupedSelectionTest PRIVATE TLLM_FMHA_TEST_HOOKS=1)