From 5ed497f1d2f1dc91c3e7fdbca5345788da43a60b Mon Sep 17 00:00:00 2001 From: linquanh Date: Wed, 9 Sep 2026 07:40:13 +0000 Subject: [PATCH] [None][fix] Remove CGA smem-reduction clamp from trtllm-gen FMHA kernel selection PR #18518 wrapped every FmhaAutoTuner::selectKernel() call in selectKernelWithCgaSmemReductionLimit(), which clamps FmhaOptions::mMaxNumCtasPerSeqKv from INT_MAX to 8 for any GmemReduction + Static-scheduler request on ALL architectures. That request pattern is exactly the trtllm-gen generation (decode) path (xqaDispatcher sets multi_block_mode -> GmemReduction + Static). On B200 (SM100) MLA decode this changes the tuned kernel selection: - computeNumCtas caps the split-KV CTA count at 8 instead of filling the machine in one wave (e.g. ~18 splits for 8-9k KV at batch 1); - the smaller grid flips selectMlaGenerationKernel's headDimPerCtaV fine-tuning from 256 to 128, duplicating the 576-dim QK^T work 4x instead of 2x; - the one-wave check in enableCgaReduction now succeeds and promotes GmemReduction to CgaSmemReduction with far fewer CTAs. Net effect: decode attention loses most of its split-KV parallelism at low batch, regressing disagg gen-only DeepSeek-R1 FP4 MTP3 at con1 (disagg-gen_only-b200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_ eplb0_mtp3_ccb-NIXL). Remove the wrapper and restore the direct selectKernel() calls, returning SM100-family kernel selection to the pre-PR behavior. The clamp existed to keep a Rubin promotion path legal; that constraint belongs in the trtllm-gen autotuner (or an SM107-gated check), not in arch-agnostic host code. Signed-off-by: linquanh --- .../trtllmGenKernels/fmha/fmhaKernels.h | 49 +++---------------- 1 file changed, 6 insertions(+), 43 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index b7844a3fe323..7569b678836f 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -311,8 +311,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); // Check if the options are valid or not. checkFmhaOptions(options, optionsFromArgs); @@ -409,8 +409,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); checkFmhaOptions(options, optionsFromArgs); updateFmhaOptions(options, optionsFromArgs); @@ -527,8 +527,8 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - std::tie(options, optionsFromArgs, ctaDim) - = selectKernelWithCgaSmemReductionLimit(options, optionsFromArgs, params.mMultiProcessorCount); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); // Overwrite AutoTuner decision: SageAttention with SfsPV is known to cause regression to persistent scheduler. // Remove this overwritten once we refresh the cubin kernels that containing the related fix. @@ -663,43 +663,6 @@ class TllmGenFmhaKernel } private: - std::tuple selectKernelWithCgaSmemReductionLimit( - FmhaOptions options, FmhaOptionsFromArgs optionsFromArgs, int32_t multiProcessorCount) const - { - if (isGmemReduction(options.mMultiCtasKvMode) && options.mTileScheduler == TileScheduler::Static) - { - constexpr int kMaxCgaClusterDimX = 16; - constexpr int kMinPreSelectCgaClusterDimX = 2; - // selectKernel may promote this candidate to CgaSmemReduction and call computeNumCtas - // before returning the selected options. Keep the candidate legal for both 1-CTA and - // 2-CTA FMHA kernels; the exact selected clusterDimX is applied below. - options.mMaxNumCtasPerSeqKv = std::min(options.mMaxNumCtasPerSeqKv, - kMaxCgaClusterDimX / std::max(options.mClusterDimX, kMinPreSelectCgaClusterDimX)); - } - - int32_t ctaDim = 512; - FmhaAutoTuner autoTuner(options, optionsFromArgs, multiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - limitCgaSmemReductionCtasKv(options); - return {options, optionsFromArgs, ctaDim}; - } - - void limitCgaSmemReductionCtasKv(FmhaOptions& options) const - { - if (!isCgaSmemReduction(options.mMultiCtasKvMode)) - { - return; - } - constexpr int kMaxCgaClusterDimX = 16; - if (options.mClusterDimX * options.mMaxNumCtasPerSeqKv > kMaxCgaClusterDimX) - { - TLLM_LOG_WARNING( - "CGA reduction is not supported when numCtasPerSeqKv * clusterDimX > 16. Set mMaxNumCtasPerSeqKv to " - "16 / clusterDimX"); - options.mMaxNumCtasPerSeqKv = kMaxCgaClusterDimX / options.mClusterDimX; - } - } - inline uint64_t hashID(int qkvLayout, int maskType, int kernelType, int scheduler, int multiCtasKvMode, int headDimPerCtaV, int headDimQk, int headDimV, int tileSizeQ, int tileSizeKv, int numTokensPerPage, bool reuseSmemKForV, bool uses2CtaMma, int sparseAttention, bool skipsSoftmax, bool groupsHeadsQ,