Skip to content

[Fix] the bug in the trtllm-gen heurisitcf for MLA kernels. #6284

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PerkzZheng merged 1 commit into from
Jul 24, 2025
Merged

[Fix] the bug in the trtllm-gen heurisitcf for MLA kernels. #6284

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 11 additions & 6 deletions cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,9 +413,13 @@ class TllmGenFmhaKernel
return std::make_tuple(numCtasPerSeqQ, numCtasPerSeqKv, numCtasX, numCtasY, numCtasZ, clusterDimX);
}

// Compute the seqLenPerCtaKv for selecting the MLA generation kernel.
int computeSeqLenPerCtaKv(RunnerParams const& params) const
// Determine if we should use the SwapsMmaAbForGeneration kernel for MLA generation.
bool useSwapsMmaAbMlaGenKernel(RunnerParams const& params) const
{
// Use the SwapsMmaAbForGeneration kernel for MLA generation when the following conditions are met:
// 1. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later).
// 2. The numCtas (after splitting the heads across multiple CTAs) <= params.mMultiProcessorCount.

// The maximum number Ctas per Kv sequence, which makes sure that each CtaKv has work to do.
// Here we assume the stepKv is 256.
int const maxNumCtasPerSeqKv = (params.mMaxSeqLenKv + 256 - 1) / 256;
Expand All @@ -427,8 +431,8 @@ class TllmGenFmhaKernel
= std::min(maxNumCtasPerSeqKv, std::max(1, int32_t(params.mMultiProcessorCount / numCtas)));
// Compute the seqLenPerCtaKv.
int const seqLenPerCtaKv = (params.mMaxSeqLenKv + numCtasPerSeqKv - 1) / numCtasPerSeqKv;
// Return the seqLenPerCtaKv.
return seqLenPerCtaKv;
// Whether we should use the SwapsMmaAbForGeneration kernel for MLA generation.
return seqLenPerCtaKv <= 1024 && numCtas <= params.mMultiProcessorCount;
}

std::pair<uint64_t, std::string> hashFromRunnerParams(
Expand All @@ -442,10 +446,11 @@ class TllmGenFmhaKernel
// We use the low-latency kernel (SwapsMmaAbForGeneration with tileSizeQ = 16) when any of the following
// conditions are met:
// 1. The number of headsQPerKv is <= 32.
// 2. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later).
// 2. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned later) and
// the numCtas (after splitting the heads across multiple CTAs) <= params.mMultiProcessorCount.

// Check the conditions.
if (params.mNumHeadsQPerKv <= 32 || computeSeqLenPerCtaKv(params) <= 1024)
if (params.mNumHeadsQPerKv <= 32 || useSwapsMmaAbMlaGenKernel(params))
{
kernelType = FmhaKernelType::SwapsMmaAbForGeneration;
}
Expand Down