KV Split Oversubscription for Mixed Sequence Lengths

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

csrc/fmhaReduction.cu

@@ -371,7 +371,7 @@ void runFmhaReduction(TllmGenFmhaKernelMetaInfo const& kernelMeta, KernelParams
   int32_t const maxNumCtasForReduction{(multiProcessorCount * 2) /
                                        static_cast<int32_t>(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;
 

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 <typename T>
 inline __device__ __host__ T divUp(T m, T n) {
   return (m + n - 1) / n;

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.

include/flashinfer/trtllm/fmha/kernelParams.h

@@ -647,9 +647,9 @@ struct KernelParams {
   }
 
   // Setup the kernel parameters.
-  template <class FmhaOptions_, class KernelMeta>
+  template <class FmhaOptions_, class KernelMeta, class CtaLaunchParams>
   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<int64_t>(ctaLaunchParams.mNumCtasX) *
+                                     ctaLaunchParams.mNumCtasY * ctaLaunchParams.mNumCtasZ *
+                                     kernelMeta.mStepQ;
     params.ptrMultiCtasKvCounter = options.multiCtasKvCounterPtr;
     params.ptrPartialStats = reinterpret_cast<float2*>(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;