change to storing seqlen_q so combine kernel can do coalesced reads for virtual batch metadata

Author: jayhshah

hopper/flash.h

@@ -151,7 +151,8 @@ struct Flash_fwd_params : public Qkv_params {
     bool pack_gqa;
 
     int * __restrict__ tile_count_semaphore;
-    int * __restrict__ num_m_blocks_ptr;
+    // int * __restrict__ num_m_blocks_ptr;
+    int * __restrict__ prepare_seqlen_q_ptr;
     // int * __restrict__ num_n_blocks_ptr;
     int * __restrict__ num_splits_dynamic_ptr;
     int * __restrict__ varlen_batch_idx_ptr; // virtual -> actual

hopper/flash_api.cpp

@@ -656,8 +656,8 @@ mha_fwd_get_scheduler_metadata(
         tile_count_semaphore = torch::empty(
             {int(scheduler_needs_semaphore) + tile_count_semaphore_offset},
             opts.dtype(torch::kInt32));
-        // ORDER: {num_m_blocks, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2}
-        params.num_m_blocks_ptr =  use_prepare_varlen ? tile_count_semaphore.data_ptr<int>() : nullptr;
+        // ORDER: {prepare_seqlen_q, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2}
+        params.prepare_seqlen_q_ptr =  use_prepare_varlen ? tile_count_semaphore.data_ptr<int>() : nullptr;
         params.num_splits_dynamic_ptr = use_prepare_varlen && use_dynamic_split ? tile_count_semaphore.data_ptr<int>() + b_rounded : nullptr;
         params.varlen_batch_idx_ptr =  use_prepare_varlen && params.varlen_sort_batches ? tile_count_semaphore.data_ptr<int>() + sort_offset : nullptr;
         // params.num_n_blocks_ptr  = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr<int>() + head_swizzle_offset : nullptr;
@@ -1058,8 +1058,8 @@ mha_fwd(at::Tensor &q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq
         if (scheduler_needs_semaphore && !use_prepare_varlen) {
             tile_count_semaphore.zero_();  // If varlen we'll manually do the zero-ing
         }
-        // ORDER: {num_m_blocks, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2}
-        params.num_m_blocks_ptr =  use_prepare_varlen ? tile_count_semaphore.data_ptr<int>() : nullptr;
+        // ORDER: {prepare_seqlen_q, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2}
+        params.prepare_seqlen_q_ptr =  use_prepare_varlen ? tile_count_semaphore.data_ptr<int>() : nullptr;
         params.num_splits_dynamic_ptr = use_prepare_varlen && use_dynamic_split ? tile_count_semaphore.data_ptr<int>() + b_rounded : nullptr;
         params.varlen_batch_idx_ptr =  use_prepare_varlen && params.varlen_sort_batches ? tile_count_semaphore.data_ptr<int>() + sort_offset : nullptr;
         // params.num_n_blocks_ptr  = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr<int>() + head_swizzle_offset : nullptr;

hopper/flash_fwd_combine_kernel.h

@@ -209,10 +209,12 @@ class FlashAttnFwdCombine {
         int seqlen_q;
         int total_q;
         int num_heads;
+        int num_heads_kv;
         int dv;
+        bool pack_gqa;
         int const* cu_seqlens_q;
         int const* seqused_q;
-        int const* varlen_batch_idx_ptr;
+        int const* prepare_seqlen_q_ptr;
     };
 
     struct StaticTileScheduler {
@@ -257,12 +259,14 @@ class FlashAttnFwdCombine {
         };
 
         struct Params {
-            int b;
-            int num_heads;
+            int const b;
+            int const num_heads;
+            int const num_heads_kv;
+            bool const pack_gqa;
             int const* const cu_seqlens_q;
             int const* const seqused_q;
+            int const* const prepare_seqlen_q_ptr;
             SchedulingAlgo algo;
-            int const* const varlen_batch_idx_ptr = nullptr;
         };
 
         SharedStorage& shared_storage;
@@ -286,10 +290,12 @@ class FlashAttnFwdCombine {
             return {
                 args.b,
                 args.num_heads,
+                args.num_heads_kv,
+                args.pack_gqa,
                 args.cu_seqlens_q,
                 args.seqused_q,
-                choose_scheduling_algo(args),
-                args.varlen_batch_idx_ptr
+                args.prepare_seqlen_q_ptr,
+                choose_scheduling_algo(args)
             }; 
         }
 
@@ -315,7 +321,6 @@ class FlashAttnFwdCombine {
         }
 
         CUTE_DEVICE BlockCoord get_block_coord_linearized_m_and_batch(Params const& params) {
-            int num_heads = params.num_heads;
             int curr_tile_id = blockIdx.x;
 
             // Scan through the batches find the batch that contains the current
@@ -338,9 +343,13 @@ class FlashAttnFwdCombine {
 
                     auto get_num_m_blocks = [&](int bidb) {
                         if (bidb >= params.b) return 0;
-                        int actual_bidb = params.varlen_batch_idx_ptr ? params.varlen_batch_idx_ptr[bidb] : bidb;
-                        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{actual_bidb, 0, params.cu_seqlens_q, params.seqused_q};
-                        return cute::ceil_div(seqlen_info.seqlen * num_heads, Int<kBlockM>{}());
+                        if (params.prepare_seqlen_q_ptr) {
+                            int length = params.prepare_seqlen_q_ptr[bidb] * (!params.pack_gqa ? params.num_heads : params.num_heads_kv);
+                            return cute::ceil_div(length, Int<kBlockM>{});
+                        } else {
+                            flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{bidb, 0, params.cu_seqlens_q, params.seqused_q};
+                            return cute::ceil_div(seqlen_info.seqlen * params.num_heads, Int<kBlockM>{});
+                        }
                     };
 
                     // Cumulative number of blocks for the next 31 batches

hopper/flash_fwd_combine_launch_template.h

@@ -40,9 +40,8 @@ void run_flash_fwd_combine(Flash_fwd_params &params, cudaStream_t stream, bool e
     };
 
     typename CombineKernel::SchedulerArguments scheduler_args  {
-        params.b, params.seqlen_q, params.total_q, params.h, params.dv,
-        params.cu_seqlens_q, params.seqused_q,
-        params.varlen_batch_idx_ptr
+        params.b, params.seqlen_q, params.total_q, params.h, params.h_k, params.dv, params.pack_gqa,
+        params.cu_seqlens_q, params.seqused_q, params.prepare_seqlen_q_ptr
     };
 
     typename CombineKernel::Params kernel_params = {

hopper/flash_fwd_launch_template.h

@@ -157,7 +157,7 @@ void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
         params.seqlen_k, params.d, params.dv, sizeof(Element),
         params.tile_count_semaphore, params.cu_seqlens_q, params.seqused_q,
         params.num_splits_dynamic_ptr,
-        params.num_m_blocks_ptr,
+        params.prepare_seqlen_q_ptr,
         params.varlen_batch_idx_ptr,
         params.num_nheads_in_l2_ptr
     };

hopper/flash_prepare_scheduler.cu

@@ -47,7 +47,7 @@ __global__ void prepare_varlen_num_blocks_kernel(
         int num_batch, int num_head, int qhead_per_khead, int num_sm, int num_splits_static,
         cutlass::FastDivmod blockm_divmod, cutlass::FastDivmod blockn_divmod,
         int* const tile_count_semaphore,
-        int* const num_m_blocks_ptr,
+        int* const prepare_seqlen_q_ptr,
         int* const num_splits_dynamic_ptr,
         int* const varlen_batch_idx_ptr,
         // int* const num_n_blocks_ptr,
@@ -78,7 +78,7 @@ __global__ void prepare_varlen_num_blocks_kernel(
 
     int lane = threadIdx.x % cutlass::NumThreadsPerWarp;
 
-    auto get_num_m_blocks = [&](int batch_idx) {
+    auto get_num_m_blocks_and_seqlen = [&](int batch_idx) {
         int seqlen;
         if (seqused_q) {
             seqlen = batch_idx < num_batch ? seqused_q[batch_idx] : 0;
@@ -91,7 +91,8 @@ __global__ void prepare_varlen_num_blocks_kernel(
         }
         if(packgqa) { seqlen *= qhead_per_khead; }
         return batch_idx < num_batch && lane < kNumBatchPerWarp
-            ? blockm_divmod.div(seqlen + blockm_divmod.divisor - 1) : 0;
+            ? cute::make_tuple(blockm_divmod.div(seqlen + blockm_divmod.divisor - 1), seqlen)
+            : cute::make_tuple(0, 0);
     };
 
     auto get_num_n_blocks = [&](int batch_idx) {
@@ -124,7 +125,10 @@ __global__ void prepare_varlen_num_blocks_kernel(
     int batch_cta_idx_offset = int(blockIdx.x) * 992;
     int bidb_start = batch_cta_idx_offset + kNumBatchPerWarp * warp_idx;
     int batch_idx = lane + bidb_start;
-    int num_m_blocks = get_num_m_blocks(batch_idx);
+    // int num_m_blocks = get_num_m_blocks(batch_idx);
+    auto seqlen_q_info = get_num_m_blocks_and_seqlen(batch_idx);
+    int num_m_blocks = cute::get<0>(seqlen_q_info);
+    int seqlen_q = cute::get<1>(seqlen_q_info);
     int num_n_blocks = get_num_n_blocks(batch_idx);
 
     auto get_nheads_in_l2 = [&](int n_blocks) {
@@ -165,47 +169,35 @@ __global__ void prepare_varlen_num_blocks_kernel(
             num_n_blocks = INT_MIN; // sort last
         } else if (is_causal) {
             // sort by shortest member to process
-            num_n_blocks = num_n_blocks * blockn_divmod.divisor - num_m_blocks * blockm_divmod.divisor;
+            num_n_blocks = num_n_blocks * blockn_divmod.divisor - seqlen_q;
         }
         int4 batch_coords[ITEMS_PER_THREAD]; // 1 item per thread
-        batch_coords[0] = make_int4(num_n_blocks, num_m_blocks, num_splits_dynamic, batch_idx);
-
-        // if (threadIdx.x == 0) {
-        //     printf("Unsorted: num_n_blocks - num_m_blocks = %d, num_m_blocks = %d, num_splits = %d, batch_idx = %d.\n", 
-        //         batch_coords[0].x, batch_coords[0].y, batch_coords[0].z, batch_coords[0].w);
-        // } __syncthreads();
+        batch_coords[0] = make_int4(num_n_blocks, seqlen_q, num_splits_dynamic, batch_idx);
 
         // Sort batches by num_n_blocks in descending order
         BlockMergeSort(temp_storage).Sort(batch_coords, PrepareSortOp<int4>());
 
-        // if (threadIdx.x == 0) {
-        //     printf("Sorted: num_n_blocks - num_m_blocks = %d, num_m_blocks = %d, num_splits = %d, batch_idx = %d.\n", 
-        //         batch_coords[0].x, batch_coords[0].y, batch_coords[0].z, batch_coords[0].w);
-        // } __syncthreads();
-
         if (is_causal) {
             // reset value to num_n_blocks
-            batch_coords[0].x = blockn_divmod.div(batch_coords[0].x + batch_coords[0].y * blockm_divmod.divisor);
+            batch_coords[0].x = blockn_divmod.div(batch_coords[0].x + batch_coords[0].y);
         }
 
         // When sorting, we re-index some metadata by 'virtual batch index'
         // and also store the vbidx -> bidx mapping.
         // 1. num_nheads_in_l2_ptr: virtual_batch_idx -> num_nheads_in_l2[batch_idx]
         // 2. num_splits_dynamic_ptr: virtual_batch_idx -> num_splits[batch_idx]
-        // 3. num_m_blocks_ptr: virtual_batch_idx -> num_m_blocks[batch_idx]
+        // 3. prepare_seqlen_q_ptr: virtual_batch_idx -> seqlen_q[batch_idx] * (packgqa ? qhead_per_khead : 1)
         // 4. varlen_batch_idx_ptr: virtual_batch_idx -> batch_idx      
         batch_idx = batch_cta_idx_offset + threadIdx.x;
         if (batch_idx < num_batch && threadIdx.x < 992) {
-            // num_n_blocks_ptr[threadIdx.x] = max(batch_coords[0].x, 1);
             if(num_nheads_in_l2_ptr) { num_nheads_in_l2_ptr[batch_idx] = get_nheads_in_l2(max(batch_coords[0].x, 1)); }
-            num_m_blocks_ptr[batch_idx] = batch_coords[0].y;
+            prepare_seqlen_q_ptr[batch_idx] = batch_coords[0].y;
             if(num_splits_dynamic_ptr) { num_splits_dynamic_ptr[batch_idx] = batch_coords[0].z; }
             varlen_batch_idx_ptr[batch_idx] = batch_coords[0].w;
         }  
     } else {
         if (batch_idx < num_batch && lane < kNumBatchPerWarp) {
-            num_m_blocks_ptr[batch_idx] = num_m_blocks;
-            // num_n_blocks_ptr[batch_idx] = max(num_n_blocks, 1);
+            prepare_seqlen_q_ptr[batch_idx] = seqlen_q;
             if(num_splits_dynamic_ptr) { num_splits_dynamic_ptr[batch_idx] = num_splits_dynamic; }
             if(num_nheads_in_l2_ptr) { num_nheads_in_l2_ptr[batch_idx] = get_nheads_in_l2(max(num_n_blocks, 1)); }
             // printf("idx = %d, num_m_blocks = %d, num_n_blocks = %d, num_split_static = %d, num_splits_dynamic = %d\n", bidb_start + lane, num_m_blocks_ptr[bidb_start + lane], num_n_blocks, num_splits_static, num_splits_dynamic);
@@ -236,7 +228,7 @@ void prepare_varlen_num_blocks(Flash_fwd_params &params, cudaStream_t stream, bo
                 params.b, !packgqa ? params.h : params.h_k, qhead_per_khead, params.num_sm, params.num_splits,
                 cutlass::FastDivmod(blockM), cutlass::FastDivmod(blockN),
                 params.tile_count_semaphore,
-                params.num_m_blocks_ptr,
+                params.prepare_seqlen_q_ptr,
                 params.num_splits_dynamic_ptr,
                 params.varlen_batch_idx_ptr,
                 // params.num_n_blocks_ptr,

hopper/tile_scheduler.hpp

@@ -25,7 +25,7 @@ struct TileSchedulerArguments {
     int const* const cu_seqlens = nullptr;
     int const* const seqused = nullptr;
     int const* const num_splits_dynamic_ptr = nullptr;
-    int const* const num_m_blocks_ptr = nullptr;
+    int const* const prepare_seqlen_q_ptr = nullptr;
     int const* const varlen_batch_idx_ptr = nullptr;
     // int const* const num_n_blocks_ptr = nullptr;
     int const* const num_nheads_in_l2_ptr = nullptr;
@@ -385,7 +385,7 @@ class VarlenDynamicPersistentTileScheduler {
         int const* const cu_seqlens;
         int const* const seqused;
         int const* const num_splits_dynamic_ptr;
-        int const* const num_m_blocks_ptr;
+        int const* const prepare_seqlen_q_ptr;
         int const* const varlen_batch_idx_ptr;
         // int const* const num_n_blocks_ptr;
         int const* const num_nheads_in_l2_ptr;
@@ -408,7 +408,7 @@ class VarlenDynamicPersistentTileScheduler {
                 cutlass::FastDivmod(!Split ? 1 : args.num_splits),
                 args.tile_count_semaphore, args.cu_seqlens, args.seqused,
                 args.num_splits_dynamic_ptr,
-                args.num_m_blocks_ptr,
+                args.prepare_seqlen_q_ptr,
                 args.varlen_batch_idx_ptr,
                 // aras.num_n_blocks_ptr,
                 args.num_nheads_in_l2_ptr};
@@ -470,7 +470,7 @@ class VarlenDynamicPersistentTileScheduler {
             int batch_idx = lane + bidb_start;
             if constexpr (Prepared) {
                 return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1
-                    ? params.num_m_blocks_ptr[batch_idx] : 0;
+                    ? cute::ceil_div(params.prepare_seqlen_q_ptr[batch_idx], kBlockM) : 0;
             } else {
                 int seqlen = params.seqlen * (!PackGQA ? 1 : params.qhead_per_khead);
                 if (seqlen > kBlockM) {
@@ -487,7 +487,6 @@ class VarlenDynamicPersistentTileScheduler {
                 }
                 return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1
                     ? cute::ceil_div(seqlen, kBlockM) : 0;
-                    // ? params.num_m_blocks_ptr[batch_idx] : 0;
             }
         };