diff --git a/benchmarks/bench_trtllm_gen_mla.py b/benchmarks/bench_trtllm_gen_mla.py index 5856571327..66f10cd74e 100644 --- a/benchmarks/bench_trtllm_gen_mla.py +++ b/benchmarks/bench_trtllm_gen_mla.py @@ -18,6 +18,7 @@ def bench_trtllm_mla( dtype, backend="auto", with_sinks=False, + seq_lens_list=None, ): """Benchmark a single (config, backend, sinks?) cell. @@ -38,16 +39,20 @@ def bench_trtllm_mla( device=device, ).to(dtype) - num_tokens = seq_len * batch_size - num_blocks = (num_tokens + page_size - 1) // page_size + 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: + 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) 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 @@ -80,7 +85,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(1024 * 1024 * 1024, dtype=torch.int8, device=device) sinks = ( torch.randn(num_q_heads, dtype=torch.float32, device=device) @@ -155,6 +160,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 @@ -238,3 +253,27 @@ def bench_trtllm_mla( f"{type(e).__name__}: {e}" ) print() + + # Mixed sequence length benchmark (production-like distribution) + 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 dbe9eda1eb..85befca21e 100644 --- a/csrc/fmhaReduction.cu +++ b/csrc/fmhaReduction.cu @@ -371,7 +371,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/common.h b/include/flashinfer/trtllm/common.h index ae4b832a8e..94de512455 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). +// 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", 1); + 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 121bc15cf9..015d0511c5 100644 --- a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -259,8 +259,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 { @@ -318,8 +317,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); } @@ -492,10 +490,26 @@ class TllmGenFmhaKernel { // benefits of a shorter mainloop. int const maxNumCtasPerSeqKv = (maxAttentionWindow + 2 * kernelMeta.mStepKv - 1) / (2 * kernelMeta.mStepKv); - // Compute numCtasPerSeqKv. - numCtasPerSeqKv = std::min( - maxNumCtasPerSeqKv, - std::max(1, int32_t(params.mMultiProcessorCount / (numCtasX * numCtasY * numCtasZ)))); + + int const baseCtas = numCtasX * numCtasY * numCtasZ; + 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 + // where long sequences get insufficient KV parallelism. + int const kMinTokensPerCta = getEnvKvOversubMinTokensPerCta(); + int const kMaxOccupancyWaves = getEnvKvOversubMaxWaves(); + int const kMaxSplits = getEnvKvOversubMaxSplits(); + int const desiredSplits = (params.mMaxSeqLenKv + kMinTokensPerCta - 1) / kMinTokensPerCta; + if (numCtasPerSeqKv < desiredSplits && desiredSplits > getEnvKvOversubMinDesiredSplits()) { + 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 3d63dd5cee..f0de2d60f2 100644 --- a/include/flashinfer/trtllm/fmha/kernelParams.h +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -647,9 +647,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. @@ -841,7 +841,11 @@ 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 + // 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; @@ -876,8 +880,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;