From 18ea66e13515cd60ea78940dd4bbbe90848ff2a0 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 21 May 2026 09:34:27 +0000 Subject: [PATCH 1/7] Attention multiwave --- benchmarks/bench_trtllm_gen_mla.py | 51 ++++++++++++++++--- csrc/fmhaReduction.cu | 2 +- .../flashinfer/trtllm/fmha/fmhaKernels.cuh | 22 +++++++- include/flashinfer/trtllm/fmha/kernelParams.h | 11 +++- 4 files changed, 74 insertions(+), 12 deletions(-) diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py index a739ccc21b..8aac7c3753 100644 --- a/benchmarks/bench_trtllm_gen_mla.py +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -11,7 +11,8 @@ def bench_trtllm_mla( - batch_size, q_len_per_request, seq_len, page_size, dtype, backend="auto" + batch_size, q_len_per_request, seq_len, page_size, dtype, backend="auto", + seq_lens_list=None, ): torch.manual_seed(42) device = "cuda:0" @@ -25,13 +26,17 @@ def bench_trtllm_mla( device=device, ).to(dtype) - num_tokens = seq_len * batch_size - num_blocks = (num_tokens + page_size - 1) // page_size + # Sequence lengths: use provided list or generate random + if seq_lens_list is not None: + seq_lens = list(seq_lens_list) + max_seq_len = max(seq_lens) + else: + seq_lens = [torch.randint(1, seq_len, (1,)).item() for _ in range(batch_size)] + seq_lens[-1] = seq_len + max_seq_len = max(seq_lens) - # Sequence lengths and block tables - seq_lens = [torch.randint(1, seq_len, (1,)).item() for _ in range(batch_size)] - seq_lens[-1] = seq_len - max_seq_len = max(seq_lens) + num_tokens = max_seq_len * batch_size + num_blocks = (num_tokens + page_size - 1) // page_size seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size @@ -67,7 +72,7 @@ def bench_trtllm_mla( # Allocate workspace buffer # todo(Yingyi): calculate the actual size of workspace buffer - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=device) + workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=device) # Run decode-MLA # warmup @@ -137,6 +142,16 @@ def bench_trtllm_mla( print(f"FLOPs: {flops / ms / 1e9:.2f} TFLOPs/s") +def sample_prod_distribution(batch_size: int, seed: int = 42) -> list[int]: + """Sample sequence lengths from a production-like distribution.""" + import random + + rng = random.Random(seed) + seq_lens = [4096, 8192, 16384, 32768, 65536, 131072] + weights = [35, 20, 15, 12, 10, 8] + return rng.choices(seq_lens, weights=weights, k=batch_size) + + if __name__ == "__main__": import argparse @@ -178,3 +193,23 @@ def bench_trtllm_mla( f"backend={args.backend}: {type(e).__name__}: {e}" ) print() + + # Mixed sequence length benchmark (production-like distribution) + print("\n===== Mixed sequence length benchmark =====") + for batch_size in [16, 32, 64, 128, 256]: + sl = sample_prod_distribution(batch_size, seed=batch_size) + try: + bench_trtllm_mla( + batch_size, 1, max(sl), 32, torch.bfloat16, + backend=args.backend, seq_lens_list=sl, + ) + except ValueError as e: + print(f"SKIPPED: {e}") + print() + except Exception as e: + print( + f"ERROR: batch_size={batch_size}, q_len=1, " + f"seq_len=mixed, page_size=32, dtype=torch.bfloat16, " + f"backend={args.backend}: {type(e).__name__}: {e}" + ) + print() diff --git a/csrc/fmhaReduction.cu b/csrc/fmhaReduction.cu index 486fb75f8b..058fa864a2 100644 --- a/csrc/fmhaReduction.cu +++ b/csrc/fmhaReduction.cu @@ -334,7 +334,7 @@ void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams int32_t const maxNumCtasForReduction{(multiProcessorCount * 2) / static_cast(gridDim.x * gridDim.y * gridDim.z)}; // The number of Ctas for the reduction work. - int32_t const numCtasForReduction{std::min(maxNumCtasForReduction, numSlices)}; + int32_t const numCtasForReduction{std::max(1, std::min(maxNumCtasForReduction, numSlices))}; // Launch more CTAs to split the reduction work if needed. gridDim.x *= numCtasForReduction; diff --git a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh index 165cc4e2d3..076c19f769 100644 --- a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -477,10 +477,28 @@ class TllmGenFmhaKernel { // benefits of a shorter mainloop. int const maxNumCtasPerSeqKv = (maxAttentionWindow + 2 * kernelMeta.mStepKv - 1) / (2 * kernelMeta.mStepKv); - // Compute numCtasPerSeqKv. + + int const baseCtas = numCtasX * numCtasY * numCtasZ; numCtasPerSeqKv = std::min( maxNumCtasPerSeqKv, - std::max(1, int32_t(params.mMultiProcessorCount / (numCtasX * numCtasY * numCtasZ)))); + std::max(1, int32_t(params.mMultiProcessorCount / baseCtas))); + + // When the longest sequence needs more KV splits than the standard + // heuristic provides, oversubscribe the SMs. This helps mixed-length batches + // where long sequences get insufficient KV parallelism. + int constexpr kMinTokensPerCta = 2048; + int constexpr kMaxOccupancyWaves = 16; + int constexpr kMaxSplits = 32; + int const desiredSplits = + (params.mMaxSeqLenKv + kMinTokensPerCta - 1) / kMinTokensPerCta; + if (numCtasPerSeqKv < desiredSplits && desiredSplits > 4) { + int const maxSplitsFromSMs = + std::max(1, int32_t(kMaxOccupancyWaves * params.mMultiProcessorCount / baseCtas)); + numCtasPerSeqKv = + std::max(numCtasPerSeqKv, std::min({desiredSplits, maxSplitsFromSMs, + maxNumCtasPerSeqKv, kMaxSplits})); + } + // Update the numCtasX. numCtasX *= numCtasPerSeqKv; // The current total number of CTAs. diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h index 7067243f29..3d6143c0ea 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -825,7 +825,16 @@ struct KernelParams { params.ptrAttentionSinks = options.ptrAttentionSinks; // The partial buffers' pointers when the multiCtasKv mode is enabled. - int64_t partialStatsBufferSize = options.mMultiProcessorCount * kernelMeta.mStepQ; + // partialStats and partialO are packed sequentially in the workspace. Each + // (batch, head, q, kv) tile writes to a unique slot, so the offset between + // them must accommodate the total number of tiles. + int const headDimPerCtaV_ = + kernelMeta.m2CtaMma ? kernelMeta.mHeadDimPerCtaV * 2 : kernelMeta.mHeadDimPerCtaV; + int const numCtasForAllHeads_ = options.mNumHeadsQ / kernelMeta.mStepQ; + int const numHeadDimCtasV_ = kernelMeta.mHeadDimV / headDimPerCtaV_; + int64_t partialStatsBufferSize = static_cast(options.mBatchSize) * + numCtasForAllHeads_ * numHeadDimCtasV_ * maxNumCtasQ * maxNumCtasKv * + kernelMeta.mStepQ; params.ptrMultiCtasKvCounter = options.multiCtasKvCounterPtr; params.ptrPartialStats = reinterpret_cast(options.multiCtasKvScratchPtr); params.ptrPartialO = params.ptrPartialStats + partialStatsBufferSize; From 82a38d1d6c7153f61c11e04a6fa0d100a4e4a302 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 21 May 2026 10:33:21 +0000 Subject: [PATCH 2/7] small fixes Signed-off-by: root --- benchmarks/bench_trtllm_gen_mla.py | 2 -- include/flashinfer/trtllm/fmha/kernelParams.h | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py index 8aac7c3753..ce6263bfe5 100644 --- a/benchmarks/bench_trtllm_gen_mla.py +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -26,7 +26,6 @@ def bench_trtllm_mla( device=device, ).to(dtype) - # Sequence lengths: use provided list or generate random if seq_lens_list is not None: seq_lens = list(seq_lens_list) max_seq_len = max(seq_lens) @@ -195,7 +194,6 @@ def sample_prod_distribution(batch_size: int, seed: int = 42) -> list[int]: print() # Mixed sequence length benchmark (production-like distribution) - print("\n===== Mixed sequence length benchmark =====") for batch_size in [16, 32, 64, 128, 256]: sl = sample_prod_distribution(batch_size, seed=batch_size) try: diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h index 3d6143c0ea..2f05856f4f 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -826,8 +826,7 @@ struct KernelParams { // The partial buffers' pointers when the multiCtasKv mode is enabled. // partialStats and partialO are packed sequentially in the workspace. Each - // (batch, head, q, kv) tile writes to a unique slot, so the offset between - // them must accommodate the total number of tiles. + // (batch, head, q, kv) tile writes to a unique slot. int const headDimPerCtaV_ = kernelMeta.m2CtaMma ? kernelMeta.mHeadDimPerCtaV * 2 : kernelMeta.mHeadDimPerCtaV; int const numCtasForAllHeads_ = options.mNumHeadsQ / kernelMeta.mStepQ; From d61224745277183f3ca13611a60782b4a81f801a Mon Sep 17 00:00:00 2001 From: root Date: Thu, 21 May 2026 10:38:32 +0000 Subject: [PATCH 3/7] precommit --- benchmarks/bench_trtllm_gen_mla.py | 16 +++++++++++++--- include/flashinfer/trtllm/fmha/fmhaKernels.cuh | 12 +++++------- include/flashinfer/trtllm/fmha/kernelParams.h | 4 ++-- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py index ce6263bfe5..135c8c8d5e 100644 --- a/benchmarks/bench_trtllm_gen_mla.py +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -11,7 +11,12 @@ def bench_trtllm_mla( - batch_size, q_len_per_request, seq_len, page_size, dtype, backend="auto", + batch_size, + q_len_per_request, + seq_len, + page_size, + dtype, + backend="auto", seq_lens_list=None, ): torch.manual_seed(42) @@ -198,8 +203,13 @@ def sample_prod_distribution(batch_size: int, seed: int = 42) -> list[int]: sl = sample_prod_distribution(batch_size, seed=batch_size) try: bench_trtllm_mla( - batch_size, 1, max(sl), 32, torch.bfloat16, - backend=args.backend, seq_lens_list=sl, + batch_size, + 1, + max(sl), + 32, + torch.bfloat16, + backend=args.backend, + seq_lens_list=sl, ) except ValueError as e: print(f"SKIPPED: {e}") diff --git a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh index 076c19f769..06f783ad0b 100644 --- a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -479,9 +479,8 @@ class TllmGenFmhaKernel { (maxAttentionWindow + 2 * kernelMeta.mStepKv - 1) / (2 * kernelMeta.mStepKv); int const baseCtas = numCtasX * numCtasY * numCtasZ; - numCtasPerSeqKv = std::min( - maxNumCtasPerSeqKv, - std::max(1, int32_t(params.mMultiProcessorCount / baseCtas))); + numCtasPerSeqKv = std::min(maxNumCtasPerSeqKv, + std::max(1, int32_t(params.mMultiProcessorCount / baseCtas))); // When the longest sequence needs more KV splits than the standard // heuristic provides, oversubscribe the SMs. This helps mixed-length batches @@ -489,14 +488,13 @@ class TllmGenFmhaKernel { int constexpr kMinTokensPerCta = 2048; int constexpr kMaxOccupancyWaves = 16; int constexpr kMaxSplits = 32; - int const desiredSplits = - (params.mMaxSeqLenKv + kMinTokensPerCta - 1) / kMinTokensPerCta; + int const desiredSplits = (params.mMaxSeqLenKv + kMinTokensPerCta - 1) / kMinTokensPerCta; if (numCtasPerSeqKv < desiredSplits && desiredSplits > 4) { int const maxSplitsFromSMs = std::max(1, int32_t(kMaxOccupancyWaves * params.mMultiProcessorCount / baseCtas)); numCtasPerSeqKv = - std::max(numCtasPerSeqKv, std::min({desiredSplits, maxSplitsFromSMs, - maxNumCtasPerSeqKv, kMaxSplits})); + std::max(numCtasPerSeqKv, + std::min({desiredSplits, maxSplitsFromSMs, maxNumCtasPerSeqKv, kMaxSplits})); } // Update the numCtasX. diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h index 2f05856f4f..b78858c4f5 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -832,8 +832,8 @@ struct KernelParams { int const numCtasForAllHeads_ = options.mNumHeadsQ / kernelMeta.mStepQ; int const numHeadDimCtasV_ = kernelMeta.mHeadDimV / headDimPerCtaV_; int64_t partialStatsBufferSize = static_cast(options.mBatchSize) * - numCtasForAllHeads_ * numHeadDimCtasV_ * maxNumCtasQ * maxNumCtasKv * - kernelMeta.mStepQ; + numCtasForAllHeads_ * numHeadDimCtasV_ * maxNumCtasQ * + maxNumCtasKv * kernelMeta.mStepQ; params.ptrMultiCtasKvCounter = options.multiCtasKvCounterPtr; params.ptrPartialStats = reinterpret_cast(options.multiCtasKvScratchPtr); params.ptrPartialO = params.ptrPartialStats + partialStatsBufferSize; From bb416d7eaab9d22bb8db7850adcab0c51a90346e Mon Sep 17 00:00:00 2001 From: root Date: Thu, 21 May 2026 11:25:16 +0000 Subject: [PATCH 4/7] review comments --- benchmarks/bench_trtllm_gen_mla.py | 9 +++++---- include/flashinfer/trtllm/fmha/kernelParams.h | 4 +++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py index 135c8c8d5e..2cd35eac06 100644 --- a/benchmarks/bench_trtllm_gen_mla.py +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -32,6 +32,9 @@ def bench_trtllm_mla( ).to(dtype) if seq_lens_list is not None: + assert len(seq_lens_list) == batch_size, ( + f"seq_lens_list length {len(seq_lens_list)} != batch_size {batch_size}" + ) seq_lens = list(seq_lens_list) max_seq_len = max(seq_lens) else: @@ -39,11 +42,9 @@ def bench_trtllm_mla( seq_lens[-1] = seq_len max_seq_len = max(seq_lens) - num_tokens = max_seq_len * batch_size - num_blocks = (num_tokens + page_size - 1) // page_size seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device) - blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size + num_blocks = int(blocks_per_seq.sum().item()) max_num_blocks_per_seq = blocks_per_seq.max().item() # Generate random but unique block IDs for all sequences @@ -76,7 +77,7 @@ def bench_trtllm_mla( # Allocate workspace buffer # todo(Yingyi): calculate the actual size of workspace buffer - workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=device) + workspace_buffer = torch.empty(1024 * 1024 * 1024, dtype=torch.int8, device=device) # Run decode-MLA # warmup diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h index b78858c4f5..9649daac1a 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -829,7 +829,9 @@ struct KernelParams { // (batch, head, q, kv) tile writes to a unique slot. int const headDimPerCtaV_ = kernelMeta.m2CtaMma ? kernelMeta.mHeadDimPerCtaV * 2 : kernelMeta.mHeadDimPerCtaV; - int const numCtasForAllHeads_ = options.mNumHeadsQ / kernelMeta.mStepQ; + int const numHeadsPerCta_ = + kernelMeta.mGroupsHeadsQ ? std::min(options.mNumHeadsQPerKv, kernelMeta.mStepQ) : 1; + int const numCtasForAllHeads_ = options.mNumHeadsQ / numHeadsPerCta_; int const numHeadDimCtasV_ = kernelMeta.mHeadDimV / headDimPerCtaV_; int64_t partialStatsBufferSize = static_cast(options.mBatchSize) * numCtasForAllHeads_ * numHeadDimCtasV_ * maxNumCtasQ * From 120714014494bf39e931a5aca899dc5989e60962 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 21 May 2026 14:10:10 +0000 Subject: [PATCH 5/7] address comments --- .../flashinfer/trtllm/fmha/fmhaKernels.cuh | 6 ++--- include/flashinfer/trtllm/fmha/kernelParams.h | 22 +++++++------------ 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh index 06f783ad0b..57ac212d6e 100644 --- a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -246,8 +246,7 @@ class TllmGenFmhaKernel { } // Prepare the kernel parameters. - auto kernelParams = KernelParams::setKernelParams( - params, kernelMeta, ctaLaunchParams.mMaxNumCtasQ, ctaLaunchParams.mMaxNumCtasKv); + auto kernelParams = KernelParams::setKernelParams(params, kernelMeta, ctaLaunchParams); // Override SageAttention parameters. auto sageParamEncode = [](int blockSize) -> int32_t { @@ -305,8 +304,7 @@ class TllmGenFmhaKernel { "fallback to GmemReduction."); // Rebuild kernelParams: setKernelParams uses kernelMeta (TMA descriptors, tile shapes) // which changed when switching from CgaSmemReduction to GmemReduction kernel. - kernelParams = KernelParams::setKernelParams( - params, kernelMeta, ctaLaunchParams.mMaxNumCtasQ, ctaLaunchParams.mMaxNumCtasKv); + kernelParams = KernelParams::setKernelParams(params, kernelMeta, ctaLaunchParams); buildLaunchConfig(launch_config, launch_attribute, kernelMeta, ctaLaunchParams, params); setNonPortableClusterIfNeeded(func, ctaLaunchParams); } diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h index 9649daac1a..25a239653e 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -643,9 +643,9 @@ struct KernelParams { } // Setup the kernel parameters. - template + template static KernelParams setKernelParams(FmhaOptions_ const& options, KernelMeta const& kernelMeta, - int32_t maxNumCtasQ, int32_t maxNumCtasKv) { + CtaLaunchParams const& ctaLaunchParams) { // Create the return struct. KernelParams params; // Memset the kernel parameters to 0. @@ -826,16 +826,10 @@ struct KernelParams { // The partial buffers' pointers when the multiCtasKv mode is enabled. // partialStats and partialO are packed sequentially in the workspace. Each - // (batch, head, q, kv) tile writes to a unique slot. - int const headDimPerCtaV_ = - kernelMeta.m2CtaMma ? kernelMeta.mHeadDimPerCtaV * 2 : kernelMeta.mHeadDimPerCtaV; - int const numHeadsPerCta_ = - kernelMeta.mGroupsHeadsQ ? std::min(options.mNumHeadsQPerKv, kernelMeta.mStepQ) : 1; - int const numCtasForAllHeads_ = options.mNumHeadsQ / numHeadsPerCta_; - int const numHeadDimCtasV_ = kernelMeta.mHeadDimV / headDimPerCtaV_; - int64_t partialStatsBufferSize = static_cast(options.mBatchSize) * - numCtasForAllHeads_ * numHeadDimCtasV_ * maxNumCtasQ * - maxNumCtasKv * kernelMeta.mStepQ; + // CTA writes to a unique slot, so the offset must accommodate all CTAs. + int64_t partialStatsBufferSize = static_cast(ctaLaunchParams.mNumCtasX) * + ctaLaunchParams.mNumCtasY * ctaLaunchParams.mNumCtasZ * + kernelMeta.mStepQ; params.ptrMultiCtasKvCounter = options.multiCtasKvCounterPtr; params.ptrPartialStats = reinterpret_cast(options.multiCtasKvScratchPtr); params.ptrPartialO = params.ptrPartialStats + partialStatsBufferSize; @@ -870,8 +864,8 @@ struct KernelParams { params.mMaxSeqLenQ = options.mMaxSeqLenQ; params.mMaxSeqLenKv = options.mMaxSeqLenKv; - params.mMaxNumCtasQ = maxNumCtasQ; - params.mMaxNumCtasKv = maxNumCtasKv; + params.mMaxNumCtasQ = ctaLaunchParams.mMaxNumCtasQ; + params.mMaxNumCtasKv = ctaLaunchParams.mMaxNumCtasKv; params.mMaxNumPagesPerSeqKv = options.mMaxNumPagesPerSeqKv; // TODO: just use mMaxSeqLenQ for number of MTP tokens. params.mSumOfSeqLensQ = options.mSumOfSeqLensQ; From 531fff788897dfa1cd34d4a3a7d23b5160bd920f Mon Sep 17 00:00:00 2001 From: root Date: Mon, 25 May 2026 09:46:25 +0000 Subject: [PATCH 6/7] Tunable constants --- include/flashinfer/trtllm/common.h | 24 +++++++++++++++++++ .../flashinfer/trtllm/fmha/fmhaKernels.cuh | 8 +++---- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/include/flashinfer/trtllm/common.h b/include/flashinfer/trtllm/common.h index ae4b832a8e..c514c49615 100644 --- a/include/flashinfer/trtllm/common.h +++ b/include/flashinfer/trtllm/common.h @@ -197,10 +197,34 @@ inline static bool getBoolEnv(char const* name) { return env && env[0] == '1' && env[1] == '\0'; } +inline static int getIntEnv(char const* name, int defaultVal) { + char const* env = std::getenv(name); + return env ? std::atoi(env) : defaultVal; +} + inline bool getEnvUseTileSizeKv64ForTrtllmGen() { static bool const useTileSizeKv64 = getBoolEnv("TRTLLM_GEN_ENABLE_TILE_SIZE_KV64"); return useTileSizeKv64; } + +// KV split oversubscription tuning (for mixed-length batches). +// Set FLASHINFER_KV_OVERSUB_MAX_WAVES=1 to disable oversubscription. +inline int getEnvKvOversubMinTokensPerCta() { + static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MIN_TOKENS_PER_CTA", 2048); + return val; +} +inline int getEnvKvOversubMaxWaves() { + static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MAX_WAVES", 16); + return val; +} +inline int getEnvKvOversubMaxSplits() { + static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MAX_SPLITS", 32); + return val; +} +inline int getEnvKvOversubMinDesiredSplits() { + static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MIN_DESIRED_SPLITS", 4); + return val; +} template inline __device__ __host__ T divUp(T m, T n) { return (m + n - 1) / n; diff --git a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh index 6db6162b50..015d0511c5 100644 --- a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -498,11 +498,11 @@ class TllmGenFmhaKernel { // When the longest sequence needs more KV splits than the standard // heuristic provides, oversubscribe the SMs. This helps mixed-length batches // where long sequences get insufficient KV parallelism. - int constexpr kMinTokensPerCta = 2048; - int constexpr kMaxOccupancyWaves = 16; - int constexpr kMaxSplits = 32; + int const kMinTokensPerCta = getEnvKvOversubMinTokensPerCta(); + int const kMaxOccupancyWaves = getEnvKvOversubMaxWaves(); + int const kMaxSplits = getEnvKvOversubMaxSplits(); int const desiredSplits = (params.mMaxSeqLenKv + kMinTokensPerCta - 1) / kMinTokensPerCta; - if (numCtasPerSeqKv < desiredSplits && desiredSplits > 4) { + if (numCtasPerSeqKv < desiredSplits && desiredSplits > getEnvKvOversubMinDesiredSplits()) { int const maxSplitsFromSMs = std::max(1, int32_t(kMaxOccupancyWaves * params.mMultiProcessorCount / baseCtas)); numCtasPerSeqKv = From 92eae0a494cfc6ab66b6601b2efd8cdb1bb1714c Mon Sep 17 00:00:00 2001 From: root Date: Mon, 25 May 2026 09:55:59 +0000 Subject: [PATCH 7/7] Tunable constants off by default --- include/flashinfer/trtllm/common.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/flashinfer/trtllm/common.h b/include/flashinfer/trtllm/common.h index c514c49615..94de512455 100644 --- a/include/flashinfer/trtllm/common.h +++ b/include/flashinfer/trtllm/common.h @@ -208,13 +208,13 @@ inline bool getEnvUseTileSizeKv64ForTrtllmGen() { } // KV split oversubscription tuning (for mixed-length batches). -// Set FLASHINFER_KV_OVERSUB_MAX_WAVES=1 to disable oversubscription. +// Disabled by default. Set FLASHINFER_KV_OVERSUB_MAX_WAVES=16 to enable. inline int getEnvKvOversubMinTokensPerCta() { static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MIN_TOKENS_PER_CTA", 2048); return val; } inline int getEnvKvOversubMaxWaves() { - static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MAX_WAVES", 16); + static int const val = getIntEnv("FLASHINFER_KV_OVERSUB_MAX_WAVES", 1); return val; } inline int getEnvKvOversubMaxSplits() {