From 4da1bfa95b94b519365ee800e79603e943c16a0e Mon Sep 17 00:00:00 2001
From: Jingxin Pan <jingxin.pan@bytedance.com>
Date: Fri, 24 Jul 2026 09:59:12 -0700
Subject: [PATCH] optimize kda prepare cu_seqlens scan with prefix-sum and
 binary search

---
 csrc/flash_kda.cpp        |  5 ++++-
 csrc/smxx/fwd_kernel1.cuh | 41 +++++++++++++++++++++++++++------------
 csrc/smxx/fwd_launch.cu   |  9 ++++++++-
 3 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/csrc/flash_kda.cpp b/csrc/flash_kda.cpp
index 5ea5573..81f5483 100644
--- a/csrc/flash_kda.cpp
+++ b/csrc/flash_kda.cpp
@@ -19,7 +19,10 @@ int64_t get_workspace_size(
 
     int64_t per_tile_bytes = 3 * CHUNK * D * 2 + D * 4 + 2 * CHUNK * CHUNK * 2;
 
-    return H * total_tiles * per_tile_bytes;
+    // Trailing buffer for the tile prefix-sum (N+1 int32), 128-byte aligned.
+    int64_t tile_prefix_bytes = ((N + 1) * 4 + 127) / 128 * 128;
+
+    return H * total_tiles * per_tile_bytes + tile_prefix_bytes;
 }
 
 void fwd(
diff --git a/csrc/smxx/fwd_kernel1.cuh b/csrc/smxx/fwd_kernel1.cuh
index a49f871..354cad4 100644
--- a/csrc/smxx/fwd_kernel1.cuh
+++ b/csrc/smxx/fwd_kernel1.cuh
@@ -82,6 +82,26 @@ struct SharedStorageK1 {
     alignas(16) cutlass::arch::ClusterTransactionBarrier tma_load_barrier;
 };
 
+// Build prefix-sum of per-sequence tile counts, so Kernel 1 can map
+// global_tile_idx -> seq_idx with an O(log N) binary search instead of an
+// O(N) linear scan per CTA in varlen mode.
+__global__ void _flash_kda_build_tile_prefix(
+    int64_t const* __restrict__ cu_seqlens,
+    int N,
+    int chunk,
+    int* __restrict__ tile_prefix
+) {
+    if (threadIdx.x == 0) {
+        int acc = 0;
+        tile_prefix[0] = 0;
+        for (int i = 0; i < N; ++i) {
+            int slen = int(cu_seqlens[i + 1] - cu_seqlens[i]);
+            acc += (slen + chunk - 1) / chunk;
+            tile_prefix[i + 1] = acc;
+        }
+    }
+}
+
 // ==================== Kernel 1: Prepare ====================
 template <
     class TmaLoadQ,
@@ -115,7 +135,8 @@ __global__ void __launch_bounds__(NumThreads, 8) _flash_kda_fwd_prepare(
     int64_t const* cu_seqlens,
     int total_tiles,
     float const* A_log_ptr,
-    float gate_scale
+    float gate_scale,
+    int const* tile_prefix
 ) {
     // --- constants
     using BF16 = cutlass::bfloat16_t;
@@ -151,18 +172,14 @@ __global__ void __launch_bounds__(NumThreads, 8) _flash_kda_fwd_prepare(
     int seq_len, t_tiles_this_seq;
 
     if constexpr (IsVarlen) {
-        // Linear scan on cu_seqlens to find (seq_idx, local_t)
-        seq_idx = 0;
-        tiles_before = 0;
-        for (int i = 0; i < N; i++) {
-            int slen = int(cu_seqlens[i + 1] - cu_seqlens[i]);
-            int n_tiles = (slen + CHUNK - 1) / CHUNK;
-            if (tiles_before + n_tiles > global_tile_idx) {
-                seq_idx = i;
-                break;
-            }
-            tiles_before += n_tiles;
+        int lo = 0, hi = N;
+        while (lo + 1 < hi) {
+            int mid = (lo + hi) >> 1;
+            if (tile_prefix[mid] <= global_tile_idx) lo = mid;
+            else hi = mid;
         }
+        seq_idx = lo;
+        tiles_before = tile_prefix[lo];
         local_t = global_tile_idx - tiles_before;
         bos = cu_seqlens[seq_idx];
         eos = cu_seqlens[seq_idx + 1];
diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu
index 74f7ee4..91a67a7 100644
--- a/csrc/smxx/fwd_launch.cu
+++ b/csrc/smxx/fwd_launch.cu
@@ -68,6 +68,8 @@ void launch_fwd(
     BF16*  ws_inv = reinterpret_cast<BF16*>(ws + n_ht * (WS::kKDecayed + WS::kQDecayed + WS::kKRestored + WS::kGTotal));
     BF16*  ws_mqk = reinterpret_cast<BF16*>(ws + n_ht * (WS::kKDecayed + WS::kQDecayed + WS::kKRestored + WS::kGTotal + WS::kINV));
 
+    int* ws_tile_prefix = reinterpret_cast<int*>(ws + n_ht * WS::kPerTile);
+
     auto ws_kd_gmem_layout = make_layout(make_shape(int(n_ht), CHUNK, D), LayoutRight{});
     auto ws_qd_gmem_layout = ws_kd_gmem_layout;
     auto ws_kr_gmem_layout = ws_kd_gmem_layout;
@@ -159,6 +161,11 @@ void launch_fwd(
 
         cudaFuncSetAttribute(kernel1, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k1);
 
+        if constexpr (IsVarlen) {
+            _flash_kda_build_tile_prefix<<<1, 32, 0, stream>>>(
+                cu_seqlens_ptr, N, CHUNK, ws_tile_prefix);
+        }
+
         dim3 grid_k1(total_tiles, H);
         dim3 block_k1(kK1Threads);
 
@@ -168,7 +175,7 @@ void launch_fwd(
             tma_store_ws_kd, tma_store_ws_qd, tma_store_ws_kr,
             tma_store_ws_gt, tma_store_ws_inv, tma_store_ws_mqk,
             scale, T_total, H, N, cu_seqlens_ptr, total_tiles,
-            A_log_ptr, gate_scale
+            A_log_ptr, gate_scale, ws_tile_prefix
         );
     }
 #endif