From 10da579483e2a714189eaa4bb635c3d8c3e8279e Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Mon, 5 Oct 2026 19:31:31 -0700
Subject: [PATCH] [None][fix] Add a static tile scheduler override for context
 FMHA

Allow TRTLLM_GEN_FMHA_CONTEXT_STATIC_TILE_SCHEDULER=1 to select the static scheduler consistently for kernel support checks and launches. Keep the persistent scheduler as the default.

Port the FMHA portion of 29ccb16af0d1c66f5f7c6522457d51862224c875. The staging-buffer portion is superseded on main by automatic cold-page sizing and removal of the Python copy engine.

Signed-off-by: Jonas Li <jonasl@nvidia.com>
Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 cpp/tensorrt_llm/common/envUtils.cpp        |  6 ++++++
 cpp/tensorrt_llm/common/envUtils.h          |  3 +++
 cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp | 20 +++++++++++++++++---
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/cpp/tensorrt_llm/common/envUtils.cpp b/cpp/tensorrt_llm/common/envUtils.cpp
index aee8dfd2b99a..7bf1c480b16b 100644
--- a/cpp/tensorrt_llm/common/envUtils.cpp
+++ b/cpp/tensorrt_llm/common/envUtils.cpp
@@ -232,6 +232,12 @@ bool getEnvUseTileSizeKv64ForTrtllmGen()
     return useTileSizeKv64;
 }
 
+bool getEnvUseStaticTileSchedulerForTrtllmGenContextFmha()
+{
+    static bool const useStatic = getBoolEnv("TRTLLM_GEN_FMHA_CONTEXT_STATIC_TILE_SCHEDULER");
+    return useStatic;
+}
+
 bool getEnvEnablePDL()
 {
     static std::once_flag flag;
diff --git a/cpp/tensorrt_llm/common/envUtils.h b/cpp/tensorrt_llm/common/envUtils.h
index a81e1b362f82..b3766734d0b9 100644
--- a/cpp/tensorrt_llm/common/envUtils.h
+++ b/cpp/tensorrt_llm/common/envUtils.h
@@ -45,6 +45,9 @@ std::optional<int32_t> getEnvXqaBlocksPerSequence();
 // Whether use tileSizeKv64 for multiCtasKvMode of trtllm-gen kernels.
 bool getEnvUseTileSizeKv64ForTrtllmGen();
 
+// Whether trtllm-gen context FMHA uses the static tile scheduler instead of the persistent one.
+bool getEnvUseStaticTileSchedulerForTrtllmGenContextFmha();
+
 // Tune the number of blocks per sequence for accuracy/performance purpose.
 bool getEnvMmhaMultiblockDebug();
 
diff --git a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
index bfa2f4e278ea..91842464a00f 100644
--- a/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
+++ b/cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
@@ -17,12 +17,26 @@
 #include "fmhaDispatcher.h"
 #include "tensorrt_llm/common/config.h"
 #include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
 
 TRTLLM_NAMESPACE_BEGIN
 
 namespace kernels
 {
 
+namespace
+{
+// Persistent is the default for performance. TRTLLM_GEN_FMHA_CONTEXT_STATIC_TILE_SCHEDULER=1 selects the static
+// variant instead: the sm107 persistent context kernel can hang in a tcgen05 TMEM-dealloc guardrail trap
+// (observed with the MLA HQk192/HV128 SeparateQkv kernel on Kimi-K3 prefill). Kernel support check and launch
+// must agree on the scheduler, so both go through this helper.
+TileScheduler contextTileScheduler()
+{
+    return tensorrt_llm::common::getEnvUseStaticTileSchedulerForTrtllmGenContextFmha() ? TileScheduler::Static
+                                                                                       : TileScheduler::Persistent;
+}
+} // namespace
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout)
@@ -112,7 +126,7 @@ bool FmhaDispatcher::isSupported()
         tllmRunnerParams.mQkvLayout = qkvLayout;
         tllmRunnerParams.setAttentionMaskType(static_cast<std::int8_t>(mFixedParams.attentionMaskType));
         tllmRunnerParams.mKernelType = FmhaKernelType::Context;
-        tllmRunnerParams.mTileScheduler = TileScheduler::Persistent;
+        tllmRunnerParams.mTileScheduler = contextTileScheduler();
         tllmRunnerParams.mMultiCtasKvMode = false;
         tllmRunnerParams.mNumHeadsQ = mFixedParams.numQHeads;
         tllmRunnerParams.mNumHeadsKv = mFixedParams.numKvHeads;
@@ -200,8 +214,8 @@ void FmhaDispatcher::run(MHARunnerParams runnerParams)
         tllmRunnerParams.mQkvLayout = qkvLayout;
         tllmRunnerParams.setAttentionMaskType(static_cast<std::int8_t>(mFixedParams.attentionMaskType));
         tllmRunnerParams.mKernelType = FmhaKernelType::Context;
-        // Always use persistent scheduler for better performance.
-        tllmRunnerParams.mTileScheduler = TileScheduler::Persistent;
+        // Persistent scheduler by default for better performance; see contextTileScheduler().
+        tllmRunnerParams.mTileScheduler = contextTileScheduler();
         tllmRunnerParams.mMultiCtasKvMode = false;
 
         tllmRunnerParams.qPtr = runnerParams.qPtr;