Revert "[ex77] fix mla split; add fwd lse; add bwd varlen (#2366)" (#2370)

This reverts commit f12b1d75c9.

examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp

@@ -39,11 +39,11 @@ namespace cutlass::fmha::kernel {
 
 using namespace cute;
 
-template<class ProblemShape, class Element, class ElementAcc>
+template<class Element, class ElementAcc>
 struct FmhaKernelBwdConvert {
 
   struct Arguments {
-    ProblemShape problem_shape;
+    tuple<int, int, int, tuple<int, int>> problem_size;
 
     const ElementAcc* ptr_src_dQ;
     tuple<int, _1, tuple<int, int>> stride_src_dQ;
@@ -85,11 +85,11 @@ struct FmhaKernelBwdConvert {
   static const int kIterationsSeq = kBlockSeq / kNumThreadsSeq;
 
   static bool can_implement(Arguments const& args) {
-    return get<2>(args.problem_shape) % kElementsPerLoad == 0;
+    return get<2>(args.problem_size) % kElementsPerLoad == 0;
   }
 
   static dim3 get_grid_shape(Params const& params) {
-    dim3 grid(size<3,0>(params.problem_shape), size<3,1>(params.problem_shape), ceil_div(std::max(size<0>(params.problem_shape), size<1>(params.problem_shape)), kBlockSeq));
+    dim3 grid(size<3,0>(params.problem_size), size<3,1>(params.problem_size), ceil_div(std::max(size<0>(params.problem_size), size<1>(params.problem_size)), kBlockSeq));
     return grid;
   }
 
@@ -102,25 +102,18 @@ struct FmhaKernelBwdConvert {
     return args;
   }
 
-  template<class StrideSrc, class StrideDest, class Count>
-  CUTLASS_DEVICE void copy(Params const& params, const ElementAcc* ptr_src, StrideSrc const& stride_src, Element* ptr_dest, StrideDest const& stride_dest, Count const& count) {
+  template<class StrideSrc, class StrideDest>
+  CUTLASS_DEVICE void copy(Params const& params, const ElementAcc* ptr_src, StrideSrc const& stride_src, Element* ptr_dest, StrideDest const& stride_dest, int count) {
     auto ptr_src_bh = ptr_src + get<2,0>(stride_src) * blockIdx.x + get<2,1>(stride_src) * blockIdx.y;
     auto ptr_dest_bh = ptr_dest + get<2,0>(stride_dest) * blockIdx.x + get<2,1>(stride_dest) * blockIdx.y;
 
-    int seqlen = count;
-    if constexpr (is_variable_length_v<decltype(count)>) {
-      int offset = count.cumulative_length[blockIdx.y];
-      ptr_dest_bh += offset * get<0>(stride_dest);
-      seqlen = count.cumulative_length[blockIdx.y + 1] - offset;
-    }
-
     for (int idx_s_t = threadIdx.y; idx_s_t < kBlockSeq; idx_s_t += kNumThreadsSeq) {
       int idx_s = idx_s_t + kBlockSeq * blockIdx.z;
-      if (idx_s >= seqlen) continue;
+      if (idx_s >= count) continue;
       auto ptr_src_bhs = ptr_src_bh + idx_s * get<0>(stride_src);
       auto ptr_dest_bhs = ptr_dest_bh + idx_s * get<0>(stride_dest);
 
-      for (int idx_d = threadIdx.x * kElementsPerLoad; idx_d < get<2>(params.problem_shape); idx_d += kElementsPerLoad * kNumThreadsD) {
+      for (int idx_d = threadIdx.x * kElementsPerLoad; idx_d < get<2>(params.problem_size); idx_d += kElementsPerLoad * kNumThreadsD) {
         ElementAcc value_src[kElementsPerLoad];
         Element value_dest[kElementsPerLoad];
         
@@ -139,13 +132,13 @@ struct FmhaKernelBwdConvert {
 
   CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
     if (params.ptr_src_dQ != nullptr) {
-      copy(params, params.ptr_src_dQ, params.stride_src_dQ, params.ptr_dest_dQ, params.stride_dest_dQ, get<0>(params.problem_shape));
+      copy(params, params.ptr_src_dQ, params.stride_src_dQ, params.ptr_dest_dQ, params.stride_dest_dQ, get<0>(params.problem_size));
     }
     if (params.ptr_src_dK != nullptr) {
-      copy(params, params.ptr_src_dK, params.stride_src_dK, params.ptr_dest_dK, params.stride_dest_dK, get<1>(params.problem_shape));
+      copy(params, params.ptr_src_dK, params.stride_src_dK, params.ptr_dest_dK, params.stride_dest_dK, get<1>(params.problem_size));
     }
     if (params.ptr_src_dV != nullptr) {
-      copy(params, params.ptr_src_dV, params.stride_src_dV, params.ptr_dest_dV, params.stride_dest_dV, get<1>(params.problem_shape));
+      copy(params, params.ptr_src_dV, params.stride_src_dV, params.ptr_dest_dV, params.stride_dest_dV, get<1>(params.problem_size));
     }
   }
 };

examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_sum_OdO.hpp

@@ -39,11 +39,11 @@ namespace cutlass::fmha::kernel {
 
 using namespace cute;
 
-template<class ProblemShape, class Element, class ElementAcc>
+template<class Element, class ElementAcc>
 struct FmhaKernelBwdSumOdO {
 
   struct Arguments {
-    ProblemShape problem_shape;
+    cute::tuple<int, int, int, cute::tuple<int, int>> problem_size;
 
     const Element* ptr_O;
     cute::tuple<int, cute::_1, cute::tuple<int, int>> stride_O;
@@ -86,11 +86,11 @@ struct FmhaKernelBwdSumOdO {
   static const int kIterationsQ = kBlockQ / kNumThreadsQ;
 
   static bool can_implement(Arguments const& args) {
-    return get<2>(args.problem_shape) % kElementsPerLoad == 0;
+    return get<2>(args.problem_size) % kElementsPerLoad == 0;
   }
 
   static dim3 get_grid_shape(Params const& params) {
-    dim3 grid(ceil_div(size<0>(params.problem_shape), kBlockQ), size<3,0>(params.problem_shape), size<3,1>(params.problem_shape));
+    dim3 grid(ceil_div(size<0>(params.problem_size), kBlockQ), size<3,0>(params.problem_size), size<3,1>(params.problem_size));
     return grid;
   }
 
@@ -110,20 +110,10 @@ struct FmhaKernelBwdSumOdO {
     auto ptr_lse_bh = params.ptr_lse + blockIdx.y * get<1,0>(params.stride_lse) + blockIdx.z * get<1,1>(params.stride_lse);
     auto ptr_scaled_lse_bh = params.ptr_scaled_lse + blockIdx.y * get<1,0>(params.stride_scaled_lse) + blockIdx.z * get<1,1>(params.stride_scaled_lse);
 
-    auto problem_q = get<0>(params.problem_shape);
-    int seqlen_q = problem_q;
-    if constexpr (is_variable_length_v<decltype(problem_q)>) {
-      int offset = problem_q.cumulative_length[blockIdx.z];
-      ptr_O_bh += offset * get<0>(params.stride_O);
-      ptr_dO_bh += offset * get<0>(params.stride_dO);
-      ptr_lse_bh += offset * get<0>(params.stride_lse);
-      seqlen_q = problem_q.cumulative_length[blockIdx.z + 1] - offset;
-    }
-
     CUTLASS_PRAGMA_UNROLL
     for (int idx_q_t = threadIdx.y; idx_q_t < kBlockQ; idx_q_t += kNumThreadsQ) {
       int idx_q = idx_q_t + kBlockQ * blockIdx.x;
-      if (idx_q >= seqlen_q) continue;
+      if (idx_q >= get<0>(params.problem_size)) continue;
       ElementAcc acc = 0;
       auto ptr_O_bhq = ptr_O_bh + idx_q * get<0>(params.stride_O);
       auto ptr_dO_bhq = ptr_dO_bh + idx_q * get<0>(params.stride_dO);
@@ -131,7 +121,7 @@ struct FmhaKernelBwdSumOdO {
       auto ptr_lse_bhq = ptr_lse_bh + idx_q * get<0>(params.stride_lse);
       auto ptr_scaled_lse_bhq = ptr_scaled_lse_bh + idx_q * get<0>(params.stride_scaled_lse);
 
-      for (int idx_d = threadIdx.x * kElementsPerLoad; idx_d < get<2>(params.problem_shape); idx_d += kElementsPerLoad * kNumThreadsD) {
+      for (int idx_d = threadIdx.x * kElementsPerLoad; idx_d < get<2>(params.problem_size); idx_d += kElementsPerLoad * kNumThreadsD) {
         Element value_O[kElementsPerLoad];
         Element value_dO[kElementsPerLoad];
         

examples/77_blackwell_fmha/kernel/fmha_options.hpp

@@ -82,4 +82,4 @@ struct Option {
   using option_value = Value;
 };
 
-}  // namespace cutlass::fmha::kernel
+}  // namespace cutlass::fmha::kernel

examples/77_blackwell_fmha/kernel/fmha_tile_scheduler.hpp

@@ -90,8 +90,8 @@ struct PersistentTileScheduler {
   struct Params {
     int num_blocks;
     FastDivmod divmod_m_block;
-    FastDivmod divmod_b;
     FastDivmod divmod_h;
+    FastDivmod divmod_b;
 
     KernelHardwareInfo hw_info;
   };
@@ -146,7 +146,7 @@ struct PersistentTileScheduler {
     params.divmod_m_block(block_decode, m_block, block_decode);
     params.divmod_b(block_decode, bidb, block_decode);
     params.divmod_h(block_decode, bidh, block_decode);
-    return make_coord(m_block, _0{}, make_coord(bidb, bidh));
+    return make_coord(m_block, _0{}, make_coord(bidh, bidb));
   }
 
   CUTLASS_DEVICE

examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp

@@ -43,8 +43,6 @@
 
 #include "collective/fmha_common.hpp"
 
-#include <cmath>
-
 namespace cutlass::fmha::kernel {
 
 using namespace cutlass::fmha::collective;
@@ -52,7 +50,6 @@ using namespace cutlass::fmha::collective;
 using namespace cute;
 
 template<
-    class ProblemShape,
     class Element,
     class ElementAcc,
     class TileShape,
@@ -121,7 +118,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
   using TensorStrideContiguousK = Stride<int, _1, Stride<int, int>>;
   using TensorStrideContiguousMN = Stride<_1, int, Stride<int, int>>;
-
+  
   // compute S
   using CollectiveMmaKQ = typename cutlass::gemm::collective::CollectiveBuilder<
       cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
@@ -277,6 +274,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
   static constexpr int SharedStorageSize = offsetof(SharedStorage, tmem_base_ptr) + sizeof(uint32_t);
   static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "using too much smem");
 
+  using ProblemShape = Shape<int, int, int, Shape<int, int>>;  // Q K D (H B), eventuall D = (D_QK, D_VO)
   using TensorStride = TensorStrideContiguousK;  // S D (H B)
   using RowTensorStride = Stride<_1, Stride<int, int>>;    // S (H B)
 
@@ -362,16 +360,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
 
   static Params to_underlying_arguments(Arguments const& args, void*) {
-    auto [Q_, K_, D, HB] = args.problem_shape;
-    int Q = Q_;
-    int K = K_;
-
-    if constexpr (is_variable_length_v<decltype(Q_)>) {
-      Q = Q_.total_length;
-    }
-    if constexpr (is_variable_length_v<decltype(K_)>) {
-      K = K_.total_length;
-    }
+    auto [Q, K, D, HB] = args.problem_shape;
 
     auto params_kq = CollectiveMmaKQ::to_underlying_arguments(
       make_shape(K, Q, D, HB),
@@ -389,10 +378,10 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
     TMA_DQ tma_red_dq = make_tma_copy(
         SM90_TMA_REDUCE_ADD{},
-        make_tensor(args.mainloop.ptr_dq_acc, make_shape(Q_, D, HB), args.mainloop.stride_dq_acc),
+        make_tensor(args.mainloop.ptr_dq_acc, make_shape(Q, D, HB), args.mainloop.stride_dq_acc),
         SmemLayoutDQ{}(_, _, _0{})
     );
-
+      
     return Params{
       args.problem_shape,
       args.mainloop,
@@ -427,11 +416,10 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
   }
 
 
-  template<class BlkCoord, class BlkOffset, class ProblemShape_>
+  template<class BlkCoord>
   CUTLASS_DEVICE void load(
       BlkCoord const& blk_coord,
-      BlkOffset const& blk_offset,
-      ProblemShape_ const& problem_shape,
+      ProblemShape const& problem_shape,
       int iter_index,
       int iter_count,
       MainloopArguments const& mainloop_args,
@@ -452,15 +440,10 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
     uint16_t mcast_mask = 0;
 
-    auto mK_in = mainloop_params.tma_load_k.get_tma_tensor(make_shape(K, D, HB));
-    auto mV_in = mainloop_params.tma_load_v.get_tma_tensor(make_shape(K, D, HB));
-    auto mQ_in = mainloop_params.tma_load_q.get_tma_tensor(make_shape(Q, D, HB));
-    auto mDO_in = mainloop_params.tma_load_do.get_tma_tensor(make_shape(Q, D, HB));
-
-    auto mK = domain_offset(select<1,2,3>(blk_offset), mK_in);
-    auto mV = domain_offset(select<1,2,3>(blk_offset), mV_in);
-    auto mQ = domain_offset(select<0,2,3>(blk_offset), mQ_in);
-    auto mDO = domain_offset(select<0,2,3>(blk_offset), mDO_in);
+    auto mK = mainloop_params.tma_load_k.get_tma_tensor(make_shape(K, D, HB));
+    auto mQ = mainloop_params.tma_load_q.get_tma_tensor(make_shape(Q, D, HB));
+    auto mV = mainloop_params.tma_load_v.get_tma_tensor(make_shape(K, D, HB));
+    auto mDO = mainloop_params.tma_load_do.get_tma_tensor(make_shape(Q, D, HB));
 
     auto gK = local_tile(mK, TileShapeKQ{}, make_coord(_,_,_), Step<_1, X, _1>{});
     auto gQ = local_tile(mQ, TileShapeKQ{}, make_coord(_,_,_), Step<X, _1, _1>{});
@@ -469,7 +452,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
     ThrMMA cta_mma_kq = TiledMmaKQ{}.get_slice(_0{});
     ThrMMA cta_mma_vdo = TiledMmaVDO{}.get_slice(_0{});
-
+    
     auto tSTgK = cta_mma_kq.partition_A(gK);
     auto tSTgQ = cta_mma_kq.partition_B(gQ);
     auto tDPTgV = cta_mma_vdo.partition_A(gV);
@@ -494,8 +477,8 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
         group_modes<0,3>(sDO), group_modes<0,3>(tDPTgDO));
 
     // set up lse and sum_odo
-
-    auto [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_batch] = blk_coord;
+    
+    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;
 
     pipeline_load_mma_q.producer_acquire(pipeline_load_mma_q_producer_state);
     auto tma_barrier = pipeline_load_mma_q.producer_get_barrier(pipeline_load_mma_q_producer_state);
@@ -512,7 +495,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
 
     // load Q
-    if (cute::elect_one_sync()) {
+    if (cute::elect_one_sync()) { 
       cute::copy(
           mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask),
           tQgQ_mkl(_, iter_index, _0{}, blk_coord_batch),
@@ -532,14 +515,12 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     int smem_idx = TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + thread_idx * 4;
     int gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4;
     auto mLSE = make_tensor(mainloop_args.ptr_lse, make_shape(Q, HB), mainloop_args.stride_lse);
-    for (int i = 0; i < 4; i++) {
-      cutlass::arch::cp_async_zfill<4>(
-          shared_tensors.smem_lse.begin() + smem_idx + i,
-          &mLSE(gmem_idx + i, blk_coord_batch),
-          gmem_idx + i < Q
-      );
-    }
-
+    cutlass::arch::cp_async_zfill<16>(
+        shared_tensors.smem_lse.begin() + smem_idx,
+        &mLSE(gmem_idx, blk_coord_batch),
+         gmem_idx < Q
+    );
+    
     pipeline_load_compute_lse.producer_commit(pipeline_load_compute_lse_producer_state, cutlass::arch::cpasync_barrier_arrive);
     ++pipeline_load_compute_lse_producer_state;
 
@@ -548,7 +529,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     tma_barrier = pipeline_load_mma_do.producer_get_barrier(pipeline_load_mma_do_producer_state);
 
     pipeline_load_mma_do.producer_expect_transaction(pipeline_load_mma_do_producer_state, kTransactionsBytesLoadV);
-
+    
     // load V
     if (cute::elect_one_sync()) {
       cute::copy(
@@ -559,7 +540,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
 
     // load dO
-    if (cute::elect_one_sync()) {
+    if (cute::elect_one_sync()) { 
       cute::copy(
           mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask),
           tDOgDO_mkl(_, iter_index, _0{}, blk_coord_batch),
@@ -575,13 +556,11 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     smem_idx = TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + thread_idx * 4;
     gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4;
     auto mSumOdO = make_tensor(mainloop_args.ptr_sum_odo, make_shape(Q, HB), mainloop_args.stride_sum_odo);
-    for (int i = 0; i < 4; i++) {
-      cutlass::arch::cp_async_zfill<4>(
-          shared_tensors.smem_sum_odo.begin() + smem_idx + i,
-          &mSumOdO(gmem_idx + i, blk_coord_batch),
-          gmem_idx + i < Q
-      );
-    }
+    cutlass::arch::cp_async<16>(
+        shared_tensors.smem_sum_odo.begin() + smem_idx,
+        &mSumOdO(gmem_idx, blk_coord_batch),
+        gmem_idx < Q
+    );
 
     pipeline_load_compute_sum_odo.producer_commit(pipeline_load_compute_sum_odo_producer_state, cutlass::arch::cpasync_barrier_arrive);
     ++pipeline_load_compute_sum_odo_producer_state;
@@ -594,7 +573,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       tma_barrier = pipeline_load_mma_q.producer_get_barrier(pipeline_load_mma_q_producer_state);
 
       // load Q
-      if (cute::elect_one_sync()) {
+      if (cute::elect_one_sync()) { 
         cute::copy(
             mainloop_params.tma_load_q.with(*tma_barrier, mcast_mask),
             tQgQ_mkl(_, iter_index, _0{}, blk_coord_batch),
@@ -605,26 +584,24 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       ++pipeline_load_mma_q_producer_state;
 
       pipeline_load_compute_lse.producer_acquire(pipeline_load_compute_lse_producer_state);
-
+      
       // load LSE
       smem_idx = TileShapeQ{} * pipeline_load_compute_lse_producer_state.index() + thread_idx * 4;
       gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4;
-      for (int i = 0; i < 4; i++) {
-        cutlass::arch::cp_async_zfill<4>(
-            shared_tensors.smem_lse.begin() + smem_idx + i,
-            &mLSE(gmem_idx + i, blk_coord_batch),
-            gmem_idx + i < Q
-        );
-      }
-
+      cutlass::arch::cp_async<16>(
+          shared_tensors.smem_lse.begin() + smem_idx,
+          &mLSE(gmem_idx, blk_coord_batch),
+          gmem_idx < Q
+      );
+      
       pipeline_load_compute_lse.producer_commit(pipeline_load_compute_lse_producer_state, cutlass::arch::cpasync_barrier_arrive);
       ++pipeline_load_compute_lse_producer_state;
 
       pipeline_load_mma_do.producer_acquire(pipeline_load_mma_do_producer_state);
       tma_barrier = pipeline_load_mma_do.producer_get_barrier(pipeline_load_mma_do_producer_state);
 
-      // load dO
-      if (cute::elect_one_sync()) {
+      // load dO  
+      if (cute::elect_one_sync()) { 
         cute::copy(
             mainloop_params.tma_load_do.with(*tma_barrier, mcast_mask),
             tDOgDO_mkl(_, iter_index, _0{}, blk_coord_batch),
@@ -635,18 +612,16 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       ++pipeline_load_mma_do_producer_state;
 
       pipeline_load_compute_sum_odo.producer_acquire(pipeline_load_compute_sum_odo_producer_state);
-
+      
       // load sum_OdO
       smem_idx = TileShapeQ{} * pipeline_load_compute_sum_odo_producer_state.index() + thread_idx * 4;
       gmem_idx = TileShapeQ{} * iter_index + thread_idx * 4;
-      for (int i = 0; i < 4; i++) {
-        cutlass::arch::cp_async_zfill<4>(
-            shared_tensors.smem_sum_odo.begin() + smem_idx + i,
-            &mSumOdO(gmem_idx + i, blk_coord_batch),
-            gmem_idx + i < Q
-        );
-      }
-
+      cutlass::arch::cp_async_zfill<16>(
+          shared_tensors.smem_sum_odo.begin() + smem_idx,
+          &mSumOdO(gmem_idx, blk_coord_batch),
+          gmem_idx < Q
+      );
+      
       pipeline_load_compute_sum_odo.producer_commit(pipeline_load_compute_sum_odo_producer_state, cutlass::arch::cpasync_barrier_arrive);
       ++pipeline_load_compute_sum_odo_producer_state;
 
@@ -656,31 +631,31 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
   }
 
 
-  template<class BlkCoord, class ProblemShape_>
+  template<class BlkCoord>
   CUTLASS_DEVICE void mma(
       BlkCoord const& blk_coord,
-      ProblemShape_ const& problem_shape,
+      ProblemShape const& problem_shape,
       int iter_index,
       int iter_count,
       MainloopArguments const& mainloop_args,
       TensorStorage& shared_tensors,
-      PipelineLoadMmaQ& pipeline_load_mma_q,
-      typename PipelineLoadMmaQ::PipelineState& pipeline_load_mma_q_consumer_state,
-      PipelineLoadMmaDO& pipeline_load_mma_do,
+      PipelineLoadMmaQ& pipeline_load_mma_q, 
+      typename PipelineLoadMmaQ::PipelineState& pipeline_load_mma_q_consumer_state,        
+      PipelineLoadMmaDO& pipeline_load_mma_do, 
       typename PipelineLoadMmaDO::PipelineState& pipeline_load_mma_do_consumer_state,
-      PipelineMmaComputeS& pipeline_mma_compute_s,
+      PipelineMmaComputeS& pipeline_mma_compute_s, 
       typename PipelineMmaComputeS::PipelineState& pipeline_mma_compute_s_producer_state,
-      PipelineMmaComputeDP& pipeline_mma_compute_dp,
+      PipelineMmaComputeDP& pipeline_mma_compute_dp, 
       typename PipelineMmaComputeDP::PipelineState& pipeline_mma_compute_dp_producer_state,
-      PipelineMmaReduceDQ& pipeline_mma_reduce_dq,
+      PipelineMmaReduceDQ& pipeline_mma_reduce_dq, 
       typename PipelineMmaReduceDQ::PipelineState& pipeline_mma_reduce_dq_producer_state,
-      PipelineComputeMmaP& pipeline_compute_mma_p,
+      PipelineComputeMmaP& pipeline_compute_mma_p, 
       typename PipelineComputeMmaP::PipelineState& pipeline_compute_mma_p_consumer_state,
-      PipelineComputeMmaDS& pipeline_compute_mma_ds,
+      PipelineComputeMmaDS& pipeline_compute_mma_ds, 
       typename PipelineComputeMmaDS::PipelineState& pipeline_compute_mma_ds_consumer_state,
       PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
       typename PipelineMmaComputeDKDV::PipelineState& pipeline_mma_compute_dkdv_producer_state) {
-
+    
     auto [Q, K, D, HB] = problem_shape;
 
     auto sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{});
@@ -710,7 +685,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     Tensor tDVrP = TiledMmaPDO::make_fragment_A(sP)(_, _, _, _0{});
     tDVrP.data() = TmemAllocation::kP;
     Tensor tDVrDOT = TiledMmaPDO::make_fragment_B(sDOT);
-
+    
     TiledMmaKQ tiled_mma_kq;
     TiledMmaVDO tiled_mma_vdo;
     TiledMmaDSK tiled_mma_dsk;
@@ -948,8 +923,6 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       TensorC const& coord,
       TensorShape const& tensor_shape) {
 
-    Tensor preds = cute::lazy::transform(coord, [&](auto const& c) { return elem_less(c, tensor_shape); });
-
     auto copy_op = make_cotiled_copy(
         Copy_Atom<UniversalCopy<uint128_t>, Element>{},
         make_layout(make_shape(_1{}, Int<sizeof(uint128_t) / sizeof(Element)>{})),
@@ -957,91 +930,42 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     );
     auto thr_copy = copy_op.get_slice(_0{});
 
-    Tensor tCg = thr_copy.partition_D(gmem);
-    Tensor tCr = thr_copy.partition_S(quantize(regs));
-    Tensor tPc = thr_copy.partition_D(preds);
+    auto tCg = thr_copy.partition_D(gmem);
+    auto tCr = thr_copy.partition_S(quantize(regs));
+    auto tCc = thr_copy.partition_D(coord);
 
-    copy_if(copy_op, tPc, tCr, tCg);
-  }
+    constexpr int R = decltype(tCr.layout())::rank;
+    auto tCg_v = group_modes<1, R>(tCg);
+    auto tCr_v = group_modes<1, R>(tCr);
+    auto tCc_v = group_modes<1, R>(tCc);
+    auto tCp_v = make_tensor<bool>(shape<1>(tCc_v));
 
-
-  template<class BlkCoord, class BlkOffset, class ProblemShape_>
-  CUTLASS_DEVICE void epilogue_clear(
-      BlkCoord const& blk_coord,
-      BlkOffset const& blk_offset,
-      ProblemShape_ const& problem_shape,
-      MainloopArguments const& mainloop_args,
-      EpilogueArguments const& epilogue_args) {
-
-    auto [Q, K, D, HB] = problem_shape;
-    auto [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_batch] = blk_coord;
-
-    auto mDK_in = make_tensor(make_gmem_ptr(epilogue_args.ptr_dk), make_shape(K, TileShapeDQK{}, HB), epilogue_args.stride_dk);
-    auto mDK = domain_offset(select<1,2,3>(blk_offset), mDK_in);
-    auto gDK = local_tile(mDK, TileShapeDSQ{}, make_coord(_,_,_), Step<_1, _1, X>{})
-        (_, _, blk_coord_k, _0{}, blk_coord_batch);
-
-    Tensor cDK = domain_offset(
-        make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}),
-        make_identity_tensor(take<0,2>(TileShapeDSQ{}))
-    );
-
-    auto mDV_in = make_tensor(make_gmem_ptr(epilogue_args.ptr_dv), make_shape(K, TileShapeDVO{}, HB), epilogue_args.stride_dv);
-    auto mDV = domain_offset(select<1,2,3>(blk_offset), mDV_in);
-    auto gDV = local_tile(mDV, TileShapePDO{}, make_coord(_,_,_), Step<_1, _1, X>{})
-        (_, _, blk_coord_k, _0{}, blk_coord_batch);
-
-    Tensor cDV = domain_offset(
-        make_coord(blk_coord_k * TileShapeK{}, _0{}),
-        make_identity_tensor(take<0,2>(TileShapePDO{}))
-    );
-    
-    if (threadIdx.x >= 256) {
-      return;
+    for (int i = 0; i < size(tCp_v); ++i) {
+      tCp_v(i) = elem_less(tCc_v(_0{},i), tensor_shape);
     }
 
-    auto tiled_copy = make_cotiled_copy(
-        Copy_Atom<UniversalCopy<uint128_t>, Element>{},
-        make_ordered_layout(make_shape(_256{}, Int<sizeof(uint128_t) / sizeof(Element)>{}), Step<_1, _0>{}),
-        make_ordered_layout(make_shape(TileShapeK{}, TileShapeDQK{}), Step<_1, _0>{}));
-
-    auto thr_copy = tiled_copy.get_slice(threadIdx.x);
-    auto tCgDK = thr_copy.partition_D(gDK);
-    auto tCcDK = thr_copy.partition_S(cDK);
-    auto tCrDK = make_tensor<Element>(shape(tCcDK));
-
-    clear(tCrDK);
-    store(tCgDK, tCrDK, tCcDK, select<1,2>(problem_shape));
-
-    auto tCgDV = thr_copy.partition_D(gDV);
-    auto tCcDV = thr_copy.partition_S(cDV);
-    auto tCrDV = make_tensor<Element>(shape(tCcDV));
-
-    clear(tCrDV);
-    store(tCgDV, tCrDV, tCcDV, select<1,2>(problem_shape));
+    copy_if(copy_op, tCp_v, tCr_v, tCg_v);
   }
 
 
-  template<class BlkCoord, class BlkOffset, class ProblemShape_>
+  template<class BlkCoord>
   CUTLASS_DEVICE void epilogue(
       BlkCoord const& blk_coord,
-      BlkOffset const& blk_offset,
-      ProblemShape_ const& problem_shape,
+      ProblemShape const& problem_shape,
       MainloopArguments const& mainloop_args,
       EpilogueArguments const& epilogue_args,
       PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
       typename PipelineMmaComputeDKDV::PipelineState& pipeline_mma_compute_dkdv_consumer_state) {
 
     auto [Q, K, D, HB] = problem_shape;
-    auto [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_batch] = blk_coord;
+    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;
 
     auto load_op = SM100_TMEM_LOAD_32dp32b16x{};
 
     auto tDKtDK = partition_fragment_C(TiledMmaDSQ{}, select<0,1>(TileShapeDSQ{}))(make_coord(_,_),_0{},_0{});
     tDKtDK.data() = TmemAllocation::kDK;
 
-    auto mDK_in = make_tensor(make_gmem_ptr(epilogue_args.ptr_dk), make_shape(K, TileShapeDQK{}, HB), epilogue_args.stride_dk);
-    auto mDK = domain_offset(select<1,2,3>(blk_offset), mDK_in);
+    auto mDK = make_tensor(make_gmem_ptr(epilogue_args.ptr_dk), make_shape(K, TileShapeDQK{}, HB), epilogue_args.stride_dk);
     auto gDK = local_tile(mDK, TileShapeDSQ{}, make_coord(_,_,_), Step<_1, _1, X>{})
         (_, _, blk_coord_k, _0{}, blk_coord_batch);
 
@@ -1076,13 +1000,12 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     auto tDVtDV = partition_fragment_C(TiledMmaDSQ{}, select<0,1>(TileShapeDSQ{}))(make_coord(_,_),_0{},_0{});
     tDVtDV.data() = TmemAllocation::kDV;
 
-    auto mDV_in = make_tensor(make_gmem_ptr(epilogue_args.ptr_dv), make_shape(K, TileShapeDVO{}, HB), epilogue_args.stride_dv);
-    auto mDV = domain_offset(select<1,2,3>(blk_offset), mDV_in);
+    auto mDV = make_tensor(make_gmem_ptr(epilogue_args.ptr_dv), make_shape(K, TileShapeDVO{}, HB), epilogue_args.stride_dv);
     auto gDV = local_tile(mDV, TileShapePDO{}, make_coord(_,_,_), Step<_1, _1, X>{})
         (_, _, blk_coord_k, _0{}, blk_coord_batch);
 
     Tensor cDV = domain_offset(
-        make_coord(blk_coord_k * TileShapeK{}, _0{}),
+        make_coord(get<1>(blk_coord) * TileShapeK{}, _0{}),
         make_identity_tensor(take<0,2>(TileShapePDO{}))
     );
 
@@ -1126,11 +1049,10 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
   }
 
 
-  template<class BlkCoord, class BlkOffset, class ProblemShape_>
+  template<class BlkCoord>
   CUTLASS_DEVICE void compute(
       BlkCoord const& blk_coord,
-      BlkOffset const& blk_offset,
-      ProblemShape_ const& problem_shape,
+      ProblemShape const& problem_shape,
       int iter_index,
       int iter_count,
       MainloopArguments const& mainloop_args,
@@ -1151,7 +1073,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       PipelineMmaComputeDKDV& pipeline_mma_compute_dkdv,
       typename PipelineMmaComputeDKDV::PipelineState& pipeline_mma_compute_dkdv_consumer_state) {
 
-
+    
     auto [Q, K, D, HB] = problem_shape;
 
     // in tmem, S & P overlap
@@ -1192,7 +1114,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     Tensor tTR_cST   = split_wg(thread_t2r.partition_D(cST));
     Tensor tTR_rST = make_tensor<ElementAcc>(shape(tTR_cST));
     Tensor tTR_tST = split_wg(thread_t2r.partition_S(tSTtST));
-
+    
     Tensor tTR_cDPT_p = thread_t2r.partition_D(cDPT);
     Tensor tTR_cDPT = split_wg(tTR_cDPT_p);
     Tensor tTR_rDPT = make_tensor<ElementAcc>(shape(tTR_cDPT));
@@ -1214,9 +1136,6 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     auto tRT_tP = split_wg(thread_r2t.partition_D(tDVrP));
     auto tRT_cST = split_wg(thread_r2t.partition_S(tDVcST));
 
-    bool is_residual_k = get<1>(blk_coord) * TileShapeK{} + TileShapeK{} >= get<1>(problem_shape);
-    int last_iter = iter_count - 1 + iter_index;
-
     CUTLASS_PRAGMA_NO_UNROLL
     while (iter_count > 0) {
       // wait for S and P
@@ -1233,28 +1152,20 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
           fn(cute::false_type{});
         }
       };
-
-      bool leading_causal_masking = false;
-      if constexpr (std::is_base_of_v<cutlass::fmha::collective::CausalMask, Mask>) {
-        leading_causal_masking = warp_uniform(iter_index == get<1>(blk_coord));
-      }
-      bool trailing_residual_masking = false;
-      if constexpr (std::is_base_of_v<cutlass::fmha::collective::ResidualMaskForBackward, Mask>) {
-        trailing_residual_masking = warp_uniform((iter_index == last_iter) || is_residual_k);
-      }
-
-      dispatch_bool(leading_causal_masking || trailing_residual_masking, [&](auto is_masked_tile) {
+      
+      dispatch_bool(std::is_base_of_v<cutlass::fmha::collective::CausalMask, Mask> &&
+          warp_uniform(iter_index == get<1>(blk_coord)), [&](auto is_causal_masked_tile) {
 
         // compute P = softmax(S, LSE)
         cute::copy(tiled_t2r, tTR_tST, tTR_rST);
-
-        if constexpr (decltype(is_masked_tile)::value) {
+  
+        if constexpr (std::is_base_of_v<cutlass::fmha::collective::CausalMask, Mask> && decltype(is_causal_masked_tile)::value) {
           Mask{}.apply_mask(tTR_rST, [&](int i) {
             auto c_transpose = tTR_cST(i);
             return make_coord(get<1>(c_transpose) + iter_index * TileShapeQ{}, get<0>(c_transpose) + get<1>(blk_coord) * TileShapeK{});
           }, problem_shape);
         }
-
+  
         ElementAcc log2_e = static_cast<ElementAcc>(M_LOG2E);
         float2 softmax_scale_log2_e;
         softmax_scale_log2_e.x = mainloop_args.softmax_scale * log2_e;
@@ -1273,16 +1184,16 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
           tTR_rST(i) = ::exp2f(out.x);
           tTR_rST(i+1) = ::exp2f(out.y);
         }
-
+  
         auto tRT_rST = quantize(tTR_rST);
         auto tRT_rST_reshaped = make_tensor(tRT_rST.data(), shape(tRT_cST));
-
+        
         cutlass::arch::fence_view_async_tmem_load();
         cutlass::arch::NamedBarrier(
           kNumComputeWarps * NumThreadsPerWarp,
           cutlass::arch::ReservedNamedBarriers::TransformBarrier
         ).arrive_and_wait();
-
+        
         cute::copy(tiled_r2t, tRT_rST_reshaped, tRT_tP);
       });
 
@@ -1364,15 +1275,15 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
 
     epilogue(
-        blk_coord, blk_offset, problem_shape, mainloop_args, epilogue_args,
+        blk_coord, problem_shape, mainloop_args, epilogue_args,
         pipeline_mma_compute_dkdv, pipeline_mma_compute_dkdv_consumer_state
     );
   }
 
-  template<class BlkCoord, class ProblemShape_>
+  template<class BlkCoord>
   CUTLASS_DEVICE void reduce(
       BlkCoord const& blk_coord,
-      ProblemShape_ const& problem_shape,
+      ProblemShape const& problem_shape,
       int iter_index,
       int iter_count,
       MainloopArguments const& mainloop_args,
@@ -1382,12 +1293,12 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       typename PipelineMmaReduceDQ::PipelineState& pipeline_mma_reduce_dq_consumer_state,
       PipelineReduceTmaStore& pipeline_reduce_tma_store,
       typename PipelineReduceTmaStore::PipelineState& pipeline_reduce_tma_store_producer_state) {
-
+    
     using X = Underscore;
-
+    
     auto [Q, K, D, HB] = problem_shape;
 
-    auto [blk_coord_q, blk_coord_k, blk_coord_d, blk_coord_batch] = blk_coord;
+    auto [blk_coord_q, blk_coord_k, blk_coord_batch] = blk_coord;
 
     // must match TileShapeDQ
     auto load_op = SM100_TMEM_LOAD_32dp32b32x{};
@@ -1396,7 +1307,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     tDQtDQ.data() = TmemAllocation::kDQ;
 
     Tensor mDQ = mainloop_params.tma_red_dq.get_tma_tensor(make_shape(Q, D, HB));
-    auto gDQ = local_tile(mDQ, TileShapeKQ{}, make_coord(_,_,_), Step<X, _1, _1>{})
+    auto gDQ = local_tile(mDQ, TileShapeKQ{}, make_coord(_,_,_), Step<_1, _1, X>{})
         (_, _, _, _0{}, blk_coord_batch);
 
     Tensor cDQ = make_identity_tensor(take<0,2>(TileShapeDSK{}));
@@ -1465,7 +1376,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
       iter_index += 1;
     }
   }
-
+  
 
   CUTLASS_DEVICE void operator()(Params const& params, char* smem) {
     int warp_idx = cutlass::canonical_warp_idx_sync();
@@ -1650,7 +1561,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     typename decltype(pipeline_compute_mma_p)::PipelineState pipeline_compute_mma_p_consumer_state;
     typename decltype(pipeline_compute_mma_ds)::PipelineState pipeline_compute_mma_ds_consumer_state;
     typename decltype(pipeline_mma_compute_dkdv)::PipelineState pipeline_mma_compute_dkdv_consumer_state;
-
+    
     auto pipeline_load_mma_q_producer_state = make_producer_start_state<decltype(pipeline_load_mma_q)>();
     auto pipeline_load_mma_do_producer_state = make_producer_start_state<decltype(pipeline_load_mma_do)>();
     auto pipeline_load_compute_lse_producer_state = make_producer_start_state<decltype(pipeline_load_compute_lse)>();
@@ -1665,45 +1576,27 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
     pipeline_init_wait(size(ClusterShape{}));
 
-    auto blk_coord = make_coord(_0{}, blockIdx.x, _0{}, make_coord(blockIdx.y, blockIdx.z));
-    auto [problem_shape, blk_offset] = apply_variable_length_offset(
-        params.problem_shape,
-        blk_coord
-    );
+    auto blk_coord = make_coord(_0{}, blockIdx.x, make_coord(blockIdx.y, blockIdx.z));
+    auto problem_shape = params.problem_shape;
     int iter_count = ceil_div(get<0>(problem_shape), TileShapeQ{});
     int iter_start = 0;
     if constexpr (std::is_base_of_v<cutlass::fmha::collective::CausalMask, Mask>) {
       iter_start = (get<1>(blk_coord) * TileShapeK{}) / TileShapeQ{};
     }
-    if (get<1>(blk_coord) * TileShapeK{} >= get<1>(problem_shape)) {
-      return;
-    }
     iter_count -= iter_start;
 
-    if (iter_count <= 0) {
-      epilogue_clear(
-          blk_coord,
-          blk_offset,
-          problem_shape,
-          params.mainloop,
-          params.epilogue
-      );
-      return;
-    }
-
     if (role == WarpRole::Load) {
       warpgroup_reg_set<RegisterAllocation::kLoad>();
-
+    
       load(
           blk_coord,
-          blk_offset,
           problem_shape,
           iter_start,
           iter_count,
           params.mainloop,
           params.mainloop_params,
           shared_storage.tensors,
-          pipeline_load_mma_q, pipeline_load_mma_q_producer_state,
+          pipeline_load_mma_q, pipeline_load_mma_q_producer_state,        
           pipeline_load_mma_do, pipeline_load_mma_do_producer_state,
           pipeline_load_compute_lse, pipeline_load_compute_lse_producer_state,
           pipeline_load_compute_sum_odo, pipeline_load_compute_sum_odo_producer_state
@@ -1715,7 +1608,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
 
       tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
       __syncwarp();
-
+    
       mma(
           blk_coord,
           problem_shape,
@@ -1723,7 +1616,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
           iter_count,
           params.mainloop,
           shared_storage.tensors,
-          pipeline_load_mma_q, pipeline_load_mma_q_consumer_state,
+          pipeline_load_mma_q, pipeline_load_mma_q_consumer_state,        
           pipeline_load_mma_do, pipeline_load_mma_do_consumer_state,
           pipeline_mma_compute_s, pipeline_mma_compute_s_producer_state,
           pipeline_mma_compute_dp, pipeline_mma_compute_dp_producer_state,
@@ -1736,10 +1629,9 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
     else if (role == WarpRole::Compute) {
       warpgroup_reg_set<RegisterAllocation::kCompute>();
-
+    
       compute(
           blk_coord,
-          blk_offset,
           problem_shape,
           iter_start,
           iter_count,
@@ -1768,7 +1660,7 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
     else if (role == WarpRole::Reduce) {
       warpgroup_reg_set<RegisterAllocation::kReduce>();
-
+    
       reduce(
           blk_coord,
           problem_shape,
@@ -1785,9 +1677,9 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     }
     else {
       warpgroup_reg_set<RegisterAllocation::kEmpty>();
-
+    
       /* no-op */
-
+      
     }
   }
 

examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp

@@ -356,7 +356,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
     typename CollectiveMainloop::PipelineO::PipelineState pipeline_mma_corr_producer_state = cutlass::make_producer_start_state<typename CollectiveMainloop::PipelineO>();
 
     CollectiveMainloop mainloop;
-    CollectiveEpilogue epilogue{params.epilogue};
+    CollectiveEpilogue epilogue;
 
     if (role == WarpRole::Softmax0 || role == WarpRole::Softmax1) {
       warpgroup_reg_set<NumRegsSoftmax>();
@@ -407,8 +407,7 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
           pipeline_s0_corr, pipeline_s0_corr_consumer_state,
           pipeline_s1_corr, pipeline_s1_corr_consumer_state,
           pipeline_mma_corr, pipeline_mma_corr_consumer_state,
-          pipeline_corr_epi, pipeline_corr_epi_producer_state,
-          epilogue
+          pipeline_corr_epi, pipeline_corr_epi_producer_state
         );
 
 

examples/77_blackwell_fmha/kernel/sm100_fmha_mla_reduction.hpp

@@ -146,7 +146,7 @@ struct Sm100FmhaMlaReductionKernel {
       ElementAcc sum_lse = 0;
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < kNLsePerThread; ++i) {
-        sum_lse = sum_lse + expf(local_lse[i] - lse_max);
+        sum_lse = sum_lse + expf(local_lse[i] - params.scale * lse_max);
       }
     
       CUTLASS_PRAGMA_UNROLL
@@ -156,7 +156,7 @@ struct Sm100FmhaMlaReductionKernel {
 
       sum_lse = __shfl_sync(0xffffffff, sum_lse, 0);
 
-      ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits<ElementAcc>::infinity() : logf(sum_lse) + lse_max;
+      ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? std::numeric_limits<ElementAcc>::infinity() : logf(sum_lse) + params.scale * lse_max;
       if (threadIdx.x == 0 and params.ptr_lse != nullptr) {
         gLSE(0) = global_lse;
       }