Add Paged Attention support to FMHA FWD CUTLASS kernel for variable length (#5033)

Summary:
Pull Request resolved: #5033

X-link: https://github.com/facebookresearch/FBGEMM/pull/2046

Added Paged attention for variable sequence length case for Blackwell Cutlass FWD kernel.

Reviewed By: Aya-ZIbra

Differential Revision: D84284273

fbshipit-source-id: d1c102d225a8ebd704811fba8a2a5791e471908d

Author: sarithad-meta

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd.cu

@@ -104,7 +104,13 @@ std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd(
   };
 
   auto dispatch_mask = [&](auto varlen) {
-    int seq_k = kIsPaged ? static_cast<int>(*seqlen_k) : varlen ? k.size(0) : k.size(1);
+    int seq_k = kIsPaged
+        ? (varlen
+            ? static_cast<int>(*max_seq_len_k)
+            : static_cast<int>(*seqlen_k))
+        : (varlen
+            ? k.size(0)
+            : k.size(1));
     if (causal) {
       if (bottom_right) {
         return dispatch_head_dim(varlen, CausalMask</*kIsQBegin=*/false>{});
@@ -113,7 +119,7 @@ std::tuple<at::Tensor, at::Tensor> dispatch_fmha_fwd(
       }
     } else if (local) {
       if (bottom_right) {
-        return dispatch_head_dim(varlen, LocalMask</*kIsQBegin=*/false>{});
+          return dispatch_head_dim(varlen, LocalMask</*kIsQBegin=*/false>{});
       } else {
         return dispatch_head_dim(varlen, LocalMask</*kIsQBegin=*/true>{});
       }

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_fwd_template.cuh

@@ -91,11 +91,19 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
               StrideLSE>,
           TileScheduler>>;
 
-  if (kIsPaged && !kIsVarlen) {
-    TORCH_CHECK(
+  if (kIsPaged) {
+    if (kIsVarlen) { // Variable length
+      TORCH_CHECK(
+          q.dim() == 3,
+          "Expect Q shape to be (total_Q_seqlen, num_Q_heads, head_dim) ",
+          "Found shape ", q.sizes());
+    }
+    else { // Fixed Length
+      TORCH_CHECK(
           q.dim() == 4,
           "Expect Q shape to be (batch_size, Q_seqlen, num_Q_heads, head_dim). ",
           "Found shape ", q.sizes());
+    }
     TORCH_CHECK(
         k.dim() == 4,
         "Expect K shape to be (num_blocks, page_block_size, num_KV_heads, head_dim) ",
@@ -113,7 +121,10 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
     TORCH_CHECK((k.size(1) % tile_N) == 0, "Page Block Size should be divisible by N tile size");
     TORCH_CHECK((v.size(1) % tile_N) == 0, "Page Block Size should be divisible by N tile size");
 
-    TORCH_CHECK(seqlen_k.has_value(), "seqlen_k should be set");
+    // For fixed length sequences, seqlen_k should be set.
+    if (!kIsVarlen) {
+        TORCH_CHECK(seqlen_k.has_value(), "seqlen_k should be set");
+    }
   }
   else if (kIsVarlen) {
     TORCH_CHECK(
@@ -153,7 +164,8 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
 
   // Extract dimensions from input tensors
   int H_Q = kIsVarlen ? q.size(1) : q.size(2); // Number of Q heads
-  int H_K = kIsVarlen ? k.size(1) : k.size(2); // Number of K heads
+  int H_K = (kIsPaged && kIsVarlen) ? k.size(2)
+          : (kIsVarlen ? k.size(1) : k.size(2)); // Number of K heads
   int D = q.size(q.dim() - 1); // Head dimension (D)
 
   TORCH_CHECK(H_Q % H_K == 0);
@@ -162,14 +174,20 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
 
   // SQ represents SumB(Q) for varlen (jagged len)
   int SQ = kIsVarlen ? q.size(0) : q.size(1);
-  int SK = kIsPaged ? static_cast<int>(*seqlen_k) : kIsVarlen ? k.size(0) : k.size(1);
+  int SK = kIsPaged
+        ? (kIsVarlen
+            ? static_cast<int>(*max_seq_len_k)
+            : static_cast<int>(*seqlen_k))
+        : (kIsVarlen
+            ? k.size(0)
+            : k.size(1));
   int B = kIsVarlen ? cu_seqlens_q->size(0) - 1 : q.size(0);
 
   // Parameters for paged attention.
   int page_table_stride = kIsPaged ? page_table.value().size(1) : 0;
   int num_blocks = kIsPaged ? k.size(0) : 1; // num_blocks
   int page_block_size = kIsPaged ? k.size(1) : 1; // page_block_size
-  // num KV tiles > 1 within a page in the case of  page_block_size > TileShapeN.
+  // num KV tiles > 1 within a page in the case of page_block_size > TileShapeN.
   int num_KV_tiles_per_page = kIsPaged ? k.size(1) / (get<1>(TileShape{}).value) : 1;
 
   ProblemShapeType problem_shape;
@@ -250,8 +268,10 @@ std::tuple<at::Tensor, at::Tensor> fmha_fwd(
   typename Operation::Arguments arguments;
   if constexpr (kIsVarlen) {
     get<2, 1>(stride_Q) = 0;
-    get<2, 1>(stride_K) = 0;
-    get<2, 1>(stride_V) = 0;
+    if (!kIsPaged) {
+        get<2, 1>(stride_K) = 0;
+        get<2, 1>(stride_V) = 0;
+    }
     get<2, 1>(stride_O) = 0;
     get<1, 1>(stride_LSE) = 0;
   }

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/collective/sm100_fmha_load_tma_warpspecialized.hpp

@@ -118,7 +118,6 @@ struct Sm100FmhaLoadTmaWarpspecialized {
     auto dV = args.dV;
     bool kIsPaged = args.ptr_page_table ? true : false;
 
-
     // Local changes (D79534034)
     int get_0 = int(get<0>(problem_shape));
     int get_1 = int(get<1>(problem_shape));
@@ -128,10 +127,12 @@ struct Sm100FmhaLoadTmaWarpspecialized {
       get_0 = get<0>(problem_shape).total_length;
     }
 
-    if constexpr (is_variable_length_v<tuple_element_t<1, ProblemShape>>) {
-      get<2, 1>(dK) = 0;
-      get<2, 1>(dV) = 0;
-      get_1 = get<1>(problem_shape).total_length;
+    if (!kIsPaged) {
+      if constexpr (is_variable_length_v<tuple_element_t<1, ProblemShape>>) {
+        get<2, 1>(dK) = 0;
+        get<2, 1>(dV) = 0;
+        get_1 = get<1>(problem_shape).total_length;
+      }
     }
 
     TMA_Q tma_load_q;
@@ -141,7 +142,8 @@ struct Sm100FmhaLoadTmaWarpspecialized {
     if (kIsPaged) { // Paged Case
       //Create TMA Atom/Descriptor for Q, K, V
       //Q
-      Layout layout_Q = make_layout(select<0,2,3>(problem_shape), dQ);
+      auto problem_shape_q = make_tuple(get_0, get_1, get<2>(problem_shape), get<3>(problem_shape));
+      Layout layout_Q = make_layout(select<0,2,3>(problem_shape_q), dQ);
       Tensor mQ = make_tensor(make_gmem_ptr(ptr_Q), layout_Q);
 
       auto cluster_layout_vmnk =
@@ -152,21 +154,16 @@ struct Sm100FmhaLoadTmaWarpspecialized {
         typename CollectiveMmaQK::TiledMma{}, cluster_layout_vmnk);
 
       // K
-      auto problem_shape_paged_k =  make_tuple(get_0, get_1, get<2>(problem_shape), get<3>(problem_shape));
-      get<1> (problem_shape_paged_k) = args.page_block_size;
-      get<3, 1>(problem_shape_paged_k) = args.num_blocks;
-      Layout layout_k = make_layout(select<1,2,3>(problem_shape_paged_k), dK);
+      auto problem_shape_paged_kv =  make_tuple(get_0, args.page_block_size, get<2>(problem_shape), make_tuple(get<0>(get<3>(problem_shape)), args.num_blocks));
+      Layout layout_k = make_layout(select<1,2,3>(problem_shape_paged_kv), dK);
       Tensor mK = make_tensor(make_gmem_ptr(ptr_K), layout_k);
 
       tma_load_k = make_tma_atom_B_sm100<Element>(
         cute::SM90_TMA_LOAD{}, mK, SmemLayoutK{}(_, _, _, _0{}), TileShapeQK{},
         typename CollectiveMmaQK::TiledMma{}, cluster_layout_vmnk);
 
       // V
-      auto problem_shape_paged_v =  make_tuple(get_0, get<2>(problem_shape), get_1, get<3>(problem_shape));
-      get<2> (problem_shape_paged_v) = args.page_block_size;
-      get<3, 1>(problem_shape_paged_v) = args.num_blocks;
-      Layout layout_v = make_layout(select<1,2,3>(problem_shape_paged_v), select<1,0,2>(dV));
+      Layout layout_v = make_layout(select<2,1,3>(problem_shape_paged_kv), select<1,0,2>(dV));
       Tensor mV = make_tensor(make_gmem_ptr(ptr_V), layout_v);
 
       tma_load_v = make_tma_atom_B_sm100<Element>(
@@ -368,7 +365,7 @@ struct Sm100FmhaLoadTmaWarpspecialized {
     }
   }
 
-template<class BlkCoord, class ProblemShape, class ParamsProblemShape>
+  template<class BlkCoord, class ProblemShape, class ParamsProblemShape>
   CUTLASS_DEVICE void
   load_paged(
       BlkCoord const& blk_coord_in, ProblemShape const& problem_shape,
@@ -418,11 +415,8 @@ template<class BlkCoord, class ProblemShape, class ParamsProblemShape>
     Tensor tQgQ = tQgQ_qdl(_, _, _0{}, get<2>(blk_coord_q));
 
     // compute gK, sK
-    ProblemShapeK problem_shape_k =  problem_shape;
-    get<1> (problem_shape_k) = params.page_block_size;
-    get<3, 1>(problem_shape_k) = params.num_blocks;
-
-    Tensor mK_kdl_p = params.tma_load_k.get_tma_tensor(select<1,2,3>(problem_shape_k));
+    ProblemShapeK problem_shape_kv = make_tuple(get<0>(problem_shape), params.page_block_size, get<2>(problem_shape), make_tuple(get<0>(get<3>(problem_shape)), params.num_blocks));
+    Tensor mK_kdl_p = params.tma_load_k.get_tma_tensor(select<1,2,3>(problem_shape_kv));
 
     Tensor gK_kdl = local_tile(mK_kdl_p, TileShapeQK{}, make_coord(_, _, _), Step<X, _1, _1>{});
     Tensor tSgK_kdl = mma_qk.partition_B(gK_kdl);
@@ -437,10 +431,7 @@ template<class BlkCoord, class ProblemShape, class ParamsProblemShape>
 
     // compute gV, sV
     ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0);
-    ProblemShapeK problem_shape_v =  problem_shape;
-    get<1> (problem_shape_v) = params.page_block_size;
-    get<3, 1>(problem_shape_v) = params.num_blocks;
-    Tensor mV_dkl_p = params.tma_load_v.get_tma_tensor(select<2,1,3>(problem_shape_v));
+    Tensor mV_dkl_p = params.tma_load_v.get_tma_tensor(select<2,1,3>(problem_shape_kv));
 
     Tensor gV_dkl = local_tile(mV_dkl_p, TileShapePV{}, make_coord(_, _, _), Step<X, _1, _1>{});
     Tensor tOgV_dkl = mma_pv.partition_B(gV_dkl);

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp

@@ -82,9 +82,9 @@ struct Sm100FmhaCtxKernelWarpspecializedSchedule {
   static const int NumRegsCorrection = 96 - (kDebugUsingPrintf ? 16 : 0);
   static const int NumRegsOther = 32 + (kDebugUsingPrintf ? 16 : 0);
   static const int NumRegsEmpty = 24;
-  
+
   static const int NumWarps = 16;
-  
+
 };
 
 
@@ -148,7 +148,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
   static const int NumWarpsCorrection = KernelSchedule::NumWarpsCorrection;
   static const int NumWarpsEpilogue = KernelSchedule::NumWarpsEpilogue;
   static const int NumWarpsLoad = KernelSchedule::NumWarpsLoad;
-  
+
   static_assert(NumWarpsEpilogue == CollectiveEpilogue::NumWarpsEpilogue);
   static_assert(NumWarpsLoad == CollectiveEpilogue::NumWarpsLoad);
 
@@ -177,13 +177,13 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
     };
 
     static constexpr bool IsPersistent = std::is_same_v<TileScheduler, PersistentTileScheduler> || std::is_same_v<TileScheduler, CausalPersistentTileScheduler>;
-    using MainloopEpilogueStorage = std::conditional_t<IsPersistent, 
-                                                       std::conditional_t<IsMla, 
+    using MainloopEpilogueStorage = std::conditional_t<IsPersistent,
+                                                       std::conditional_t<IsMla,
                                                                           std::conditional_t<CollectiveMainloop::IsOrderLoadEpilogue, UnionType, StructType>,
                                                                           StructType>,
                                                        UnionType>;
 
-    MainloopEpilogueStorage mainloop_epilogue; 
+    MainloopEpilogueStorage mainloop_epilogue;
 
     struct PipelineStorage {
       alignas(16) typename CollectiveMainloop::PipelineQ::SharedStorage load_q;
@@ -305,7 +305,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
       shared_storage.pipelines.load_q,
       pipeline_load_q_params,
       ClusterShape{},  cute::true_type{}, /*mask calc*/cute::false_type{});
-    
+
     typename CollectiveMainloop::PipelineKV::Params pipeline_load_kv_params;
     if (role == WarpRole::Load) {
       pipeline_load_kv_params.role = CollectiveMainloop::PipelineKV::ThreadCategory::Producer;
@@ -565,7 +565,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
       warpgroup_reg_set<NumRegsOther>();
 
       if constexpr (IsMla && CollectiveMainloop::IsOrderLoadEpilogue) {
-        cutlass::arch::NamedBarrier::arrive((NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp, 
+        cutlass::arch::NamedBarrier::arrive((NumWarpsLoad + NumWarpsEpilogue) * NumThreadsPerWarp,
                                       cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
       }