From 9c0a3731ead1306278962522206be52dc1e9da2e Mon Sep 17 00:00:00 2001
From: 0z5a <dinadegttyarenko6@gmail.com>
Date: Fri, 14 Aug 2026 13:03:56 +0800
Subject: [PATCH] Reuse tile prefix in recurrence kernel

---
 csrc/smxx/fwd_kernel2.cuh | 7 ++-----
 csrc/smxx/fwd_launch.cu   | 2 +-
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/csrc/smxx/fwd_kernel2.cuh b/csrc/smxx/fwd_kernel2.cuh
index 26f73fe..dc8d144 100644
--- a/csrc/smxx/fwd_kernel2.cuh
+++ b/csrc/smxx/fwd_kernel2.cuh
@@ -147,6 +147,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     int H,
     int N,
     int64_t const* cu_seqlens,
+    int const* tile_prefix,
     int total_tiles
 ) {
     using BF16 = cutlass::bfloat16_t;
@@ -221,11 +222,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     if constexpr (IsVarlen) {
         bos = cu_seqlens[seq_idx];
         eos = cu_seqlens[seq_idx + 1];
-        // Compute tile_base via linear scan (no host-precomputed table)
-        tile_base = 0;
-        for (int i = 0; i < seq_idx; i++) {
-            tile_base += (int(cu_seqlens[i + 1] - cu_seqlens[i]) + CHUNK - 1) / CHUNK;
-        }
+        tile_base = tile_prefix[seq_idx];
     } else {
         int T_seq = T_total / N;
         bos = seq_idx * T_seq;
diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu
index 91a67a7..fbfed7f 100644
--- a/csrc/smxx/fwd_launch.cu
+++ b/csrc/smxx/fwd_launch.cu
@@ -210,7 +210,7 @@ void launch_fwd(
             tma_load_initial_state,
             tma_store_final_state,
             tma_store_out,
-            out_ptr, T_total, H, N, cu_seqlens_ptr, total_tiles
+            out_ptr, T_total, H, N, cu_seqlens_ptr, ws_tile_prefix, total_tiles
         );
     }
 #endif