1
0

flash_bwd_kernel.h 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303
  1. /******************************************************************************
  2. * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
  3. ******************************************************************************/
  4. #pragma once
  5. #include "cute/tensor.hpp"
  6. #include <cutlass/cutlass.h>
  7. #include <cutlass/arch/reg_reconfig.h>
  8. #include <cutlass/array.h>
  9. #include <cutlass/numeric_types.h>
  10. #include <cutlass/numeric_conversion.h>
  11. #include <cutlass/kernel_hardware_info.h>
  12. #include "cutlass/pipeline/pipeline.hpp"
  13. #include "utils.h"
  14. namespace flash {
  15. using namespace cute;
  16. template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
  17. class FlashAttnBwd {
  18. public:
  19. // Type Aliases
  20. static constexpr bool Is_causal = CollectiveMainloop_::Is_causal;
  21. static constexpr bool Is_local = CollectiveMainloop_::Is_local;
  22. static_assert(CollectiveMainloop_::Varlen == CollectiveEpilogue_::Varlen);
  23. static constexpr bool Varlen = CollectiveMainloop_::Varlen;
  24. // Mainloop derived types
  25. using CollectiveMainloop = CollectiveMainloop_;
  26. using TileShape_MNK = typename CollectiveMainloop::TileShape_MNK;
  27. using TiledMmaSdP = typename CollectiveMainloop::TiledMmaSdP;
  28. using TiledMmadKV = typename CollectiveMainloop::TiledMmadKV;
  29. using ArchTag = typename CollectiveMainloop::ArchTag;
  30. using ClusterShape = typename CollectiveMainloop::ClusterShape;
  31. using MainloopArguments = typename CollectiveMainloop::Arguments;
  32. using MainloopParams = typename CollectiveMainloop::Params;
  33. static constexpr bool dKV_swapAB = CollectiveMainloop::dKV_swapAB;
  34. // Epilogue derived types
  35. using CollectiveEpilogue = CollectiveEpilogue_;
  36. using EpilogueArguments = typename CollectiveEpilogue::Arguments;
  37. using EpilogueParams = typename CollectiveEpilogue::Params;
  38. static_assert(ArchTag::kMinComputeCapability >= 90);
  39. using TileScheduler = TileScheduler_;
  40. using TileSchedulerArguments = typename TileScheduler::Arguments;
  41. using TileSchedulerParams = typename TileScheduler::Params;
  42. static constexpr uint32_t NumLoadWarpGroups = 1;
  43. static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMmaSdP{})) / cutlass::NumThreadsPerWarpGroup;
  44. static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMmaSdP{})) + (NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup);
  45. static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
  46. static_assert(NumMmaWarpGroups == 2 || NumMmaWarpGroups == 3);
  47. /// Register requirement for Load and Math WGs
  48. static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 2 ? 24 : 32;
  49. static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 240 : 160;
  50. // If you want to print from the producer warp, you'd need to increase the number of registers
  51. // Otherwise you'll get CUDA error.
  52. // static constexpr uint32_t LoadRegisterRequirement = 40;
  53. // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152;
  54. // Kernel level shared memory storage
  55. struct SharedStorage {
  56. struct TensorStorage : cute::aligned_struct<128> {
  57. union {
  58. typename CollectiveMainloop::TensorStorage mainloop;
  59. typename CollectiveEpilogue::TensorStorage epilogue;
  60. };
  61. } tensors;
  62. struct PipelineStorage : cute::aligned_struct<16> {
  63. alignas(16) cutlass::arch::ClusterTransactionBarrier barrier_KV;
  64. alignas(16) cutlass::arch::ClusterBarrier barrier_dKV;
  65. alignas(16) typename CollectiveMainloop::MainloopPipeline::SharedStorage pipeline_q;
  66. alignas(16) typename CollectiveMainloop::MainloopPipeline_dO::SharedStorage pipeline_do;
  67. alignas(16) typename TileScheduler::SharedStorage smem_scheduler;
  68. } pipelines;
  69. };
  70. static constexpr int SharedStorageSize = sizeof(SharedStorage);
  71. // Device side arguments
  72. struct Arguments {
  73. MainloopArguments mainloop{};
  74. EpilogueArguments epilogue{};
  75. cutlass::KernelHardwareInfo hw_info{};
  76. TileSchedulerArguments scheduler{};
  77. };
  78. // Kernel entry point API
  79. struct Params {
  80. MainloopParams mainloop{};
  81. EpilogueParams epilogue{};
  82. cutlass::KernelHardwareInfo hw_info{};
  83. TileSchedulerParams scheduler{};
  84. };
  85. //
  86. // Methods
  87. //
  88. // Convert to underlying arguments. In this case, a simple copy for the aliased type.
  89. static
  90. Params
  91. to_underlying_arguments(Arguments const& args) {
  92. CUTLASS_TRACE_HOST("to_underlying_arguments():");
  93. // Get SM count if needed, otherwise use user supplied SM count
  94. int sm_count = args.hw_info.sm_count;
  95. if (sm_count <= 0) {
  96. CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
  97. " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
  98. sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
  99. }
  100. CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
  101. cutlass::KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
  102. return {
  103. CollectiveMainloop::to_underlying_arguments(args.mainloop),
  104. CollectiveEpilogue::to_underlying_arguments(args.epilogue),
  105. hw_info,
  106. TileScheduler::to_underlying_arguments(args.scheduler)
  107. };
  108. }
  109. // Computes the kernel launch grid shape based on runtime parameters
  110. static dim3
  111. get_grid_shape(Params const& params) {
  112. return TileScheduler::get_grid_shape(params.scheduler, params.hw_info.sm_count);
  113. }
  114. static dim3
  115. get_block_shape() {
  116. return dim3(MaxThreadsPerBlock, 1, 1);
  117. }
  118. CUTLASS_DEVICE
  119. void
  120. operator()(Params const& params, char* smem_buf) {
  121. static constexpr int NumMmaThreads = NumMmaWarpGroups * cutlass::NumThreadsPerWarpGroup;
  122. static constexpr int NumCopyThreads = NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup;
  123. static constexpr int kBlockM = get<0>(TileShape_MNK{});
  124. static constexpr int kBlockN = get<1>(TileShape_MNK{});
  125. using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
  126. using PipelineParams = typename MainloopPipeline::Params;
  127. using PipelineState = typename MainloopPipeline::PipelineState;
  128. using MainloopPipeline_dO = typename CollectiveMainloop::MainloopPipeline_dO;
  129. using PipelineParams_dO = typename MainloopPipeline_dO::Params;
  130. using PipelineState_dO = typename MainloopPipeline_dO::PipelineState;
  131. static constexpr bool Q_dO_same_stages = std::is_same_v<MainloopPipeline, MainloopPipeline_dO>;
  132. SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
  133. int const lane_predicate = cute::elect_one_sync();
  134. int const warp_idx = cutlass::canonical_warp_idx_sync();
  135. // Issue Tma Descriptor Prefetch from a single thread
  136. if (warp_idx == 0 && lane_predicate) {
  137. CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
  138. CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue);
  139. }
  140. // Obtain warp index
  141. int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
  142. PipelineParams pipeline_params;
  143. pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesQ + CollectiveMainloop::TmaTransactionBytesLSE;
  144. int warp_group_idx = cutlass::canonical_warp_group_idx();
  145. pipeline_params.role = warp_group_idx == 0
  146. ? MainloopPipeline::ThreadCategory::Producer
  147. : MainloopPipeline::ThreadCategory::Consumer;
  148. pipeline_params.is_leader = warp_group_thread_idx == 0;
  149. pipeline_params.num_consumers = NumMmaThreads;
  150. if (warp_idx == 0 && lane_predicate) {
  151. shared_storage.pipelines.barrier_KV.init(1 /*numThreads*/);
  152. // shared_storage.barrier_dKV.init(size(ClusterShape{}) /*numThreads*/);
  153. }
  154. // We're counting on pipeline_q to call cutlass::arch::fence_barrier_init();
  155. MainloopPipeline pipeline_q(shared_storage.pipelines.pipeline_q, pipeline_params, ClusterShape{});
  156. auto role_dO = warp_group_idx == 0
  157. ? MainloopPipeline_dO::ThreadCategory::Producer
  158. : MainloopPipeline_dO::ThreadCategory::Consumer;
  159. PipelineParams_dO pipeline_params_dO {pipeline_params.transaction_bytes, role_dO, pipeline_params.is_leader, pipeline_params.num_consumers};
  160. MainloopPipeline_dO pipeline_do(shared_storage.pipelines.pipeline_do, cute::conditional_return<Q_dO_same_stages>(pipeline_params, pipeline_params_dO), ClusterShape{});
  161. CollectiveMainloop collective_mainloop;
  162. CollectiveEpilogue collective_epilogue;
  163. // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster
  164. if constexpr (size(ClusterShape{}) > 1) {
  165. cute::cluster_arrive_relaxed();
  166. cute::cluster_wait();
  167. } else {
  168. __syncthreads();
  169. }
  170. if (warp_group_idx == 0) { // Producer
  171. cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
  172. int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
  173. if (warp_idx_in_warpgroup == 0) { // Load K, V, and do TMA on Q and dO
  174. PipelineState smem_pipe_write = cutlass::make_producer_start_state<MainloopPipeline>();
  175. PipelineState_dO smem_pipe_write_do = cutlass::make_producer_start_state<MainloopPipeline_dO>();
  176. TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.pipelines.smem_scheduler));
  177. for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/true>(params.scheduler);
  178. work_tile_info.is_valid(params.scheduler);
  179. work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/true>(params.scheduler, work_tile_info)) {
  180. auto block_coord = work_tile_info.get_block_coord(params.scheduler);
  181. auto [n_block, bidh, bidb] = block_coord;
  182. // With Varlen it's possible to have query length = 0. We want to skip the iteration.
  183. if constexpr (Is_causal || Is_local || Varlen) {
  184. int const m_block_min = collective_mainloop.get_m_block_min(params.mainloop, n_block, bidb);
  185. int const m_block_max = collective_mainloop.get_m_block_max(params.mainloop, n_block, bidb);
  186. if (m_block_min >= m_block_max) {
  187. scheduler.prefetch_next_work(params.scheduler, work_tile_info);
  188. continue;
  189. }
  190. }
  191. auto scheduler_prefetch = [&scheduler, &params, &work_tile_info]() {
  192. scheduler.prefetch_next_work(params.scheduler, work_tile_info);
  193. };
  194. collective_mainloop.load(params.mainloop, pipeline_q, pipeline_do, smem_pipe_write,
  195. smem_pipe_write_do, shared_storage, scheduler_prefetch, block_coord);
  196. }
  197. collective_mainloop.load_tail(pipeline_q, pipeline_do, smem_pipe_write, smem_pipe_write_do);
  198. } else if (warp_idx_in_warpgroup == 1) {
  199. TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.pipelines.smem_scheduler));
  200. for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
  201. work_tile_info.is_valid(params.scheduler);
  202. work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
  203. auto block_coord = work_tile_info.get_block_coord(params.scheduler);
  204. auto [n_block, bidh, bidb] = block_coord;
  205. if constexpr (Is_causal || Is_local || Varlen) {
  206. int const m_block_min = collective_mainloop.get_m_block_min(params.mainloop, n_block, bidb);
  207. int const m_block_max = collective_mainloop.get_m_block_max(params.mainloop, n_block, bidb);
  208. if (m_block_min >= m_block_max) { continue; }
  209. }
  210. collective_mainloop.store_dq(params.mainloop, shared_storage, block_coord);
  211. }
  212. }
  213. } else { // Consumer
  214. cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
  215. TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.pipelines.smem_scheduler));
  216. // Initialize matmul objects.
  217. TiledMmadKV tiled_mma_dKV;
  218. PipelineState smem_pipe_read;
  219. PipelineState_dO smem_pipe_read_do;
  220. collective_mainloop.mma_init();
  221. scheduler.init_consumer();
  222. int work_idx = 0;
  223. CUTLASS_PRAGMA_NO_UNROLL
  224. for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
  225. work_tile_info.is_valid(params.scheduler);
  226. work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
  227. auto block_coord = work_tile_info.get_block_coord(params.scheduler);
  228. auto [n_block, bidh, bidb] = block_coord;
  229. if constexpr (Is_causal || Is_local || Varlen) {
  230. int const m_block_min = collective_mainloop.get_m_block_min(params.mainloop, n_block, bidb);
  231. int const m_block_max = collective_mainloop.get_m_block_max(params.mainloop, n_block, bidb);
  232. if (m_block_min >= m_block_max) { // We exit early and write 0 to dK and dV
  233. collective_epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord);
  234. continue;
  235. }
  236. }
  237. // dK and dV output accumulator.
  238. Tensor tdKrdK = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
  239. Tensor tdVrdV = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
  240. collective_mainloop.mma(params.mainloop, pipeline_q, pipeline_do, smem_pipe_read, smem_pipe_read_do,
  241. tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage);
  242. collective_epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV,
  243. threadIdx.x - NumCopyThreads, block_coord);
  244. ++work_idx;
  245. }
  246. collective_epilogue.store_tail();
  247. }
  248. }
  249. };
  250. } // namespace flash