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 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(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(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;