Convert FA3 to Kernel Hub format

build.toml

@@ -0,0 +1,593 @@
+[general]
+name = "flash_attn3"
+universal = false
+cuda-minver = "12.4"
+cuda-maxver = "12.4"
+
+[torch]
+src = [
+  "torch-ext/pytorch_shim.h",
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.flash_attn]
+backend = "cuda"
+cuda-capabilities = ["8.0", "9.0a"]
+cuda-flags = [
+  "-O3",
+  "-std=c++17",
+  "--ftemplate-backtrace-limit=0",              # To debug template code
+  "--use_fast_math",
+  "-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
+  "-DCUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1",
+  "-DCUTLASS_ENABLE_GDC_FOR_SM90",
+  "--expt-relaxed-constexpr",
+  "--expt-extended-lambda",
+  "--use_fast_math",
+  "-DNDEBUG",
+]
+
+src = [
+  "flash-attn/cuda_check.h",
+  "flash-attn/flash_api.cpp",
+  "flash-attn/flash_fwd_combine.cu",
+  "flash-attn/flash_fwd_combine_kernel.h",
+  "flash-attn/flash_fwd_combine_launch_template.h",
+  "flash-attn/flash.h",
+  "flash-attn/flash_prepare_scheduler.cu",
+  "flash-attn/heuristics.h",
+  "flash-attn/seqlen.h",
+  "flash-attn/static_switch.h",
+  "flash-attn/tile_size.h",
+  "flash-attn/utils.h",
+]
+depends = ["torch", "cutlass_3_9"]
+
+[kernel.flash_attn_sm80]
+backend = "cuda"
+cuda-capabilities = ["8.0", "9.0a"]
+cuda-flags = [
+  "-O3",
+  "-std=c++17",
+  "--ftemplate-backtrace-limit=0",              # To debug template code
+  "--use_fast_math",
+  "-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
+  "-DCUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1",
+  "-DCUTLASS_ENABLE_GDC_FOR_SM90",
+  "--expt-relaxed-constexpr",
+  "--expt-extended-lambda",
+  "--use_fast_math",
+  "-DNDEBUG",
+]
+src = [
+  "flash-attn/block.h",
+  "flash-attn/copy_sm90_bulk_reduce.hpp",
+  "flash-attn/epilogue_bwd.hpp",
+  "flash-attn/epilogue_fwd.hpp",
+  "flash-attn/flash.h",
+  "flash-attn/flash_bwd_kernel_sm80.h",
+  "flash-attn/flash_bwd_kernel_sm90.h",
+  "flash-attn/flash_bwd_launch_template.h",
+  "flash-attn/flash_bwd_postprocess_kernel.h",
+  "flash-attn/flash_bwd_preprocess_kernel.h",
+  "flash-attn/flash_fwd_launch_template.h",
+  "flash-attn/flash_fwd_kernel_sm80.h",
+  "flash-attn/flash_fwd_kernel_sm90.h",
+  "flash-attn/heuristics.h",
+  "flash-attn/mainloop_bwd_sm80.hpp",
+  "flash-attn/mainloop_fwd_sm80.hpp",
+  "flash-attn/mainloop_bwd_sm90_tma_gmma_ws.hpp",
+  "flash-attn/mainloop_fwd_sm90_tma_gmma_ws.hpp",
+  "flash-attn/mask.h",
+  "flash-attn/named_barrier.hpp",
+  "flash-attn/pack_gqa.h",
+  "flash-attn/paged_kv.h",
+  "flash-attn/rotary.h",
+  "flash-attn/sm90_pipeline_no_cluster.hpp",
+  "flash-attn/softmax.h",
+  "flash-attn/tile_size.h",
+  "flash-attn/tile_scheduler.hpp",
+
+  "flash-attn/instantiations/flash_bwd_hdim128_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_split_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_softcap_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_split_sm80.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_split_softcap_sm80.cu"
+]
+include = ["flash-attn"]
+depends = ["torch", "cutlass_3_9"]
+
+[kernel.flash_attn_sm90]
+backend = "cuda"
+cuda-capabilities = ["8.0", "9.0a"]
+cuda-flags = [
+  "-O3",
+  "-std=c++17",
+  "--ftemplate-backtrace-limit=0",              # To debug template code
+  "--use_fast_math",
+  "-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
+  "-DCUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1",
+  "-DCUTLASS_ENABLE_GDC_FOR_SM90",
+  "--expt-relaxed-constexpr",
+  "--expt-extended-lambda",
+  "--use_fast_math",
+  "-DNDEBUG",
+]
+src = [
+  "flash-attn/block.h",
+  "flash-attn/copy_sm90_bulk_reduce.hpp",
+  "flash-attn/epilogue_bwd.hpp",
+  "flash-attn/epilogue_fwd.hpp",
+  "flash-attn/flash.h",
+  "flash-attn/flash_bwd_kernel_sm80.h",
+  "flash-attn/flash_bwd_kernel_sm90.h",
+  "flash-attn/flash_bwd_launch_template.h",
+  "flash-attn/flash_bwd_postprocess_kernel.h",
+  "flash-attn/flash_bwd_preprocess_kernel.h",
+  "flash-attn/flash_fwd_launch_template.h",
+  "flash-attn/flash_fwd_kernel_sm80.h",
+  "flash-attn/flash_fwd_kernel_sm90.h",
+  "flash-attn/heuristics.h",
+  "flash-attn/mainloop_bwd_sm80.hpp",
+  "flash-attn/mainloop_fwd_sm80.hpp",
+  "flash-attn/mainloop_bwd_sm90_tma_gmma_ws.hpp",
+  "flash-attn/mainloop_fwd_sm90_tma_gmma_ws.hpp",
+  "flash-attn/mask.h",
+  "flash-attn/named_barrier.hpp",
+  "flash-attn/pack_gqa.h",
+  "flash-attn/paged_kv.h",
+  "flash-attn/rotary.h",
+  "flash-attn/sm90_pipeline_no_cluster.hpp",
+  "flash-attn/softmax.h",
+  "flash-attn/tile_size.h",
+  "flash-attn/tile_scheduler.hpp",
+
+  "flash-attn/instantiations/flash_bwd_hdim128_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim128_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim192_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim256_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim64_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_bwd_hdim96_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim128_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_128_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim192_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim256_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_256_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_512_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim64_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdim96_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimall_fp16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_bf16_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_e4m3_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_paged_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_paged_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_paged_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_paged_split_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_softcap_packgqa_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_softcap_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_split_sm90.cu",
+  "flash-attn/instantiations/flash_fwd_hdimdiff_fp16_split_softcap_sm90.cu",
+]
+include = ["flash-attn"]
+depends = ["torch", "cutlass_3_9"]
+
+# [kernel.flash_attn_sm100]
+# backend = "cuda"
+# cuda-capabilities = ["8.0", "9.0a", "10.0"]
+# cuda-flags = [
+#   "-O3",
+#   "-std=c++17",
+#   "--ftemplate-backtrace-limit=0",              # To debug template code
+#   "--use_fast_math",
+#   "-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
+#   "-DCUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1",
+#   "-DCUTLASS_ENABLE_GDC_FOR_SM90",
+#   "--expt-relaxed-constexpr",
+#   "--expt-extended-lambda",
+#   "--use_fast_math",
+#   "-DNDEBUG",
+# ]
+# src = [
+#   "flash-attn/block.h",
+#   "flash-attn/copy_sm90_bulk_reduce.hpp",
+#   "flash-attn/epilogue_bwd.hpp",
+#   "flash-attn/epilogue_fwd.hpp",
+#   "flash-attn/flash.h",
+#   "flash-attn/flash_bwd_kernel_sm80.h",
+#   "flash-attn/flash_bwd_kernel_sm90.h",
+#   "flash-attn/flash_bwd_launch_template.h",
+#   "flash-attn/flash_bwd_postprocess_kernel.h",
+#   "flash-attn/flash_bwd_preprocess_kernel.h",
+#   "flash-attn/flash_fwd_launch_template.h",
+#   "flash-attn/flash_fwd_kernel_sm80.h",
+#   "flash-attn/flash_fwd_kernel_sm90.h",
+#   "flash-attn/heuristics.h",
+#   "flash-attn/mainloop_bwd_sm80.hpp",
+#   "flash-attn/mainloop_fwd_sm80.hpp",
+#   "flash-attn/mainloop_bwd_sm90_tma_gmma_ws.hpp",
+#   "flash-attn/mainloop_fwd_sm90_tma_gmma_ws.hpp",
+#   "flash-attn/mask.h",
+#   "flash-attn/named_barrier.hpp",
+#   "flash-attn/pack_gqa.h",
+#   "flash-attn/paged_kv.h",
+#   "flash-attn/rotary.h",
+#   "flash-attn/sm90_pipeline_no_cluster.hpp",
+#   "flash-attn/softmax.h",
+#   "flash-attn/tile_size.h",
+#   "flash-attn/tile_scheduler.hpp",
+#
+#   "flash-attn/instantiations/flash_fwd_hdim128_bf16_sm100.cu",
+# ]
+# include = ["flash-attn"]
+# depends = ["torch", "cutlass_3_9"]

flake.lock

@@ -0,0 +1,168 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1747046372,
+        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-compat_2": {
+      "locked": {
+        "lastModified": 1733328505,
+        "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "hf-nix": {
+      "inputs": {
+        "flake-compat": "flake-compat_2",
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1750234878,
+        "narHash": "sha256-q9DRC9zdpzUf88qqg1qbhP1qgJbE2cMtn8oUmosuyT8=",
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "rev": "c7132f90763d756da3e77da62e01be0a4546dc57",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "hf-nix",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "hf-nix": "hf-nix",
+        "nixpkgs": [
+          "kernel-builder",
+          "hf-nix",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1750275112,
+        "narHash": "sha256-gqAxmLLt0tYvuRYumOZHQgryMeEFdt6j3nEC8B5rT14=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "1b63210b2a1fc3cda2e3a579e7aa8f8c8532626f",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1747820358,
+        "narHash": "sha256-fTqsZsUX6M3yeEvgyQvXcbGmT2CaRVyVwsi8eK29Oj4=",
+        "owner": "danieldk",
+        "repo": "nixpkgs",
+        "rev": "d3c1681180717528068082103bf323147de6ab0b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "danieldk",
+        "ref": "cudatoolkit-12.9-kernel-builder",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for Hopper Flash Attention kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      path = ./.;
+      rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate;
+    };
+}

flash-attn/block.h

@@ -0,0 +1,139 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+namespace flash {
+
+template <class SeqlenInfo_t, int kBlockM, int kBlockN, bool Is_causal, bool Is_local, bool PackGQA=false, bool Split=false>
+struct BlockMN {
+
+    static
+    CUTLASS_DEVICE
+    cute::tuple<int, int> get_n_block_min_max(
+            SeqlenInfo_t const& seqlen_info,
+            int const m_block, int const bidb, int const split_idx, int const num_splits,
+            int const window_size_left, int const window_size_right,
+            cutlass::FastDivmod const& attention_chunk_divmod,
+            cutlass::FastDivmod const& qhead_per_khead_divmod) {
+
+        int const seqlen_k = seqlen_info.seqlen_k;
+        int const seqlen_q = seqlen_info.seqlen_q;
+        int n_block_max = cute::ceil_div(seqlen_k, kBlockN);
+        if constexpr (Is_causal || Is_local) {
+            int m_idx_max = (m_block + 1) * kBlockM;
+            // TODO: check off-by-1 error
+            if (PackGQA) { m_idx_max = qhead_per_khead_divmod.divide(m_idx_max - 1) + 1 ; }
+            int const n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q;
+            int n_idx_right = !Is_local ? n_idx : n_idx + window_size_right;
+            if (Is_local && attention_chunk_divmod.divisor > 0) {
+                n_idx_right = std::min(n_idx_right, flash::round_up(attention_chunk_divmod, n_idx));
+            }
+            n_block_max = std::min(n_block_max, cute::ceil_div(n_idx_right, kBlockN));
+        }
+        int n_block_min = 0;
+        if constexpr (Is_local) {
+            int m_idx_min = m_block * kBlockM;
+            if (PackGQA) { m_idx_min = qhead_per_khead_divmod.divide(m_idx_min); }
+            int const n_idx = m_idx_min + seqlen_k - seqlen_q;
+            int n_idx_left = n_idx - window_size_left;
+            if (attention_chunk_divmod.divisor > 0) {
+                n_idx_left = std::max(n_idx_left, flash::round_down(attention_chunk_divmod, n_idx));
+            }
+            n_block_min = std::max(int(0), n_idx_left / kBlockN);
+        }
+        // if (threadIdx.x == 128) { printf("Inside, bid.x = %d, bid.y = %d, bid.z = %d, split_idx = %d, n_block_min: %d, n_block_max: %d\n", blockIdx.x, blockIdx.y, blockIdx.z, split_idx, n_block_min, n_block_max); }
+        if constexpr (Split) {
+            uint32_t num_splits_dynamic_u = reinterpret_cast<uint32_t const&>(split_idx) >> 16; // first 16 bits are for num_splits
+            int num_splits_dynamic = reinterpret_cast<int&>(num_splits_dynamic_u);
+            int split_idx_actual = split_idx & 0x0000FFFF;
+            int num_splits_actual = num_splits_dynamic > 0 ? num_splits_dynamic : num_splits;
+            int num_n_blocks_per_split = n_block_max <= n_block_min ? 0 : cute::ceil_div(n_block_max - n_block_min, num_splits_actual);
+            n_block_min = n_block_min + split_idx_actual * num_n_blocks_per_split;
+            n_block_max = std::min(n_block_min + num_n_blocks_per_split, n_block_max);
+            // if (threadIdx.x == 128) { printf("Inside, bid.x = %d, bid.y = %d, bid.z = %d, split_idx = %d, num_splits_dynamic = %d, num_splits_actual = %d, num_n_blocks_per_split = %d, n_block_min: %d, n_block_max: %d\n", blockIdx.x, blockIdx.y, blockIdx.z, split_idx, num_splits_dynamic, num_splits_actual, num_n_blocks_per_split, n_block_min, n_block_max); }
+        }
+        // if (threadIdx.x == 128) { printf("After split, inside, bid.y = %d, bid.z = %d, split_idx = %d, n_block_min: %d, n_block_max: %d\n", blockIdx.y, blockIdx.z, split_idx, n_block_min, n_block_max); }
+        return {n_block_min, n_block_max};
+    }
+
+    static
+    CUTLASS_DEVICE
+    cute::tuple<int, int> get_n_block_k_new_min_max(
+            SeqlenInfo_t const& seqlen_info,
+            int const m_block, int const bidb, int const split_idx, int const num_splits,
+            int const window_size_left, int const window_size_right,
+            cutlass::FastDivmod const& attention_chunk_divmod,
+            cutlass::FastDivmod const& qhead_per_khead_divmod) {
+
+        auto [n_block_min, n_block_max] = get_n_block_min_max(
+            seqlen_info, m_block, bidb, split_idx, num_splits,
+            window_size_left, window_size_right, attention_chunk_divmod, qhead_per_khead_divmod);
+        int const idx_k_new_min = std::max(n_block_min * kBlockN - seqlen_info.seqlen_k_og, 0);
+        int const idx_k_new_max = std::min(n_block_max * kBlockN - seqlen_info.seqlen_k_og, seqlen_info.seqlen_k_new);
+        int const n_block_new_min = idx_k_new_min / kBlockN;
+        int const n_block_new_max = idx_k_new_max > idx_k_new_min ? cute::ceil_div(idx_k_new_max, kBlockN) : n_block_new_min;
+        // if (threadIdx.x == 128 && m_block == 0) { printf("bidb = %d, seqlen_k_new = %d, seqlen_k_og = %d, n_block_min = %d, n_block_max = %d, idx_k_new_min = %d, idx_k_new_max = %d, n_block_new_min = %d, n_block_new_max = %d\n", bidb, seqlen_k_new, seqlen_k_og, n_block_min, n_block_max, idx_k_new_min, idx_k_new_max, n_block_new_min, n_block_new_max);}
+        return {n_block_new_min, n_block_new_max};
+    }
+
+    static
+    CUTLASS_DEVICE
+    cute::tuple<int, int> get_m_block_min_max(
+            SeqlenInfo_t const& seqlen_info,
+            int const n_block, int const bidb,
+            int const window_size_left, int const window_size_right, int const sink_token_length) {
+        // TODO: support attention_chunk
+        int const seqlen_q = seqlen_info.seqlen_q;
+        int const seqlen_k = seqlen_info.seqlen_k;
+        int m_block_max = cute::ceil_div(seqlen_q, kBlockM);
+        if constexpr (Is_local) {
+            if (n_block >= cute::ceil_div(sink_token_length, kBlockN)) {
+                m_block_max = std::min(m_block_max, cute::ceil_div((n_block + 1) * kBlockN + seqlen_q - seqlen_k + window_size_left, kBlockM));
+            }
+        }
+        int m_block_min = 0;
+        if constexpr (Is_causal || Is_local) {
+            m_block_min = std::max(m_block_min, (n_block * kBlockN + seqlen_q - seqlen_k - window_size_right) / kBlockM);
+        }
+        return {m_block_min, m_block_max};
+    }
+
+    // If we have separate iterations with causal or local masking at the start, where do we stop
+    static
+    CUTLASS_DEVICE
+    int get_n_block_min_causal_local_mask(
+            SeqlenInfo_t const& seqlen_info,
+            int const m_block, int const n_block_min, int const window_size_right,
+            cutlass::FastDivmod const& attention_chunk_divmod,
+            cutlass::FastDivmod const& qhead_per_khead_divmod) {
+        int const m_idx_min = !PackGQA ? m_block * kBlockM : qhead_per_khead_divmod.divide(m_block * kBlockM);
+        int const n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q;
+        int n_idx_right = !Is_local ? n_idx : n_idx + window_size_right;
+        if (Is_local && attention_chunk_divmod.divisor > 0) {
+            n_idx_right = std::min(n_idx_right, flash::round_up(attention_chunk_divmod, n_idx));
+        }
+        return std::max(n_block_min, n_idx_right / kBlockN);
+    }
+
+    // If we have separate iterations with local masking at the end, where do we stop the non-masked iterations
+    static
+    CUTLASS_DEVICE
+    int get_n_block_min_before_local_mask(
+            SeqlenInfo_t const& seqlen_info,
+            int const m_block, int const n_block_min, int const window_size_left,
+            cutlass::FastDivmod const& attention_chunk_divmod,
+            cutlass::FastDivmod const& qhead_per_khead_divmod) {
+        int const m_idx_max = !PackGQA ? (m_block + 1) * kBlockM : qhead_per_khead_divmod.divide((m_block + 1) * kBlockM - 1) + 1;
+        int const n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q;
+        int n_idx_left = !Is_local ? n_idx : n_idx - window_size_left;
+        if (Is_local && attention_chunk_divmod.divisor > 0) {
+            n_idx_left = std::max(n_idx_left, flash::round_down(attention_chunk_divmod, n_idx));
+        }
+        return !Is_local ? n_block_min : std::max(n_block_min, cute::ceil_div(n_idx_left, kBlockN));
+    }
+
+};
+
+} // namespace flash

flash-attn/copy_sm90_bulk_reduce.hpp

@@ -0,0 +1,49 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include<cute/arch/copy_sm90_tma.hpp>
+
+namespace cute
+{
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct SM90_BULK_REDUCE_ADD
+{
+  CUTE_HOST_DEVICE static void
+  copy(float const* smem_ptr,
+       float      * gmem_ptr, int32_t store_bytes)
+  {
+#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [%0], [%1], %2;\n"
+                     :
+                     : "l"(gmem_ptr), "r"(smem_int_ptr), "r"(store_bytes)
+                     : "memory");
+#else
+    CUTE_INVALID_CONTROL_PATH("Trying to use BULK_REDUCE_ADD without CUTE_ARCH_TMA_SM90_ENABLED.");
+#endif
+  }
+
+  CUTE_HOST_DEVICE static void
+  copy(float const* smem_ptr,
+       float      * gmem_ptr, int32_t store_bytes, uint64_t cache_hint)
+  {
+#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [%0], [%1], %2, %3;\n"
+                     :
+                     : "l"(gmem_ptr), "r"(smem_int_ptr), "r"(store_bytes), "l"(cache_hint)
+                     : "memory");
+#else
+    CUTE_INVALID_CONTROL_PATH("Trying to use BULK_REDUCE_ADD without CUTE_ARCH_TMA_SM90_ENABLED.");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // end namespace cute

flash-attn/cuda_check.h

@@ -0,0 +1,19 @@
+/******************************************************************************
+ * Copyright (c) 2024, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <assert.h>
+#include <stdlib.h>
+
+#define CHECK_CUDA(call)                        \
+    do {                                                                                                  \
+        cudaError_t status_ = call;                                                                       \
+        if (status_ != cudaSuccess) {                                                                     \
+            fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \
+            exit(1);                                                                                      \
+        }                                                                                                 \
+    } while(0)
+
+#define CHECK_CUDA_KERNEL_LAUNCH() CHECK_CUDA(cudaGetLastError())

flash-attn/epilogue_bwd.hpp

@@ -0,0 +1,533 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/barrier.h"
+#include "cute/tensor.hpp"
+
+#include "cutlass/gemm/collective/builders/sm90_common.inl"
+
+#include "seqlen.h"
+#include "named_barrier.hpp"
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class TileShape_MNK_, class Element_, class ArchTag_,
+          int NumEpilogueThreads_, bool Varlen_, bool dKV_swapAB_, int AtomLayoutKdKV=1>
+struct CollectiveEpilogueBwd {
+
+    using TileShape_MNK = TileShape_MNK_;
+    using Element = Element_;
+    using ArchTag = ArchTag_;
+    static constexpr int NumEpilogueThreads = NumEpilogueThreads_;
+    static constexpr bool Varlen = Varlen_;
+    static constexpr bool dKV_swapAB = dKV_swapAB_;
+    static constexpr bool Use_TMA = !Varlen && ArchTag::kMinComputeCapability >= 90;
+
+    static_assert(ArchTag::kMinComputeCapability >= 80);
+
+    using GmemTiledCopydKVTMA = cute::SM90_TMA_STORE;
+
+    // These are for storing the output tensor without TMA (e.g., for setting output to zero)
+    static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);
+    static_assert(get<2>(TileShape_MNK{}) % kGmemElemsPerLoad == 0, "Headdim must be a multiple of kGmemElemsPerLoad");
+    static constexpr int kHeadDim = get<2>(TileShape_MNK{});
+    static constexpr int kGmemThreadsPerRow = cutlass::gcd(kHeadDim / kGmemElemsPerLoad, NumEpilogueThreads);
+    static_assert(NumEpilogueThreads % kGmemThreadsPerRow == 0, "NumEpilogueThreads must be a multiple of kGmemThreadsPerRow");
+    using GmemLayoutAtom = Layout<Shape <Int<NumEpilogueThreads / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
+                                  Stride<Int<kGmemThreadsPerRow>, _1>>;
+    using GmemTiledCopydKV = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoad>>>{}));  // Val layout, 8 or 16 vals per store
+
+    using SmemLayoutAtomdKVTMA = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
+                                          // TODO: do we have to change this if dKV_swapAB is true?
+                                          decltype(cute::get<1>(TileShape_MNK{})), Int<CUTE_STATIC_V(cute::get<2>(TileShape_MNK{})) / AtomLayoutKdKV>>());
+    using SmemLayoutdKVTMA = decltype(tile_to_shape(SmemLayoutAtomdKVTMA{}, select<1, 2>(TileShape_MNK{})));
+    using SmemLayoutdKVtTMA =
+        decltype(cute::composition(SmemLayoutdKVTMA{},
+                                   make_layout(make_shape(get<2>(TileShape_MNK{}), get<1>(TileShape_MNK{})),
+                                               make_stride(decltype(get<1>(TileShape_MNK{})){}, _1{}))));
+
+    // If we don't use TMA
+    static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : (kHeadDim % 32 == 0 ? 32 : 16);
+    static constexpr int kSwizzle = kBlockKSmem == 64 ? 3 : (kBlockKSmem == 32 ? 2 : 1);
+    using SmemLayoutAtomdKVSTG =
+        decltype(composition(Swizzle<kSwizzle, 3, 3>{},
+                             Layout<Shape<Int<8>, Int<kBlockKSmem>>,
+                             Stride<Int<kBlockKSmem>, _1>>{}));
+
+    using SmemLayoutAtomdKV = std::conditional_t<Use_TMA, SmemLayoutAtomdKVTMA, SmemLayoutAtomdKVSTG>;
+    using SmemLayoutdKV = decltype(tile_to_shape(SmemLayoutAtomdKV{}, select<1, 2>(TileShape_MNK{})));
+    using SmemLayoutdKVt =
+        decltype(cute::composition(SmemLayoutdKV{},
+                                   make_layout(make_shape(get<2>(TileShape_MNK{}), get<1>(TileShape_MNK{})),
+                                               make_stride(decltype(get<1>(TileShape_MNK{})){}, _1{}))));
+
+    using SmemCopyAtomdKV = Copy_Atom<
+        std::conditional_t<
+            ArchTag::kMinComputeCapability >= 90,
+            std::conditional_t<!dKV_swapAB, cute::SM90_U32x4_STSM_N, cute::SM90_U16x8_STSM_T>,
+            AutoVectorizingCopyWithAssumedAlignment<128>
+        >,
+        Element>;
+
+    static constexpr size_t SmemAlignmentdKV = ArchTag::kMinComputeCapability >= 90 ? cutlass::detail::alignment_for_swizzle(SmemLayoutdKV{}) : 128;
+    static_assert(SmemAlignmentdKV >= 128, "Require at least 128B alignment");
+
+    struct TensorStorage : cute::aligned_struct<SmemAlignmentdKV> {
+        cute::array_aligned<Element, cute::cosize_v<SmemLayoutdKV>, SmemAlignmentdKV> smem_dk;
+        cute::array_aligned<Element, cute::cosize_v<SmemLayoutdKV>, SmemAlignmentdKV> smem_dv;
+    };
+
+    using ShapedKV = cute::Shape<int32_t, int32_t, int32_t, int32_t>;  // (seqlen_k, d, head, batch)
+    using StridedKV = cute::Stride<int64_t, _1, int64_t, int64_t>;
+
+    using TMA_dKV = std::conditional_t<
+        Use_TMA,
+        decltype(make_tma_copy(
+            GmemTiledCopydKVTMA{},
+            make_tensor(make_gmem_ptr(static_cast<Element*>(nullptr)), ShapedKV{}, StridedKV{}),
+            SmemLayoutdKVTMA{},
+            select<1, 2>(TileShape_MNK{}),
+            _1{})),  // no mcast for dKV
+        std::nullptr_t
+        >;
+
+    // Host side kernel arguments
+    struct Arguments {
+        Element* ptr_dK;
+        ShapedKV const shape_dK;
+        StridedKV const stride_dK;
+        Element* ptr_dV;
+        ShapedKV const shape_dV;
+        StridedKV const stride_dV;
+        int const num_heads_q;
+        int* dk_semaphore;
+        int* dv_semaphore;
+        int const* cu_seqlens;
+        int const* seqused;
+    };
+
+    // Device side kernel params
+    struct Params {
+        Element* ptr_dK;
+        ShapedKV const shape_dK;
+        StridedKV const stride_dK;
+        Element* ptr_dV;
+        ShapedKV const shape_dV;
+        StridedKV const stride_dV;
+        TMA_dKV tma_store_dK, tma_store_dV;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    static Params
+    to_underlying_arguments(Arguments const& args) {
+        Tensor mdK = make_tensor(make_gmem_ptr(args.ptr_dK), args.shape_dK, args.stride_dK);
+        Tensor mdV = make_tensor(make_gmem_ptr(args.ptr_dV), args.shape_dV, args.stride_dV);
+        TMA_dKV tma_store_dK = [&] {
+            if constexpr (Use_TMA) {
+                return make_tma_copy(GmemTiledCopydKVTMA{}, mdK, SmemLayoutdKVTMA{}, select<1, 2>(TileShape_MNK{}), _1{}); // no mcast for dKV
+            } else {
+                return nullptr;
+            }
+        }();
+        TMA_dKV tma_store_dV = [&] {
+            if constexpr (Use_TMA) {
+                return make_tma_copy(GmemTiledCopydKVTMA{}, mdV, SmemLayoutdKVTMA{}, select<1, 2>(TileShape_MNK{}), _1{}); // no mcast for dKV
+            } else {
+                return nullptr;
+            }
+        }();
+        return {args.ptr_dK, args.shape_dK, args.stride_dK, args.ptr_dV, args.shape_dV, args.stride_dV,
+                tma_store_dK, tma_store_dV, args.cu_seqlens, args.seqused};
+    }
+
+    /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
+    CUTLASS_DEVICE
+    static void prefetch_tma_descriptors(Params const& params) {
+        if constexpr (Use_TMA) {
+            cute::prefetch_tma_descriptor(params.tma_store_dK.get_tma_descriptor());
+            cute::prefetch_tma_descriptor(params.tma_store_dV.get_tma_descriptor());
+        }
+    }
+
+    template <typename SharedStorage, typename FrgTensorO, typename TiledMma>
+    CUTLASS_DEVICE void
+    store(Params const& params,
+          FrgTensorO const& tdKrdK,
+          FrgTensorO const& tdVrdV,
+          SharedStorage& shared_storage,
+          TiledMma tiled_mma,
+          int thread_idx,
+          cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+          ) {
+
+        auto [n_block, bidh, bidb] = block_coord;
+        Tensor sdK = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dk.data()), SmemLayoutdKV{}));
+        Tensor sdV = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dv.data()), SmemLayoutdKV{}));
+        Tensor sdKt = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dk.data()), SmemLayoutdKVt{}));
+        Tensor sdVt = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dv.data()), SmemLayoutdKVt{}));
+        auto smem_tiled_copy_dKV = make_tiled_copy_C(SmemCopyAtomdKV{}, tiled_mma);
+        auto smem_thr_copy_dKV = smem_tiled_copy_dKV.get_thread_slice(thread_idx);
+
+        Tensor tdVrdV_out = make_tensor_like<Element>(tdVrdV);
+        flash::convert_type_out(tdVrdV, tdVrdV_out);
+        Tensor tdKrdK_out = make_tensor_like<Element>(tdKrdK);
+        flash::convert_type_out(tdKrdK, tdKrdK_out);
+        Tensor taccdKrdK = smem_thr_copy_dKV.retile_S(tdKrdK_out);        // ((Atom,AtomNum), MMA_M, MMA_N)
+        Tensor taccdVrdV = smem_thr_copy_dKV.retile_S(tdVrdV_out);        // ((Atom,AtomNum), MMA_M, MMA_N)
+        // if (blockIdx.x == 0 && threadIdx.x == 128) { print(smem_thr_copy_dKV); print(sdK); printf("\n"); print(sdKt); printf("\n"); }
+        Tensor taccdKsdK = smem_thr_copy_dKV.partition_D(cute::conditional_return<!dKV_swapAB>(sdK, sdKt));     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+        Tensor taccdVsdV = smem_thr_copy_dKV.partition_D(cute::conditional_return<!dKV_swapAB>(sdV, sdVt));     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+
+        // Make sure all WGs have finished reading K and V
+        flash::named_barrier_sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+        cute::copy(smem_tiled_copy_dKV, taccdVrdV, taccdVsdV);
+        cute::copy(smem_tiled_copy_dKV, taccdKrdK, taccdKsdK);
+        if constexpr (Use_TMA) {
+            cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
+            cutlass::arch::NamedBarrier::arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
+                                                cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+
+            Tensor mdK = params.tma_store_dK.get_tma_tensor(params.shape_dK);
+            Tensor mdV = params.tma_store_dV.get_tma_tensor(params.shape_dV);
+            Tensor gdK = local_tile(mdK(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+            Tensor gdV = local_tile(mdV(_, _, bidh, bidb), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+            auto block_tma_dK = params.tma_store_dK.get_slice(_0{});
+            auto block_tma_dV = params.tma_store_dV.get_slice(_0{});
+            Tensor tdKgdK = block_tma_dK.partition_D(gdK);  // (TMA, TMA_M, TMA_K)
+            Tensor tdKsdK = block_tma_dK.partition_S(sdK); // (TMA, TMA_M, TMA_K)
+            Tensor tdVgdV = block_tma_dV.partition_D(gdV);  // (TMA, TMA_M, TMA_K)
+            Tensor tdVsdV = block_tma_dV.partition_S(sdV); // (TMA, TMA_M, TMA_K)
+            int warp_idx_sync = __shfl_sync(0xffffffff, thread_idx / cutlass::NumThreadsPerWarp, 0);
+            if (warp_idx_sync == NumEpilogueThreads / cutlass::NumThreadsPerWarp - 1) {
+                cutlass::arch::NamedBarrier::sync(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
+                                                cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+                if (cute::elect_one_sync()) {
+                    cute::copy(params.tma_store_dV, tdVsdV, tdVgdV);
+                    cute::copy(params.tma_store_dK, tdKsdK, tdKgdK);
+                    tma_store_arrive();
+                }
+            }
+            tma_store_wait<0>();
+            // // Tell warp 0 that smem_k and smem_v are ready
+            // cutlass::arch::NamedBarrier::arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast<uint32_t>(BwdNamedBarriers::KVEmpty) /*id*/);
+
+        } else {
+            flash::named_barrier_sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            static constexpr int kBlockN = get<1>(TileShape_MNK{});
+            flash::SeqlenInfo<Varlen, kBlockN> seqlen_info{bidb, size<0>(params.shape_dK), params.cu_seqlens, params.seqused};
+            bool const is_varlen = Varlen && params.cu_seqlens;
+            Tensor mdK = make_tensor(make_gmem_ptr(params.ptr_dK), params.shape_dK, params.stride_dK)(_, _, bidh, !is_varlen ? bidb : 0);
+            Tensor gdK = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mdK), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+            Tensor mdV = make_tensor(make_gmem_ptr(params.ptr_dV), params.shape_dV, params.stride_dV)(_, _, bidh, !is_varlen ? bidb : 0);
+            Tensor gdV = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mdV), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+
+            GmemTiledCopydKV gmem_tiled_copy_dKV;
+            auto gmem_thr_copy_dKV = gmem_tiled_copy_dKV.get_thread_slice(thread_idx);
+            Tensor tdKVgdV = gmem_thr_copy_dKV.partition_D(gdV);
+            Tensor tdKVsdV = gmem_thr_copy_dKV.partition_S(sdV); // (TMA, TMA_M, TMA_K)
+            Tensor tdKVgdK = gmem_thr_copy_dKV.partition_D(gdK);
+            Tensor tdKVsdK = gmem_thr_copy_dKV.partition_S(sdK); // (TMA, TMA_M, TMA_K)
+            Tensor tdKVrdV = make_fragment_like(tdKVgdV);
+            Tensor tdKVrdK = make_fragment_like(tdKVgdK);
+            Tensor cdKV = cute::make_identity_tensor(select<1, 2>(TileShape_MNK{}));  // (BLK_N,BLK_K) -> (blk_n,blk_k)
+            // Repeat the partitioning with identity layouts
+            Tensor tdKVcdKV = gmem_thr_copy_dKV.partition_D(cdKV);
+            Tensor tdKVpdV = make_tensor<bool>(make_shape(size<2>(tdKVgdV)));
+            Tensor tdKVpdK = make_tensor<bool>(make_shape(size<2>(tdKVgdK)));
+            #pragma unroll
+            for (int k = 0; k < size(tdKVpdV); ++k) { tdKVpdV(k) = get<1>(tdKVcdKV(_0{}, _0{}, k)) < get<1>(params.shape_dV); }
+            #pragma unroll
+            for (int k = 0; k < size(tdKVpdK); ++k) { tdKVpdK(k) = get<1>(tdKVcdKV(_0{}, _0{}, k)) < get<1>(params.shape_dK); }
+            // Need to check OOB when reading from smem if kBlockN isn't evenly tiled
+            static constexpr bool EvenN = kBlockN % CUTE_STATIC_V(size<0>(GmemLayoutAtom{})) == 0;
+            flash::copy</*Is_even_MN=*/EvenN, /*Is_even_K=*/true, /*Clear_OOB_MN=*/false>(
+                gmem_tiled_copy_dKV, tdKVsdV, tdKVrdV, tdKVcdKV, tdKVpdV, kBlockN);
+            flash::copy</*Is_even_MN=*/EvenN, /*Is_even_K=*/true, /*Clear_OOB_MN=*/false>(
+                gmem_tiled_copy_dKV, tdKVsdK, tdKVrdK, tdKVcdKV, tdKVpdK, kBlockN);
+            // // Tell warp 0 that smem_k and smem_v are ready
+            // cutlass::arch::fence_view_async_shared(); // ensure smem reads are done before next TMA to smem_k/v
+            // flash::named_barrier_arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast<uint32_t>(BwdNamedBarriers::KVEmpty) /*id*/);
+            // Construct identity layout for gdKV
+            // Clear_OOB_K must be false since we don't want to write zeros to gmem
+            flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+                gmem_tiled_copy_dKV, tdKVrdV, tdKVgdV, tdKVcdKV, tdKVpdV, std::min(seqlen_info.seqlen - n_block * kBlockN, kBlockN)
+            );
+            flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+                gmem_tiled_copy_dKV, tdKVrdK, tdKVgdK, tdKVcdKV, tdKVpdK, std::min(seqlen_info.seqlen - n_block * kBlockN, kBlockN)
+            );
+        }
+    }
+
+    CUTLASS_DEVICE void
+    store_tail() {
+        // if constexpr (Use_TMA) { tma_store_wait<0>(); }
+    }
+
+    // Write 0 to dK and dV
+    CUTLASS_DEVICE void
+    store_zero(
+         Params const& params,
+         int thread_idx,
+         cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+         ) {
+        static constexpr int kBlockN = get<1>(TileShape_MNK{});
+        auto [n_block, bidh, bidb] = block_coord;
+        flash::SeqlenInfo<Varlen, kBlockN> seqlen_info{bidb, size<0>(params.shape_dK), params.cu_seqlens, params.seqused};
+        bool const is_varlen = Varlen && params.cu_seqlens;
+        Tensor mdK = make_tensor(make_gmem_ptr(params.ptr_dK), params.shape_dK, params.stride_dK)(_, _, bidh, !is_varlen ? bidb : 0);
+        Tensor gdK = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mdK), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+        Tensor mdV = make_tensor(make_gmem_ptr(params.ptr_dV), params.shape_dV, params.stride_dV)(_, _, bidh, !is_varlen ? bidb : 0);
+        Tensor gdV = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mdV), select<1, 2>(TileShape_MNK{}), make_coord(n_block, _0{}));  // (M, K)
+
+        GmemTiledCopydKV gmem_tiled_copy_dKV;
+        auto gmem_thr_copy_dKV = gmem_tiled_copy_dKV.get_thread_slice(thread_idx);
+        Tensor tdKVgdK = gmem_thr_copy_dKV.partition_D(gdK);
+        Tensor tdKVgdV = gmem_thr_copy_dKV.partition_D(gdV);
+        Tensor tdKVrdKV = make_fragment_like(tdKVgdK);
+        clear(tdKVrdKV);
+        // Construct identity layout for gdKV
+        Tensor cdKV = cute::make_identity_tensor(select<1, 2>(TileShape_MNK{}));  // (BLK_M,BLK_K) -> (blk_m,blk_k)
+        // Repeat the partitioning with identity layouts
+        Tensor tdKVcdKV = gmem_thr_copy_dKV.partition_D(cdKV);
+        Tensor tdKVpdK = make_tensor<bool>(make_shape(size<2>(tdKVgdK)));
+        Tensor tdKVpdV = make_tensor<bool>(make_shape(size<2>(tdKVgdV)));
+        #pragma unroll
+        for (int k = 0; k < size(tdKVpdK); ++k) { tdKVpdK(k) = get<1>(tdKVcdKV(_0{}, _0{}, k)) < get<1>(params.shape_dK); }
+        #pragma unroll
+        for (int k = 0; k < size(tdKVpdV); ++k) { tdKVpdV(k) = get<1>(tdKVcdKV(_0{}, _0{}, k)) < get<1>(params.shape_dV); }
+        // Clear_OOB_K must be false since we don't want to write zeros to gmem
+        flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+            gmem_tiled_copy_dKV, tdKVrdKV, tdKVgdK, tdKVcdKV, tdKVpdK, seqlen_info.seqlen - n_block * kBlockN
+        );
+        flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+            gmem_tiled_copy_dKV, tdKVrdKV, tdKVgdV, tdKVcdKV, tdKVpdV, seqlen_info.seqlen - n_block * kBlockN
+        );
+    }
+
+};
+
+template <class TileShape_MNK_, class ElementAccum, class ArchTag_,
+          int NumEpilogueThreads_, bool Varlen_, bool Deterministic>
+struct CollectiveEpilogueBwdGQA {
+
+    using TileShape_MNK = TileShape_MNK_;
+    using Element = ElementAccum;
+    using ArchTag = ArchTag_;
+    static constexpr int NumEpilogueThreads = NumEpilogueThreads_;
+    static constexpr bool Varlen = Varlen_;
+    static constexpr bool Use_TMA = ArchTag::kMinComputeCapability >= 90;
+
+    static_assert(ArchTag::kMinComputeCapability >= 80);
+
+    static constexpr int kBlockN = get<1>(TileShape_MNK{});
+    static constexpr int kHeadDim = get<2>(TileShape_MNK{});
+    static_assert(NumEpilogueThreads % cutlass::NumThreadsPerWarp == 0, "NumEpilogueThreads must be a multiple of NumThreadsPerWarp");
+    static constexpr int NumWarpGroups = NumEpilogueThreads / cutlass::NumThreadsPerWarpGroup;
+    // Thread layout, 256 or 384 threads per row
+    // We split into NumWarpGroups so that we can use the same postprocessing kernel as dQ
+    using R2SLayoutAtomdKVaccum = Layout<Shape<Int<cutlass::NumThreadsPerWarpGroup>, Int<NumWarpGroups>>>;
+    using R2STiledCopydKVaccum = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{}, R2SLayoutAtomdKVaccum{},
+                                                         Layout<Shape < _4>>{}));  // Val layout, 4 vals per store
+    // For Sm80
+    using R2GLayoutAtomdKVaccum = Layout<Shape<Int<NumEpilogueThreads>>>;
+    using R2GTiledCopydKVaccum = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{}, R2GLayoutAtomdKVaccum{},
+                                                         Layout<Shape < _1>>{}));  // Val layout, 1 vals per store
+
+    using SmemLayoutdKVaccum = Layout<Shape<Int<kBlockN * kHeadDim / NumWarpGroups>, Int<NumWarpGroups>>>;
+    using SmemLayoutdKVaccumFlat = Layout<Shape<Int<kBlockN * kHeadDim>>>;
+
+    // Strangely without this SmemAlignment, the total smem for hdim 128 (80 x 128) is 228KB even though we
+    // only need 227KB. We use the same alignment as the non-GQA epilogue to avoid this issue.
+    static constexpr int SmemAlignment = kHeadDim % 64 == 0 ? 1024 : (kHeadDim % 32 == 0 ? 512 : 256);
+    struct TensorStorageTMA : cute::aligned_struct<SmemAlignment> {
+        cute::array_aligned<ElementAccum, cute::cosize_v<SmemLayoutdKVaccum>, SmemAlignment> smem_dkv;
+    };
+    struct TensorStorageSTG {
+        cute::array<ElementAccum, 0> smem_dkv;
+    };
+    using TensorStorage = std::conditional_t<Use_TMA, TensorStorageTMA, TensorStorageSTG>;
+
+    using ShapedKV = cute::Shape<int32_t, int32_t, int32_t>;  // (seqlen_k_rounded * d, head, batch)
+    using StridedKV = cute::Stride<_1, int64_t, int64_t>;
+
+    // Host side kernel arguments
+    struct Arguments {
+        ElementAccum* ptr_dKaccum;
+        ShapedKV const shape_dKaccum;
+        StridedKV const stride_dKaccum;
+        ElementAccum* ptr_dVaccum;
+        ShapedKV const shape_dVaccum;
+        StridedKV const stride_dVaccum;
+        int num_heads_q;
+        int* dk_semaphore;
+        int* dv_semaphore;
+        int const* cu_seqlens;
+        int const* seqused;
+    };
+
+    // Device side kernel params
+    struct Params {
+        ElementAccum* ptr_dKaccum;
+        ShapedKV const shape_dKaccum;
+        StridedKV const stride_dKaccum;
+        ElementAccum* ptr_dVaccum;
+        ShapedKV const shape_dVaccum;
+        StridedKV const stride_dVaccum;
+        cutlass::FastDivmod qhead_per_khead_divmod;
+        int* dk_semaphore;
+        int* dv_semaphore;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    static Params
+    to_underlying_arguments(Arguments const& args) {
+        if constexpr (Deterministic) {
+            assert(args.dk_semaphore != nullptr);
+            assert(args.dv_semaphore != nullptr);
+        }
+        return {args.ptr_dKaccum, args.shape_dKaccum, args.stride_dKaccum, args.ptr_dVaccum, args.shape_dVaccum, args.stride_dVaccum,
+                cutlass::FastDivmod(cute::ceil_div(args.num_heads_q, get<1>(args.shape_dKaccum))),
+                args.dk_semaphore, args.dv_semaphore,
+                args.cu_seqlens, args.seqused};
+    }
+
+    /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
+    CUTLASS_DEVICE
+    static void prefetch_tma_descriptors(Params const& params) {
+    }
+
+    template <typename SharedStorage, typename FrgTensorO, typename TiledMma>
+    CUTLASS_DEVICE void
+    store(Params const& params,
+          FrgTensorO const& tdKrdK,
+          FrgTensorO const& tdVrdV,
+          SharedStorage& shared_storage,
+          TiledMma tiled_mma,
+          int thread_idx,
+          cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+          ) {
+
+        auto [n_block, bidh, bidb] = block_coord;
+        int bidh_idx_in_group;
+        int bidh_kv = params.qhead_per_khead_divmod.divmod(bidh_idx_in_group, bidh);
+        Tensor sdKV = make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dkv.data()), SmemLayoutdKVaccum{});
+        Tensor sdKV_flat = make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_dkv.data()), SmemLayoutdKVaccumFlat{});
+        static constexpr int dKV_TMA_num_bytes = CUTE_STATIC_V(size(sdKV_flat)) * sizeof(ElementAccum);
+
+        flash::SeqlenInfo<Varlen, kBlockN> seqlen_info{bidb, size<0>(params.shape_dKaccum), params.cu_seqlens, params.seqused};
+        bool const is_varlen = Varlen && params.cu_seqlens;
+        Tensor mdKaccum = make_tensor(make_gmem_ptr(params.ptr_dKaccum), params.shape_dKaccum, params.stride_dKaccum)(_, bidh_kv, !is_varlen ? bidb : 0);
+        Tensor mdVaccum = make_tensor(make_gmem_ptr(params.ptr_dVaccum), params.shape_dVaccum, params.stride_dVaccum)(_, bidh_kv, !is_varlen ? bidb : 0);
+        Tensor gdKaccum = local_tile(domain_offset(make_coord(seqlen_info.offset_padded * kHeadDim), mdKaccum), Shape<Int<kBlockN * kHeadDim>>{}, make_coord(n_block));  // (M * K)
+        Tensor gdVaccum = local_tile(domain_offset(make_coord(seqlen_info.offset_padded * kHeadDim), mdVaccum), Shape<Int<kBlockN * kHeadDim>>{}, make_coord(n_block));  // (M * K)
+
+        R2STiledCopydKVaccum r2s_tiled_copy_dKVaccum;
+        auto r2s_thr_copy_dKVaccum = r2s_tiled_copy_dKVaccum.get_thread_slice(thread_idx);
+        Tensor tdKVsdKVaccum = r2s_thr_copy_dKVaccum.partition_D(sdKV);
+
+        // Only used if !Use_TMA
+        R2GTiledCopydKVaccum r2g_tiled_copy_dKVaccum;
+        auto r2g_thr_copy_dKVaccum = r2g_tiled_copy_dKVaccum.get_thread_slice(thread_idx);
+
+        // Make sure all WGs have finished reading K and V, otherwise we get racy dQ
+        // because smem_q could be changed.
+        flash::named_barrier_sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+        if constexpr (Use_TMA) {
+            Tensor taccdKVrdV = r2s_thr_copy_dKVaccum.retile_S(tdVrdV); // ((Atom,AtomNum), MMA_M, MMA_N)
+            cute::copy(r2s_tiled_copy_dKVaccum, taccdKVrdV, tdKVsdKVaccum);
+        }
+
+        // int const num_batch = params.num_batch;
+        int const num_batch = get<2>(params.shape_dKaccum);
+        int const num_head_kv = get<1>(params.shape_dKaccum);
+        int *lock_ptr = !Deterministic ? nullptr : params.dv_semaphore + bidb * num_head_kv + bidh_kv;
+        using Barrier = cutlass::GenericBarrier<cutlass::detail::SyncwarpSync>;
+
+        // if (thread_idx == 0) { printf("blockIdx.x = %d, blockIdx.y = %d, blockIdx.z = %d, bidb = %d, bidh_kv = %d, lock_ptr = %p, dv_semaphore = %p, num_batch = %d, num_head_kv = %d, n_block = %d, bihd_idx_in_group = %d\n", blockIdx.x, blockIdx.y, blockIdx.z, bidb, bidh_kv, lock_ptr, params.dv_semaphore, num_batch, num_head_kv, n_block, bidh_idx_in_group);}
+
+        if constexpr (Deterministic) {
+            Barrier::wait_eq(lock_ptr, thread_idx, n_block * num_batch * num_head_kv, bidh_idx_in_group);
+        }
+        // if (thread_idx == 0) { printf("After barrier blockIdx.x = %d, blockIdx.y = %d, blockIdx.z = %d, bidb = %d, bidh_kv = %d, lock_ptr = %p, dv_semaphore = %p\n", blockIdx.x, blockIdx.y, blockIdx.z, bidb, bidh_kv, lock_ptr, params.dv_semaphore);}
+        if constexpr (Use_TMA) {
+            cutlass::arch::fence_view_async_shared();
+            cutlass::arch::NamedBarrier::sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            if (thread_idx == 0) {
+                SM90_BULK_REDUCE_ADD::copy(raw_pointer_cast(sdKV_flat.data()), raw_pointer_cast(gdVaccum.data()), dKV_TMA_num_bytes, static_cast<uint64_t>(TMA::CacheHintSm90::EVICT_LAST));
+                tma_store_arrive();
+                tma_store_wait<0>();
+            }
+        } else {
+            Tensor tdVrdV_atomic = r2g_thr_copy_dKVaccum.retile_S(tdVrdV);
+            Tensor tdVgdV_atomic = r2g_thr_copy_dKVaccum.partition_D(gdVaccum);
+            static_assert(CUTE_STATIC_V(size(tdVrdV_atomic)) == CUTE_STATIC_V(size(tdVgdV_atomic)));
+            #pragma unroll
+            for (int i = 0; i < size(tdVrdV_atomic); ++i) { atomicAdd(&tdVgdV_atomic(i), tdVrdV_atomic(i)); }
+        }
+        if constexpr (Deterministic) {
+            Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv);
+        }
+
+        if constexpr (Use_TMA) {
+            cutlass::arch::NamedBarrier::sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            Tensor taccdKVrdK = r2s_thr_copy_dKVaccum.retile_S(tdKrdK); // ((Atom,AtomNum), MMA_M, MMA_N)
+            cute::copy(r2s_tiled_copy_dKVaccum, taccdKVrdK, tdKVsdKVaccum);
+        }
+        lock_ptr = !Deterministic ? nullptr : params.dk_semaphore + bidb * num_head_kv + bidh_kv;
+        // if (thread_idx == 0) { printf("blockIdx.x = %d, blockIdx.y = %d, blockIdx.z = %d, bidb = %d, bidh_kv = %d, lock_ptr = %p, dk_semaphore = %p, num_batch = %d, num_head_kv = %d, n_block = %d, bihd_idx_in_group = %d\n", blockIdx.x, blockIdx.y, blockIdx.z, bidb, bidh_kv, lock_ptr, params.dk_semaphore, num_batch, num_head_kv, n_block, bidh_idx_in_group);}
+
+        if constexpr (Deterministic) {
+            Barrier::wait_eq(lock_ptr, thread_idx, n_block * num_batch * num_head_kv, bidh_idx_in_group);
+        }
+        // if (thread_idx == 0) { printf("After barrier blockIdx.x = %d, blockIdx.y = %d, blockIdx.z = %d, bidb = %d, bidh_kv = %d, lock_ptr = %p, dk_semaphore = %p\n", blockIdx.x, blockIdx.y, blockIdx.z, bidb, bidh_kv, lock_ptr, params.dk_semaphore);}
+        if constexpr (Use_TMA) {
+            cutlass::arch::fence_view_async_shared();
+            cutlass::arch::NamedBarrier::sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            if (thread_idx == 0) {
+                SM90_BULK_REDUCE_ADD::copy(raw_pointer_cast(sdKV_flat.data()), raw_pointer_cast(gdKaccum.data()), dKV_TMA_num_bytes, static_cast<uint64_t>(TMA::CacheHintSm90::EVICT_LAST));
+                tma_store_arrive();
+                tma_store_wait<0>();
+            }
+        } else {
+            Tensor tdKrdK_atomic = r2g_thr_copy_dKVaccum.retile_S(tdKrdK);
+            Tensor tdKgdK_atomic = r2g_thr_copy_dKVaccum.partition_D(gdKaccum);
+            static_assert(CUTE_STATIC_V(size(tdKrdK_atomic)) == CUTE_STATIC_V(size(tdKgdK_atomic)));
+            #pragma unroll
+            for (int i = 0; i < size(tdKrdK_atomic); ++i) { atomicAdd(&tdKgdK_atomic(i), tdKrdK_atomic(i)); }
+        }
+        if constexpr (Deterministic) {
+            Barrier::arrive_inc(lock_ptr, thread_idx, n_block * num_batch * num_head_kv);
+        }
+        // // Tell warp 0 that smem_k and smem_v are ready
+        // flash::named_barrier_arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp, static_cast<uint32_t>(BwdNamedBarriers::KVEmpty) /*id*/);
+    }
+
+    CUTLASS_DEVICE void
+    store_tail() {
+    }
+
+    // Write 0 to dK and dV
+    CUTLASS_DEVICE void
+    store_zero(
+         Params const& params,
+         int thread_idx,
+         cute::tuple<int32_t, int32_t, int32_t> const& block_coord
+         ) {
+        // Don't need to do anything since dKaccum and dVaccum are already zero-initialized
+    }
+
+};
+
+} // namespace flash

flash-attn/epilogue_fwd.hpp

@@ -0,0 +1,484 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cutlass/cutlass.h>
+#include <cutlass/fast_math.h>  // For FastDivMod
+#include "cute/tensor.hpp"
+
+#include "cutlass/gemm/collective/builders/sm90_common.inl"
+#include "cutlass/epilogue/collective/builders/sm90_common.inl"
+
+#include "seqlen.h"
+#include "named_barrier.hpp"
+#include "pack_gqa.h"
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class TileShape_MNK_PV_, class ClusterShape_, class Element_, class ArchTag_,
+          int NumEpilogueThreads_, bool Varlen_, bool PackGQA_, bool Split_, bool FP8PermuteCol=false>
+struct CollectiveEpilogueFwd {
+
+    using TileShape_MNK_PV = TileShape_MNK_PV_;
+    using ClusterShape = ClusterShape_;
+    using Element = Element_;
+    using ElementPartial = float;
+    using ArchTag = ArchTag_;
+    static constexpr int NumEpilogueThreads = NumEpilogueThreads_;
+    static constexpr bool Varlen = Varlen_;
+    static constexpr bool PackGQA = PackGQA_;
+    static constexpr bool Split = Split_;
+    static constexpr bool Use_smem = !(Split && !Varlen);
+    static constexpr bool Use_TMA_O = ArchTag::kMinComputeCapability >= 90 && !Varlen && !Split && !PackGQA;
+
+    static_assert(ArchTag::kMinComputeCapability >= 80);
+    static_assert(ArchTag::kMinComputeCapability >= 90 || CUTE_STATIC_V(size(ClusterShape{})) == 1);
+    static_assert(sizeof(Element) <= 2);
+
+    static constexpr int kBlockM = get<0>(TileShape_MNK_PV{});
+    static constexpr int kHeadDimV = get<1>(TileShape_MNK_PV{});
+
+    static constexpr bool LargeHeadDimV = kHeadDimV > 256;
+
+    using GmemTiledCopyOTMA = cute::SM90_TMA_STORE;
+
+    // These are for storing the output tensor without TMA (e.g., for setting output to zero)
+    static constexpr int kGmemElemsPerStore = sizeof(cute::uint128_t) / sizeof(Element);
+    static_assert(kHeadDimV % kGmemElemsPerStore == 0, "Headdim must be a multiple of kGmemElemsPerStore");
+    // We want each "row" to have 64 elements (128 bytes, i.e. 1 cache line). We want each thread to have 4 elements
+    // in the M direction and 2 elements in the K direction. In the case of PackGQA, this reduces the number of times
+    // we need to call divmod.
+    static constexpr int kBytePerRow = kHeadDimV * sizeof(Element);
+    static constexpr int kBlockKGmem = (kBytePerRow % 128 == 0 ? 128 : (kBytePerRow % 64 == 0 ? 64 : 32)) / sizeof(Element);
+    static constexpr int kGmemThreadsPerRow = kBlockKGmem / kGmemElemsPerStore;
+    // If PackGQA, we split the work of compute O_ptr among threads in the same row, so we need this to within a warp
+    static_assert(cutlass::NumThreadsPerWarp % kGmemThreadsPerRow == 0);
+    static_assert(NumEpilogueThreads % kGmemThreadsPerRow == 0, "NumEpilogueThreads must be a multiple of kGmemThreadsPerRow");
+    using GmemLayoutAtom = Layout<Shape <Int<NumEpilogueThreads / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
+                                  Stride<Int<kGmemThreadsPerRow>, _1>>;
+    static_assert(kBlockM % CUTE_STATIC_V(shape<0>(GmemLayoutAtom{})) == 0, "kBlockM must be a multiple of NumEpilogueThreads / kGmemThreadsPerRow");
+    using GmemTiledCopyO = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerStore>>>{}));  // Val layout, 8 or 16 vals per store
+
+    using SmemLayoutAtomOTMA = decltype(cutlass::gemm::collective::detail::ss_smem_selector<GMMA::Major::K, Element,
+        decltype(cute::get<0>(TileShape_MNK_PV{})), decltype(cute::get<1>(TileShape_MNK_PV{}))>());
+    using SmemLayoutOTMA = decltype(tile_to_shape(SmemLayoutAtomOTMA{}, select<0, 1>(TileShape_MNK_PV{})));
+    static constexpr int kSwizzle = kBlockKGmem == 128 ? 4 : (kBlockKGmem == 64 ? 3 : (kBlockKGmem == 32 ? 2 : 1));
+    static constexpr int kSwizzleBase = sizeof(Element) == 4 ? 2 : (sizeof(Element) == 2 ? 3 : 4);
+    using SmemLayoutAtomO = decltype(
+        composition(Swizzle<kSwizzle, kSwizzleBase, kSwizzleBase>{},
+                    Layout<Shape<_8, Int<kBlockKGmem>>,
+                           Stride<Int<kBlockKGmem>, _1>>{}));
+    using SmemLayoutOSTS = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_MNK_PV{})));
+    using SmemLayoutO = std::conditional_t<ArchTag::kMinComputeCapability >= 90, SmemLayoutOTMA, SmemLayoutOSTS>;
+
+    using ShapeO = cute::Shape<int32_t, int32_t, int32_t, int32_t, int32_t>;  // (seqlen_q, d, head, batch, num_splits)
+    using StrideO = cute::Stride<int64_t, _1, int64_t, int64_t, int64_t>;
+    using StrideLSE = cute::Stride<_1, int64_t, int64_t, int64_t>;            // (seqlen_q, head, batch, num_splits)
+    // ((qhead_per_khead, seqlen_q), d, nheads_kv, batch, num_splits)
+    using ShapeOPacked = std::conditional_t<!PackGQA, ShapeO, cute::Shape<cute::Shape<int32_t, int32_t>, int32_t, int32_t, int32_t, int32_t>>;
+    using StrideOPacked = std::conditional_t<!PackGQA, StrideO, cute::Stride<cute::Stride<int64_t, int64_t>, _1, int64_t, int64_t, int64_t>>;
+    // ((qhead_per_khead, seqlen_q), nheads_kv, batch, num_splits)
+    using ShapeLSEPacked = std::conditional_t<!PackGQA, cute::Shape<int32_t, int32_t, int32_t, int32_t>, cute::Shape<cute::Shape<int32_t, int32_t>, int32_t, int32_t, int32_t>>;
+    using StrideLSEPacked = std::conditional_t<!PackGQA, StrideLSE, cute::Stride<cute::Stride<int64_t, _1>, int64_t, int64_t, int64_t>>;
+
+    using CopyOpR2S = std::conditional_t<
+        ArchTag::kMinComputeCapability >= 90,
+        // cute::SM90_U32x4_STSM_N if Element size is 2 bytes (fp16, bf16)
+        decltype(cutlass::epilogue::collective::detail::sm90_get_smem_store_op_for_accumulator<StrideO, Element>()),
+        AutoVectorizingCopyWithAssumedAlignment<128>
+    >;
+    using SmemCopyAtomO = Copy_Atom<CopyOpR2S, Element>;
+
+    // static constexpr size_t SmemAlignmentO = cutlass::detail::alignment_for_swizzle(SmemLayoutO{});
+    // static_assert(SmemAlignmentO >= 128, "Require at least 128B alignment");
+    // struct TensorStorage : cute::aligned_struct<SmemAlignmentO> {
+    //     cute::array_aligned<Element, Use_smem ? cute::cosize_v<SmemLayoutO> : 0, SmemAlignmentO> smem_o;
+    // };
+    struct TensorStorage : cute::aligned_struct<128> {
+        cute::array_aligned<Element, Use_smem ? cute::cosize_v<SmemLayoutO> : 0> smem_o;
+    };
+
+    using TMA_O = std::conditional_t<
+        Use_TMA_O,
+        decltype(make_tma_copy(
+            GmemTiledCopyOTMA{},
+            make_tensor(make_gmem_ptr(static_cast<Element*>(nullptr)), ShapeO{}, StrideO{}),
+            SmemLayoutOTMA{},
+            select<0, 1>(TileShape_MNK_PV{}),
+            _1{})),  // no mcast for O
+        std::nullptr_t
+    >;
+
+    // Host side kernel arguments
+    struct Arguments {
+        Element* ptr_O;
+        ShapeO const shape_O;
+        StrideO const stride_O;
+        ElementPartial* ptr_O_partial;
+        StrideO const stride_O_partial;
+        float* ptr_LSE;
+        StrideLSE const stride_LSE;
+        float* ptr_LSE_partial;
+        StrideLSE const stride_LSE_partial;
+        int32_t const nheads_kv;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    // Device side kernel params
+    struct Params {
+        Element* ptr_O;
+        ShapeO const shape_O;
+        StrideO const stride_O;
+        ShapeOPacked const shape_O_packed;
+        StrideOPacked const stride_O_packed;
+        ElementPartial* ptr_O_partial;
+        StrideO const stride_O_partial;
+        StrideOPacked const stride_O_partial_packed;
+        float* ptr_LSE;
+        StrideLSE const stride_LSE;
+        ShapeLSEPacked const shape_LSE_packed;
+        StrideLSEPacked const stride_LSE_packed;
+        float* ptr_LSE_partial;
+        StrideLSE const stride_LSE_partial;
+        StrideLSEPacked const stride_LSE_partial_packed;
+        cutlass::FastDivmod qhead_per_khead_divmod;
+        TMA_O tma_store_O;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    static Params
+    to_underlying_arguments(Arguments const& args) {
+        Tensor mO = make_tensor(make_gmem_ptr(args.ptr_O), args.shape_O, args.stride_O);
+        TMA_O tma_store_O = [&]{
+            if constexpr (Use_TMA_O) {
+                return make_tma_copy(GmemTiledCopyOTMA{}, mO, SmemLayoutO{}, select<0, 1>(TileShape_MNK_PV{}), _1{}); // no mcast
+            } else {
+                return nullptr;
+            }
+        }();
+        // If PackGQA, reshape O to be ((qhead_per_khead, seqlen_q), head_size, nhead_k, batch_size, num_splits)
+        int const qhead_per_khead = !PackGQA ? 1 : cute::ceil_div(get<2>(args.shape_O), args.nheads_kv);
+        auto const shape_O_packed = cute::conditional_return<!PackGQA>(
+            args.shape_O,
+            make_shape(make_shape(qhead_per_khead, get<0>(args.shape_O)), get<1>(args.shape_O), args.nheads_kv, get<3>(args.shape_O), get<4>(args.shape_O))
+        );
+        auto const stride_O_packed = cute::conditional_return<!PackGQA>(
+            args.stride_O,
+            make_stride(make_stride(get<2>(args.stride_O), get<0>(args.stride_O)), get<1>(args.stride_O), get<2>(args.stride_O) * qhead_per_khead, get<3>(args.stride_O), get<4>(args.stride_O))
+        );
+        auto const stride_O_partial_packed = cute::conditional_return<!PackGQA>(
+            args.stride_O_partial,
+            make_stride(make_stride(get<2>(args.stride_O_partial), get<0>(args.stride_O_partial)), get<1>(args.stride_O_partial), get<2>(args.stride_O_partial) * qhead_per_khead, get<3>(args.stride_O_partial), get<4>(args.stride_O_partial))
+        );
+        // If PackGQA, Reshape LSE to be ((qhead_per_khead, seqlen_q), nhead_k, batch_size, num_splits)
+        auto const shape_LSE_packed = cute::conditional_return<!PackGQA>(
+            select<0, 2, 3, 4>(args.shape_O),
+            make_shape(make_shape(qhead_per_khead, get<0>(args.shape_O)), args.nheads_kv, get<3>(args.shape_O), get<4>(args.shape_O))
+        );
+        auto const stride_LSE_packed = cute::conditional_return<!PackGQA>(
+            args.stride_LSE,
+            make_stride(make_stride(get<1>(args.stride_LSE), get<0>(args.stride_LSE)), get<1>(args.stride_LSE) * qhead_per_khead, get<2>(args.stride_LSE), get<3>(args.stride_LSE))
+        );
+        auto const stride_LSE_partial_packed = cute::conditional_return<!PackGQA>(
+            args.stride_LSE_partial,
+            make_stride(make_stride(get<1>(args.stride_LSE_partial), get<0>(args.stride_LSE_partial)), get<1>(args.stride_LSE_partial) * qhead_per_khead, get<2>(args.stride_LSE_partial), get<3>(args.stride_LSE_partial))
+        );
+        return {args.ptr_O, args.shape_O, args.stride_O, shape_O_packed, stride_O_packed,
+                args.ptr_O_partial, args.stride_O_partial, stride_O_partial_packed,
+                args.ptr_LSE, args.stride_LSE, shape_LSE_packed, stride_LSE_packed,
+                args.ptr_LSE_partial, args.stride_LSE_partial, stride_LSE_partial_packed,
+                cutlass::FastDivmod(qhead_per_khead),
+                tma_store_O, args.cu_seqlens, args.seqused};
+    }
+
+    /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
+    CUTLASS_DEVICE
+    static void prefetch_tma_descriptors(Params const& params) {
+        if constexpr (Use_TMA_O) {
+            cute::prefetch_tma_descriptor(params.tma_store_O.get_tma_descriptor());
+        }
+    }
+
+    template <typename SharedStorage, typename FrgTensorO, typename FrgTensorLSE, typename TiledMma>
+    CUTLASS_DEVICE void
+    store(Params const& params,
+          FrgTensorO& tOrO,
+          FrgTensorLSE const& lse,
+          SharedStorage& shared_storage,
+          TiledMma tiled_mma,
+          int thread_idx,
+          cute::tuple<int32_t, int32_t, int32_t, int32_t> const& block_coord
+          ) {
+
+        auto [m_block, bidh, bidb, split_idx] = block_coord;
+        int num_splits = get<4>(params.shape_O_packed);
+        if constexpr (Split && Varlen) {
+            uint32_t num_splits_dynamic_u = reinterpret_cast<uint32_t const&>(split_idx) >> 16; // first 16 bits are for num_splits
+            int num_splits_dynamic = reinterpret_cast<int&>(num_splits_dynamic_u);
+            num_splits = num_splits_dynamic > 0 ? num_splits_dynamic : num_splits;
+            split_idx &= 0x0000FFFF;  // Only use the lower 16 bits of split_idx
+        }
+        bool const is_split = !Split ? false : (!Varlen ? true : num_splits > 1);
+
+        Tensor sO = make_tensor(make_smem_ptr(shared_storage.tensors.epilogue.smem_o.data()), SmemLayoutO{});
+        // Tensor sO_pi = cute::as_position_independent_swizzle_tensor(sO);
+
+        static constexpr bool NeedFP8Permute = FP8PermuteCol && (sizeof(Element) == 2 || sizeof(Element) == 4);
+        // If we will possibly need tOrO in FP32, we'd want to permute tOrO before type conversion.
+        // Otherwise we can permute after conversion.
+        if constexpr (NeedFP8Permute && Split) { flash::permute_output_fp8_Vcolmajor(tOrO); }
+        Tensor tOrO_out = make_tensor_like<Element>(tOrO);
+        flash::convert_type_out(tOrO, tOrO_out);
+        if constexpr (NeedFP8Permute && !Split) { flash::permute_output_fp8_Vcolmajor(tOrO_out); }
+
+        // Make sure all WGs have finished reading V
+        // Technically we don't need this if we're not using smem, but the mainloop makes the assumption that
+        // all epilogue threads sync at least once during the epilogue (so that we can start loading Q with
+        // cp.async if we need).
+        flash::named_barrier_sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+
+        // Step 1: Write O from rmem -> smem
+        if constexpr (Use_smem) {
+            auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
+            auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
+            Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);        // ((Atom,AtomNum), MMA_M, MMA_N)
+            Tensor taccOsO = smem_thr_copy_O.partition_D(sO);     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+            // Tensor taccOsO = smem_thr_copy_O.partition_D(sO_pi);     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+            cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
+            if constexpr (Use_TMA_O) {
+                cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
+                cutlass::arch::NamedBarrier::arrive(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
+                                                    cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            } else {
+                flash::named_barrier_sync(NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+            }
+        } else {
+            if constexpr (ArchTag::kMinComputeCapability >= 90) {
+                #pragma unroll
+                for (uint32_t cta_id = 0; cta_id < size(ClusterShape{}); ++cta_id) {
+                    shared_storage.pipelines.barrier_O.arrive(cta_id);
+                }
+            }
+        }
+
+        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{bidb, size<0>(params.shape_O), params.cu_seqlens, params.seqused};
+        bool is_varlen = Varlen && params.cu_seqlens;
+        int offset_o = seqlen_info.offset;
+        int seqlen_o = seqlen_info.seqlen;
+        int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / cutlass::NumThreadsPerWarpGroup, 0);
+
+        // Step 2: Write LSE from rmem -> gmem
+        auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
+        // (MMA,MMA_M,MMA_K)
+        Tensor taccOcO = thread_mma.partition_C(cute::make_identity_tensor(select<0, 1>(TileShape_MNK_PV{})));
+        static_assert(decltype(size<0, 0>(taccOcO))::value == 2);
+        static_assert(decltype(size<0, 1>(taccOcO))::value == 2);
+        Tensor taccOcO_rowcol = make_tensor(taccOcO.data(), flash::convert_layout_acc_rowcol(taccOcO.layout()));
+        Tensor taccOcO_row = taccOcO_rowcol(_, _0{});
+        CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row));                     // MMA_M
+
+        using PackGQA_t = flash::PackGQAManager<get<0>(TileShape_MNK_PV{}), get<1>(TileShape_MNK_PV{}), NumEpilogueThreads, Element>;
+        using PackGQApartial_t = flash::PackGQAManager<get<0>(TileShape_MNK_PV{}), get<1>(TileShape_MNK_PV{}), NumEpilogueThreads, ElementPartial>;
+
+        Tensor mLSE = make_tensor(make_gmem_ptr((!is_split ? params.ptr_LSE : params.ptr_LSE_partial) + offset_o * get<0>(!is_split ? params.stride_LSE : params.stride_LSE_partial)),
+                                  params.shape_LSE_packed,
+                                  !is_split ? params.stride_LSE_packed : params.stride_LSE_partial_packed)(_, bidh, !is_varlen ? bidb : 0, !is_split ? 0 : split_idx);
+        // if (thread_idx == 0) { printf("Before LSE write, m_block: %d, bidh: %d, bidb: %d, split_idx: %d, offset_o: %d, seqlen_o: %d\n", m_block, bidh, bidb, split_idx, offset_o, seqlen_o); print(mLSE); printf("\n"); }
+        if (!LargeHeadDimV || warp_group_idx == 0) {
+            if constexpr (!PackGQA) {
+                #pragma unroll
+                for (int mi = 0; mi < size(lse); ++mi) {
+                    int const row = m_block * kBlockM + get<0>(taccOcO_row(mi));
+                    if (get<1>(taccOcO_row(_0{})) == 0 && row < seqlen_o) { mLSE(row) = lse(mi); }
+                }
+            } else {
+                PackGQA_t::store_LSE(mLSE, lse, tiled_mma, params.qhead_per_khead_divmod, thread_idx, seqlen_o, m_block);
+            }
+        }
+
+        // Step 3: Write O from smem -> gmem
+        if constexpr (Use_TMA_O) {
+            Tensor mO = params.tma_store_O.get_tma_tensor(params.shape_O)(_, _, bidh, bidb, split_idx);
+            Tensor gO = local_tile(mO, select<0, 1>(TileShape_MNK_PV{}), make_coord(m_block, _0{}));  // (M, K)
+            auto block_tma_O = params.tma_store_O.get_slice(_0{});
+            Tensor tOgO = block_tma_O.partition_D(gO);  // (TMA, TMA_M, TMA_K)
+            Tensor tOsO = block_tma_O.partition_S(sO); // (TMA, TMA_M, TMA_K)
+            int warp_idx_sync = __shfl_sync(0xffffffff, thread_idx / cutlass::NumThreadsPerWarp, 0);
+            if (warp_idx_sync == NumEpilogueThreads / cutlass::NumThreadsPerWarp - 1) {
+                cutlass::arch::NamedBarrier::sync(NumEpilogueThreads + cutlass::NumThreadsPerWarp,
+                                                  cutlass::arch::ReservedNamedBarriers::EpilogueBarrier);
+                if (cute::elect_one_sync()) {
+                    cute::copy(params.tma_store_O, tOsO, tOgO);
+                    tma_store_arrive();
+                    tma_store_wait<0>();
+                    #pragma unroll
+                    for (uint32_t cta_id = 0; cta_id < size(ClusterShape{}); ++cta_id) {
+                        shared_storage.pipelines.barrier_O.arrive(cta_id);
+                    }
+                }
+            }
+        } else {  // Don't use TMA in Varlen case since we don't want to overwrite the output of another sequence
+            if (!is_split) {
+                Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + offset_o * get<0>(params.stride_O)), params.shape_O_packed, params.stride_O_packed)(_, _, bidh, !is_varlen ? bidb : 0, _0{});
+                Tensor gO = local_tile(mO, select<0, 1>(TileShape_MNK_PV{}), make_coord(m_block, _0{}));  // (M, K)
+                // if (thread_idx == 0) { printf("Before O write, m_block: %d, bidh: %d, bidb: %d, split_idx: %d, offset_o: %d, seqlen_o: %d, mO_addr = %p, addr diff = %d\n", m_block, bidh, bidb, split_idx, offset_o, seqlen_o, mO.data(), reinterpret_cast<int>(&mO(0)) - reinterpret_cast<int>(params.ptr_O)); }
+                GmemTiledCopyO gmem_tiled_copy_O;
+                auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
+                Tensor tOsO = gmem_thr_copy_O.partition_S(sO);        // ((Atom,AtomNum),ATOM_M,ATOM_N)
+                // Tensor tOsO = gmem_thr_copy_O.partition_S(sO_pi);        // ((Atom,AtomNum),ATOM_M,ATOM_N)
+                Tensor tOrO = make_fragment_like(tOsO);
+                cute::copy(gmem_tiled_copy_O, tOsO, tOrO);
+                if constexpr (ArchTag::kMinComputeCapability >= 90) {
+                    cutlass::arch::fence_view_async_shared(); // ensure smem reads are done before next TMA to smem_v
+                    #pragma unroll
+                    for (uint32_t cta_id = 0; cta_id < size(ClusterShape{}); ++cta_id) {
+                        shared_storage.pipelines.barrier_O.arrive(cta_id);
+                    }
+                }
+                if constexpr (!PackGQA) {
+                    // (BLK_M,BLK_K) -> (blk_m,blk_k)
+                    Tensor tOcO = gmem_thr_copy_O.partition_D(cute::make_identity_tensor(select<0, 1>(TileShape_MNK_PV{})));
+                    Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOsO)));
+                    #pragma unroll
+                    for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(params.shape_O); }
+                    Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
+                    // Clear_OOB_K must be false since we don't want to write zeros to gmem
+                    flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+                        gmem_tiled_copy_O, tOrO, tOgO, tOcO, tOpO, seqlen_o - m_block * kBlockM
+                    );
+                } else {
+                    // If PackGQA, we split the work of compute O_ptr among threads in the same row
+                    PackGQA_t::store_O(mO, tOrO, params.qhead_per_khead_divmod, thread_idx, seqlen_o, m_block);
+                }
+            } else {
+                Tensor mOpartial = make_tensor(make_gmem_ptr(params.ptr_O_partial + offset_o * get<0>(params.stride_O_partial)), params.shape_O_packed, params.stride_O_partial_packed)(_, _, bidh, !is_varlen ? bidb : 0, split_idx);
+                Tensor gOpartial = local_tile(mOpartial, select<0, 1>(TileShape_MNK_PV{}), make_coord(m_block, _0{}));  // (M, K)
+                // We already arrived on barrier_O earlier if !Use_smem
+                if constexpr (Use_smem) {
+                    if constexpr (ArchTag::kMinComputeCapability >= 90) {
+                        #pragma unroll
+                        for (uint32_t cta_id = 0; cta_id < size(ClusterShape{}); ++cta_id) {
+                            shared_storage.pipelines.barrier_O.arrive(cta_id);
+                        }
+                    }
+                }
+                if constexpr (!PackGQA) {
+                    static constexpr int kGmemElemsPerStoreDirect = 2;
+                    cute::Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementPartial> gmem_copy_direct;
+                    // Reshape acc from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
+                    Tensor tOrO_rowcol = make_tensor(tOrO.data(), flash::convert_layout_acc_rowcol(tOrO.layout()));
+                    Tensor tOrO_copy = cute::tiled_divide(tOrO_rowcol, Shape<_1, Int<kGmemElemsPerStoreDirect>>{});
+                    Tensor tOgO = thread_mma.partition_C(gOpartial);
+                    Tensor tOgO_rowcol = make_tensor(tOgO.data(), flash::convert_layout_acc_rowcol(tOgO.layout()));
+                    Tensor tOgO_copy = cute::tiled_divide(tOgO_rowcol, Shape<_1, Int<kGmemElemsPerStoreDirect>>{});
+                    Tensor taccOcO_col = taccOcO_rowcol(_0{}, _);
+                    #pragma unroll
+                    for (int m = 0; m < size(taccOcO_row); ++m) {
+                        if (get<0>(taccOcO_row(m)) < seqlen_o - m_block * kBlockM) {
+                            #pragma unroll
+                            for (int k = 0; k < size(taccOcO_col) / kGmemElemsPerStoreDirect; ++k) {
+                                if (get<1>(taccOcO_col(k * kGmemElemsPerStoreDirect)) < get<1>(params.shape_O)) {
+                                    cute::copy(gmem_copy_direct, tOrO_copy(_, m, k), tOgO_copy(_, m, k));
+                                }
+                            }
+                        }
+                    }
+                } else {
+                    PackGQApartial_t::store_O_direct(mOpartial, tOrO, tiled_mma, params.qhead_per_khead_divmod, thread_idx, seqlen_o, m_block);
+                }
+            }
+        }
+    }
+
+    CUTLASS_DEVICE void
+    store_tail() {
+        // Don't need to do tma_store_wait<0>() here since we already did in @store
+    }
+
+    // Write 0 to output and -inf to LSE
+    CUTLASS_DEVICE void
+    store_zero(
+         Params const& params,
+         int thread_idx,
+         cute::tuple<int32_t, int32_t, int32_t, int32_t> const& block_coord
+         ) {
+        static constexpr int kBlockM = get<0>(TileShape_MNK_PV{});
+        auto [m_block, bidh, bidb, split_idx] = block_coord;
+        int num_splits = get<4>(params.shape_O_packed);
+        if constexpr (Split && Varlen) {
+            uint32_t num_splits_dynamic_u = reinterpret_cast<uint32_t const&>(split_idx) >> 16; // first 16 bits are for num_splits
+            int num_splits_dynamic = reinterpret_cast<int&>(num_splits_dynamic_u);
+            num_splits = num_splits_dynamic > 0 ? num_splits_dynamic : num_splits;
+            split_idx &= 0x0000FFFF;  // Only use the lower 16 bits of split_idx
+        }
+        bool const is_split = !Split ? false : (!Varlen ? true : num_splits > 1);
+
+        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{bidb, size<0>(params.shape_O), params.cu_seqlens, params.seqused};
+        bool const is_varlen = Varlen && params.cu_seqlens;
+        int offset_o = seqlen_info.offset;
+        int seqlen_o = seqlen_info.seqlen;
+        int qhead_per_khead = !PackGQA ? 1 : params.qhead_per_khead_divmod.divisor;
+        Tensor mLSE = make_tensor(make_gmem_ptr((!is_split ? params.ptr_LSE : params.ptr_LSE_partial) + offset_o * get<0>(!is_split ? params.stride_LSE : params.stride_LSE_partial)),
+                                  params.shape_LSE_packed,
+                                  !is_split ? params.stride_LSE_packed : params.stride_LSE_partial_packed)(_, bidh, !is_varlen ? bidb : 0, !is_split ? 0 : split_idx);
+        Tensor gLSE = local_tile(mLSE, Shape<Int<kBlockM>>{}, make_coord(m_block));
+
+        static_assert(kBlockM <= NumEpilogueThreads);
+        if (thread_idx < kBlockM) {
+            const int row = m_block * kBlockM + thread_idx;
+            if constexpr (!PackGQA) {
+                if (row < seqlen_o) { mLSE(row) = -INFINITY; }
+            } else {
+                if (row < seqlen_o * qhead_per_khead) {
+                    int m_idx, h_idx;
+                    m_idx = params.qhead_per_khead_divmod.divmod(h_idx, row);
+                    // mLSE has shape ((qhead_per_khead, seqlen_q)) and it's unhappy with just 1 "make_coord"
+                    mLSE(make_coord(make_coord(h_idx, m_idx))) = -INFINITY;
+                }
+            }
+        }
+
+        // If split, we don't have to write 0 to mOpartial if the mha_combine kernel is used,
+        // since it will not use the value of O if LSE is -inf.
+        if (!is_split) {
+            Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + offset_o * get<0>(params.stride_O)), params.shape_O_packed, params.stride_O_packed)(_, _, bidh, !is_varlen ? bidb : 0, _0{});
+
+            GmemTiledCopyO gmem_tiled_copy_O;
+            auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
+            Tensor tOcO = gmem_thr_copy_O.partition_D(cute::make_identity_tensor(select<0, 1>(TileShape_MNK_PV{})));
+            if constexpr (!PackGQA) {
+                Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOcO)));
+                #pragma unroll
+                for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(params.shape_O); }
+                Tensor gO = local_tile(mO, select<0, 1>(TileShape_MNK_PV{}), make_coord(m_block, _0{}));  // (M, K)
+                Tensor tOgO = gmem_thr_copy_O.partition_D(gO);
+                Tensor tOrO = make_fragment_like(tOgO);
+                cute::clear(tOrO);
+                // Clear_OOB_K must be false since we don't want to write zeros to gmem
+                flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+                    gmem_tiled_copy_O, tOrO, tOgO, tOcO, tOpO, seqlen_o - m_block * kBlockM
+                );
+            } else {
+                // If PackGQA, we split the work of compute O_ptr among threads in the same row
+                using PackGQA_t = flash::PackGQAManager<get<0>(TileShape_MNK_PV{}), get<1>(TileShape_MNK_PV{}), NumEpilogueThreads, Element>;
+                Tensor tOrO = make_tensor<Element>(make_shape(Shape<_1, Int<kGmemElemsPerStore>>{}, size<1>(tOcO), size<2>(tOcO)));
+                cute::clear(tOrO);
+                PackGQA_t::store_O(mO, tOrO, params.qhead_per_khead_divmod, thread_idx, seqlen_o, m_block);
+            }
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash.h

@@ -0,0 +1,218 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda.h>
+#include <vector>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Qkv_params {
+    using index_t = int64_t;
+    // The QKV matrices.
+    void *__restrict__ q_ptr;
+    void *__restrict__ k_ptr;
+    void *__restrict__ v_ptr;
+
+    // The stride between rows of the Q, K and V matrices.
+    index_t q_batch_stride;
+    index_t k_batch_stride;
+    index_t v_batch_stride;
+    index_t q_row_stride;
+    index_t k_row_stride;
+    index_t v_row_stride;
+    index_t q_head_stride;
+    index_t k_head_stride;
+    index_t v_head_stride;
+    index_t v_dim_stride;
+
+    // The number of heads.
+    int h, h_k;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Flash_fwd_params : public Qkv_params {
+    using index_t = int64_t;
+
+    // The O matrix (output).
+    void * __restrict__ o_ptr;
+    void * __restrict__ oaccum_ptr;
+
+    // The stride between rows of O.
+    index_t o_batch_stride;
+    index_t o_row_stride;
+    index_t o_head_stride;
+
+    // The pointer to the softmax sum.
+    void * __restrict__ softmax_lse_ptr;
+    void * __restrict__ softmax_lseaccum_ptr;
+
+    // For FP8 scaling
+    float * __restrict__ q_descale_ptr;
+    float * __restrict__ k_descale_ptr;
+    float * __restrict__ v_descale_ptr;
+    index_t q_descale_batch_stride;
+    index_t q_descale_head_stride;
+    index_t k_descale_batch_stride;
+    index_t k_descale_head_stride;
+    index_t v_descale_batch_stride;
+    index_t v_descale_head_stride;
+
+    // The dimensions.
+    int b, seqlen_q, seqlen_k, seqlen_knew, d, seqlen_q_rounded, seqlen_k_rounded, d_rounded, rotary_dim;
+    int total_q, total_k, total_knew;
+    int b_k;  // When having KV cache and with cache_batch_idx, K & V might have larger batch size than Q
+    int dv, dv_rounded;  // For the case where V headdim is different from Q/K headdim
+
+    // The scaling factors for the kernel.
+    float scale_softmax;
+    float softcap;
+
+    // array of length b+1 holding starting offset of each sequence.
+    int * __restrict__ cu_seqlens_q;
+    int * __restrict__ cu_seqlens_k;
+    int * __restrict__ cu_seqlens_knew;
+    int * __restrict__ leftpad_k;
+
+    // If provided, the actual length of each q/k sequence.
+    int *__restrict__ seqused_q;
+    int *__restrict__ seqused_k;
+
+    // The stride between rows of Oaccum.
+    index_t oaccum_split_stride;
+    index_t oaccum_batch_stride;
+    index_t oaccum_row_stride;
+    index_t oaccum_head_stride;
+
+    // The stride between rows of LSEaccum.
+    index_t lseaccum_split_stride;
+    index_t lseaccum_batch_stride;
+    index_t lseaccum_head_stride;
+
+    // The K_new and V_new matrices.
+    void * __restrict__ knew_ptr;
+    void * __restrict__ vnew_ptr;
+
+    // The stride between rows of the Q, K and V matrices.
+    index_t knew_batch_stride;
+    index_t vnew_batch_stride;
+    index_t knew_row_stride;
+    index_t vnew_row_stride;
+    index_t knew_head_stride;
+    index_t vnew_head_stride;
+
+    void *__restrict__ qv_ptr;
+    index_t qv_batch_stride;
+    index_t qv_row_stride;
+    index_t qv_head_stride;
+
+    // The cos and sin matrices for rotary embedding.
+    void * __restrict__ rotary_cos_ptr;
+    void * __restrict__ rotary_sin_ptr;
+    int *__restrict__ seqlens_rotary;
+
+    // The indices to index into the KV cache.
+    int * __restrict__ kv_batch_idx;
+
+    // Paged KV cache
+    int * __restrict__ page_table;
+    index_t page_table_batch_stride;
+    int page_size;
+    int num_pages;
+    bool pagedkv_tma;
+
+    // The dropout probability (probability of keeping an activation).
+    float p_dropout;
+    // uint32_t p_dropout_in_uint;
+    // uint16_t p_dropout_in_uint16_t;
+    uint8_t p_dropout_in_uint8_t;
+
+    // Scale factor of 1 / (1 - p_dropout).
+    float rp_dropout;
+
+    // Local window size
+    int window_size_left, window_size_right;
+    int attention_chunk;
+
+    // Pointer to the RNG seed (idx 0) and offset (idx 1).
+    uint64_t * rng_state;
+
+    bool is_bf16;
+    bool is_fp32;
+    bool is_e4m3;
+    bool is_causal;
+    bool is_local;
+
+    bool is_rotary_interleaved;
+
+    int num_splits;  // For split-KV version
+    bool pack_gqa;
+
+    int * __restrict__ tile_count_semaphore;
+    // int * __restrict__ num_m_blocks_ptr;
+    // int * __restrict__ num_n_blocks_ptr;
+    int * __restrict__ num_splits_dynamic_ptr;
+    bool skip_scheduler_metadata_computation;
+
+    int arch;
+    int num_sm;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Flash_bwd_params : public Flash_fwd_params {
+    using index_t = int64_t;
+
+    // The dO and dQKV matrices.
+    void *__restrict__ do_ptr;
+    void *__restrict__ dq_ptr;
+    void *__restrict__ dk_ptr;
+    void *__restrict__ dv_ptr;
+
+    // To accumulate dQ
+    void *__restrict__ dq_accum_ptr;
+    void *__restrict__ dk_accum_ptr;
+    void *__restrict__ dv_accum_ptr;
+
+    // // To accumulate dK and dV in case we're splitting the bwd along seqlen_q
+    // dimension void *__restrict__ dk_accum_ptr; void *__restrict__
+    // dv_accum_ptr;
+
+    // The stride between rows of the dO, dQ, dK and dV matrices.
+    index_t do_batch_stride;
+    index_t do_row_stride;
+    index_t do_head_stride;
+    index_t dq_batch_stride;
+    index_t dk_batch_stride;
+    index_t dv_batch_stride;
+    index_t dq_row_stride;
+    index_t dk_row_stride;
+    index_t dv_row_stride;
+    index_t dq_head_stride;
+    index_t dk_head_stride;
+    index_t dv_head_stride;
+
+    // The pointer to the softmax d sum.
+    void *__restrict__ dsoftmax_sum;
+    void *__restrict__ softmax_lse_log2_ptr;
+
+    int *__restrict__ dq_semaphore;
+    int *__restrict__ dk_semaphore;
+    int *__restrict__ dv_semaphore;
+
+    bool deterministic;
+    index_t dq_accum_split_stride;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Arch, typename T, int kHeadDim, int kHeadDimV, bool Split, bool PagedKVNonTMA, bool Has_softcap, bool PackGQA>
+void run_mha_fwd_(Flash_fwd_params &params, cudaStream_t stream);
+void prepare_varlen_num_blocks(Flash_fwd_params &params, cudaStream_t stream, bool packgqa, int blockM, int blockN, bool enable_pdl);
+template <int Arch, typename T, int kHeadDim, bool Has_softcap>
+void run_mha_bwd_(Flash_bwd_params &params, cudaStream_t stream);
+template <typename T, typename Tpartial, int kBlockK>
+void run_mha_fwd_combine_(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);

flash-attn/flash_api.cpp

@@ -0,0 +1,1720 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#include <Python.h>
+#include <torch/nn/functional/padding.h>
+#include <ATen/cuda/CUDAContextLight.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include <cutlass/numeric_types.h>
+
+#include "flash.h"
+#include "static_switch.h"
+#include "tile_size.h"
+#include "heuristics.h"
+#include "cuda_check.h"
+
+
+extern "C" {
+/* Creates a dummy empty _C module that can be imported from Python.
+    The import from Python will load the .so consisting of this file
+    in this extension, so that the TORCH_LIBRARY static initializers
+    below are run. */
+PyObject* PyInit__C(void)
+{
+    static struct PyModuleDef module_def = {
+        PyModuleDef_HEAD_INIT,
+        "_C",   /* name of module */
+        NULL,   /* module documentation, may be NULL */
+        -1,     /* size of per-interpreter state of the module,
+                    or -1 if the module keeps state in global variables. */
+        NULL,   /* methods */
+    };
+    return PyModule_Create(&module_def);
+}
+}
+
+#define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA")
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+void set_params_fprop(Flash_fwd_params &params,
+                      // sizes
+                      const size_t b,
+                      const size_t seqlen_q,
+                      const size_t seqlen_k,
+                      const size_t seqlen_q_rounded,
+                      const size_t seqlen_k_rounded,
+                      const size_t h,
+                      const size_t h_k,
+                      const size_t d,
+                      const size_t d_rounded,
+                      // device pointers
+                      const at::Tensor q,
+                      const at::Tensor k,
+                      const at::Tensor v,
+                      at::Tensor out,
+                      void *cu_seqlens_q_d,
+                      void *cu_seqlens_k_d,
+                      void *seqused_q,
+                      void *seqused_k,
+                      void *softmax_lse_d,
+                      float p_dropout,
+                      float softmax_scale,
+                      int window_size_left,
+                      int window_size_right,
+                      int attention_chunk,
+                      const float softcap=0.f,
+                      const int sm_margin=0) {
+
+    // Reset the parameters
+    params = {};
+
+    params.is_bf16 = q.dtype() == torch::kBFloat16;
+    params.is_e4m3 = q.dtype() == torch::kFloat8_e4m3fn;
+
+    // Set the pointers and strides.
+    params.q_ptr = q.data_ptr();
+    params.k_ptr = k.data_ptr();
+    params.v_ptr = v.data_ptr();
+    // All stride are in elements, not bytes.
+    params.q_row_stride = q.stride(-3);
+    params.k_row_stride = k.stride(-3);
+    params.v_row_stride = v.stride(-3);
+    params.q_head_stride = q.stride(-2);
+    params.k_head_stride = k.stride(-2);
+    params.v_head_stride = v.stride(-2);
+    params.v_dim_stride = v.stride(-1);
+    params.o_ptr = out.data_ptr();
+    params.o_row_stride = out.stride(-3);
+    params.o_head_stride = out.stride(-2);
+
+    if (cu_seqlens_q_d == nullptr) {
+        params.q_batch_stride = q.stride(0);
+        params.o_batch_stride = out.stride(0);
+    }
+    if (cu_seqlens_k_d == nullptr) {
+        params.k_batch_stride = k.stride(0);
+        params.v_batch_stride = v.stride(0);
+    }
+
+    params.cu_seqlens_q = static_cast<int *>(cu_seqlens_q_d);
+    params.cu_seqlens_k = static_cast<int *>(cu_seqlens_k_d);
+    params.seqused_q = static_cast<int *>(seqused_q);
+    params.seqused_k = static_cast<int *>(seqused_k);
+
+    // Softmax sum
+    params.softmax_lse_ptr = softmax_lse_d;
+
+    // Set the dimensions.
+    params.b = b;
+    params.h = h;
+    params.h_k = h_k;
+    params.seqlen_q = seqlen_q;
+    params.seqlen_k = seqlen_k;
+    params.seqlen_q_rounded = seqlen_q_rounded;
+    params.seqlen_k_rounded = seqlen_k_rounded;
+    params.d = d;
+    params.d_rounded = d_rounded;
+
+    // Set the different scale values.
+    params.scale_softmax = softmax_scale;
+    params.softcap = softcap;
+
+    // Set this to probability of keeping an element to simplify things.
+    params.p_dropout = 1.f - p_dropout;
+    // Convert p from float to int so we don't have to convert the random uint to float to compare.
+    // [Minor] We want to round down since when we do the comparison we use <= instead of <
+    // params.p_dropout_in_uint = uint32_t(std::floor(params.p_dropout * 4294967295.0));
+    // params.p_dropout_in_uint16_t = uint16_t(std::floor(params.p_dropout * 65535.0));
+    params.p_dropout_in_uint8_t = uint8_t(std::floor(params.p_dropout * 255.0));
+    params.rp_dropout = 1.f / params.p_dropout;
+    TORCH_CHECK(p_dropout < 1.f);
+    #ifdef FLASHATTENTION_DISABLE_DROPOUT
+        TORCH_CHECK(p_dropout == 0.0f, "This flash attention build does not support dropout.");
+    #endif
+
+    // Causal is the special case where window_size_right == 0 and window_size_left < 0.
+    // Local is the more general case where window_size_right >= 0 or window_size_left >= 0.
+    params.is_causal = window_size_left < 0 && window_size_right == 0 && attention_chunk == 0;
+    params.is_local = (window_size_left >= 0 || window_size_right >= 0 || attention_chunk >= 1) && !params.is_causal;
+
+    // TODO: check this
+    if (window_size_left < 0) { window_size_left = seqlen_k - 1; }
+    if (window_size_right < 0) { window_size_right = seqlen_q - 1; }
+    if (attention_chunk > 0) {
+        window_size_left = std::min(window_size_left, attention_chunk - 1);
+        window_size_right = std::min(window_size_right, attention_chunk - 1);
+    }
+    params.window_size_left = window_size_left;
+    params.window_size_right = window_size_right;
+    params.attention_chunk = attention_chunk;
+
+    params.arch = at::cuda::getCurrentDeviceProperties()->major * 10 + at::cuda::getCurrentDeviceProperties()->minor;
+    params.num_sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount - sm_margin;
+
+    #ifdef FLASHATTENTION_DISABLE_LOCAL
+        TORCH_CHECK(!params.is_local, "This flash attention build does not support local attention.");
+    #endif
+}
+
+void set_params_dgrad(Flash_bwd_params &params,
+                      // sizes
+                      const size_t b,
+                      const size_t seqlen_q,
+                      const size_t seqlen_k,
+                      const size_t seqlen_q_rounded,
+                      const size_t seqlen_k_rounded,
+                      const size_t h,
+                      const size_t h_k,
+                      const size_t d,
+                      const size_t d_rounded,
+                      // device pointers
+                      const at::Tensor q,
+                      const at::Tensor k,
+                      const at::Tensor v,
+                      const at::Tensor out,
+                      const at::Tensor dout,
+                      at::Tensor dq,
+                      at::Tensor dk,
+                      at::Tensor dv,
+                      void *cu_seqlens_q_d,
+                      void *cu_seqlens_k_d,
+                      void *seqused_q,
+                      void *seqused_k,
+                      void *dq_accum_d,
+                      void *dk_accum_d,
+                      void *dv_accum_d,
+                      void *softmax_lse_d,
+                      void *dsoftmax_sum_d,
+                      float p_dropout,
+                      float softmax_scale,
+                      int window_size_left,
+                      int window_size_right,
+                      int attention_chunk,
+                      const float softcap=0.f,
+                      bool deterministic=false,
+                      int const sm_margin=0) {
+
+    set_params_fprop(params,
+                     b, seqlen_q, seqlen_k, seqlen_q_rounded, seqlen_k_rounded, h, h_k, d, d_rounded,
+                     q, k, v, out,
+                     cu_seqlens_q_d,
+                     cu_seqlens_k_d,
+                     seqused_q,
+                     seqused_k,
+                     softmax_lse_d,
+                     p_dropout,
+                     softmax_scale,
+                     window_size_left,
+                     window_size_right,
+                     attention_chunk,
+                     softcap,
+                     sm_margin);
+
+    // Set the pointers and strides.
+    params.do_ptr = dout.data_ptr();
+    params.do_row_stride = dout.stride(-3);
+    params.do_head_stride = dout.stride(-2);
+    params.dq_ptr = dq.data_ptr();
+    params.dk_ptr = dk.data_ptr();
+    params.dv_ptr = dv.data_ptr();
+    params.dq_row_stride = dq.stride(-3);
+    params.dk_row_stride = dk.stride(-3);
+    params.dv_row_stride = dv.stride(-3);
+    params.dq_head_stride = dq.stride(-2);
+    params.dk_head_stride = dk.stride(-2);
+    params.dv_head_stride = dv.stride(-2);
+
+    if (cu_seqlens_q_d == nullptr) {
+        params.do_batch_stride = dout.stride(0);
+        params.dq_batch_stride = dq.stride(0);
+        params.dk_batch_stride = dk.stride(0);
+        params.dv_batch_stride = dv.stride(0);
+    }
+
+    params.dq_accum_ptr = dq_accum_d;
+    params.dk_accum_ptr = dk_accum_d;
+    params.dv_accum_ptr = dv_accum_d;
+
+    // Softmax sum
+    params.dsoftmax_sum = dsoftmax_sum_d;
+
+    params.deterministic = deterministic;
+}
+
+template <int Arch, int Split, bool PagedKVNonTMA, bool PackGQA, bool Has_softcap>
+void run_mha_fwd_constexpr(Flash_fwd_params &params, cudaStream_t stream) {
+    if (!params.is_e4m3) {
+        if (params.is_bf16) {
+            #ifndef FLASHATTENTION_DISABLE_HDIM64
+            if (params.d <= 64) {
+                if constexpr (Arch == 90) {
+                    if (params.dv > 256) {
+                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 64, 512, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    } else if (params.dv > 64) {
+                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 64, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    }
+                }
+                return run_mha_fwd_<Arch, cutlass::bfloat16_t, 64, 64, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+            }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM96
+            if (params.d <= 96) { return run_mha_fwd_<Arch, cutlass::bfloat16_t, 96, 96, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM128
+            if (params.d <= 128) { return run_mha_fwd_<Arch, cutlass::bfloat16_t, 128, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM192
+            if (params.d <= 192) {
+                if constexpr (Arch == 90) {
+                    if (params.dv <= 128) {
+                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    }
+                }
+                return run_mha_fwd_<Arch, cutlass::bfloat16_t, 192, 192, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+            }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM256
+            if (params.d <= 256) { return run_mha_fwd_<Arch, cutlass::bfloat16_t, 256, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+        } else {
+            #ifndef FLASHATTENTION_DISABLE_FP16
+            #ifndef FLASHATTENTION_DISABLE_HDIM64
+            if (params.d <= 64) {
+                if constexpr (Arch == 90) {
+                    if (params.dv > 256) {
+                        return run_mha_fwd_<Arch, cutlass::half_t, 64, 512, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    } else if (params.dv > 64) {
+                        return run_mha_fwd_<Arch, cutlass::half_t, 64, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    }
+                }
+                return run_mha_fwd_<Arch, cutlass::half_t, 64, 64, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+            }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM96
+            if (params.d <= 96) { return run_mha_fwd_<Arch, cutlass::half_t, 96, 96, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM128
+            if (params.d <= 128) { return run_mha_fwd_<Arch, cutlass::half_t, 128, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM192
+            if (params.d <= 192) {
+                if constexpr (Arch == 90) {
+                    if (params.dv <= 128) {
+                        return run_mha_fwd_<Arch, cutlass::half_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                    }
+                }
+                return run_mha_fwd_<Arch, cutlass::half_t, 192, 192, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+            }
+            #endif
+            #ifndef FLASHATTENTION_DISABLE_HDIM256
+            if (params.d <= 256) { return run_mha_fwd_<Arch, cutlass::half_t, 256, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+            #endif
+            #else
+            TORCH_CHECK(false, "This flash attention build does not support FP16.");
+            #endif
+        }
+    } else {
+        #ifndef FLASHATTENTION_DISABLE_FP8
+        #ifndef FLASHATTENTION_DISABLE_HDIM64
+        if (params.d <= 64) { return run_mha_fwd_<90, cutlass::float_e4m3_t, 64, 64, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM96
+        if (params.d <= 96) { return run_mha_fwd_<90, cutlass::float_e4m3_t, 96, 96, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM128
+        if (params.d <= 128) { return run_mha_fwd_<90, cutlass::float_e4m3_t, 128, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM192
+        if (params.d <= 192) {
+            if constexpr (Arch == 90) {
+                if (params.dv <= 128) {
+                    return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+                }
+            }
+            return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 192, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
+        }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM256
+        if (params.d <= 256) { return run_mha_fwd_<90, cutlass::float_e4m3_t, 256, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); }
+        #endif
+        #else
+        TORCH_CHECK(false, "This flash attention build does not support FP8.");
+        #endif
+    }
+}
+
+void run_mha_fwd(Flash_fwd_params &params, cudaStream_t stream) {
+    // HEADDIM_SWITCH(params.d, [&] {
+    //     run_mha_fwd_<cutlass::half_t, kHeadSize>(params, stream);
+    // });
+    TORCH_CHECK(params.num_splits >= 1);
+    ARCH_SWITCH(params.arch, Arch, [&] {
+        SPLIT_SWITCH(params.num_splits > 1, Split, [&] {
+            PAGEDKV_SWITCH(params.page_table && !params.pagedkv_tma, PagedKVNonTMA, [&] {
+                PACKGQA_SWITCH(params.pack_gqa, PackGQA_, [&] {
+                    // Always enable PackGQA for Sm8x or PagedKVNonTMA or Split to reduce compilation
+                    static constexpr bool PackGQA = PackGQA_ || Arch < 90 || PagedKVNonTMA || Split;
+                    SOFTCAP_SWITCH(params.softcap > 0.0, Has_softcap, [&] {
+                        run_mha_fwd_constexpr<Arch, Split, PagedKVNonTMA, PackGQA, Has_softcap>(params, stream);
+                    });
+                });
+            });
+        });
+    });
+}
+
+void run_mha_fwd_combine(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl=false) {
+    #ifndef FLASHATTENTION_DISABLE_SPLIT
+    // If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively
+    // so that kBlockM is smaller and we have more parallelism.
+    if (params.is_fp32) {
+        if (params.dv <= 64) {
+            run_mha_fwd_combine_<float, float, 64>(params, stream, enable_pdl);
+        } else {
+            run_mha_fwd_combine_<float, float, 128>(params, stream, enable_pdl);
+        }
+    } else if (params.is_bf16) {
+        if (params.dv <= 64) {
+            run_mha_fwd_combine_<cutlass::bfloat16_t, float, 64>(params, stream, enable_pdl);
+        } else {
+            run_mha_fwd_combine_<cutlass::bfloat16_t, float, 128>(params, stream, enable_pdl);
+        }
+    } else {
+        if (params.dv <= 64) {
+            run_mha_fwd_combine_<cutlass::half_t, float, 64>(params, stream, enable_pdl);
+        } else {
+            run_mha_fwd_combine_<cutlass::half_t, float, 128>(params, stream, enable_pdl);
+        }
+    }
+    #else
+    TORCH_CHECK(false, "This flash attention build does not support combine kernels.");
+    #endif
+}
+
+inline bool get_pagedkv_tma(Flash_fwd_params const& params) {
+    if (params.arch < 90 || !params.page_table || params.leftpad_k || params.knew_ptr) { return false; }
+    // This needs to match the kernel configs
+    auto kBlockMN_kernel_args_sm90 = tile_size_fwd_sm90(params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, false /*v_colmajor*/, false /*paged_kv_non_TMA*/, params.softcap > 0.f);
+    int const kBlockM = std::get<0>(kBlockMN_kernel_args_sm90);
+    int const kBlockN = std::get<1>(kBlockMN_kernel_args_sm90);
+    // Heuristic: when seqlen_q <= kBlockM, we're not compute bound, and somehow using TMA is slower,
+    // at least for MLA.
+    return params.page_size % kBlockN == 0 && params.seqlen_q * (params.h / params.h_k) > kBlockM;
+}
+
+inline bool get_pack_gqa(Flash_fwd_params const& params) {
+    // Always enable PackGQA for Sm8x or PagedKVNonTMA or Split to reduce compilation and binary size.
+    // Has little effect on speed.
+    if (params.arch < 90 || (params.page_table && !params.pagedkv_tma) || params.num_splits > 1) { return true; }
+    #ifdef FLASHATTENTION_DISABLE_PACKGQA
+    return false;
+    #else
+    // params.page_table must already be set
+    if (params.h == params.h_k) { return false; }
+    // This needs to match the kernel configs
+    auto kBlockMN_kernel_args_sm90 = tile_size_fwd_sm90(params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, false /*v_colmajor*/, params.page_table && !params.pagedkv_tma, params.softcap > 0.f);
+    int const kBlockM = std::get<0>(kBlockMN_kernel_args_sm90);
+    return should_pack_gqa(params.cu_seqlens_q || params.seqused_q, params.seqlen_q, params.h / params.h_k, kBlockM);
+    #endif
+}
+
+inline int get_num_splits(Flash_fwd_params const& params) {
+    #ifdef FLASHATTENTION_DISABLE_SPLIT
+    return 1;
+    #else
+    // Always enable PackGQA for Split
+    // params.page_table must already be set
+    // This needs to match the kernel configs
+    bool varlen = params.cu_seqlens_q || params.cu_seqlens_k || params.seqused_q || params.seqused_k || params.leftpad_k;
+    auto kBlockMN_kernel_args_sm90 = tile_size_fwd_sm90(params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, false /*v_colmajor*/, params.page_table && !params.pagedkv_tma, params.softcap > 0.f);
+    // Strictly speaking we need to pass in (varlen && params.num_splits > 1) but num_splits
+    // has not been set here. It's OK though because we might just underestimate kBlockN a bit
+    auto kBlockMN_kernel_args_sm8x = tile_size_fwd_sm8x(params.arch == 86 || params.arch == 89, params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, params.page_table, varlen, params.softcap > 0.f, params.knew_ptr);
+    int const kBlockM = params.arch >= 90 ? std::get<0>(kBlockMN_kernel_args_sm90) : std::get<0>(kBlockMN_kernel_args_sm8x);
+    int const kBlockN = params.arch >= 90 ? std::get<1>(kBlockMN_kernel_args_sm90) : std::get<1>(kBlockMN_kernel_args_sm8x);
+    int seqlen_q_packgqa = params.seqlen_q * (params.h / params.h_k);
+    // If is_local, we're not going to load all of seqlen_k
+    int const seqlen_k_loaded = !params.is_local
+        ? params.seqlen_k
+        : std::max(0, std::min(params.seqlen_k, params.window_size_right + params.window_size_left + 1 + kBlockM));
+    int const num_n_blocks = (seqlen_k_loaded + kBlockN - 1) / kBlockN;
+    int const num_m_blocks = (seqlen_q_packgqa + kBlockM - 1) / kBlockM;
+    int const size_one_kv_head = params.seqlen_k * (params.d + params.dv) * (params.is_e4m3 ? 1 : 2);
+    // Always enable PackGQA for Split
+    // If varlen, we use dynamic split, so this heuristic just needs to get an upper bound on num_splits.
+    // We assume the case where there's 1 long sequence and the rest are short, i.e. pretending
+    // that batch = 1.
+    int total_mblocks = (params.num_splits_dynamic_ptr ? 1 : params.b) * params.h_k * num_m_blocks;
+    return num_splits_heuristic(total_mblocks, params.num_sm, num_n_blocks, num_m_blocks, size_one_kv_head, params.is_causal || params.is_local, 128);
+    #endif
+}
+
+inline int get_max_headdim() {
+    #ifndef FLASHATTENTION_DISABLE_HDIM256
+    return 256;
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM192
+    return 192;
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM128
+    return 128;
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM96
+    return 96;
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM64
+    return 64;
+    #endif
+    return 0;
+}
+
+inline int round_up_headdim(int head_size) {
+    #ifndef FLASHATTENTION_DISABLE_HDIM64
+    if (head_size <= 64) { return 64; }
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM96
+    if (head_size <= 96) { return 96; }
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM128
+    if (head_size <= 128) { return 128; }
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM192
+    if (head_size <= 192) { return 192; }
+    #endif
+    #ifndef FLASHATTENTION_DISABLE_HDIM256
+    if (head_size <= 256) { return 256; }
+    #endif
+    return 256;
+}
+
+inline int round_up_headdimv(int head_size) {
+    if (head_size <= 64) { return 64; }
+    if (head_size <= 96) { return 96; }
+    if (head_size <= 128) { return 128; }
+    if (head_size <= 192) { return 192; }
+    if (head_size <= 256) { return 256; }
+    return 512;
+}
+
+// Only applicable to the case where seqused_k (i.e. cache_seqlens) is available
+at::Tensor
+mha_fwd_get_scheduler_metadata(
+        int64_t batch_size,
+        int64_t max_seqlen_q,
+        int64_t max_seqlen_k,
+        int64_t num_heads,
+        int64_t num_heads_k,
+        int64_t headdim,
+        int64_t headdim_v,
+        at::ScalarType qkv_dtype,
+        at::Tensor seqused_k, // b
+        std::optional<at::Tensor> cu_seqlens_q_,  // b+1
+        std::optional<at::Tensor> cu_seqlens_k_,  // b+1
+        std::optional<at::Tensor> cu_seqlens_k_new_,  // b+1
+        std::optional<at::Tensor> seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
+        std::optional<at::Tensor> leftpad_k_, // b
+        std::optional<int64_t> page_size,
+        int64_t max_seqlen_k_new,  // 0 means we're not appending new KV
+        bool is_causal,
+        int64_t window_size_left,
+        int64_t window_size_right,
+        int64_t attention_chunk,
+        bool has_softcap,
+        int64_t num_splits,
+        std::optional<bool> pack_gqa_,
+        int64_t sm_margin
+        ) {
+
+    TORCH_CHECK(qkv_dtype == at::ScalarType::Half || qkv_dtype == at::ScalarType::BFloat16 || qkv_dtype == at::ScalarType::Float8_e4m3fn,
+                "FlashAttention only supports fp16, bf16, and fp8_e4m3 data type");
+    TORCH_CHECK(num_heads % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
+
+    // Reset the parameters
+    Flash_fwd_params params{};
+    params.is_bf16 = qkv_dtype == at::ScalarType::BFloat16;
+    params.is_e4m3 = qkv_dtype == at::ScalarType::Float8_e4m3fn;
+    params.b = batch_size;
+    params.seqlen_q = max_seqlen_q;
+    params.seqlen_k = max_seqlen_k;
+    params.h = num_heads;
+    params.h_k = num_heads_k;
+    params.d = headdim;
+    params.dv = headdim_v;
+    params.d_rounded = round_up_headdim(headdim);
+    params.dv_rounded = headdim_v == headdim ? params.d_rounded : round_up_headdimv(headdim_v);
+    params.seqlen_knew = max_seqlen_k_new;
+
+    bool const is_varlen_q = cu_seqlens_q_.has_value();
+    params.cu_seqlens_q = is_varlen_q ? cu_seqlens_q_.value().data_ptr<int>() : nullptr;
+    bool const is_varlen_k = cu_seqlens_k_.has_value();
+    params.cu_seqlens_k = is_varlen_k ? cu_seqlens_k_.value().data_ptr<int>() : nullptr;
+    params.cu_seqlens_knew = cu_seqlens_k_new_.has_value() ? cu_seqlens_k_new_.value().data_ptr<int>() : nullptr;
+    params.seqused_q = seqused_q_.has_value() ? seqused_q_.value().data_ptr<int>() : nullptr;
+    params.seqused_k = seqused_k.data_ptr<int>();
+    params.leftpad_k = leftpad_k_.has_value() ? leftpad_k_.value().data_ptr<int>() : nullptr;
+    params.knew_ptr = params.seqlen_knew > 0 ? reinterpret_cast<int*>(1) : nullptr;
+    if (window_size_left >= max_seqlen_k - 1) { window_size_left = -1; }
+    if (window_size_right >= max_seqlen_q - 1) { window_size_right = -1; }
+    // causal=true is the same as causal=false in this case
+    if (max_seqlen_q == 1 && window_size_left == -1 && window_size_right == -1 && attention_chunk == 0) {
+        // Special case of hdim 128 where we want causal to have kBlockN=128, better for pagedKV and TMA
+        if ((headdim <= 64 || headdim > 128) || !page_size.has_value()) {
+            is_causal = false;
+        }
+    }
+    if (is_causal) { window_size_right = 0; }
+
+    params.is_causal = window_size_left < 0 && window_size_right == 0 && attention_chunk == 0;
+    params.is_local = (window_size_left >= 0 || window_size_right >= 0 || attention_chunk >= 1) && !params.is_causal;
+    if (window_size_left < 0) { window_size_left = max_seqlen_k - 1; }
+    if (window_size_right < 0) { window_size_right = max_seqlen_q - 1; }
+    if (attention_chunk > 0) {
+        window_size_left = std::min(window_size_left, attention_chunk - 1);
+        window_size_right = std::min(window_size_right, attention_chunk - 1);
+    }
+    params.window_size_left = window_size_left;
+    params.window_size_right = window_size_right;
+    params.attention_chunk = attention_chunk;
+    params.arch = at::cuda::getCurrentDeviceProperties()->major * 10 + at::cuda::getCurrentDeviceProperties()->minor;
+    params.num_sm = at::cuda::getCurrentDeviceProperties()->multiProcessorCount - sm_margin;
+    params.softcap = has_softcap ? 1.0f : 0.0f;
+
+    params.page_size = page_size.has_value() ? page_size.value() : 1;
+    params.page_table = !page_size.has_value() ? nullptr : reinterpret_cast<int*>(1);
+
+    bool const use_dynamic_split = params.b <= 992;
+    params.num_splits_dynamic_ptr = !use_dynamic_split ? nullptr : reinterpret_cast<int*>(1);
+
+    params.pagedkv_tma = get_pagedkv_tma(params);
+    params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits;
+    // Always enable PackGQA for Split, and get_pack_gqa requires params.num_splits to decide
+    params.pack_gqa = pack_gqa_.has_value() ? pack_gqa_.value() : get_pack_gqa(params);
+
+    bool is_varlen = true;
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)seqused_k.get_device()};
+
+    auto opts = seqused_k.options();
+    // This needs to be set after get_num_splits
+    at::Tensor tile_count_semaphore;  // Contains the semaphore and optionally num_splits_dynamic
+    bool const scheduler_needs_semaphore = params.arch >= 90 || params.num_splits > 1;
+    if (scheduler_needs_semaphore || use_dynamic_split) {
+        tile_count_semaphore = torch::empty({int(scheduler_needs_semaphore) + int(use_dynamic_split) * params.b}, opts.dtype(torch::kInt32));
+        if (scheduler_needs_semaphore) {
+            if (!use_dynamic_split) { tile_count_semaphore.zero_(); }  // If varlen we'll manually do the zero-ing
+            params.tile_count_semaphore = tile_count_semaphore.data_ptr<int>();
+        } else {
+            params.tile_count_semaphore = nullptr;
+        }
+        params.num_splits_dynamic_ptr = use_dynamic_split ? tile_count_semaphore.data_ptr<int>() + 1 : nullptr;
+    }
+
+    if (params.num_splits_dynamic_ptr) {
+        auto kBlockMN_kernel_args_sm90 = tile_size_fwd_sm90(params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, false /*v_colmajor*/, params.page_table && !params.pagedkv_tma, params.softcap > 0.f);
+        auto kBlockMN_kernel_args_sm8x = tile_size_fwd_sm8x(params.arch == 86 || params.arch == 89, params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, params.page_table, is_varlen && params.num_splits > 1, params.softcap > 0.f, params.knew_ptr);
+        int const kBlockM = params.arch >= 90 ? std::get<0>(kBlockMN_kernel_args_sm90) : std::get<0>(kBlockMN_kernel_args_sm8x);
+        int const kBlockN = params.arch >= 90 ? std::get<1>(kBlockMN_kernel_args_sm90) : std::get<1>(kBlockMN_kernel_args_sm8x);
+        auto stream = at::cuda::getCurrentCUDAStream().stream();
+        prepare_varlen_num_blocks(params, stream, params.pack_gqa, kBlockM, kBlockN, false /*enable_pdl*/);
+        CHECK_CUDA_KERNEL_LAUNCH();
+    }
+    return tile_count_semaphore;
+}
+
+// b: batch_size
+// b_k: batch_size_k
+// s_q: seqlen_q
+// s_k: seqlen_k
+// s_k_new: seqlen_k_new
+// h: num_heads
+// h_k: num_heads_k
+// d: head_size
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>
+mha_fwd(at::Tensor q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+        at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, h_k, d) if there is page_table.
+        at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, page_size, h_k, dv) if there is page_table.
+        std::optional<at::Tensor> k_new_,  // (b, s_k_new, h_k, d) or (total_k_new, h_k, d) if there is cu_seqlens_k_new
+        std::optional<at::Tensor> v_new_,  // (b, s_k_new, h_k, dv) or (total_k_new, h_k, dv) if there is cu_seqlens_k_new
+        std::optional<at::Tensor> q_v_,  // (b, s_q, h, dv) or (total_q_new, h, dv) if there is cu_seqlens_q
+        std::optional<at::Tensor> out_,  // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
+        std::optional<at::Tensor> cu_seqlens_q_,  // b+1
+        std::optional<at::Tensor> cu_seqlens_k_,  // b+1
+        std::optional<at::Tensor> cu_seqlens_k_new_,  // b+1
+        std::optional<at::Tensor> seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
+        std::optional<at::Tensor> seqused_k_, // b. If given, only this many elements of each batch element's keys are used.
+        std::optional<int64_t> max_seqlen_q_,
+        // TODO: check if we need max_seqlen_k
+        std::optional<int64_t> max_seqlen_k_,
+        std::optional<at::Tensor> page_table_, // (b_k, max_num_pages_per_seq)
+        std::optional<at::Tensor> kv_batch_idx_, // b. indices to index into the KV cache
+        std::optional<at::Tensor> leftpad_k_, // b
+        std::optional<at::Tensor> rotary_cos_, // seqlen_ro x (rotary_dim / 2)
+        std::optional<at::Tensor> rotary_sin_, // seqlen_ro x (rotary_dim / 2)
+        std::optional<at::Tensor> seqlens_rotary_, // b
+        std::optional<at::Tensor> q_descale_,  // (b, h_k), not (b, h)
+        std::optional<at::Tensor> k_descale_,  // (b, h_k)
+        std::optional<at::Tensor> v_descale_,  // (b, h_k)
+        std::optional<double> softmax_scale_,
+        bool is_causal,
+        int64_t window_size_left,
+        int64_t window_size_right,
+        int64_t attention_chunk,
+        double softcap,
+        bool is_rotary_interleaved,   // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
+        std::optional<at::Tensor> scheduler_metadata_,  // (b + 1)
+        int64_t num_splits,
+        std::optional<bool> pack_gqa_,
+        int64_t sm_margin
+        ) {
+
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    bool is_sm8x = dprops->major >= 8;
+    TORCH_CHECK(is_sm8x, "FlashAttention only supports Ampere GPUs or newer.");
+
+    auto q_type = q.scalar_type();
+    TORCH_CHECK(q_type == at::ScalarType::Half || q_type == at::ScalarType::BFloat16 || q_type == at::ScalarType::Float8_e4m3fn,
+                "FlashAttention only supports fp16, bf16, and fp8_e4m3 data type");
+    if (dprops->major < 9) {
+        TORCH_CHECK(q_type == at::ScalarType::Half || q_type == at::ScalarType::BFloat16,
+                    "FlashAttention on Ampere/Ada cards only supports fp16 and bf16 data type");
+    }
+    TORCH_CHECK(k.scalar_type() == q_type, "query and key must have the same dtype");
+    TORCH_CHECK(v.scalar_type() == q_type, "query and value must have the same dtype");
+
+    CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v);
+
+    TORCH_CHECK(q.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(k.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+
+    at::Tensor page_table;
+    const bool paged_KV = page_table_.has_value();
+    if (paged_KV) {
+        page_table = page_table_.value();
+        CHECK_DEVICE(page_table);
+        TORCH_CHECK(page_table.dtype() == torch::kInt32, "page_table must have dtype torch.int32");
+        TORCH_CHECK(page_table.stride(-1) == 1, "page_table must have contiguous last dimension");
+    }
+
+    at::Tensor cu_seqlens_q;
+    bool const is_varlen_q = cu_seqlens_q_.has_value();
+    if (is_varlen_q) {
+        cu_seqlens_q = cu_seqlens_q_.value();
+        CHECK_DEVICE(cu_seqlens_q); CHECK_CONTIGUOUS(cu_seqlens_q);
+        TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32, "cu_seqlens_q must have dtype torch.int32");
+        TORCH_CHECK(max_seqlen_q_.has_value(), "max_seqlen_q must be provided if cu_seqlens_q is provided");
+    }
+    at::Tensor cu_seqlens_k;
+    bool const is_varlen_k = cu_seqlens_k_.has_value();
+    if (is_varlen_k) {
+        cu_seqlens_k = cu_seqlens_k_.value();
+        CHECK_DEVICE(cu_seqlens_k); CHECK_CONTIGUOUS(cu_seqlens_k);
+        TORCH_CHECK(cu_seqlens_k.dtype() == torch::kInt32, "cu_seqlens_k must have dtype torch.int32");
+        TORCH_CHECK(max_seqlen_k_.has_value(), "max_seqlen_k must be provided if cu_seqlens_k is provided");
+        TORCH_CHECK(!paged_KV, "If cu_seqlens_k is passed in, then page table is not supported");
+        TORCH_CHECK(!kv_batch_idx_.has_value(), "If cu_seqlens_k is passed in, then page table is not supported");
+    }
+
+    auto const sizes = q.sizes();
+    const int batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.size(0) - 1;
+    int seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_.value();
+    int total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
+    int num_heads = q.size(-2);
+    int const head_size = q.size(-1);
+    int const head_size_v = v.size(-1);
+    int const max_num_pages_per_seq = !paged_KV ? 0 : page_table.size(1);
+    int const num_pages = !paged_KV ? 0 : k.size(0);
+    int const page_size = !paged_KV ? 1 : k.size(1);
+    int const seqlen_k = !is_varlen_k ? (!paged_KV ? k.size(1) : max_num_pages_per_seq * page_size) : max_seqlen_k_.value();
+    int const total_k = !is_varlen_k ? batch_size * k.size(1) : k.size(0);
+    int const num_heads_k = k.size(-2);
+    int const batch_size_k = !paged_KV ? (!is_varlen_k ? k.size(0) : cu_seqlens_k.size(0) - 1) : page_table.size(0);
+    double softmax_scale = 1.0 / sqrt(double(head_size));
+    if (softmax_scale_.has_value()) {
+        softmax_scale = softmax_scale_.value();
+    }
+    if (!kv_batch_idx_.has_value()) {
+        TORCH_CHECK(batch_size == batch_size_k, "batch_size must be equal to batch_size_k");
+    }
+    int const max_headdim = get_max_headdim();
+    TORCH_CHECK(head_size <= max_headdim, "FlashAttention forward only supports head dimension at most " + std::to_string(max_headdim));
+    TORCH_CHECK(num_heads % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
+    if (head_size_v != head_size) {
+        TORCH_CHECK((head_size > 128 && head_size <= 192 && head_size_v > 96 && head_size_v <= 128) ||
+                   (head_size <= 64 && head_size_v <= 512),
+                   "If V headdim is different from Q/K dim, we only support Q/K headdim in (128, 192] and V headdim in (96, 128], "
+                   "or (Q/K <= 64 and V <= 512).");
+        TORCH_CHECK(dprops->major == 9, "Only Hopper supports different V headdim");
+        if (head_size_v > 256) {
+            TORCH_CHECK(q_type == at::ScalarType::Half || q_type == at::ScalarType::BFloat16,
+                        "HeaddimV > 256 requires fp16 and bf16 data type");
+        }
+    }
+
+    // This needs to go before kBlockM & kBlockN since we rely on the correct window_size and is_causal to set kBlockM
+    // TODO: check this
+    if (window_size_left >= seqlen_k - 1) { window_size_left = -1; }
+    if (window_size_right >= seqlen_q - 1) { window_size_right = -1; }
+    // causal=true is the same as causal=false in this case
+    if (seqlen_q == 1 && window_size_left == -1 && window_size_right == -1 && attention_chunk == 0) {
+        // Special case of hdim 128 where we want causal to have kBlockN=128, better for pagedKV and TMA
+        if ((head_size <= 64 || head_size > 128) || !paged_KV) {
+            is_causal = false;
+        }
+    }
+    if (is_causal) { window_size_right = 0; }
+
+    if (!is_varlen_q) {
+        CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
+    } else {
+        CHECK_SHAPE(q, total_q, num_heads, head_size);
+        CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
+    }
+    if (!paged_KV) {
+        if (!is_varlen_k) {
+            CHECK_SHAPE(k, batch_size_k, seqlen_k, num_heads_k, head_size);
+            CHECK_SHAPE(v, batch_size_k, seqlen_k, num_heads_k, head_size_v);
+        } else {
+            CHECK_SHAPE(k, total_k, num_heads_k, head_size);
+            CHECK_SHAPE(v, total_k, num_heads_k, head_size_v);
+            CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
+        }
+    } else {
+        CHECK_SHAPE(k, num_pages, page_size, num_heads_k, head_size);
+        CHECK_SHAPE(v, num_pages, page_size, num_heads_k, head_size_v);
+        CHECK_SHAPE(page_table, batch_size_k, max_num_pages_per_seq);
+    }
+
+    if (seqused_q_.has_value()){
+        auto seqused_q = seqused_q_.value();
+        TORCH_CHECK(seqused_q.dtype() == torch::kInt32, "seqused_q must have dtype int32");
+        CHECK_DEVICE(seqused_q); CHECK_CONTIGUOUS(seqused_q);
+        CHECK_SHAPE(seqused_q, batch_size);
+    }
+    if (seqused_k_.has_value()) {
+        auto seqused_k = seqused_k_.value();
+        TORCH_CHECK(seqused_k.dtype() == torch::kInt32, "seqused_k must have dtype int32");
+        CHECK_DEVICE(seqused_k); CHECK_CONTIGUOUS(seqused_k);
+        CHECK_SHAPE(seqused_k, batch_size);
+    }
+
+    if (leftpad_k_.has_value()) {
+        auto leftpad_k = leftpad_k_.value();
+        TORCH_CHECK(leftpad_k.dtype() == torch::kInt32, "leftpad_k must have dtype int32");
+        CHECK_DEVICE(leftpad_k); CHECK_CONTIGUOUS(leftpad_k);
+        CHECK_SHAPE(leftpad_k, batch_size);
+    }
+
+    // This is what we will template on
+    bool const is_varlen = is_varlen_q || is_varlen_k || seqused_q_.has_value() || seqused_k_.has_value() || leftpad_k_.has_value();
+    #ifdef FLASHATTENTION_DISABLE_VARLEN
+        TORCH_CHECK(!is_varlen, "This flash attention build does not support varlen.");
+    #endif
+
+    int const alignment = q_type == torch::kFloat8_e4m3fn ? 16 : 8;
+    TORCH_CHECK(head_size % alignment == 0, "head_size should be a multiple of " + std::to_string(alignment));
+    TORCH_CHECK(head_size_v % alignment == 0, "head_size_v should be a multiple of " + std::to_string(alignment));
+
+    auto opts = q.options();
+    auto out_type = q_type == at::ScalarType::Float8_e4m3fn ? at::ScalarType::BFloat16 : q_type;
+    at::Tensor out;
+    if (out_.has_value()) {
+        out = out_.value();
+        TORCH_CHECK(out.scalar_type() == out_type, "For FP16/BF16 input, output must have the same dtype as inputs. For FP8 input, output must have dtype BF16");
+        CHECK_DEVICE(out);
+        TORCH_CHECK(out.stride(-1) == 1, "Output tensor must have contiguous last dimension");
+        if (!is_varlen_q) {
+            CHECK_SHAPE(out, batch_size, seqlen_q, num_heads, head_size_v);
+        } else {
+            CHECK_SHAPE(out, total_q, num_heads, head_size_v);
+        }
+    } else {
+        out = !is_varlen_q
+            ? torch::empty({batch_size, seqlen_q, num_heads, head_size_v}, opts.dtype(out_type))
+            : torch::empty({total_q, num_heads, head_size_v}, opts.dtype(out_type));
+    }
+
+    auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
+    int const head_size_rounded = round_up_headdim(head_size);
+    int const head_size_v_rounded = head_size_v == head_size ? head_size_rounded : round_up_headdimv(head_size_v);
+    int const seqlen_q_rounded = round_multiple(seqlen_q, 128);
+    int const seqlen_k_rounded = round_multiple(seqlen_k, 128);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)q.get_device()};
+
+    at::Tensor softmax_lse;
+    if (!is_varlen_q) {
+        softmax_lse = torch::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
+    } else {
+        softmax_lse = torch::empty({num_heads, total_q}, opts.dtype(at::kFloat));
+    }
+
+    Flash_fwd_params params;
+    set_params_fprop(params,
+                     batch_size,
+                     seqlen_q, seqlen_k,
+                     seqlen_q_rounded, seqlen_k_rounded,
+                     num_heads, num_heads_k,
+                     head_size, head_size_rounded,
+                     q, k, v, out,
+                     !is_varlen_q ? nullptr : cu_seqlens_q.data_ptr(),
+                     !is_varlen_k ? nullptr : cu_seqlens_k.data_ptr(),
+                     seqused_q_.has_value() ? seqused_q_.value().data_ptr() : nullptr,
+                     seqused_k_.has_value() ? seqused_k_.value().data_ptr() : nullptr,
+                     softmax_lse.data_ptr(),
+                     /*p_dropout=*/0.f,
+                     softmax_scale,
+                     window_size_left,
+                     window_size_right,
+                     attention_chunk,
+                     softcap,
+                     sm_margin);
+    params.total_q = total_q;
+    params.total_k = total_k;
+    params.b_k = batch_size_k;
+    params.dv = head_size_v;
+    params.dv_rounded = head_size_v_rounded;
+    if (leftpad_k_.has_value()) {  // This needs to be set before get_pagedkv_tma
+        params.leftpad_k = static_cast<int *>(leftpad_k_.value().data_ptr());
+    }
+    if (paged_KV) {
+        params.page_table = page_table.data_ptr<int>();
+        params.page_table_batch_stride = page_table.stride(0);
+    }
+    params.page_size = page_size;
+    params.num_pages = num_pages;
+
+    if (k_new_.has_value()) {  // This needs to be set before get_pagedkv_tma
+        at::Tensor k_new, v_new;
+        TORCH_CHECK(v_new_.has_value(), "If k_new is supplied, v_new must also be passed in");
+        TORCH_CHECK(seqused_k_.has_value(), "If k_new is supplied, seqlens_k must also be passed in");
+        TORCH_CHECK(seqlen_q <= seqlen_k, "If k_new is supplied, it must have seqlen <= the seqlen of the KV cache");
+        at::Tensor cu_seqlens_k_new;
+        bool const is_varlen_k_new = cu_seqlens_k_new_.has_value();
+        if (is_varlen_k_new) {
+            cu_seqlens_k_new = cu_seqlens_k_new_.value();
+            CHECK_DEVICE(cu_seqlens_k_new); CHECK_CONTIGUOUS(cu_seqlens_k_new);
+            TORCH_CHECK(cu_seqlens_k_new.dtype() == torch::kInt32, "cu_seqlens_k_new must have dtype torch.int32");
+        }
+        k_new = k_new_.value();
+        v_new = v_new_.value();
+        TORCH_CHECK(k_new.dtype() == q_type, "k_new must have the same dtype as query");
+        TORCH_CHECK(v_new.dtype() == q_type, "v_new must have the same dtype as query");
+        CHECK_DEVICE(k_new); CHECK_DEVICE(v_new);
+        TORCH_CHECK(k_new.stride(-1) == 1, "k_new tensor must have contiguous last dimension");
+        TORCH_CHECK(v_new.stride(-1) == 1, "v_new tensor must have contiguous last dimension");
+        // We don't need max_seqlen_k_new, so seqlen_k_new can be whatever when is_varlen_k_new
+        int seqlen_k_new = !is_varlen_k_new ? k_new.size(1) : 0;
+        int total_k_new = !is_varlen_k_new ? batch_size * k_new.size(1): k_new.size(0);
+        if (!is_varlen_k_new) {
+            CHECK_SHAPE(k_new, batch_size, seqlen_k_new, num_heads_k, head_size);
+            CHECK_SHAPE(v_new, batch_size, seqlen_k_new, num_heads_k, head_size_v);
+        } else {
+            CHECK_SHAPE(k_new, total_k_new, num_heads_k, head_size);
+            CHECK_SHAPE(v_new, total_k_new, num_heads_k, head_size_v);
+            CHECK_SHAPE(cu_seqlens_k_new, batch_size + 1);
+        }
+        params.seqlen_knew = seqlen_k_new;
+        params.total_knew = total_k_new;
+        params.knew_ptr = k_new.data_ptr();
+        params.vnew_ptr = v_new.data_ptr();
+        // All stride are in elements, not bytes.
+        params.knew_row_stride = k_new.stride(-3);
+        params.vnew_row_stride = v_new.stride(-3);
+        params.knew_head_stride = k_new.stride(-2);
+        params.vnew_head_stride = v_new.stride(-2);
+        if (!is_varlen_k_new) {
+            params.knew_batch_stride = k_new.stride(0);
+            params.vnew_batch_stride = v_new.stride(0);
+        }
+        if (is_varlen_k_new) {
+            params.cu_seqlens_knew = static_cast<int*>(cu_seqlens_k_new.data_ptr());
+        }
+    }
+
+    // 992 = 32 * 31 is the max supported batch in prepare_varlen_num_blocks kernel
+    bool const use_dynamic_split = is_varlen && params.b <= 992;
+    // Temporarily set num_splits_dynamic_ptr to 1 since get_num_splits checks it
+    params.num_splits_dynamic_ptr = !use_dynamic_split ? nullptr : reinterpret_cast<int*>(1);
+
+    params.pagedkv_tma = get_pagedkv_tma(params);
+    params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits;
+    // Always enable PackGQA for Split, and get_pack_gqa requires params.num_splits to decide
+    params.pack_gqa = pack_gqa_.has_value() ? pack_gqa_.value() : get_pack_gqa(params);
+
+    // This needs to be set after get_num_splits
+    at::Tensor tile_count_semaphore;  // Contains the semaphore and optionally num_splits_dynamic
+    // We don't use the persistent scheduler if Split and not Varlen
+    bool const scheduler_needs_semaphore = params.arch >= 90
+        ? (((params.is_causal || params.is_local) && (params.num_splits == 1)) || is_varlen)
+        : ((params.is_causal && !is_varlen) || (is_varlen && params.num_splits > 1));
+    if (scheduler_needs_semaphore || use_dynamic_split) {
+        int metadata_size = int(scheduler_needs_semaphore) + int(use_dynamic_split) * params.b;
+        params.skip_scheduler_metadata_computation = scheduler_metadata_.has_value();
+        if (scheduler_metadata_.has_value()) {
+            at::Tensor scheduler_metadata = scheduler_metadata_.value();
+            CHECK_DEVICE(scheduler_metadata);
+            CHECK_SHAPE(scheduler_metadata, metadata_size);
+            CHECK_CONTIGUOUS(scheduler_metadata);
+            TORCH_CHECK(scheduler_metadata.dtype() == torch::kInt32, "scheduler_metadata must have dtype int32");
+            tile_count_semaphore = scheduler_metadata;
+        } else {
+            tile_count_semaphore = torch::empty({metadata_size}, opts.dtype(torch::kInt32));
+        }
+        if (scheduler_needs_semaphore && !use_dynamic_split) {
+            tile_count_semaphore.zero_();  // If varlen we'll manually do the zero-ing
+        }
+        params.tile_count_semaphore = scheduler_needs_semaphore ? tile_count_semaphore.data_ptr<int>() : nullptr;
+        params.num_splits_dynamic_ptr = use_dynamic_split ? tile_count_semaphore.data_ptr<int>() + 1 : nullptr;
+    }
+
+    if (q_v_.has_value()) {
+        TORCH_CHECK(head_size <= 64, "q_v is only supported for head_size <= 64");
+        TORCH_CHECK(q_type == at::ScalarType::Half || q_type == at::ScalarType::BFloat16,
+                    "q_v is only supported for fp16 and bf16 data type");
+        TORCH_CHECK(params.arch == 90, "q_v is only supported for Hopper GPUs");
+        at::Tensor q_v = q_v_.value();
+        TORCH_CHECK(q_v.dtype() == q_type, "q_v must have the same dtype as query");
+        CHECK_DEVICE(q_v);
+        TORCH_CHECK(q_v.stride(-1) == 1, "q_v tensor must have contiguous last dimension");
+        if (!is_varlen_q) {
+            CHECK_SHAPE(q_v, batch_size, seqlen_q, num_heads, head_size_v);
+        } else {
+            CHECK_SHAPE(q_v, total_q, num_heads, head_size_v);
+        }
+        params.qv_ptr = q_v.data_ptr();
+        // All stride are in elements, not bytes.
+        params.qv_row_stride = q_v.stride(-3);
+        params.qv_head_stride = q_v.stride(-2);
+        if (!is_varlen_q) {
+            params.qv_batch_stride = q_v.stride(0);
+        }
+    }
+
+    if (rotary_cos_.has_value()) {
+        TORCH_CHECK(k_new_.has_value(), "If rotary cos/sin are provided, new key / value to be appended to KV cache must also be provided");
+        auto rotary_cos = rotary_cos_.value();
+        CHECK_DEVICE(rotary_cos); CHECK_CONTIGUOUS(rotary_cos);
+        params.rotary_dim = rotary_cos.size(1) * 2;
+        TORCH_CHECK(params.rotary_dim <= head_size, "rotary_dim must be <= headdim");
+        TORCH_CHECK(params.rotary_dim % 16 == 0, "Only rotary dimensions divisible by 16 are currently supported");
+        const int seqlen_ro = rotary_cos.size(0);
+        if (paged_KV) {
+            TORCH_CHECK(seqlen_ro >= seqlen_k, "cos/sin seqlen must be at least the seqlen of KV cache");
+        }
+        CHECK_SHAPE(rotary_cos, seqlen_ro, params.rotary_dim / 2);
+        TORCH_CHECK(rotary_cos.scalar_type() == q_type, "rotary_cos must have the same dtype as query");
+
+        TORCH_CHECK(rotary_sin_.has_value(), "If rotary cos is provided, rotary sin must also be provided");
+        auto rotary_sin = rotary_sin_.value();
+        CHECK_DEVICE(rotary_sin); CHECK_CONTIGUOUS(rotary_sin);
+        CHECK_SHAPE(rotary_sin, seqlen_ro, params.rotary_dim / 2);
+        TORCH_CHECK(rotary_sin.scalar_type() == q_type, "rotary_cos must have the same dtype as query");
+        params.rotary_cos_ptr = rotary_cos.data_ptr();
+        params.rotary_sin_ptr = rotary_sin.data_ptr();
+        params.is_rotary_interleaved = is_rotary_interleaved;
+        if (seqlens_rotary_.has_value()) {
+            at::Tensor seqlens_rotary = seqlens_rotary_.value();
+            CHECK_DEVICE(seqlens_rotary); CHECK_CONTIGUOUS(seqlens_rotary);
+            TORCH_CHECK(seqlens_rotary.dtype() == torch::kInt32, "seqlens_rotary must have dtype torch.int32");
+            CHECK_SHAPE(seqlens_rotary, batch_size);
+            params.seqlens_rotary = seqlens_rotary.data_ptr<int>();
+        }
+    } else {
+        params.rotary_dim = 0;
+    }
+
+    if (kv_batch_idx_.has_value()) {
+        auto kv_batch_idx = kv_batch_idx_.value();
+        CHECK_DEVICE(kv_batch_idx); CHECK_CONTIGUOUS(kv_batch_idx);
+        TORCH_CHECK(kv_batch_idx.scalar_type() == torch::kInt32, "kv_batch_idx must have dtype int32");
+        params.kv_batch_idx = reinterpret_cast<int *>(kv_batch_idx.data_ptr());
+    }
+
+    at::Tensor out_accum, softmax_lse_accum;
+    auto outaccum_type = at::ScalarType::Float;
+    if (params.num_splits > 1) {
+        TORCH_CHECK(params.num_splits <= 256, "num_splits > 256 not supported");
+        if (!is_varlen_q) {
+            out_accum = torch::empty({params.num_splits, batch_size, num_heads, seqlen_q, head_size_v}, opts.dtype(outaccum_type));
+            softmax_lse_accum = torch::empty({params.num_splits, batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
+            params.oaccum_batch_stride = out_accum.stride(1);
+            params.lseaccum_batch_stride = softmax_lse_accum.stride(1);
+        } else {
+            out_accum = torch::empty({params.num_splits, num_heads, total_q, head_size_v}, opts.dtype(outaccum_type));
+            softmax_lse_accum = torch::empty({params.num_splits, num_heads, total_q}, opts.dtype(at::kFloat));
+        }
+        params.is_fp32 = false;
+        params.oaccum_ptr = out_accum.data_ptr();
+        params.softmax_lseaccum_ptr = softmax_lse_accum.data_ptr();
+        params.oaccum_split_stride = out_accum.stride(0);
+        params.oaccum_row_stride = out_accum.stride(-2);
+        params.oaccum_head_stride = out_accum.stride(-3);
+        params.lseaccum_split_stride = softmax_lse_accum.stride(0);
+        params.lseaccum_head_stride = softmax_lse_accum.stride(-2);
+    }
+
+    if (q_type == at::ScalarType::Float8_e4m3fn) {
+        if (q_descale_.has_value()) {
+            auto q_descale = q_descale_.value();
+            CHECK_DEVICE(q_descale);
+            CHECK_SHAPE(q_descale, batch_size, num_heads_k);
+            params.q_descale_ptr = q_descale.data_ptr<float>();
+            params.q_descale_batch_stride = q_descale.stride(0);
+            params.q_descale_head_stride = q_descale.stride(1);
+        } else {
+            params.q_descale_ptr = nullptr;
+        }
+        if (k_descale_.has_value()) {
+            auto k_descale = k_descale_.value();
+            CHECK_DEVICE(k_descale);
+            CHECK_SHAPE(k_descale, batch_size, num_heads_k);
+            params.k_descale_ptr = k_descale.data_ptr<float>();
+            params.k_descale_batch_stride = k_descale.stride(0);
+            params.k_descale_head_stride = k_descale.stride(1);
+        } else {
+            params.k_descale_ptr = nullptr;
+        }
+        if (v_descale_.has_value()) {
+            auto v_descale = v_descale_.value();
+            CHECK_DEVICE(v_descale);
+            CHECK_SHAPE(v_descale, batch_size, num_heads_k);
+            params.v_descale_ptr = v_descale.data_ptr<float>();
+            params.v_descale_batch_stride = v_descale.stride(0);
+            params.v_descale_head_stride = v_descale.stride(1);
+        } else {
+            params.v_descale_ptr = nullptr;
+        }
+    }
+
+    #ifdef FLASHATTENTION_DISABLE_LOCAL
+    TORCH_CHECK(!params.is_local, "This flash attention build does not support local attention.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_SOFTCAP
+    TORCH_CHECK(params.softcap == 0.0, "This flash attention build does not support tanh softcapping.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_SPLIT
+    TORCH_CHECK(params.num_splits == 1, "This flash attention build does not support splits.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_PACKGQA
+    TORCH_CHECK(!params.pack_gqa || params.arch < 90 || (params.page_table && !params.pagedkv_tma) || params.num_splits > 1, "This flash attention build does not support pack_gqa.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_PAGEDKV
+    TORCH_CHECK(!(params.page_table && !params.pagedkv_tma), "This flash attention build does not support paged KV.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_APPENDKV
+    TORCH_CHECK(!k_new_.has_value(), "This flash attention build does not support appending KV.");
+    #endif
+
+    if (total_q > 0 && (total_k + params.total_knew) > 0 && num_heads_k > 0) {
+        auto stream = at::cuda::getCurrentCUDAStream().stream();
+        run_mha_fwd(params, stream);
+        if (params.num_splits > 1) {
+            if (out_type == at::ScalarType::BFloat16) {
+                // Since we want output in BF16. Otherwise fwd_combine will output to FP16
+                params.is_bf16 = true;
+            }
+            // Unless there's seqused_q, for the purpose of attn_combine, we can just treat it as batch=1
+            // and seqlen = total_q, and don't need to dispatch to Varlen there.
+            // However, with dynamic split, each row needs to know which batch it belongs to
+            // to read the number of splits, so we just use the varlen version of combine kernel.
+            // if (is_varlen_q && !seqused_q_.has_value()) {
+            // if (is_varlen_q) {
+            //     params.b = 1;
+            //     params.seqlen_q = total_q;
+            // }
+            // This will zero out the semaphore if needed
+            run_mha_fwd_combine(params, stream, true /*enable_pdl*/);
+        } else if (scheduler_needs_semaphore && params.skip_scheduler_metadata_computation) {
+            // need to zero out the semaphore in this case
+            tile_count_semaphore.index({torch::indexing::Slice(0, 1)}).zero_();
+        }
+    } else if (total_q > 0 && num_heads_k > 0) {
+        // If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0.
+        out.zero_();
+        softmax_lse.fill_(std::numeric_limits<float>::infinity());
+    }
+
+    // return {out, softmax_lse};
+    return {out, softmax_lse, out_accum, softmax_lse_accum};
+}
+
+#ifdef FLASHATTENTION_DISABLE_BACKWARD
+void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
+    TORCH_CHECK(false, "Flash-Attention was built with backward disabled");
+}
+#else
+template <int Arch, bool Has_softcap>
+void run_mha_bwd_constexpr(Flash_bwd_params &params, cudaStream_t stream) {
+    if (!params.is_bf16) {
+        #ifndef FLASHATTENTION_DISABLE_FP16
+        #ifndef FLASHATTENTION_DISABLE_HDIM64
+        if (params.d_rounded == 64) { return run_mha_bwd_<Arch, cutlass::half_t, 64, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM96
+        if (params.d_rounded == 96) { return run_mha_bwd_<Arch, cutlass::half_t, 96, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM128
+        if (params.d_rounded == 128) { return run_mha_bwd_<Arch, cutlass::half_t, 128, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM192
+        if (params.d_rounded == 192) { return run_mha_bwd_<Arch, cutlass::half_t, 192, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM256
+        if (params.d_rounded == 256) { return run_mha_bwd_<Arch, cutlass::half_t, 256, Has_softcap>(params, stream); }
+        #endif
+        #else
+        TORCH_CHECK(false, "This flash attention build does not support FP16.");
+        #endif
+    } else {
+        #ifndef FLASHATTENTION_DISABLE_HDIM64
+        if (params.d_rounded == 64) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 64, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM96
+        if (params.d_rounded == 96) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 96, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM128
+        if (params.d_rounded == 128) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 128, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM192
+        if (params.d_rounded == 192) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 192, Has_softcap>(params, stream); }
+        #endif
+        #ifndef FLASHATTENTION_DISABLE_HDIM256
+        if (params.d_rounded == 256) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 256, Has_softcap>(params, stream); }
+        #endif
+    }
+}
+
+void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
+        // FP16_SWITCH(!params.is_bf16, [&] {
+        //     HEADDIM_SWITCH(params.d, [&] {
+        //         run_mha_bwd_<elem_type, kHeadDim>(params, stream);
+        //     });
+        // });
+    ARCH_SWITCH(params.arch, Arch, [&] {
+        SOFTCAP_SWITCH(params.softcap > 0.f, Has_softcap, [&] {
+            run_mha_bwd_constexpr<Arch, Has_softcap>(params, stream);
+        });
+    });
+}
+#endif
+
+
+// b: batch_size
+// s_q: seqlen_q
+// s_k: seqlen_k
+// h: num_heads
+// h_k: num_heads_k
+// d: head_size
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_bwd(
+    at::Tensor dout,  // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
+    at::Tensor q,     // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    at::Tensor k,     // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
+    at::Tensor v,     // (b, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k
+    at::Tensor out,   // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
+    at::Tensor softmax_lse,    // (b, h, s_q) or (h, total_q) if there is cu_seqlens_q
+    std::optional<at::Tensor> dq_,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    std::optional<at::Tensor> dk_,   // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
+    std::optional<at::Tensor> dv_,   // (b, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k
+    std::optional<at::Tensor> cu_seqlens_q_,   // b+1
+    std::optional<at::Tensor> cu_seqlens_k_,   // b+1
+    std::optional<at::Tensor> seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
+    std::optional<at::Tensor> seqused_k_, // b. If given, only this many elements of each batch element's keys are used.
+    std::optional<int64_t> max_seqlen_q_,
+    std::optional<int64_t> max_seqlen_k_,
+    std::optional<double> softmax_scale_,
+    bool is_causal,
+    int64_t window_size_left,
+    int64_t window_size_right,
+    double softcap,
+    bool deterministic,
+    int64_t sm_margin
+) {
+
+    #ifdef FLASHATTENTION_DISABLE_BACKWARD
+        TORCH_CHECK(false, "This flash attention build does not support backward.");
+    #endif
+
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    bool is_sm8x = dprops->major >= 8;
+    TORCH_CHECK(is_sm8x, "FlashAttention only supports Ampere GPUs or newer.");
+
+    auto q_type = q.dtype();
+    TORCH_CHECK(q_type == torch::kFloat16 || q_type == torch::kBFloat16,
+                "FlashAttention only support fp16 and bf16 data type");
+    TORCH_CHECK(k.dtype() == q_type, "query and key must have the same dtype");
+    TORCH_CHECK(v.dtype() == q_type, "query and value must have the same dtype");
+    TORCH_CHECK(out.dtype() == q_type, "query and out must have the same dtype");
+    TORCH_CHECK(dout.dtype() == q_type, "query and dout must have the same dtype");
+
+    CHECK_DEVICE(q); CHECK_DEVICE(k); CHECK_DEVICE(v);
+    CHECK_DEVICE(out); CHECK_DEVICE(dout); CHECK_DEVICE(softmax_lse);
+
+    TORCH_CHECK(q.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(k.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(out.stride(-1) == 1, "out tensor must have contiguous last dimension");
+    TORCH_CHECK(dout.stride(-1) == 1, "dout tensor must have contiguous last dimension");
+
+    at::Tensor cu_seqlens_q;
+    bool const is_varlen_q = cu_seqlens_q_.has_value();
+    if (is_varlen_q) {
+        cu_seqlens_q = cu_seqlens_q_.value();
+        CHECK_DEVICE(cu_seqlens_q); CHECK_CONTIGUOUS(cu_seqlens_q);
+        TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32, "cu_seqlens_q must have dtype torch.int32");
+        TORCH_CHECK(max_seqlen_q_.has_value(), "max_seqlen_q must be provided if cu_seqlens_q is provided");
+    }
+    at::Tensor cu_seqlens_k;
+    bool const is_varlen_k = cu_seqlens_k_.has_value();
+    if (is_varlen_k) {
+        cu_seqlens_k = cu_seqlens_k_.value();
+        CHECK_DEVICE(cu_seqlens_k); CHECK_CONTIGUOUS(cu_seqlens_k);
+        TORCH_CHECK(cu_seqlens_k.dtype() == torch::kInt32, "cu_seqlens_k must have dtype torch.int32");
+        TORCH_CHECK(max_seqlen_k_.has_value(), "max_seqlen_k must be provided if cu_seqlens_k is provided");
+    }
+    // This is what we will template on
+    bool const is_varlen = is_varlen_q || is_varlen_k || seqused_q_.has_value() || seqused_k_.has_value();
+    #ifdef FLASHATTENTION_DISABLE_VARLEN
+        TORCH_CHECK(!is_varlen, "This flash attention build does not support varlen.");
+    #endif
+
+    auto const sizes = q.sizes();
+    int const batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.size(0) - 1;
+    int const seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_.value();
+    int const total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
+    int const num_heads = q.size(-2);
+    int const head_size = q.size(-1);
+    int const head_size_v = v.size(-1);
+    int const seqlen_k = !is_varlen_k ? k.size(1) : max_seqlen_k_.value();
+    int const total_k = !is_varlen_k ? batch_size * k.size(1) : k.size(0);
+    int const num_heads_k = k.size(-2);
+    TORCH_CHECK(head_size % 8 == 0, "head_size should be a multiple of 8");
+    TORCH_CHECK(head_size_v % 8 == 0, "head_size_v should be a multiple of 8");
+    int const max_headdim = get_max_headdim();
+    TORCH_CHECK(std::max(head_size, head_size_v) <= max_headdim, "FlashAttention forward only supports head dimension at most " + std::to_string(max_headdim));
+    TORCH_CHECK(num_heads % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
+    double softmax_scale = 1.0 / sqrt(double(head_size));
+    if (softmax_scale_.has_value()) {
+        softmax_scale = softmax_scale_.value();
+    }
+
+    // This needs to go before kBlockM & kBlockN since we rely on the correct window_size and is_causal to set kBlockM
+    if (window_size_left >= seqlen_k - 1) { window_size_left = -1; }
+    if (window_size_right >= seqlen_q - 1) { window_size_right = -1; }
+    if (is_causal) { window_size_right = 0; }
+    // There's a case where is_causal=false, window_size=(-1, 0). Then set_params_bprop will set params.is_causal=true.
+    // If we don't have is_causal here matching params.is_causal, we might get the wrong kBlockM (and cause IMA).
+    is_causal = window_size_left < 0 && window_size_right == 0;
+
+    int const arch = at::cuda::getCurrentDeviceProperties()->major * 10 + at::cuda::getCurrentDeviceProperties()->minor;
+    int const head_size_rounded = round_up_headdim(std::max(head_size, head_size_v));
+    int const head_size_v_rounded = head_size_rounded;
+    // Very important that these match the kernel configs
+    bool const is_local = (window_size_left >= 0 || window_size_right >= 0) && !is_causal;
+    int const kBlockM_sm90 = head_size_rounded <= 64 ? (is_causal && softcap > 0.0 ? 96 : 128)
+        : (head_size_rounded <= 96 ? 64
+           : (head_size_rounded <= 128 ? (is_causal || is_local || softcap > 0.0 ? 64 : 80)
+              : 64));
+    int const kBlockM_sm80 = head_size_rounded <= 64 ? 128 : 64;
+    int const kBlockM_sm86 = head_size_rounded <= 192 ? 64 : 32;
+    int const kBlockM = arch >= 90 ? kBlockM_sm90 : (arch == 86 || arch == 89 ? kBlockM_sm86 : kBlockM_sm80);
+    int const kBlockN_sm90 = head_size_rounded <= 128
+        ? 128
+        : (head_size_rounded <= 192 ? 96 : 80);
+    int const kBlockN_sm80 = head_size_rounded <= 128
+        ? 128
+        : (head_size_rounded <= 192 ? 80 : 64);
+    int const kBlockN_sm86 = head_size_rounded <= 64 ? 128
+        : (head_size_rounded <= 96 ? 128
+           : (head_size_rounded <= 128 ? 96
+              : (head_size_rounded <= 192 ? 64 : 64)));
+    int const kBlockN = arch >= 90 ? kBlockN_sm90 : (arch == 86 || arch == 89 ? kBlockN_sm86 : kBlockN_sm80);
+    auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
+    int const seqlen_q_rounded = round_multiple(seqlen_q, kBlockM);
+    int const seqlen_k_rounded = round_multiple(seqlen_k, kBlockN);
+    int const total_q_padded_rounded = round_multiple(total_q + batch_size * kBlockM, kBlockM);
+    int const total_k_padded_rounded = round_multiple(total_k + batch_size * kBlockN, kBlockN);
+
+    if (!is_varlen_q) {
+        CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
+        CHECK_SHAPE(out, batch_size, seqlen_q, num_heads, head_size_v);
+        CHECK_SHAPE(dout, batch_size, seqlen_q, num_heads, head_size_v);
+    } else {
+        CHECK_SHAPE(q, total_q, num_heads, head_size);
+        CHECK_SHAPE(out, total_q, num_heads, head_size_v);
+        CHECK_SHAPE(dout, total_q, num_heads, head_size_v);
+        CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
+    }
+    if (!is_varlen_k) {
+        CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size);
+        CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size_v);
+    } else {
+        CHECK_SHAPE(k, total_k, num_heads_k, head_size);
+        CHECK_SHAPE(v, total_k, num_heads_k, head_size_v);
+        CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
+    }
+
+    if (seqused_q_.has_value()){
+        auto seqused_q = seqused_q_.value();
+        TORCH_CHECK(seqused_q.dtype() == torch::kInt32, "seqused_q must have dtype int32");
+        CHECK_DEVICE(seqused_q); CHECK_CONTIGUOUS(seqused_q);
+        CHECK_SHAPE(seqused_q, batch_size);
+    }
+    if (seqused_k_.has_value()){
+        auto seqused_k = seqused_k_.value();
+        TORCH_CHECK(seqused_k.dtype() == torch::kInt32, "seqused_k must have dtype int32");
+        CHECK_DEVICE(seqused_k); CHECK_CONTIGUOUS(seqused_k);
+        CHECK_SHAPE(seqused_k, batch_size);
+    }
+
+    at::Tensor dq, dk, dv;
+    if (dq_.has_value()) {
+        dq = dq_.value();
+        TORCH_CHECK(dq.dtype() == q_type, "dq must have the same dtype as q");
+        CHECK_DEVICE(dq);
+        TORCH_CHECK(dq.stride(-1) == 1, "dq must have contiguous last dimension");
+        if (!is_varlen_q) {
+            CHECK_SHAPE(dq, batch_size, seqlen_q, num_heads, head_size);
+        } else {
+            CHECK_SHAPE(dq, total_q, num_heads, head_size);
+        }
+    } else {
+        dq = torch::empty_like(q);
+    }
+    if (dk_.has_value()) {
+        dk = dk_.value();
+        TORCH_CHECK(dk.dtype() == q_type, "dk must have the same dtype as q");
+        CHECK_DEVICE(dk);
+        TORCH_CHECK(dk.stride(-1) == 1, "dk must have contiguous last dimension");
+        if (!is_varlen_k) {
+            CHECK_SHAPE(dk, batch_size, seqlen_k, num_heads_k, head_size);
+        } else {
+            CHECK_SHAPE(dk, total_k, num_heads_k, head_size);
+        }
+    } else {
+        dk = torch::empty_like(k);
+    }
+    if (dv_.has_value()) {
+        dv = dv_.value();
+        TORCH_CHECK(dv.dtype() == q_type, "dv must have the same dtype as q");
+        CHECK_DEVICE(dv);
+        TORCH_CHECK(dv.stride(-1) == 1, "dv must have contiguous last dimension");
+        if (!is_varlen_k) {
+            CHECK_SHAPE(dv, batch_size, seqlen_k, num_heads_k, head_size_v);
+        } else {
+            CHECK_SHAPE(dv, total_k, num_heads_k, head_size_v);
+        }
+    } else {
+        dv = torch::empty_like(v);
+    }
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)q.get_device()};
+
+    auto opts = q.options();
+    // Need softmax_d to have total_q_padded_rounded since we want its address to be aligned by 16/8 bytes for TMA / LDG.64
+    at::Tensor softmax_d, softmax_lse_log2;
+    if (!is_varlen) {
+        // Need softmax_d to have seqlen_q_rounded since we want its address to be aligned by 16/8 bytes for TMA / LDG.64
+        softmax_d = torch::empty({batch_size, num_heads, seqlen_q_rounded}, opts.dtype(at::kFloat));
+        softmax_lse_log2 = torch::empty({batch_size, num_heads, seqlen_q_rounded}, opts.dtype(at::kFloat));
+    } else {
+        softmax_d = torch::empty({num_heads, total_q_padded_rounded}, opts.dtype(at::kFloat));
+        softmax_lse_log2 = torch::empty({num_heads, total_q_padded_rounded}, opts.dtype(at::kFloat));
+    }
+    at::Tensor dq_accum, dk_accum, dv_accum;
+    if (!is_varlen) {
+        dq_accum = torch::empty({batch_size, num_heads, seqlen_q_rounded * head_size_rounded}, opts.dtype(at::kFloat));
+    } else {
+        dq_accum = torch::empty({num_heads, total_q_padded_rounded * head_size_rounded}, opts.dtype(at::kFloat));
+    }
+    if (num_heads_k != num_heads) {  // MQA / GQA
+        if (!is_varlen) {
+            dk_accum = torch::zeros({batch_size, num_heads_k, seqlen_k_rounded * head_size_rounded}, opts.dtype(at::kFloat));
+            dv_accum = torch::zeros({batch_size, num_heads_k, seqlen_k_rounded * head_size_v_rounded}, opts.dtype(at::kFloat));
+        } else {
+            dk_accum = torch::zeros({num_heads_k, total_k_padded_rounded, head_size_rounded}, opts.dtype(at::kFloat));
+            dv_accum = torch::zeros({num_heads_k, total_k_padded_rounded, head_size_v_rounded}, opts.dtype(at::kFloat));
+        }
+    }
+
+    Flash_bwd_params params;
+    set_params_dgrad(params,
+                     batch_size,
+                     seqlen_q, seqlen_k,
+                     seqlen_q_rounded, seqlen_k_rounded,
+                     num_heads, num_heads_k,
+                     head_size, head_size_rounded,
+                     q, k, v, out,
+                     dout, dq, dk, dv,
+                     !is_varlen_q ? nullptr : cu_seqlens_q.data_ptr(),
+                     !is_varlen_k ? nullptr : cu_seqlens_k.data_ptr(),
+                     seqused_q_.has_value() ? seqused_q_.value().data_ptr() : nullptr,
+                     seqused_k_.has_value() ? seqused_k_.value().data_ptr() : nullptr,
+                     dq_accum.data_ptr(),
+                     num_heads_k != num_heads ? dk_accum.data_ptr() : nullptr,
+                     num_heads_k != num_heads ? dv_accum.data_ptr() : nullptr,
+                     softmax_lse.data_ptr(),
+                     softmax_d.data_ptr(),
+                     /*p_dropout=*/0.f,
+                     softmax_scale,
+                     window_size_left,
+                     window_size_right,
+                     0,  // attention_chunk
+                     softcap,
+                     deterministic,
+                     sm_margin);
+    params.total_q = total_q;
+    params.total_k = total_k;
+    params.softmax_lse_log2_ptr = softmax_lse_log2.data_ptr();
+    params.dv = head_size_v;
+    params.dv_rounded = head_size_v_rounded;
+
+    // auto tile_count_semaphore = (params.is_causal || params.is_local) ? torch::zeros({1}, opts.dtype(torch::kInt32)) : torch::empty({1}, opts.dtype(torch::kInt32));
+    // params.tile_count_semaphore = tile_count_semaphore.data_ptr<int>();
+    // Will be zero'ed out in the backward preprocess kernel
+    at::Tensor dq_semaphore = torch::empty({(seqlen_q + kBlockM - 1) / kBlockM, batch_size, num_heads}, opts.dtype(torch::kInt32));
+    params.dq_semaphore = dq_semaphore.data_ptr<int>();
+    if (num_heads_k != num_heads && params.deterministic) {
+        // TODO: do we need to zero them out?
+        at::Tensor dk_semaphore = torch::empty({(seqlen_k + kBlockN - 1) / kBlockN, batch_size, num_heads_k}, opts.dtype(torch::kInt32));
+        at::Tensor dv_semaphore = torch::empty({(seqlen_k + kBlockN - 1) / kBlockN, batch_size, num_heads_k}, opts.dtype(torch::kInt32));
+        params.dk_semaphore = dk_semaphore.data_ptr<int>();
+        params.dv_semaphore = dv_semaphore.data_ptr<int>();
+    }
+
+    #ifdef FLASHATTENTION_DISABLE_LOCAL
+    TORCH_CHECK(!params.is_local, "This flash attention build does not support local attention.");
+    #endif
+    #ifdef FLASHATTENTION_DISABLE_SOFTCAP
+    TORCH_CHECK(params.softcap == 0.0, "This flash attention build does not support tanh softcapping.");
+    #endif
+
+    if (total_q > 0 && total_k > 0 && num_heads_k > 0) {
+        auto stream = at::cuda::getCurrentCUDAStream().stream();
+        run_mha_bwd(params, stream);
+    } else if (total_k > 0 && num_heads_k > 0) {
+        // If seqlen_q == 0, then we have an empty tensor. We need to set the output to 0.
+        dk.zero_();
+        dv.zero_();
+        softmax_d.zero_();
+    } else if (total_q > 0 && num_heads_k > 0) {
+        dq.zero_();
+        softmax_d.zero_();
+    }
+
+    return { dq, dk, dv, softmax_d, softmax_lse_log2, dq_accum, dk_accum, dv_accum };
+}
+
+std::tuple<at::Tensor, at::Tensor>
+mha_combine(at::Tensor out_partial,         // num_splits x batch_size x seqlen x num_heads x head_size
+            at::Tensor lse_partial,         // num_splits x batch_size x seqlen x num_heads
+            std::optional<at::Tensor> out_,        // batch_size x seqlen x num_heads x head_size
+            std::optional<at::ScalarType> out_dtype_
+            ) {
+
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    bool is_sm8x = dprops->major >= 8;
+    TORCH_CHECK(is_sm8x, "Attention combine function only supports Ampere GPUs or newer.");
+
+    auto out_partial_type = out_partial.scalar_type();
+    TORCH_CHECK(out_partial_type == at::ScalarType::Float, "Attention combine function only support fp32 data type");
+    TORCH_CHECK(lse_partial.scalar_type() == at::ScalarType::Float, "Attention combine function only support fp32 data type");
+
+    CHECK_DEVICE(out_partial); CHECK_DEVICE(lse_partial);
+
+    TORCH_CHECK(out_partial.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(lse_partial.stride(-2) == 1, "LSE tensor must be contiguous in the seqlen dimension");
+
+    const auto sizes = out_partial.sizes();
+
+    const int num_splits = sizes[0];
+    const int batch_size = sizes[1];
+    const int seqlen = sizes[2];
+    const int num_heads = sizes[3];
+    const int head_size_og = sizes[4];
+    TORCH_CHECK(num_splits <= 256, "FlashAttention combine only supports num_splits at most 256");
+
+    CHECK_SHAPE(out_partial, num_splits, batch_size, seqlen, num_heads, head_size_og);
+    CHECK_SHAPE(lse_partial, num_splits, batch_size, seqlen, num_heads);
+
+    int const alignment = 4;
+    at::Tensor out_partial_padded;
+    auto pad = [](at::Tensor x, int alignment) {
+        return x.size(-1) % alignment == 0 ? x : torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, alignment - x.size(-1) % alignment}));
+    };
+    out_partial_padded = pad(out_partial, alignment);
+
+    auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
+    const int head_size = round_multiple(head_size_og, alignment);
+
+    auto opts = out_partial.options();
+    at::ScalarType out_type = out_dtype_.value_or(out_partial.scalar_type());
+    TORCH_CHECK(out_type == at::ScalarType::Float || out_type == at::ScalarType::BFloat16 || out_type == at::ScalarType::Half, "Output type must be FP32, FP16 or BF16");
+    at::Tensor out;
+    if (out_.has_value()) {
+        out = out_.value();
+        TORCH_CHECK(out.scalar_type() == out_type);
+        CHECK_DEVICE(out);
+        TORCH_CHECK(out.stride(-1) == 1, "Output tensor must have contiguous last dimension");
+        CHECK_SHAPE(out, batch_size, seqlen, num_heads, head_size_og);
+        if (head_size_og % alignment != 0) {
+            out = torch::empty({batch_size, seqlen, num_heads, head_size}, opts.dtype(out_type));
+        }
+    } else {
+        out = torch::empty({batch_size, seqlen, num_heads, head_size}, opts.dtype(out_type));
+    }
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)out_partial.get_device()};
+
+    auto softmax_lse = torch::empty({batch_size, num_heads, seqlen}, opts.dtype(at::kFloat)).transpose(1, 2);
+
+    Flash_fwd_params params {};  // Need to reset the params to set everything to zero
+    params.is_fp32 = out_type == at::ScalarType::Float;
+    params.is_bf16 = out_type == at::ScalarType::BFloat16;
+    params.oaccum_ptr = out_partial_padded.data_ptr();
+    params.softmax_lseaccum_ptr = lse_partial.data_ptr();
+    params.o_ptr = out.data_ptr();
+    params.softmax_lse_ptr = softmax_lse.data_ptr();
+    params.b = batch_size;
+    params.h = num_heads;
+    params.seqlen_q = seqlen;
+    params.dv = head_size;
+    params.num_splits = num_splits;
+    params.oaccum_split_stride = out_partial_padded.stride(0);
+    params.oaccum_row_stride = out_partial_padded.stride(2);
+    params.oaccum_head_stride = out_partial_padded.stride(3);
+    params.oaccum_batch_stride = out_partial_padded.stride(1);
+    params.lseaccum_split_stride = lse_partial.stride(0);
+    params.lseaccum_head_stride = lse_partial.stride(3);
+    params.lseaccum_batch_stride = lse_partial.stride(1);
+    params.o_row_stride = out.stride(1);
+    params.o_head_stride = out.stride(2);
+    params.o_batch_stride = out.stride(0);
+    params.arch = at::cuda::getCurrentDeviceProperties()->major * 10 + at::cuda::getCurrentDeviceProperties()->minor;
+
+    if (seqlen > 0 && batch_size > 0) {
+        auto stream = at::cuda::getCurrentCUDAStream().stream();
+        run_mha_fwd_combine(params, stream, false /*enable_pdl*/);
+    }
+
+    at::Tensor out_padded = out;
+    if (head_size_og % alignment != 0) {
+        out = out.index({"...", torch::indexing::Slice(torch::indexing::None, head_size_og)});
+        // if (out_.has_value()) { out_.value().copy_(out); }
+    }
+
+    return {out, softmax_lse};
+}
+
+#ifdef false
+
+TORCH_LIBRARY(flash_attn_3, m) {
+    m.def("fwd("
+        "Tensor q,"
+        "Tensor k,"
+        "Tensor v,"
+        "Tensor(k_new!)? k_new = None,"
+        "Tensor(v_new!)? v_new = None,"
+        "Tensor? q_v = None,"
+        "Tensor(out!)? out = None,"
+        "Tensor? cu_seqlens_q = None,"
+        "Tensor? cu_seqlens_k = None,"
+        "Tensor? cu_seqlens_k_new = None,"
+        "Tensor? seqused_q = None,"
+        "Tensor? seqused_k = None,"
+        "int? max_seqlen_q = None,"
+        "int? max_seqlen_k = None,"
+        "Tensor? page_table = None,"
+        "Tensor? kv_batch_idx = None,"
+        "Tensor? leftpad_k = None,"
+        "Tensor? rotary_cos = None,"
+        "Tensor? rotary_sin = None,"
+        "Tensor? seqlens_rotary = None,"
+        "Tensor? q_descale = None,"
+        "Tensor? k_descale = None,"
+        "Tensor? v_descale = None,"
+        "float? softmax_scale = None,"
+        "bool is_causal = False,"
+        "int window_size_left = -1,"
+        "int window_size_right = -1,"
+        "int attention_chunk = 0,"
+        "float softcap = 0.0,"
+        "bool is_rotary_interleaved = False,"
+        "Tensor? scheduler_metadata = None,"
+        "int num_splits = 0,"
+        "bool? pack_gqa = None,"
+        "int sm_margin = 0) -> (Tensor(out!), Tensor, Tensor, Tensor)");
+    m.def("bwd("
+        "Tensor dout,"
+        "Tensor q,"
+        "Tensor k,"
+        "Tensor v,"
+        "Tensor out,"
+        "Tensor softmax_lse,"
+        "Tensor(dq!)? dq = None,"
+        "Tensor(dk!)? dk = None,"
+        "Tensor(dv!)? dv = None,"
+        "Tensor? cu_seqlens_q = None,"
+        "Tensor? cu_seqlens_k = None,"
+        "Tensor? seqused_q = None,"
+        "Tensor? seqused_k = None,"
+        "int? max_seqlen_q = None,"
+        "int? max_seqlen_k = None,"
+        "float? softmax_scale = None,"
+        "bool is_causal = False,"
+        "int window_size_left = -1,"
+        "int window_size_right = -1,"
+        "float softcap = 0.0,"
+        "bool deterministic = False,"
+        "int sm_margin = 0) -> (Tensor(dq!), Tensor(dk!), Tensor(dv!), Tensor, Tensor, Tensor, Tensor, Tensor)");
+    m.def("fwd_combine("
+        "Tensor out_partial,"
+        "Tensor lse_partial,"
+        "Tensor(out!)? out = None,"
+        "ScalarType? out_dtype = None) -> (Tensor(out!), Tensor)");
+    m.def("get_scheduler_metadata("
+        "int batch_size,"
+        "int max_seqlen_q,"
+        "int max_seqlen_k,"
+        "int num_heads,"
+        "int num_heads_k,"
+        "int headdim,"
+        "int headdim_v,"
+        "ScalarType qkv_dtype,"
+        "Tensor seqused_k,"
+        "Tensor? cu_seqlens_q = None,"
+        "Tensor? cu_seqlens_k = None,"
+        "Tensor? cu_seqlens_k_new = None,"
+        "Tensor? seqused_q = None,"
+        "Tensor? leftpad_k = None,"
+        "int? page_size = None,"
+        "int max_seqlen_k_new = 0,"
+        "bool is_causal = False,"
+        "int window_size_left = -1,"
+        "int window_size_right = -1,"
+        "int attention_chunk = 0,"
+        "bool has_softcap = False,"
+        "int num_splits = 0,"
+        "bool? pack_gqa = None,"
+        "int sm_margin = 0) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(flash_attn_3, CUDA, m) {
+    m.impl("fwd", &mha_fwd);
+    m.impl("bwd", &mha_bwd);
+    m.impl("fwd_combine", &mha_combine);
+    m.impl("get_scheduler_metadata", &mha_fwd_get_scheduler_metadata);
+}
+
+#endif

flash-attn/flash_bwd_kernel_sm80.h

@@ -0,0 +1,173 @@
+/******************************************************************************
+ * Copyright (c) 2024, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/kernel_hardware_info.h>
+
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
+class FlashAttnBwdSm80 {
+
+public:
+
+    // Type Aliases
+    static constexpr bool Is_causal = CollectiveMainloop_::Is_causal;
+    static constexpr bool Is_local = CollectiveMainloop_::Is_local;
+    static_assert(CollectiveMainloop_::Varlen == CollectiveEpilogue_::Varlen);
+    static constexpr bool Varlen = CollectiveMainloop_::Varlen;
+
+    // Mainloop derived types
+    using CollectiveMainloop = CollectiveMainloop_;
+    using TileShape_MNK = typename CollectiveMainloop::TileShape_MNK;
+    using TiledMmaSdP = typename CollectiveMainloop::TiledMmaSdP;
+    using TiledMmadKV = typename CollectiveMainloop::TiledMmadKV;
+    using ArchTag = typename CollectiveMainloop::ArchTag;
+    using MainloopArguments = typename CollectiveMainloop::Arguments;
+    using MainloopParams = typename CollectiveMainloop::Params;
+    static constexpr bool dKV_swapAB = CollectiveMainloop::dKV_swapAB;
+
+    // Epilogue derived types
+    using CollectiveEpilogue = CollectiveEpilogue_;
+    using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+    using EpilogueParams = typename CollectiveEpilogue::Params;
+
+    static_assert(ArchTag::kMinComputeCapability >= 80);
+
+    using TileScheduler = TileScheduler_;
+    using TileSchedulerArguments = typename flash::TileSchedulerArguments;
+    using TileSchedulerParams = typename TileScheduler::Params;
+
+    static constexpr uint32_t NumThreads = CUTE_STATIC_V(size(TiledMmaSdP{}));
+    static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMmaSdP{}));
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+
+    // Kernel level shared memory storage
+    struct SharedStorage {
+        struct TensorStorage : cute::aligned_struct<128> {
+            union {
+                typename CollectiveMainloop::TensorStorage mainloop;
+                typename CollectiveEpilogue::TensorStorage epilogue;
+            };
+        } tensors;
+
+        alignas(16) typename TileScheduler::SharedStorage smem_scheduler;
+
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    // Device side arguments
+    struct Arguments {
+        MainloopArguments mainloop{};
+        EpilogueArguments epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerArguments scheduler{};
+    };
+
+    // Kernel entry point API
+    struct Params {
+        MainloopParams mainloop{};
+        EpilogueParams epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerParams scheduler{};
+    };
+
+    //
+    // Methods
+    //
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        CUTLASS_TRACE_HOST("to_underlying_arguments():");
+
+        // Get SM count if needed, otherwise use user supplied SM count
+        int sm_count = args.hw_info.sm_count;
+        if (sm_count <= 0) {
+            CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+                "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+            sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+        }
+
+        CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
+
+        cutlass::KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
+        return {
+            CollectiveMainloop::to_underlying_arguments(args.mainloop),
+            CollectiveEpilogue::to_underlying_arguments(args.epilogue),
+            hw_info,
+            TileScheduler::to_underlying_arguments(args.scheduler)
+        };
+    }
+
+    // Computes the kernel launch grid shape based on runtime parameters
+    static dim3
+    get_grid_shape(Params const& params) {
+        return TileScheduler::get_grid_shape(params.scheduler, params.hw_info.sm_count);
+    }
+
+    static dim3
+    get_block_shape() {
+        return dim3(MaxThreadsPerBlock, 1, 1);
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        static constexpr int kBlockM = get<0>(TileShape_MNK{});
+        static constexpr int kBlockN = get<1>(TileShape_MNK{});
+
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+        CollectiveMainloop mainloop;
+        CollectiveEpilogue epilogue;
+
+        TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.smem_scheduler));
+        // Initialize matmul objects.
+        TiledMmadKV tiled_mma_dKV;
+
+        scheduler.init_consumer();
+
+        int warp_idx = cutlass::canonical_warp_idx_sync();
+        CUTLASS_PRAGMA_NO_UNROLL
+        for (auto work_tile_info = warp_idx == 0 ? scheduler.template get_initial_work</*IsProducerWarp=*/true>(params.scheduler) : scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+             work_tile_info.is_valid(params.scheduler);
+             work_tile_info = warp_idx == 0 ? scheduler.template get_next_work</*IsProducerWarp=*/true>(params.scheduler, work_tile_info) : scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
+
+            auto block_coord_ = work_tile_info.get_block_coord(params.scheduler);
+            auto [n_block, bidh, bidb, _ /*split_idx*/] = block_coord_;
+            cute::tuple<int32_t, int32_t, int32_t> block_coord = {n_block, bidh, bidb};
+
+            // dK and dV output accumulator.
+            Tensor tdKrdK = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
+            Tensor tdVrdV = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
+            bool tile_valid = mainloop.mma(params.mainloop, tdKrdK, tdVrdV, threadIdx.x,
+                                           block_coord, shared_storage);
+            scheduler.prefetch_next_work(params.scheduler, work_tile_info);
+            if (tile_valid) {
+                epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV,
+                               threadIdx.x, block_coord);
+            } else {
+                epilogue.store_zero(params.epilogue, threadIdx.x, block_coord);
+            }
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_bwd_kernel_sm90.h

@@ -0,0 +1,282 @@
+
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/kernel_hardware_info.h>
+#include "cutlass/pipeline/pipeline.hpp"
+
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
+class FlashAttnBwdSm90 {
+
+public:
+
+    // Type Aliases
+    static constexpr bool Is_causal = CollectiveMainloop_::Is_causal;
+    static constexpr bool Is_local = CollectiveMainloop_::Is_local;
+    static_assert(CollectiveMainloop_::Varlen == CollectiveEpilogue_::Varlen);
+    static constexpr bool Varlen = CollectiveMainloop_::Varlen;
+
+    // Mainloop derived types
+    using CollectiveMainloop = CollectiveMainloop_;
+    using TileShape_MNK = typename CollectiveMainloop::TileShape_MNK;
+    using TiledMmaSdP = typename CollectiveMainloop::TiledMmaSdP;
+    using TiledMmadKV = typename CollectiveMainloop::TiledMmadKV;
+    using ArchTag = typename CollectiveMainloop::ArchTag;
+    using ClusterShape = typename CollectiveMainloop::ClusterShape;
+    using MainloopArguments = typename CollectiveMainloop::Arguments;
+    using MainloopParams = typename CollectiveMainloop::Params;
+    static constexpr bool dKV_swapAB = CollectiveMainloop::dKV_swapAB;
+
+    // Epilogue derived types
+    using CollectiveEpilogue = CollectiveEpilogue_;
+    using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+    using EpilogueParams = typename CollectiveEpilogue::Params;
+
+    static_assert(ArchTag::kMinComputeCapability >= 90);
+
+    using TileScheduler = TileScheduler_;
+    using TileSchedulerArguments = typename flash::TileSchedulerArguments;
+    using TileSchedulerParams = typename TileScheduler::Params;
+
+    static constexpr uint32_t NumLoadWarpGroups = 1;
+    static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMmaSdP{})) / cutlass::NumThreadsPerWarpGroup;
+    static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMmaSdP{})) + (NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup);
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+    static_assert(NumMmaWarpGroups == 2 || NumMmaWarpGroups == 3);
+
+    /// Register requirement for Load and Math WGs
+    static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 2 ? 24 : 32;
+    static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 240 : 160;
+    // If you want to print from the producer warp, you'd need to increase the number of registers
+    // Otherwise you'll get CUDA error.
+    // static constexpr uint32_t LoadRegisterRequirement = 40;
+    // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152;
+
+    // Kernel level shared memory storage
+    struct SharedStorage {
+        struct TensorStorage : cute::aligned_struct<128> {
+            union {
+                typename CollectiveMainloop::TensorStorage mainloop;
+                typename CollectiveEpilogue::TensorStorage epilogue;
+            };
+        } tensors;
+
+        struct PipelineStorage : cute::aligned_struct<16> {
+            alignas(16) cutlass::arch::ClusterTransactionBarrier barrier_KV;
+            alignas(16) typename CollectiveMainloop::MainloopPipeline::SharedStorage pipeline_q;
+            alignas(16) typename CollectiveMainloop::MainloopPipeline_dO::SharedStorage pipeline_do;
+            alignas(16) typename TileScheduler::SharedStorage smem_scheduler;
+        } pipelines;
+
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    // Device side arguments
+    struct Arguments {
+        MainloopArguments mainloop{};
+        EpilogueArguments epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerArguments scheduler{};
+    };
+
+    // Kernel entry point API
+    struct Params {
+        MainloopParams mainloop{};
+        EpilogueParams epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerParams scheduler{};
+    };
+
+    //
+    // Methods
+    //
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        CUTLASS_TRACE_HOST("to_underlying_arguments():");
+
+        // Get SM count if needed, otherwise use user supplied SM count
+        int sm_count = args.hw_info.sm_count;
+        if (sm_count <= 0) {
+            CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+                "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+            sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+        }
+
+        CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
+
+        cutlass::KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
+        return {
+            CollectiveMainloop::to_underlying_arguments(args.mainloop),
+            CollectiveEpilogue::to_underlying_arguments(args.epilogue),
+            hw_info,
+            TileScheduler::to_underlying_arguments(args.scheduler)
+        };
+    }
+
+    // Computes the kernel launch grid shape based on runtime parameters
+    static dim3
+    get_grid_shape(Params const& params) {
+        return TileScheduler::get_grid_shape(params.scheduler, params.hw_info.sm_count);
+    }
+
+    static dim3
+    get_block_shape() {
+        return dim3(MaxThreadsPerBlock, 1, 1);
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        static constexpr int NumMmaThreads = NumMmaWarpGroups * cutlass::NumThreadsPerWarpGroup;
+        static constexpr int NumCopyThreads = NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup;
+        static constexpr int kBlockM = get<0>(TileShape_MNK{});
+        static constexpr int kBlockN = get<1>(TileShape_MNK{});
+
+        using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
+        using PipelineParams = typename MainloopPipeline::Params;
+        using PipelineState = typename MainloopPipeline::PipelineState;
+        using MainloopPipeline_dO = typename CollectiveMainloop::MainloopPipeline_dO;
+        using PipelineParams_dO = typename MainloopPipeline_dO::Params;
+        using PipelineState_dO = typename MainloopPipeline_dO::PipelineState;
+        static constexpr bool Q_dO_same_stages = std::is_same_v<MainloopPipeline, MainloopPipeline_dO>;
+
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+        int const lane_predicate = cute::elect_one_sync();
+        int const warp_idx = cutlass::canonical_warp_idx_sync();
+
+        // Issue Tma Descriptor Prefetch from a single thread
+        if (warp_idx == 0 && lane_predicate) {
+            CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
+            CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue);
+        }
+
+        // Obtain warp index
+        int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+
+        PipelineParams pipeline_params;
+        pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesQ + CollectiveMainloop::TmaTransactionBytesLSE;
+        int warp_group_idx = cutlass::canonical_warp_group_idx();
+        pipeline_params.role = warp_group_idx == 0
+            ? MainloopPipeline::ThreadCategory::Producer
+            : MainloopPipeline::ThreadCategory::Consumer;
+        pipeline_params.is_leader = warp_group_thread_idx == 0;
+        pipeline_params.num_consumers = NumMmaThreads;
+
+        if (warp_idx == 0 && lane_predicate) {
+            shared_storage.pipelines.barrier_KV.init(1 /*numThreads*/);
+        }
+        // We're counting on pipeline_q to call cutlass::arch::fence_barrier_init();
+        MainloopPipeline pipeline_q(shared_storage.pipelines.pipeline_q, pipeline_params, ClusterShape{});
+        auto role_dO = warp_group_idx == 0
+            ? MainloopPipeline_dO::ThreadCategory::Producer
+            : MainloopPipeline_dO::ThreadCategory::Consumer;
+        PipelineParams_dO pipeline_params_dO {pipeline_params.transaction_bytes, role_dO, pipeline_params.is_leader, pipeline_params.num_consumers};
+        MainloopPipeline_dO pipeline_do(shared_storage.pipelines.pipeline_do, cute::conditional_return<Q_dO_same_stages>(pipeline_params, pipeline_params_dO), ClusterShape{});
+
+        CollectiveMainloop mainloop;
+        CollectiveEpilogue epilogue;
+
+        // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster
+        if constexpr (size(ClusterShape{}) > 1) {
+            cute::cluster_arrive_relaxed();
+            cute::cluster_wait();
+        } else {
+            __syncthreads();
+        }
+
+        TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.pipelines.smem_scheduler));
+
+        if (warp_group_idx == 0) {  // Producer
+            cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
+
+            int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
+            if (warp_idx_in_warpgroup == 0) {  // Load K, V, and do TMA on Q and dO
+                PipelineState smem_pipe_write = cutlass::make_producer_start_state<MainloopPipeline>();
+                PipelineState_dO smem_pipe_write_do = cutlass::make_producer_start_state<MainloopPipeline_dO>();
+                for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/true>(params.scheduler);
+                     work_tile_info.is_valid(params.scheduler);
+                     work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/true>(params.scheduler, work_tile_info)) {
+                    auto block_coord_ = work_tile_info.get_block_coord(params.scheduler);
+                    auto [n_block, bidh, bidb, _ /*split_idx*/] = block_coord_;
+                    cute::tuple<int32_t, int32_t, int32_t> block_coord = {n_block, bidh, bidb};
+                    auto scheduler_prefetch = [&scheduler, &params, &work_tile_info]() {
+                        scheduler.prefetch_next_work(params.scheduler, work_tile_info);
+                    };
+                    mainloop.load(params.mainloop, pipeline_q, pipeline_do, smem_pipe_write,
+                                  smem_pipe_write_do, shared_storage, scheduler_prefetch, block_coord);
+                }
+                mainloop.load_tail(pipeline_q, pipeline_do, smem_pipe_write, smem_pipe_write_do);
+            } else if (warp_idx_in_warpgroup == 1) {
+                for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+                     work_tile_info.is_valid(params.scheduler);
+                     work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
+                    auto block_coord_ = work_tile_info.get_block_coord(params.scheduler);
+                    auto [n_block, bidh, bidb, _ /*split_idx*/] = block_coord_;
+                    cute::tuple<int32_t, int32_t, int32_t> block_coord = {n_block, bidh, bidb};
+                    mainloop.store_dq(params.mainloop, shared_storage, block_coord);
+                }
+            }
+        } else {  // Consumer
+            cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
+            // Initialize matmul objects.
+            TiledMmadKV tiled_mma_dKV;
+
+            PipelineState smem_pipe_read;
+            PipelineState_dO smem_pipe_read_do;
+
+            mainloop.mma_init();
+            scheduler.init_consumer();
+
+            int work_idx = 0;
+            CUTLASS_PRAGMA_NO_UNROLL
+            for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+                 work_tile_info.is_valid(params.scheduler);
+                 work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
+                auto block_coord_ = work_tile_info.get_block_coord(params.scheduler);
+                auto [n_block, bidh, bidb, _ /*split_idx*/] = block_coord_;
+                cute::tuple<int32_t, int32_t, int32_t> block_coord = {n_block, bidh, bidb};
+
+                // dK and dV output accumulator.
+                Tensor tdKrdK = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
+                Tensor tdVrdV = partition_fragment_C(tiled_mma_dKV, select<!dKV_swapAB ? 1 : 2, !dKV_swapAB? 2 : 1>(TileShape_MNK{}));
+                bool tile_valid = mainloop.mma(
+                    params.mainloop, pipeline_q, pipeline_do, smem_pipe_read, smem_pipe_read_do,
+                    tdKrdK, tdVrdV, threadIdx.x - NumCopyThreads, work_idx, block_coord, shared_storage);
+                if (tile_valid) {
+                    epilogue.store(params.epilogue, tdKrdK, tdVrdV, shared_storage, tiled_mma_dKV,
+                                   threadIdx.x - NumCopyThreads, block_coord);
+                } else {
+                    epilogue.store_zero(params.epilogue, threadIdx.x - NumCopyThreads, block_coord);
+                }
+
+            }
+            epilogue.store_tail();
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_bwd_launch_template.h

@@ -0,0 +1,390 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/device_kernel.h"  // For device_kernel
+#include "cutlass/kernel_launch.h"  // For kernel_launch
+#include "cutlass/cluster_launch.hpp"  // For ClusterLauncher
+
+#include "static_switch.h"
+#include "flash.h"
+#include "flash_bwd_preprocess_kernel.h"
+#include "flash_bwd_postprocess_kernel.h"
+#include "tile_scheduler.hpp"
+#include "mainloop_bwd_sm90_tma_gmma_ws.hpp"
+#include "mainloop_bwd_sm80.hpp"
+#include "epilogue_bwd.hpp"
+#include "flash_bwd_kernel_sm90.h"
+#include "flash_bwd_kernel_sm80.h"
+
+using namespace cute;
+
+template <int Arch, int kHeadDim, int kBlockM, int kBlockN, typename Element,
+          bool Is_causal, bool Is_local, bool Has_softcap, bool Varlen, bool Deterministic, bool GQA,
+          int Stages_dO=2, int Stages_dS_or_QSm80=2,
+          bool SdP_swapAB=true, bool dKV_swapAB=false, bool dQ_swapAB=false,
+          int NumMmaWarpGroups=2, int AtomLayoutMSdP=1, int AtomLayoutNdKV=2, int AtomLayoutMdQ=1,
+          bool V_in_regs=false>
+void run_flash_bwd(Flash_bwd_params &params, cudaStream_t stream) {
+    static_assert(!(Is_causal && Is_local), "Is_causal and Is_local cannot be true at the same time.");
+    using ElementAccum = float;
+    using ArchTag = std::conditional_t<Arch >= 90, cutlass::arch::Sm90, cutlass::arch::Sm80>;
+
+    int const total_q_padded_rounded = cute::round_up(params.total_q + params.b * kBlockM, kBlockM);
+    int const total_k_padded_rounded = cute::round_up(params.total_k + params.b * kBlockN, kBlockN);
+    bool const is_varlen_q = params.cu_seqlens_q;
+    bool const is_varlen_k = params.cu_seqlens_k;
+    int seqlen_q = !is_varlen_q ? params.seqlen_q : params.total_q;
+    int seqlen_k = !is_varlen_k ? params.seqlen_k : params.total_k;
+    int seqlen_q_rounded = !is_varlen_q ? params.seqlen_q_rounded : total_q_padded_rounded;
+    int seqlen_k_rounded = !is_varlen_k ? params.seqlen_k_rounded : total_k_padded_rounded;
+    int batch_q = !is_varlen_q ? params.b : 1;
+    int batch_k = !is_varlen_k ? params.b : 1;
+
+    using TileShape_MK = cute::Shape<Int<kBlockM>, Int<kHeadDim>>;
+    using PreprocessKernel = flash::FlashAttnBwdPreprocess<TileShape_MK, Element, ElementAccum, ArchTag, /*Clear_dQaccum=*/true, Varlen>;
+    typename PreprocessKernel::Arguments preprocess_args {
+        static_cast<Element const*>(params.o_ptr),
+        {seqlen_q, params.dv, params.h, batch_q},  // shape_O
+        {params.o_row_stride, _1{}, params.o_head_stride, !is_varlen_q ? params.o_batch_stride : 0},  // stride_O
+        static_cast<Element const*>(params.do_ptr),
+        {params.do_row_stride, _1{}, params.do_head_stride, !is_varlen_q ? params.do_batch_stride : 0},  // stride_dO
+        static_cast<float*>(params.dsoftmax_sum),
+        {seqlen_q_rounded, params.h, batch_q},  // shape_dPsum
+        {_1{}, seqlen_q_rounded, !is_varlen_q ? params.h * params.seqlen_q_rounded : 0},  // stride_dPsum
+        static_cast<float*>(params.softmax_lse_ptr),
+        {_1{}, seqlen_q, !is_varlen_q ? params.h * params.seqlen_q : 0},  // stride_LSE
+        static_cast<float*>(params.softmax_lse_log2_ptr),
+        {_1{}, seqlen_q_rounded, !is_varlen_q ? params.h * params.seqlen_q_rounded : 0},  // stride_LSE_log2
+        static_cast<ElementAccum*>(params.dq_accum_ptr),
+        {seqlen_q_rounded * params.d_rounded, params.h, batch_q},  // shape_dQaccum
+        {_1{}, seqlen_q_rounded * params.d_rounded, !is_varlen_q ? params.d_rounded * seqlen_q_rounded * params.h : 0},  // stride_dQaccum
+        params.b,
+        params.dq_semaphore,
+        params.cu_seqlens_q,
+        params.seqused_q
+    };
+    typename PreprocessKernel::Params preprocess_params = PreprocessKernel::to_underlying_arguments(preprocess_args);
+    int num_m_block = cute::ceil_div(params.seqlen_q, kBlockM);
+    dim3 grid_m(num_m_block, params.h, params.b);
+    cutlass::kernel_launch<PreprocessKernel>(grid_m, PreprocessKernel::MaxThreadsPerBlock, PreprocessKernel::SharedStorageSize, stream, preprocess_params, false /*launch_with_pdl*/);
+    CHECK_CUDA_KERNEL_LAUNCH();
+
+    using TileShape_MNK = cute::Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
+    using ClusterShape = cute::Shape<_1, Int<1>, _1>;  // Currently doesn't not support cluster
+    // Stages_dS_or_QSm80 is Stages_dS if Sm90 and Stages if Sm80
+    static constexpr int Stages = Arch >= 90 ? 2 : Stages_dS_or_QSm80;
+    static constexpr int Stages_dS = Arch >= 90 ? Stages_dS_or_QSm80 : 1;
+    using CollectiveMainloop = std::conditional_t<
+        Arch >= 90,
+        flash::CollectiveMainloopBwdSm90<Stages, Stages_dO, Stages_dS, ClusterShape, TileShape_MNK, Element, ElementAccum, cutlass::arch::Sm90,
+            Is_causal, Is_local, Has_softcap, Varlen, Deterministic,
+            SdP_swapAB, dKV_swapAB, dQ_swapAB, NumMmaWarpGroups, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs>,
+        flash::CollectiveMainloopBwdSm80<Stages, Stages_dO, TileShape_MNK, Element, ElementAccum, cutlass::arch::Sm80,
+            Is_causal, Is_local, Has_softcap, Varlen, Deterministic,
+            SdP_swapAB, dKV_swapAB, dQ_swapAB, NumMmaWarpGroups, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs>
+    >;
+    using CollectiveEpilogue = std::conditional_t<
+        !GQA,
+        flash::CollectiveEpilogueBwd<TileShape_MNK, Element, ArchTag, CollectiveMainloop::NumMmaThreads, Varlen, dKV_swapAB, NumMmaWarpGroups * (Arch >= 90 ? 1 : cutlass::NumWarpsPerWarpGroup) / AtomLayoutNdKV>,
+        flash::CollectiveEpilogueBwdGQA<TileShape_MNK, ElementAccum, ArchTag, CollectiveMainloop::NumMmaThreads, Varlen, Deterministic>
+    >;
+    using Scheduler = std::conditional_t<
+        Is_causal && !Varlen,
+        flash::SingleTileBwdLPTScheduler,
+        flash::SingleTileScheduler<Varlen, false /*Split*/, false /*PackGQA*/, kBlockN>
+    >;
+    using AttnKernel = std::conditional_t<
+        Arch >= 90,
+        flash::enable_sm90_or_later<flash::FlashAttnBwdSm90<CollectiveMainloop, CollectiveEpilogue, Scheduler>>,
+        flash::enable_sm80_to_sm89<flash::FlashAttnBwdSm80<CollectiveMainloop, CollectiveEpilogue, Scheduler>>
+    >;
+
+    typename CollectiveMainloop::Arguments mainloop_args {
+        static_cast<Element const*>(params.q_ptr),
+        {seqlen_q, params.d, params.h, batch_q},  // shape_Q
+        {params.q_row_stride, _1{}, params.q_head_stride, !is_varlen_q ? params.q_batch_stride : 0},  // stride_Q
+        static_cast<Element const*>(params.k_ptr),
+        {seqlen_k, params.d, params.h_k, batch_k},  // shape_K
+        {params.k_row_stride, _1{}, params.k_head_stride, !is_varlen_k ? params.k_batch_stride : 0},  // stride_K
+        static_cast<Element const*>(params.v_ptr),
+        {seqlen_k, params.dv, params.h_k, batch_k},  // shape_V
+        {params.v_row_stride, _1{}, params.v_head_stride, !is_varlen_k ? params.v_batch_stride : 0},  // stride_V
+        static_cast<Element const*>(params.do_ptr),
+        {seqlen_q, params.dv, params.h, batch_q},  // shape_dO
+        {params.do_row_stride, _1{}, params.do_head_stride, !is_varlen_q ? params.do_batch_stride : 0},  // stride_dO
+        static_cast<ElementAccum*>(params.dq_accum_ptr),
+        {seqlen_q_rounded * params.d_rounded, params.h, batch_q},  // shape_dQaccum
+        {_1{}, seqlen_q_rounded * params.d_rounded, !is_varlen_q ? params.d_rounded * params.seqlen_q_rounded * params.h : 0}, // stride_dQaccum
+        static_cast<float*>(params.softmax_lse_log2_ptr),
+        {seqlen_q_rounded, params.h, batch_q},  // shape_LSE
+        {_1{}, seqlen_q_rounded, !is_varlen_q ? params.h * params.seqlen_q_rounded : 0},  // stride_LSE_log2
+        static_cast<float*>(params.dsoftmax_sum),
+        {_1{}, seqlen_q_rounded, !is_varlen_q ? params.h * params.seqlen_q_rounded : 0},  // stride_dPsum
+        params.scale_softmax,
+        params.window_size_left, params.window_size_right, 0 /*attention_chunk*/,
+        params.softcap,
+        params.b,
+        params.dq_semaphore,
+        params.cu_seqlens_q, params.cu_seqlens_k,
+        params.seqused_q, params.seqused_k
+    };
+    // The case work with GQA is ugly but idk how to fix it.
+    typename CollectiveEpilogue::Arguments epilogue_args {
+        static_cast<typename CollectiveEpilogue::Element*>(!GQA ? params.dk_ptr : params.dk_accum_ptr),
+        [&] {
+            if constexpr (!GQA) {
+                return typename CollectiveEpilogue::ShapedKV {seqlen_k, params.d, params.h, batch_k};  // shape_dK
+            } else {
+                return typename CollectiveEpilogue::ShapedKV {seqlen_k_rounded * params.d_rounded, params.h_k, batch_k};  // shape_dKaccum
+            }
+        }(),
+        [&] {
+            if constexpr (!GQA) {
+                return typename CollectiveEpilogue::StridedKV {params.dk_row_stride, _1{}, params.dk_head_stride, !is_varlen_k ? params.dk_batch_stride : 0};  // stride_dK
+            } else {
+                return typename CollectiveEpilogue::StridedKV {_1{}, params.d_rounded * seqlen_k_rounded, !is_varlen_k ? params.h_k * params.d_rounded * params.seqlen_k_rounded : 0};  // stride_dKaccum
+            }
+        }(),
+        static_cast<typename CollectiveEpilogue::Element*>(!GQA ? params.dv_ptr : params.dv_accum_ptr),
+        [&] {
+            if constexpr (!GQA) {
+                return typename CollectiveEpilogue::ShapedKV {seqlen_k, params.dv, params.h, batch_k};  // shape_dV
+            } else {
+                return typename CollectiveEpilogue::ShapedKV {seqlen_k_rounded * params.dv_rounded, params.h_k, batch_k};  // shape_dVaccum
+            }
+        }(),
+        [&] {
+            if constexpr (!GQA) {
+                return typename CollectiveEpilogue::StridedKV {params.dv_row_stride, _1{}, params.dv_head_stride, !is_varlen_k ? params.dv_batch_stride : 0};  // stride_dV
+            } else {
+                return typename CollectiveEpilogue::StridedKV {_1{}, params.dv_rounded * seqlen_k_rounded, !is_varlen_k ? params.h_k * params.dv_rounded * params.seqlen_k_rounded : 0};  // stride_dVaccum
+            }
+        }(),
+        params.h,
+        params.dk_semaphore,
+        params.dv_semaphore,
+        params.cu_seqlens_k,
+        params.seqused_k,
+    };
+
+    int num_blocks_n = cutlass::ceil_div(params.seqlen_k, get<1>(TileShape_MNK{}));
+    num_blocks_n = cutlass::round_up(num_blocks_n, size<1>(ClusterShape{}));
+    typename flash::TileSchedulerArguments scheduler_args {
+        num_blocks_n, params.h, params.b, 1 /*num_splits*/,
+        params.h / params.h_k,
+        params.seqlen_k,
+        params.seqlen_q, params.d, params.dv, sizeof(Element),
+        params.tile_count_semaphore, params.cu_seqlens_k, params.seqused_k
+    };
+
+    int device;
+    cudaGetDevice(&device);
+    typename AttnKernel::Params kernel_params = AttnKernel::to_underlying_arguments({
+        mainloop_args, epilogue_args, {device, params.num_sm}, scheduler_args
+    });
+
+    dim3 grid_dims = AttnKernel::get_grid_shape(kernel_params);
+    dim3 block_dims = AttnKernel::get_block_shape();
+    int smem_size = AttnKernel::SharedStorageSize;
+    // int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q));
+    // int smem_size_do = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_do));
+    // int smem_size_ds = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_ds));
+    // int smem_size_dqacc = [&] {
+    //     if constexpr (Arch >= 90) {
+    //         return sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_dqacc));
+    //     } else {
+    //         return 0;
+    //     }
+    // }();
+    // int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k));
+    // int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v));
+    // int smem_size_lse = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_lse));
+    // int smem_size_dpsum = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_dpsum));
+    // printf("smem_size = %d, q = %d, k = %d, v = %d, do = %d, ds = %d, dqacc = %d, lse = %d, dpsum = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v, smem_size_do, smem_size_ds, smem_size_dqacc, smem_size_lse, smem_size_dpsum);
+    if constexpr (size(ClusterShape{}) > 1) {
+        void const* kernel = (void const*) cutlass::device_kernel<AttnKernel>;
+        if (smem_size >= 48 * 1024) {
+            CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+        }
+        dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
+        cutlass::ClusterLauncher::launch(
+            grid_dims, cluster_dims, block_dims, smem_size, stream, kernel, kernel_params, false /*launch_with_pdl*/);
+    } else {
+        if (smem_size >= 48 * 1024) {
+            CHECK_CUDA(cudaFuncSetAttribute(cutlass::device_kernel<AttnKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+        }
+        cutlass::kernel_launch<AttnKernel>(grid_dims, block_dims, smem_size, stream, kernel_params, false /*launch_with_pdl*/);
+    }
+    CHECK_CUDA_KERNEL_LAUNCH();
+
+    using PostprocessKernel = flash::FlashAttnBwdPostprocessConvertdQ<TileShape_MK, Element, ElementAccum, ArchTag,
+        AttnKernel::CollectiveMainloop::NumMmaThreads,
+        typename AttnKernel::CollectiveMainloop::TiledMmadQ,
+        AttnKernel::CollectiveMainloop::dQ_swapAB
+        >;
+    typename PostprocessKernel::Arguments postprocess_args {
+        static_cast<ElementAccum const*>(params.dq_accum_ptr),
+        {seqlen_q_rounded * params.d_rounded, params.h, batch_q},  // shape_dQaccum
+        {_1{}, seqlen_q_rounded * params.d_rounded, !is_varlen_q ? params.d_rounded * params.seqlen_q_rounded * params.h : 0}, // stride_dQaccum
+        static_cast<Element*>(params.dq_ptr),
+        {seqlen_q, params.d, params.h, batch_q},  // shape_dQ
+        {params.dq_row_stride, _1{}, params.dq_head_stride, params.dq_batch_stride},  // stride_dQ
+        params.scale_softmax,
+        params.cu_seqlens_q,
+        params.seqused_q
+    };
+    typename PostprocessKernel::Params postprocess_params = PostprocessKernel::to_underlying_arguments(postprocess_args);
+    int num_m_block_postprocess = cute::ceil_div(params.seqlen_q, get<0>(TileShape_MK{}));
+    dim3 grid_m_postprocess(num_m_block_postprocess, params.h, params.b);
+    int smem_size_postprocess = PostprocessKernel::SharedStorageSize;
+    if (smem_size_postprocess >= 48 * 1024) {
+        CHECK_CUDA(cudaFuncSetAttribute(cutlass::device_kernel<PostprocessKernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_postprocess));
+    }
+    cutlass::kernel_launch<PostprocessKernel>(grid_m_postprocess, PostprocessKernel::MaxThreadsPerBlock, smem_size_postprocess, stream, postprocess_params, false /*launch_with_pdl*/);
+    CHECK_CUDA_KERNEL_LAUNCH();
+
+    if constexpr (GQA) {
+        using TileShape_NK = cute::Shape<Int<kBlockN>, Int<kHeadDim>>;
+        using PostprocessKerneldKV = flash::FlashAttnBwdPostprocessConvertdQ<TileShape_NK, Element, ElementAccum, ArchTag,
+            AttnKernel::CollectiveEpilogue::NumEpilogueThreads,
+            typename AttnKernel::CollectiveMainloop::TiledMmadKV,
+            AttnKernel::CollectiveMainloop::dKV_swapAB
+            >;
+        typename PostprocessKerneldKV::Arguments postprocess_dK_args {
+            static_cast<ElementAccum const*>(params.dk_accum_ptr),
+            {seqlen_k_rounded * params.d_rounded, params.h_k, batch_k},  // shape_dKaccum
+            {_1{}, seqlen_k_rounded * params.d_rounded, !is_varlen_k ? params.d_rounded * params.seqlen_k_rounded * params.h_k : 0},  // stride_dKaccum
+            static_cast<Element*>(params.dk_ptr),
+            {seqlen_k, params.d, params.h_k, batch_k},  // shape_dK
+            {params.dk_row_stride, _1{}, params.dk_head_stride, params.dk_batch_stride},  // stride_dK
+            1.f,
+            params.cu_seqlens_k,
+            params.seqused_k
+        };
+        typename PostprocessKerneldKV::Params postprocess_dK_params = PostprocessKerneldKV::to_underlying_arguments(postprocess_dK_args);
+        typename PostprocessKerneldKV::Arguments postprocess_dV_args {
+            static_cast<ElementAccum const*>(params.dv_accum_ptr),
+            {seqlen_k_rounded * params.dv_rounded, params.h_k, batch_k},  // shape_dVaccum
+            {_1{}, seqlen_k_rounded * params.dv_rounded, !is_varlen_k ? params.dv_rounded * params.seqlen_k_rounded * params.h_k : 0},  // stride_dVaccum
+            static_cast<Element*>(params.dv_ptr),
+            {seqlen_k, params.dv, params.h_k, batch_k},  // shape_dV
+            {params.dv_row_stride, _1{}, params.dv_head_stride, params.dv_batch_stride},  // stride_dV
+            1.f,
+            params.cu_seqlens_k,
+            params.seqused_k
+        };
+        typename PostprocessKerneldKV::Params postprocess_dV_params = PostprocessKerneldKV::to_underlying_arguments(postprocess_dV_args);
+        int num_n_block_postprocess = cute::ceil_div(params.seqlen_k, get<0>(TileShape_NK{}));
+        dim3 grid_n_postprocess(num_n_block_postprocess, params.h_k, params.b);
+        int smem_size_postprocess = PostprocessKerneldKV::SharedStorageSize;
+        if (smem_size_postprocess >= 48 * 1024) {
+            CHECK_CUDA(cudaFuncSetAttribute(cutlass::device_kernel<PostprocessKerneldKV>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_postprocess));
+        }
+        cutlass::kernel_launch<PostprocessKerneldKV>(grid_n_postprocess, PostprocessKerneldKV::MaxThreadsPerBlock, smem_size_postprocess, stream, postprocess_dK_params, false /*launch_with_pdl*/);
+        CHECK_CUDA_KERNEL_LAUNCH();
+        cutlass::kernel_launch<PostprocessKerneldKV>(grid_n_postprocess, PostprocessKerneldKV::MaxThreadsPerBlock, smem_size_postprocess, stream, postprocess_dV_params, false /*launch_with_pdl*/);
+        CHECK_CUDA_KERNEL_LAUNCH();
+    }
+
+}
+
+template<int Arch, typename T, int kBlockM, int kBlockN, int kHeadDim, bool Is_causal, bool Is_local, bool Has_softcap,
+         int Stages_dO=2, int Stages_dS_or_QSm80=2,
+         bool SdP_swapAB=true, bool dKV_swapAB=false, bool dQ_swapAB=false,
+         int NumMmaWarpGroups=2, int AtomLayoutMSdP=1, int AtomLayoutNdKV=2, int AtomLayoutMdQ=1,
+         bool V_in_regs=false>
+void run_mha_bwd_dispatch(Flash_bwd_params &params, cudaStream_t stream) {
+    VARLEN_SWITCH(params.cu_seqlens_q != nullptr || params.cu_seqlens_k != nullptr, Varlen, [&] {
+        BOOL_SWITCH(params.h != params.h_k, GQA, [&] {
+//             BOOL_SWITCH(params.deterministic, Deterministic, [&] {
+            // run_flash_bwd<kHeadDim, kBlockM, kBlockN, T, Is_causal, Is_local, Has_softcap, Varlen, false, GQA, Stages_dO, Stages_dS_or_QSm80, SdP_swapAB, dKV_swapAB, dQ_swapAB, NumMmaWarpGroups, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ>(params, stream);
+            run_flash_bwd<Arch, kHeadDim, kBlockM, kBlockN, T, Is_causal, Is_local, Has_softcap, Varlen /*Varlen*/, false /*Deterministic*/, GQA, Stages_dO, Stages_dS_or_QSm80, SdP_swapAB, dKV_swapAB, dQ_swapAB, NumMmaWarpGroups, AtomLayoutMSdP, AtomLayoutNdKV, AtomLayoutMdQ, V_in_regs>(params, stream);
+//             });
+        });
+    });
+}
+
+
+template<int Arch, typename T, bool Has_softcap>
+void run_mha_bwd_hdim64(Flash_bwd_params &params, cudaStream_t stream) {
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        if constexpr (Arch >= 90) {
+            if constexpr (Is_causal && Has_softcap) {
+                // register spill with 128 x 128
+                run_mha_bwd_dispatch<Arch, T, 96, 128, 64, Is_causal, Is_local, Has_softcap, 2, 2, true, false, true, 2, 1, 2, 2, false>(params, stream);
+            } else {
+                // With ShuffleStats we no longer have register spilling when Has_softcap and using 128 x 128 block.
+                run_mha_bwd_dispatch<Arch, T, 128, 128, 64, Is_causal, Is_local, Has_softcap, 2, 2, true, false, false, 2, 1, 2, 2, false>(params, stream);
+            }
+        } else if constexpr (Arch == 86 || Arch == 89) {
+            run_mha_bwd_dispatch<Arch, T, 64, 128, 64, Is_causal, Is_local, Has_softcap, 2, 2, false, false, false, 2, 2, 4, 2, true>(params, stream);
+            // run_mha_bwd_dispatch<Arch, T, 96, 96, 64, Is_causal, Is_local, Has_softcap, 1, 2, false, true, true, 2, 2, 4, 4, false>(params, stream);
+            // run_mha_bwd_dispatch<Arch, T, 80, 128, 64, Is_causal, Is_local, Has_softcap, 1, 2, true, false, true, 2, 2, 4, 2, true>(params, stream);
+            // run_mha_bwd_dispatch<Arch, T, 96, 128, 64, Is_causal, Is_local, Has_softcap, 1, 2, true, false, true, 2, 1, 8, 4, false>(params, stream);
+        } else {
+            run_mha_bwd_dispatch<Arch, T, 128, 128, 64, Is_causal, Is_local, Has_softcap, 2, 2, false, false, false, 2, 4, 4, 4, false>(params, stream);
+        }
+    });
+}
+
+template<int Arch, typename T, bool Has_softcap>
+void run_mha_bwd_hdim96(Flash_bwd_params &params, cudaStream_t stream) {
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        if constexpr (Arch >= 90) {
+            run_mha_bwd_dispatch<Arch, T, 64, 128, 96, Is_causal, Is_local, Has_softcap, 2, 2, true, false, false, 2, 1, 2, 1, true>(params, stream);
+        } else if constexpr (Arch == 86 || Arch == 89) {
+            run_mha_bwd_dispatch<Arch, T, 64, 128, 96, Is_causal, Is_local, Has_softcap, 1, 2, false, false, false, 2, 2, 4, 2, true>(params, stream);
+        } else {
+            run_mha_bwd_dispatch<Arch, T, 64, 128, 96, Is_causal, Is_local, Has_softcap, 2, 2, false, false, false, 2, 2, 4, 2, false>(params, stream);
+        }
+    });
+}
+
+template<int Arch, typename T, bool Has_softcap>
+void run_mha_bwd_hdim128(Flash_bwd_params &params, cudaStream_t stream) {
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        if constexpr (Arch >= 90) {
+            if constexpr (Is_causal || Is_local || Has_softcap) {
+                run_mha_bwd_dispatch<Arch, T, 64, 128, 128, Is_causal, Is_local, Has_softcap, 2, 2, true, false, false, 2, 1, 2, 1, false>(params, stream);
+            } else {
+                run_mha_bwd_dispatch<Arch, T, 80, 128, 128, Is_causal, Is_local, Has_softcap, 2, 2, true, false, true, 2, 1, 2, 1, false>(params, stream);
+            }
+        } else if constexpr (Arch == 86 || Arch == 89) {
+            run_mha_bwd_dispatch<Arch, T, 64, 96, 128, Is_causal, Is_local, Has_softcap, 1, 2, false, false, false, 2, 2, 2, 2, true>(params, stream);
+        } else {
+            run_mha_bwd_dispatch<Arch, T, 64, 128, 128, Is_causal, Is_local, Has_softcap, 2, 2, false, false, false, 2, 2, 2, 2, false>(params, stream);
+        }
+    });
+}
+
+template<int Arch, typename T, bool Has_softcap>
+void run_mha_bwd_hdim192(Flash_bwd_params &params, cudaStream_t stream) {
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        if constexpr (Arch >= 90) {
+            run_mha_bwd_dispatch<Arch, T, 64, 96, 192, Is_causal, Is_local, Has_softcap, 1, 1, false, true, false, 3, 1, 1, 1, false>(params, stream);
+        } else if constexpr (Arch == 86 || Arch == 89) {
+            run_mha_bwd_dispatch<Arch, T, 64, 64, 192, Is_causal, Is_local, Has_softcap, 1, 1, false, false, false, 2, 2, 2, 2, true>(params, stream);
+        } else {
+            run_mha_bwd_dispatch<Arch, T, 64, 80, 192, Is_causal, Is_local, Has_softcap, 1, 2, false, true, false, 2, 4, 2, 2, false>(params, stream);
+        }
+    });
+}
+
+template<int Arch, typename T, bool Has_softcap>
+void run_mha_bwd_hdim256(Flash_bwd_params &params, cudaStream_t stream) {
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        if constexpr (Arch >= 90) {
+            run_mha_bwd_dispatch<Arch, T, 64, 80, 256, Is_causal, Is_local, Has_softcap, 1, 1, false, true, true, 2, 1, 1, 1, false>(params, stream);
+        } else if constexpr (Arch == 86 || Arch == 89) {
+            run_mha_bwd_dispatch<Arch, T, 32, 64, 256, Is_causal, Is_local, Has_softcap, 1, 1, false, false, false, 2, 2, 2, 1, true>(params, stream);
+            // run_mha_bwd_dispatch<Arch, T, 64, 32, 256, Is_causal, Is_local, Has_softcap, 1, 1, false, false, false, 2, 4, 1, 2, true>(params, stream);
+        } else {
+            run_mha_bwd_dispatch<Arch, T, 64, 64, 256, Is_causal, Is_local, Has_softcap, 1, 1, false, false, false, 2, 4, 2, 2, false>(params, stream);
+        }
+    });
+}

flash-attn/flash_bwd_postprocess_kernel.h

@@ -0,0 +1,256 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/numeric_conversion.h>
+#include "cutlass/arch/barrier.h"
+
+#include "seqlen.h"
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class TileShape_MK_, class Element, class ElementAccum, class ArchTag_, int kNThreads, class TiledMma, bool dQ_swapAB>
+class FlashAttnBwdPostprocessConvertdQ {
+
+public:
+
+    // Type Aliases
+    using TileShape_MK = TileShape_MK_;
+    using ArchTag = ArchTag_;
+
+    static_assert(ArchTag::kMinComputeCapability >= 75);
+    static constexpr bool IsSm90 = ArchTag::kMinComputeCapability >= 90;
+
+    static constexpr uint32_t MaxThreadsPerBlock = kNThreads;
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 2;
+
+    static constexpr int kBlockM = get<0>(TileShape_MK{});
+    static constexpr int kHeadDim = get<1>(TileShape_MK{});
+    static_assert(!IsSm90 || kNThreads % cutlass::NumThreadsPerWarpGroup == 0, "kNThreads must be a multiple of NumThreadsPerWarpGroup");
+    static constexpr int NumdQWarpGgroups = kNThreads / cutlass::NumThreadsPerWarpGroup;
+    using R2SLayoutAtomdQaccum = std::conditional_t<
+        IsSm90,
+        Layout<Shape<Int<cutlass::NumThreadsPerWarpGroup>, Int<NumdQWarpGgroups>>>,
+        Layout<Shape<Int<kNThreads>>>
+    >;
+    using R2STiledCopydQaccum = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{}, R2SLayoutAtomdQaccum{},
+                                                         Layout<Shape<Int<IsSm90 ? 4 : 1>>>{}));  // Val layout, 1 or 4 vals per read
+    using G2SLayoutAtomdQaccum = Layout<Shape<Int<kNThreads>>>;
+    // UniversalCopy instead of AutoVectorizingCopyWithAssumedAlignment as the latter generates cp.async instructions
+    using G2STiledCopydQaccum = decltype(make_tiled_copy(Copy_Atom<UniversalCopy<uint128_t>, ElementAccum>{}, G2SLayoutAtomdQaccum{},
+                                                         Layout<Shape<_4>>{}));  // Val layout, 4 vals per read
+    // We don't do bound checking for the gmem -> smem load so we just assert here.
+    static_assert(IsSm90 || (kBlockM * kHeadDim) % (kNThreads * 4) == 0);
+    static constexpr int SmemdQaccumSize = size(TileShape_MK{});
+    using SmemLayoutdQaccumFlat = Layout<Shape<Int<SmemdQaccumSize>>>;
+    using SmemLayoutdQaccum = std::conditional_t<
+        IsSm90,
+        Layout<Shape<Int<kBlockM * kHeadDim / NumdQWarpGgroups>, Int<NumdQWarpGgroups>>>,
+        Layout<Shape<Int<kBlockM * kHeadDim>>>
+    >;
+
+    // We can't just use kHeadDim here. E.g. if MMA shape is 64 x 96 but split across 2 WGs,
+    // then setting kBlockKSmem to 32 will cause "Static shape_div failure".
+    // We want to treat it as 64 x 48, so kBlockKSmem should be 16.
+    static constexpr int MmaShapeN = get<1>(typename TiledMma::AtomShape_MNK{});
+    static constexpr int kBlockKSmem = MmaShapeN % 64 == 0 ? 64 : (MmaShapeN % 32 == 0 ? 32 : 16);
+    static constexpr int kSwizzle = kBlockKSmem == 64 ? 3 : (kBlockKSmem == 32 ? 2 : 1);
+    using SmemLayoutAtomdQ =
+        decltype(composition(Swizzle<kSwizzle, 3, 3>{},
+                 Layout<Shape<Int<8>, Int<kBlockKSmem>>,
+                 Stride<Int<kBlockKSmem>, _1>>{}));
+    using SmemLayoutdQ = decltype(tile_to_shape(SmemLayoutAtomdQ{}, TileShape_MK{}));
+    using SmemLayoutdQt =
+        decltype(cute::composition(SmemLayoutdQ{},
+                                   make_layout(make_shape(get<1>(TileShape_MK{}), get<0>(TileShape_MK{})),
+                                               make_stride(Int<get<0>(TileShape_MK{})>{}, _1{}))));
+
+    using SmemCopyAtomdQ = Copy_Atom<
+        std::conditional_t<
+            IsSm90,
+            std::conditional_t<!dQ_swapAB, cute::SM90_U32x4_STSM_N, cute::SM90_U16x8_STSM_T>,
+            AutoVectorizingCopyWithAssumedAlignment<128>
+        >,
+        Element>;
+
+    static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);
+    static_assert(kHeadDim % kGmemElemsPerLoad == 0, "Headdim must be a multiple of kGmemElemsPerLoad");
+    static constexpr int kGmemThreadsPerRow = cutlass::gcd(kHeadDim / kGmemElemsPerLoad, int(MaxThreadsPerBlock));
+    static_assert(MaxThreadsPerBlock % kGmemThreadsPerRow == 0, "MaxThreadsPerBlock must be a multiple of kGmemThreadsPerRow");
+    using GmemLayoutAtom = Layout<Shape <Int<MaxThreadsPerBlock / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
+                                  Stride<Int<kGmemThreadsPerRow>, _1>>;
+    using GmemTiledCopy = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoad>>>{}));  // Val layout, 8 or 16 vals per load
+
+    struct SharedStorage : cute::aligned_struct<128> {
+        cute::array_aligned<ElementAccum, cute::cosize_v<SmemLayoutdQaccum>> smem_dqacc;
+        cute::array_aligned<Element, cute::cosize_v<SmemLayoutdQ>> smem_dq;
+        alignas(16) cutlass::arch::ClusterTransactionBarrier barrier_dQaccum;
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    using ShapedQ = cute::Shape<int32_t, int32_t, int32_t, int32_t>;   // (seqlen_q, d, head, batch)
+    using StridedQ = cute::Stride<int64_t, _1, int64_t, int64_t>;
+    using ShapedQaccum = cute::Shape<int32_t, int32_t, int32_t>;  // (seqlen_q * d, head, batch)
+    using StridedQaccum = cute::Stride<_1, int64_t, int64_t>;
+
+    // Device side arguments
+    struct Arguments {
+        ElementAccum const* ptr_dQaccum;
+        ShapedQaccum const shape_dQaccum;
+        StridedQaccum const stride_dQaccum;
+        Element* ptr_dQ;
+        ShapedQ const shape_dQ;
+        StridedQ const stride_dQ;
+        float const softmax_scale;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    // Kernel entry point API
+    struct Params {
+        ElementAccum const* ptr_dQaccum;
+        ShapedQaccum const shape_dQaccum;
+        StridedQaccum const stride_dQaccum;
+        Element* ptr_dQ;
+        ShapedQ const shape_dQ;
+        StridedQ const stride_dQ;
+        float const softmax_scale;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        return {
+            args.ptr_dQaccum,
+            args.shape_dQaccum,
+            args.stride_dQaccum,
+            args.ptr_dQ,
+            args.shape_dQ,
+            args.stride_dQ,
+            args.softmax_scale,
+            args.cu_seqlens,
+            args.seqused
+        };
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        static constexpr int kBlockM = get<0>(TileShape_MK{});
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+        Tensor sdQaccum = make_tensor(make_smem_ptr(shared_storage.smem_dqacc.data()), SmemLayoutdQaccum{});
+        Tensor sdQaccum_flat = make_tensor(make_smem_ptr(shared_storage.smem_dqacc.data()), SmemLayoutdQaccumFlat{});
+        Tensor sdQ = make_tensor(make_smem_ptr(shared_storage.smem_dq.data()), SmemLayoutdQ{});
+        Tensor sdQt = make_tensor(make_smem_ptr(shared_storage.smem_dq.data()), SmemLayoutdQt{});
+
+        int const thread_idx = threadIdx.x;
+        int const m_block = blockIdx.x;
+        int const bidh = blockIdx.y;
+        int const bidb = blockIdx.z;
+
+        flash::SeqlenInfo<true /*Varlen*/, kBlockM> seqlen_info(bidb, size<0>(params.shape_dQ), params.cu_seqlens, params.seqused);
+        bool const is_varlen = params.cu_seqlens;
+        if (is_varlen && m_block * kBlockM >= seqlen_info.seqlen) { return; }
+
+        // Step 1: load dQaccum from gmem to smem
+        Tensor mdQaccum = make_tensor(make_gmem_ptr(reinterpret_cast<ElementAccum const*>(params.ptr_dQaccum)),
+                                      params.shape_dQaccum, params.stride_dQaccum)(_, bidh, !is_varlen ? bidb : 0);
+        Tensor gdQaccum = local_tile(domain_offset(make_coord(seqlen_info.offset_padded * kHeadDim), mdQaccum), Shape<Int<kBlockM * kHeadDim>>{}, make_coord(m_block));  // (M * K)
+        if constexpr (IsSm90) {  // Use BulkCopy
+            static constexpr uint32_t TmaTransactionBytesdQaccum = static_cast<uint32_t>(size(SmemLayoutdQaccumFlat{}) * cute::sizeof_bits_v<ElementAccum> / 8);
+            auto bulk_copy = Copy_Traits<SM90_BULK_COPY_AUTO>{};
+            // if (thread0()) { print(gdQaccum); printf("\n"); print(sdQaccum_flat); printf("\n"); }
+            if (thread_idx == 0) {
+                shared_storage.barrier_dQaccum.init(1 /*numThreads*/);
+                shared_storage.barrier_dQaccum.arrive_and_expect_tx(TmaTransactionBytesdQaccum);
+                copy(bulk_copy.with(*reinterpret_cast<uint64_t*>(&shared_storage.barrier_dQaccum)), gdQaccum, sdQaccum_flat);
+            }
+            __syncthreads();
+            shared_storage.barrier_dQaccum.wait(0);
+        } else {
+            G2STiledCopydQaccum g2s_tiled_copy_dQaccum;
+            auto g2s_thr_copy_dQaccum = g2s_tiled_copy_dQaccum.get_thread_slice(thread_idx);
+            Tensor tdQgdQaccumg2s = g2s_thr_copy_dQaccum.partition_S(gdQaccum);
+            Tensor tdQsdQaccumg2s = g2s_thr_copy_dQaccum.partition_D(sdQaccum);
+            cute::copy(g2s_tiled_copy_dQaccum, tdQgdQaccumg2s, tdQsdQaccumg2s);
+            __syncthreads();
+        }
+
+        // __syncthreads(); if (cute::thread0()) { print_tensor(sdQaccum); }
+
+        // Step 2: Load dQaccum from smem to register, then convert fp32 -> fp16/bf16
+        R2STiledCopydQaccum s2r_tiled_copy_dQaccum;
+        auto s2r_thr_copy_dQaccum = s2r_tiled_copy_dQaccum.get_thread_slice(thread_idx);
+        Tensor tdQsdQaccum = s2r_thr_copy_dQaccum.partition_S(sdQaccum);
+        TiledMma tiled_mma_dQ;
+        Tensor taccdQrdQaccum = partition_fragment_C(tiled_mma_dQ, select<!dQ_swapAB ? 0 : 1, !dQ_swapAB ? 1 : 0>(TileShape_MK{}));
+        // if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 1) { print(tiled_mma_dQ); printf("\n"); }
+        // if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 1) { print(tdQsdQaccum); }
+        // if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 1) { print(taccdQrdQaccum); }
+        CUTE_STATIC_ASSERT_V(size(taccdQrdQaccum) == size(tdQsdQaccum));
+        Tensor tdQrdQaccum = s2r_thr_copy_dQaccum.retile_D(taccdQrdQaccum);
+        cute::copy(s2r_tiled_copy_dQaccum, tdQsdQaccum, tdQrdQaccum);
+        #pragma unroll
+        for (int i = 0; i < size(taccdQrdQaccum); ++i) { taccdQrdQaccum(i) *= params.softmax_scale; }
+        // Convert tdQrdQ from fp32 to fp16
+        Tensor rdQ = make_tensor_like<Element>(taccdQrdQaccum);
+        flash::convert_type_out(taccdQrdQaccum, rdQ);
+
+        // Step 3: Copy dQ from register to smem
+        auto smem_tiled_copy_dQ = make_tiled_copy_C(SmemCopyAtomdQ{}, tiled_mma_dQ);
+        auto smem_thr_copy_dQ = smem_tiled_copy_dQ.get_thread_slice(thread_idx);
+        Tensor taccdQrdQ = smem_thr_copy_dQ.retile_S(rdQ);  // ((Atom,AtomNum), MMA_N, MMA_N)
+        // if (cute::thread0()) { print(smem_tiled_copy_dQ); }
+        // if (cute::thread0()) { print(smem_thr_copy_dQ); }
+        // if (cute::thread0()) { print(sdQ); }
+        Tensor taccdQsdQ = smem_thr_copy_dQ.partition_D(cute::conditional_return<!dQ_swapAB>(sdQ, sdQt));  // ((Atom,AtomNum),PIPE_M,PIPE_N)
+        cute::copy(smem_tiled_copy_dQ, taccdQrdQ, taccdQsdQ);
+        __syncthreads();
+
+        // Step 4: Copy dQ from smem to register to prepare for coalesced write to gmem
+        Tensor mdQ = make_tensor(make_gmem_ptr(params.ptr_dQ), params.shape_dQ, params.stride_dQ)(_, _, bidh, !is_varlen ? bidb : 0);
+        Tensor gdQ = local_tile(domain_offset(make_coord(seqlen_info.offset, _0{}), mdQ), TileShape_MK{}, make_coord(m_block, _0{}));  // (M, K)
+        GmemTiledCopy gmem_tiled_copy_dQ;
+        auto gmem_thr_copy_dQ = gmem_tiled_copy_dQ.get_thread_slice(thread_idx);
+        Tensor tdQsdQ = gmem_thr_copy_dQ.partition_S(sdQ);    // ((Atom,AtomNum),ATOM_M,ATOM_N)
+        Tensor tdQgdQ = gmem_thr_copy_dQ.partition_D(gdQ);
+
+        Tensor tdQrdQ = make_fragment_like(tdQsdQ);
+        Tensor tdQcdQ = gmem_thr_copy_dQ.partition_D(cute::make_identity_tensor(TileShape_MK{}));
+        Tensor tdQpdQ = make_tensor<bool>(make_shape(size<2>(tdQgdQ)));
+        #pragma unroll
+        for (int k = 0; k < size(tdQpdQ); ++k) { tdQpdQ(k) = get<1>(tdQcdQ(_0{}, _0{}, k)) < get<1>(params.shape_dQ); }
+        // Need to check OOB when reading from smem if kBlockM isn't evenly tiled
+        static constexpr bool EvenM = kBlockM % CUTE_STATIC_V(size<0>(GmemLayoutAtom{})) == 0;
+        flash::copy</*Is_even_MN=*/EvenM, /*Is_even_K=*/true, /*Clear_OOB_MN=*/false>(
+            gmem_tiled_copy_dQ, tdQsdQ, tdQrdQ, tdQcdQ, tdQpdQ, kBlockM);
+
+        // Step 5: Copy dQ from register to gmem
+        // Clear_OOB_K must be false since we don't want to write zeros to gmem
+        flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/false, /*Clear_OOB_K=*/false>(
+            gmem_tiled_copy_dQ, tdQrdQ, tdQgdQ, tdQcdQ, tdQpdQ, std::min(seqlen_info.seqlen - m_block * kBlockM, kBlockM)
+        );
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_bwd_preprocess_kernel.h

@@ -0,0 +1,252 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/numeric_conversion.h>
+
+#include "seqlen.h"
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class TileShape_MK_, class Element, class ElementAccum, class ArchTag_, bool Clear_dQaccum, bool Varlen>
+class FlashAttnBwdPreprocess {
+
+public:
+
+    // Type Aliases
+    using TileShape_MK = TileShape_MK_;
+    using ArchTag = ArchTag_;
+
+    static_assert(std::is_same_v<Element, cutlass::half_t> && ArchTag::kMinComputeCapability >= 75 ||
+                  std::is_same_v<Element, cutlass::bfloat16_t> && ArchTag::kMinComputeCapability >= 80 ||
+                  std::is_same_v<Element, cutlass::float_e4m3_t> && ArchTag::kMinComputeCapability >= 89);
+
+    static constexpr uint32_t MaxThreadsPerBlock = 256;
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 2;
+    static constexpr int SharedStorageSize = 0;
+
+    static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);
+    static_assert(get<1>(TileShape_MK{}) % kGmemElemsPerLoad == 0, "Headdim must be a multiple of kGmemElemsPerLoad");
+    static constexpr int kBlockM = get<0>(TileShape_MK{});
+    static constexpr int kHeadDim = get<1>(TileShape_MK{});
+    // We want kBlockKGmem to be a power of 2 so that when we do the summing,
+    // it's just between threads in the same warp
+    static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32);
+    static constexpr int kGmemThreadsPerRow = kBlockKGmem / kGmemElemsPerLoad;
+    static_assert(MaxThreadsPerBlock % kGmemThreadsPerRow == 0, "MaxThreadsPerBlock must be a multiple of kGmemThreadsPerRow");
+    using GmemLayoutAtom = Layout<Shape <Int<MaxThreadsPerBlock / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
+                                  Stride<Int<kGmemThreadsPerRow>, _1>>;
+    using GmemTiledCopy = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoad>>>{}));  // Val layout, 8 or 16 vals per load
+
+    static constexpr int kGmemElemsPerLoadAccum = sizeof(cute::uint128_t) / sizeof(ElementAccum);
+    static_assert((kBlockM * kHeadDim / kGmemElemsPerLoadAccum) % MaxThreadsPerBlock == 0, "MaxThreadsPerBlock must divide kBlockM * kHeadDim / kGmemElemsPerLoadAccum");
+    using GmemLayoutAtomAccum = Layout<Shape<Int<MaxThreadsPerBlock>>>;
+    using GmemTiledCopyAccum = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{},
+                        GmemLayoutAtomAccum{},
+                        Layout<Shape<Int<kGmemElemsPerLoadAccum>>>{}));  // Val layout, 4 vals per store
+
+    using ShapeO = cute::Shape<int32_t, int32_t, int32_t, int32_t>;  // (seqlen_q, d, head, batch)
+    using StrideO = cute::Stride<int64_t, _1, int64_t, int64_t>;
+    using ShapedPsum = cute::Shape<int32_t, int32_t, int32_t>;  // (seqlen_q, head, batch)
+    using StridedPsum = cute::Stride<_1, int64_t, int64_t>;
+    using ShapedQaccum = cute::Shape<int32_t, int32_t, int32_t>;  // (seqlen_q * d, head, batch)
+    using StridedQaccum = cute::Stride<_1, int64_t, int64_t>;
+
+    // Device side arguments
+    struct Arguments {
+        Element const* ptr_O;
+        ShapeO const shape_O;
+        StrideO const stride_O;
+        Element const* ptr_dO;
+        StrideO const stride_dO;
+        float* ptr_dPsum;
+        ShapedPsum const shape_dPsum;
+        StridedPsum const stride_dPsum;
+        float const* ptr_LSE;
+        StridedPsum const stride_LSE;
+        float *ptr_LSE_log2;
+        StridedPsum const stride_LSE_log2;
+        ElementAccum* ptr_dQaccum;
+        ShapedQaccum const shape_dQaccum;
+        StridedQaccum const stride_dQaccum;
+        int num_batch;  // We need this to know the size of dq_semaphore in case of varlen
+        int* dq_semaphore;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    // Kernel entry point API
+    struct Params {
+        Element const* ptr_O;
+        ShapeO const shape_O;
+        StrideO const stride_O;
+        Element const* ptr_dO;
+        StrideO const stride_dO;
+        float* ptr_dPsum;
+        ShapedPsum const shape_dPsum;
+        StridedPsum const stride_dPsum;
+        float const* ptr_LSE;
+        StridedPsum const stride_LSE;
+        float* ptr_LSE_log2;
+        StridedPsum const stride_LSE_log2;
+        ElementAccum* ptr_dQaccum;
+        ShapedQaccum const shape_dQaccum;
+        StridedQaccum const stride_dQaccum;
+        int num_batch;
+        int* dq_semaphore;
+        int const* cu_seqlens = nullptr;
+        int const* seqused = nullptr;
+    };
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        return {
+            args.ptr_O,
+            args.shape_O,
+            args.stride_O,
+            args.ptr_dO,
+            args.stride_dO,
+            args.ptr_dPsum,
+            args.shape_dPsum,
+            args.stride_dPsum,
+            args.ptr_LSE,
+            args.stride_LSE,
+            args.ptr_LSE_log2,
+            args.stride_LSE_log2,
+            args.ptr_dQaccum,
+            args.shape_dQaccum,
+            args.stride_dQaccum,
+            args.num_batch,
+            args.dq_semaphore,
+            args.cu_seqlens,
+            args.seqused
+        };
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, [[maybe_unused]] char* smem_buf) {
+
+        static constexpr int kBlockM = get<0>(TileShape_MK{});
+
+        int const thread_idx = threadIdx.x;
+        int const m_block = blockIdx.x;
+        int const bidh = blockIdx.y;
+        int const bidb = blockIdx.z;
+
+        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info(bidb, size<0>(params.shape_O), params.cu_seqlens, params.seqused);
+        bool const is_varlen = Varlen && params.cu_seqlens;
+        int const seqlen_o = seqlen_info.seqlen;
+        if (is_varlen && m_block * kBlockM >= seqlen_o) { return; }
+
+        Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O), params.shape_O, params.stride_O)(_, _, bidh, !is_varlen ? bidb : 0);
+        Tensor gO = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mO), TileShape_MK{}, make_coord(m_block, _0{}));  // (M, K)
+        Tensor mdO = make_tensor(make_gmem_ptr(params.ptr_dO), params.shape_O, params.stride_dO)(_, _, bidh, !is_varlen ? bidb : 0);
+        Tensor gdO = local_tile(cute::domain_offset(make_coord(seqlen_info.offset, _0{}), mdO), TileShape_MK{}, make_coord(m_block, _0{}));  // (M, K)
+
+        auto shape_LSE = select<0, 2, 3>(params.shape_O);
+        Tensor mLSE = make_tensor(make_gmem_ptr(params.ptr_LSE), shape_LSE, params.stride_LSE)(_, bidh, !is_varlen ? bidb : 0);
+        Tensor gLSE = local_tile(cute::domain_offset(make_coord(seqlen_info.offset), mLSE), Shape<Int<kBlockM>>{}, make_coord(m_block));
+        static_assert(kBlockM <= MaxThreadsPerBlock);
+        float lse = thread_idx < seqlen_o - m_block * kBlockM && thread_idx < kBlockM ? gLSE(thread_idx) : INFINITY;
+
+        GmemTiledCopy gmem_tiled_copy_O;
+        auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
+
+        Tensor tOgO = gmem_thr_copy_O.partition_S(gO);
+        Tensor tOgdO = gmem_thr_copy_O.partition_S(gdO);
+        // Construct identity layout for gO
+        Tensor cO = cute::make_identity_tensor(TileShape_MK{});  // (BLK_M,BLK_K) -> (blk_m,blk_k)
+        // Repeat the partitioning with identity layouts
+        Tensor tOcO = gmem_thr_copy_O.partition_D(cO);
+        Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOgO)));
+        #pragma unroll
+        for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(params.shape_O); }
+
+        // (8, kBlockM / 32, kHeadDim / 64) or (8, kBlockM / 16, kHeadDim / 128)
+        Tensor tOrO = make_fragment_like(tOgO);
+        Tensor tOrdO = make_fragment_like(tOgdO);
+        flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/true, /*Clearn_OOB_K=*/true>(
+            gmem_tiled_copy_O, tOgO, tOrO, tOcO, tOpO, seqlen_o - m_block * kBlockM
+        );
+        flash::copy</*Is_even_MN=*/false, /*Is_even_K=*/false, /*Clear_OOB_MN=*/true, /*Clearn_OOB_K=*/true>(
+            gmem_tiled_copy_O, tOgdO, tOrdO, tOcO, tOpO, seqlen_o - m_block * kBlockM
+        );
+        // if (threadIdx.x == 222) { printf("bidx = %d, bidy = %d, bidz = %d, seqlen_o = %d, m_block = %d, seqlen_o - m_block * kBlockM = %d, tOgO addr = %p\n", blockIdx.x, blockIdx.y, blockIdx.z, seqlen_o, m_block, seqlen_o - m_block * kBlockM, &tOgO(0));}
+
+        // Reshape from e.g. (8, kBlockM / 32, kHeadDim / 64) to (kBlockM / 32, (8, kHeadDim / 64))
+        Layout l = make_layout(get<1>(tOrO.layout()), make_layout(get<0>(tOrO.layout()), get<2>(tOrO.layout())));
+        Tensor tOrO_l = make_tensor(tOrO.data(), l);
+        Tensor o_fp32 = make_tensor_like<float>(tOrO_l);
+        flash::convert_type_out(tOrO_l, o_fp32);
+        Tensor tOrdO_l = make_tensor(tOrdO.data(), l);
+        Tensor do_fp32 = make_tensor_like<float>(tOrdO_l);
+        flash::convert_type_out(tOrdO_l, do_fp32);
+        // Sum across the last dimension
+        Tensor dP_sum = make_tensor<float>(make_shape(size<0>(o_fp32)));
+        #pragma unroll
+        for (int mi = 0; mi < size<0>(o_fp32); ++mi) {
+            float dP_sum_cur = do_fp32(mi, 0) * o_fp32(mi, 0);
+            #pragma unroll
+            for (int ni = 1; ni < size<1>(o_fp32); ni++) {
+                dP_sum_cur += do_fp32(mi, ni) * o_fp32(mi, ni);
+            }
+            flash::SumOp<float> sum_op;
+            dP_sum(mi) = flash::Allreduce<kGmemThreadsPerRow>::run(dP_sum_cur, sum_op);
+        }
+
+        Tensor mdPsum = make_tensor(make_gmem_ptr(params.ptr_dPsum), params.shape_dPsum, params.stride_dPsum)(_, bidh, !is_varlen ? bidb : 0);
+        Tensor gdPsum = local_tile(cute::domain_offset(make_coord(seqlen_info.offset_padded), mdPsum), Shape<Int<kBlockM>>{}, make_coord(m_block));
+        if (get<1>(tOcO(_0{}, _0{}, _0{})) == 0) {
+            #pragma unroll
+            for (int mi = 0; mi < size(dP_sum); ++mi) {
+                int const row = get<0>(tOcO(_0{}, mi, _0{}));
+                gdPsum(row) = row < seqlen_o - m_block * kBlockM ? dP_sum(mi) : 0;
+            }
+        }
+
+        int const seqlen_rounded = cute::round_up(seqlen_o, kBlockM);
+        Tensor mLSElog2 = make_tensor(make_gmem_ptr(params.ptr_LSE_log2), params.shape_dPsum, params.stride_LSE_log2)(_, bidh, !is_varlen ? bidb : 0);
+        Tensor gLSElog2 = local_tile(cute::domain_offset(make_coord(seqlen_info.offset_padded), mLSElog2), Shape<Int<kBlockM>>{}, make_coord(m_block));
+        if (thread_idx < seqlen_rounded - m_block * kBlockM && thread_idx < kBlockM) {
+            gLSElog2(thread_idx) = lse == -INFINITY ? 0.f : lse * float(M_LOG2E);
+        }
+
+        if constexpr (Clear_dQaccum) {
+            Tensor mdQaccum = make_tensor(make_gmem_ptr(params.ptr_dQaccum), params.shape_dQaccum, params.stride_dQaccum)(_, bidh, !is_varlen ? bidb : 0);
+            Tensor gdQaccum = local_tile(cute::domain_offset(make_coord(seqlen_info.offset_padded * kHeadDim), mdQaccum), Shape<Int<kBlockM * kHeadDim>>{}, make_coord(m_block));
+            GmemTiledCopyAccum gmem_tiled_copy_dQaccum;
+            auto gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_thread_slice(thread_idx);
+            Tensor tdQgdQaccum = gmem_thr_copy_dQaccum.partition_D(gdQaccum);
+            Tensor zero = make_fragment_like(tdQgdQaccum);
+            clear(zero);
+            cute::copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementAccum>{}, zero, tdQgdQaccum);
+        }
+
+        if (params.dq_semaphore != nullptr && thread_idx == 0) {
+            int const num_batch = params.num_batch;
+            int const num_head = get<2>(params.shape_O);
+            params.dq_semaphore[bidh + bidb * num_head + m_block * num_head * num_batch] = 0;
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_fwd_combine.cu

@@ -0,0 +1,13 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+
+#include "flash_fwd_combine_launch_template.h"
+
+template void run_mha_fwd_combine_<float, float, 64>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);
+template void run_mha_fwd_combine_<float, float, 128>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);
+
+template void run_mha_fwd_combine_<cutlass::half_t, float, 64>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);
+template void run_mha_fwd_combine_<cutlass::half_t, float, 128>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);
+
+template void run_mha_fwd_combine_<cutlass::bfloat16_t, float, 64>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);
+template void run_mha_fwd_combine_<cutlass::bfloat16_t, float, 128>(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl);

flash-attn/flash_fwd_combine_kernel.h

@@ -0,0 +1,482 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/arch/memory.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/numeric_conversion.h>
+
+#include "cutlass/arch/grid_dependency_control.h"
+
+#include "seqlen.h"
+#include "utils.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class TileShape_MK_, int kLogMaxSplits_, int kNThreads, int AlignmentLSE_,
+          bool Is_even_K, bool Varlen, class Element, class ElementPartial, class ArchTag_>
+class FlashAttnFwdCombine {
+
+public:
+
+    // Type Aliases
+    using TileShape_MK = TileShape_MK_;
+    using ArchTag = ArchTag_;
+    static constexpr int kMaxSplits = 1 << kLogMaxSplits_;
+    static constexpr int AlignmentLSE = std::min(AlignmentLSE_, int(128 / 8 / sizeof(float)));
+    static_assert(AlignmentLSE >= 1);
+    static constexpr int kStages = 4;
+
+    static_assert(ArchTag::kMinComputeCapability >= 75);
+    static constexpr bool Has_cp_async = ArchTag::kMinComputeCapability >= 80;
+
+    static constexpr uint32_t MaxThreadsPerBlock = kNThreads;
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 2;
+
+    static constexpr int kBlockM = get<0>(TileShape_MK{});
+    static constexpr int kBlockK = get<1>(TileShape_MK{});
+
+    static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(ElementPartial);
+    static_assert(kBlockK % kGmemElemsPerLoad == 0, "kBlockK must be a multiple of kGmemElemsPerLoad");
+    static constexpr int kBlockKGmem = kBlockK % 128 == 0 ? 128 : (kBlockK % 64 == 0 ? 64 : 32);
+    static constexpr int kGmemThreadsPerRow = kBlockKGmem / kGmemElemsPerLoad;
+    static_assert(MaxThreadsPerBlock % kGmemThreadsPerRow == 0, "MaxThreadsPerBlock must be a multiple of kGmemThreadsPerRow");
+    using GmemCopyAtom = std::conditional_t<
+        Has_cp_async,
+        cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<uint128_t>, ElementPartial>,
+        cute::Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementPartial>
+    >;
+    using GmemLayoutAtom = Layout<Shape <Int<MaxThreadsPerBlock / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
+                                  Stride<Int<kGmemThreadsPerRow>, _1>>;
+    static_assert(kBlockM % CUTE_STATIC_V(shape<0>(GmemLayoutAtom{})) == 0);
+    using GmemTiledCopyAccum = decltype(
+        make_tiled_copy(GmemCopyAtom{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoad>>>{}));  // Val layout, 4 vals per load
+    using GmemTiledCopy = decltype(
+        make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                        GmemLayoutAtom{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoad>>>{}));  // Val layout, 4 vals per load
+
+    using AlignmentTypeLSE = cute::uint_byte_t<static_cast<int>(sizeof(float)) * AlignmentLSE>;
+    static constexpr int kGmemElemsPerLoadLSE = sizeof(AlignmentTypeLSE) / sizeof(float);
+    static_assert(kBlockM % kGmemElemsPerLoadLSE == 0, "kBlockM must be a multiple of kGmemElemsPerLoadLSE");
+    static_assert(kBlockM % 8 == 0, "kBlockM must be a multiple of 8");
+    static constexpr int kBlockMSmem = kBlockM % 128 == 0 ? 128 : (kBlockM % 64 == 0 ? 64 : (kBlockM % 32 == 0 ? 32 : (kBlockM % 16 == 0 ? 16 : 8)));
+    static constexpr int kGmemThreadsPerRowLSE = kBlockMSmem / kGmemElemsPerLoadLSE;
+    static_assert(MaxThreadsPerBlock % kGmemThreadsPerRowLSE == 0, "MaxThreadsPerBlock must be a multiple of kGmemThreadsPerRowLSE");
+    using GmemLayoutAtomLSE = Layout<Shape <Int<MaxThreadsPerBlock / kGmemThreadsPerRowLSE>, Int<kGmemThreadsPerRowLSE>>,
+                                     Stride<Int<kGmemThreadsPerRowLSE>, _1>>;
+    static_assert(kMaxSplits % CUTE_STATIC_V(shape<0>(GmemLayoutAtomLSE{})) == 0);
+    using GmemCopyAtomLSE = std::conditional_t<
+        Has_cp_async,
+        cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeLSE>, float>,
+        cute::Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<AlignmentLSE * sizeof(float) * 8>, float>
+    >;
+    using GmemTiledCopyLSE = decltype(
+        make_tiled_copy(GmemCopyAtomLSE{},
+                        GmemLayoutAtomLSE{},
+                        Layout<Shape<_1, Int<kGmemElemsPerLoadLSE>>>{}));  // Val layout, 4 vals per load
+
+    // Otherwise we get IMA when some threads access sLSE, as we're not doing any masking
+    static_assert((kBlockM * kMaxSplits * AlignmentLSE) % kNThreads == 0, "kNThreads must divide kBlockM * kMaxSplits * AlignmentLSE");
+    // This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts
+    using SmemLSESwizzle = std::conditional_t<
+        kBlockMSmem == 8,
+        Swizzle<5, 0, 5>,
+        std::conditional_t<kBlockMSmem == 16, Swizzle<4, 0, 4>, Swizzle<3, 2, 3>>
+    >;
+    using SmemLayoutAtomLSE =
+        decltype(composition(SmemLSESwizzle{},
+                 Layout<Shape<Int<8>, Int<kBlockMSmem>>,
+                 Stride<Int<kBlockMSmem>, _1>>{}));
+    using SmemLayoutLSE = decltype(tile_to_shape(SmemLayoutAtomLSE{}, Shape<Int<kMaxSplits>, Int<kBlockM>>{}));
+
+    using SmemLayoutO = Layout<Shape<Int<kBlockM>, Int<kBlockK>, Int<kStages>>,
+                               Stride<Int<kBlockK>, _1, Int<kBlockM * kBlockK>>>;
+
+    // We want each column (kMaxSplits) to be processed by threads in the same warp.
+    // To reduce the number of shuffles, we want as few threads on the same column as possible.
+    // E.g., if kBlockM is divisible by 64, and there are 256 threads, we want 4 threads (0, 1, 2, 4) per column
+    // have have 64 such quads.
+    static_assert(MaxThreadsPerBlock % kBlockMSmem == 0, "MaxThreadsPerBlock must be a multiple of kBlockMSmem");
+    static constexpr int kSmemThreadsPerColLSEt = MaxThreadsPerBlock / kBlockMSmem;
+    static_assert(cutlass::NumThreadsPerWarp % kSmemThreadsPerColLSEt == 0, "kSmemThreadsPerColLSEt must divide NumThreadsPerWarp");
+    using S2RLayoutAtomLSE = Layout<Shape<Int<kSmemThreadsPerColLSEt>, Int<MaxThreadsPerBlock / kSmemThreadsPerColLSEt>>>;
+    using S2RTiledCopyLSE = decltype(make_tiled_copy(cute::Copy_Atom<cute::DefaultCopy, float>{}, S2RLayoutAtomLSE{}, Layout<_1>{}));
+
+    using ShapeOPartial = cute::Shape<int32_t, int32_t, int32_t, int32_t, int32_t>;  // (seqlen, d, num_splits, head, batch)
+    using StrideOPartial = cute::Stride<int64_t, _1, int64_t, int64_t, int64_t>;
+    using ShapeLSEPartial = cute::Shape<int32_t, int32_t, int32_t, int32_t>;  // (seqlen, num_splits, head, batch)
+    using StrideLSEPartial = cute::Stride<_1, int64_t, int64_t, int64_t>;  // (seqlen, num_splits, head, batch)
+    using ShapeO = cute::Shape<int32_t, int32_t, int32_t, int32_t>;  // (seqlen, d, head, batch)
+    using StrideO = cute::Stride<int64_t, _1, int64_t, int64_t>;
+    using ShapeLSE = cute::Shape<int32_t, int32_t, int32_t>;  // (seqlen, head, batch)
+    using StrideLSE = cute::Stride<_1, int64_t, int64_t>;  // (seqlen, head, batch)
+
+    struct SharedStorage : cute::aligned_struct<128> {
+        cute::array_aligned<float, cute::cosize_v<SmemLayoutLSE>> smem_lse_partial;
+        cute::array_aligned<int, kBlockM> smem_max_valid_split;
+        cute::array_aligned<ElementPartial, cute::cosize_v<SmemLayoutO>> smem_o_partial;
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    // Device side arguments
+    struct Arguments {
+        ElementPartial const* const ptr_O_partial;
+        ShapeOPartial const shape_O_partial;
+        StrideOPartial const stride_O_partial;
+        float const* const ptr_LSE_partial;
+        ShapeLSEPartial const shape_LSE_partial;
+        StrideLSEPartial const stride_LSE_partial;
+        Element* const ptr_O;
+        StrideO const stride_O;
+        float* const ptr_LSE;
+        StrideLSE const stride_LSE;
+        int const* const cu_seqlens = nullptr;
+        int const* const seqused = nullptr;
+        int const* const num_splits_dynamic_ptr = nullptr;
+        int* const semaphore_to_reset = nullptr;
+    };
+
+    // Kernel entry point API
+    struct Params {
+        ElementPartial const* const ptr_O_partial;
+        ShapeOPartial const shape_O_partial;
+        StrideOPartial const stride_O_partial;
+        float const* const ptr_LSE_partial;
+        ShapeLSEPartial const shape_LSE_partial;
+        StrideLSEPartial const stride_LSE_partial;
+        Element* const ptr_O;
+        StrideO const stride_O;
+        float* const ptr_LSE;
+        StrideLSE const stride_LSE;
+        cutlass::FastDivmod seqlen_divmod, head_divmod;
+        int const* const cu_seqlens = nullptr;
+        int const* const seqused = nullptr;
+        int const* const num_splits_dynamic_ptr = nullptr;
+        int* const semaphore_to_reset = nullptr;
+    };
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        assert(get<1>(args.shape_LSE_partial) <= kMaxSplits);
+        return {
+            args.ptr_O_partial,
+            args.shape_O_partial,
+            args.stride_O_partial,
+            args.ptr_LSE_partial,
+            args.shape_LSE_partial,
+            args.stride_LSE_partial,
+            args.ptr_O,
+            args.stride_O,
+            args.ptr_LSE,
+            args.stride_LSE,
+            cutlass::FastDivmod(get<0>(args.shape_LSE_partial)), cutlass::FastDivmod(get<2>(args.shape_LSE_partial)),
+            args.cu_seqlens,
+            args.seqused,
+            args.num_splits_dynamic_ptr,
+            args.semaphore_to_reset
+        };
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+        Tensor sLSE = make_tensor(make_smem_ptr(shared_storage.smem_lse_partial.data()), SmemLayoutLSE{});
+        Tensor sMaxValidSplit = make_tensor(make_smem_ptr(shared_storage.smem_max_valid_split.data()), Shape<Int<kBlockM>>{});
+        Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o_partial.data()), SmemLayoutO{});
+
+        int const thread_idx = threadIdx.x;
+        int const m_block = blockIdx.x;
+        int const k_block = blockIdx.y;
+        int const batch = blockIdx.z;
+        int const num_splits = params.num_splits_dynamic_ptr ? params.num_splits_dynamic_ptr[batch] : get<1>(params.shape_LSE_partial);
+
+        if (params.semaphore_to_reset && threadIdx.x == 0 && blockIdx.x == gridDim.x - 1 && blockIdx.y == gridDim.y - 1 && blockIdx.z == gridDim.z - 1) {
+            cutlass::arch::wait_on_dependent_grids();
+            *params.semaphore_to_reset = 0;
+        }
+        if (num_splits <= 1) { return; }
+        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{batch, size<0>(params.shape_LSE_partial), params.cu_seqlens, params.seqused};
+        int const offset = seqlen_info.offset;
+        int const seqlen = seqlen_info.seqlen;
+        int max_idx = seqlen * get<2>(params.shape_LSE_partial);
+        if constexpr (Varlen) {
+            if (m_block * kBlockM >= max_idx) { return; }
+        }
+
+        cutlass::FastDivmod seqlen_divmod_dynamic(seqlen);
+
+        // Step 1: load LSE_partial from gmem -> smem
+        Tensor mLSEpartial = make_tensor(make_gmem_ptr(params.ptr_LSE_partial + offset * get<0>(params.stride_LSE_partial)),
+                                         select<1, 0, 2, 3>(params.shape_LSE_partial),
+                                         select<1, 0, 2, 3>(params.stride_LSE_partial))(_, _, _, !Varlen ? batch : 0);  // (num_splits, seqlen, head)
+        Tensor mLSEpartial_copy = cute::tiled_divide(mLSEpartial, Shape<_1, Int<kGmemElemsPerLoadLSE>>{});
+        GmemTiledCopyLSE gmem_tiled_copy_LSE;
+        auto gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_thread_slice(thread_idx);
+        Tensor tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE);
+
+        // Construct identity layout for sLSE
+        Tensor cLSE = make_identity_tensor(make_shape(size<0>(sLSE), size<1>(sLSE)));    // (NUM_SPLITS, BLK_M) -> (num_splits, blk_m)
+        // Repeat the partitioning with identity layouts
+        Tensor tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE);
+
+        cutlass::arch::wait_on_dependent_grids();
+
+        #pragma unroll
+        for (int m = 0; m < size<2>(tLSEcLSE); ++m) {
+            int mi = int(get<1>(tLSEcLSE(_0{}, _0{}, m)));
+            int idx = m_block * kBlockM + mi;
+            if (idx < max_idx) {
+                int m_idx, bidh;
+                if constexpr (!Varlen) {
+                    bidh = params.seqlen_divmod.divmod(m_idx, idx);
+                } else {
+                    bidh = seqlen_divmod_dynamic.divmod(m_idx, idx);
+                }
+                Tensor mLSEpartial_cur_copy = mLSEpartial_copy(_, _, m_idx, bidh);
+                #pragma unroll
+                for (int s = 0; s < size<1>(tLSEcLSE); ++s) {
+                    int si = get<0>(tLSEcLSE(_0{}, s, _0{}));
+                    // if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && thread_idx < 32) { printf("thread_idx = %d, m = %d, s = %d, addr = %p, bank = %d\n", thread_idx, m, s, reinterpret_cast<float *>(&(tLSEsLSE(_0{}, s, m))), reinterpret_cast<int>(&(tLSEsLSE(_0{}, s, m))) / 4 % 32);}
+                    if (si < num_splits) {
+                        cute::copy(gmem_tiled_copy_LSE, mLSEpartial_cur_copy(_, si), tLSEsLSE(_, s, m));
+                    } else {
+                        cute::fill(tLSEsLSE(_, s, m), -INFINITY);
+                    }
+                }
+            } else {
+                // We don't need to zero out the rest of the LSEs, as we will not write the output to gmem
+                // cute::fill(tLSEsLSE(_, _, m), -INFINITY);
+            }
+        }
+        if constexpr (Has_cp_async) { cute::cp_async_fence(); }
+
+        // Step 2: Load O_partial from gmem -> smem for split = 0, 1, ..., kStages - 2.
+        // We want these async loads to be in flight as we compute the LSE.
+        GmemTiledCopyAccum gmem_tiled_copy_O_partial;
+        auto gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_thread_slice(thread_idx);
+        // Construct identity layout for gO
+        Tensor cO = cute::make_identity_tensor(TileShape_MK{});  // (BLK_M,BLK_K) -> (blk_m,blk_k)
+        // Repeat the partitioning with identity layouts
+        Tensor tOcO = gmem_thr_copy_O_partial.partition_D(cO);
+        Tensor mOpartial = make_tensor(make_gmem_ptr(params.ptr_O_partial + offset * get<0>(params.stride_O_partial)),
+                                       params.shape_O_partial, params.stride_O_partial)(_, _, _, _, !Varlen ? batch : 0);  // (seqlen, d, num_splits, head)
+
+        // Precompute these values to avoid recomputing them in the loop
+        Tensor tOmidx = make_tensor<int>(make_shape(size<1>(tOcO)));
+        Tensor tObidh = make_tensor<int>(make_shape(size<1>(tOcO)));
+        Tensor tOrOptr = make_tensor<ElementPartial const*>(make_shape(size<1>(tOcO)));
+        #pragma unroll
+        for (int m = 0; m < size<1>(tOcO); ++m) {
+            int mi = get<0>(tOcO(_0{}, m, _0{}));
+            int idx = m_block * kBlockM + mi;
+            if constexpr (!Varlen) {
+                tObidh(m) = params.seqlen_divmod.divmod(tOmidx(m), idx);
+            } else {
+                tObidh[m] = seqlen_divmod_dynamic.divmod(tOmidx(m), idx);
+            }
+            tOrOptr[m] = &mOpartial(tOmidx(m), k_block * kBlockK, _0{}, tObidh(m));
+            if (idx >= max_idx) {
+                tObidh[m] = -1;
+            }
+        }
+
+        Tensor tOpO = make_tensor<bool>(make_shape(size<2>(tOcO)));
+        if constexpr (!(Is_even_K)) {
+            #pragma unroll
+            for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(_0{}, _0{}, k)) < get<1>(params.shape_O_partial) - k_block * kBlockK; }
+        }
+
+        Tensor tOsOpartial = gmem_thr_copy_O_partial.partition_D(sO);
+
+        auto load_O_partial = [&] (int split, int stage) {
+            Tensor tOsOpartial_cur = tOsOpartial(_, _, _, stage);
+            #pragma unroll
+            for (int m = 0; m < size<1>(tOcO); ++m) {
+                if (tObidh(m) >= 0)  {
+                    Tensor mOpartial_cur = make_tensor(make_gmem_ptr(tOrOptr[m]), mOpartial(_0{}, _, _, _0{}).layout());
+                    Tensor mOpartial_cur_copy = cute::tiled_divide(mOpartial_cur, Shape<Int<kGmemElemsPerLoad>>{});
+                    #pragma unroll
+                    for (int k = 0; k < size<2>(tOcO); ++k) {
+                        int k_idx = get<1>(tOcO(_0{}, _0{}, k)) / kGmemElemsPerLoad;
+                        if (Is_even_K || tOpO(k)) {
+                            cute::copy(gmem_tiled_copy_O_partial, mOpartial_cur_copy(_, k_idx, split), tOsOpartial_cur(_, m, k));
+                        }
+                    }
+                }
+            }
+        };
+
+        for (int s = 0; s < kStages - 1; ++s) {
+            if (s < num_splits) { load_O_partial(s, s); }
+            if constexpr (Has_cp_async) { cute::cp_async_fence(); }
+        }
+
+        // Step 3: load and transpose LSE_partial from smem -> rmem
+        if constexpr (Has_cp_async) { cutlass::arch::cp_async_wait<kStages - 1>(); }
+        __syncthreads();
+
+        S2RTiledCopyLSE s2r_tiled_copy_LSE;
+        auto s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_thread_slice(thread_idx);
+        Tensor ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE);
+        Tensor ts2rrLSE = make_fragment_like(ts2rsLSE);
+        cute::copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE);
+
+        // Step 4: compute the final LSE along the split dimension
+        Tensor lse_sum = make_tensor<float>(make_shape(size<2>(ts2rrLSE)));
+        Tensor ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE);
+        // We compute the max valid split for each row to short-circuit the computation later
+        Tensor max_valid_split = make_tensor<int>(make_shape(size<2>(ts2rrLSE)));
+        static_assert(CUTE_STATIC_V(size<0>(ts2rrLSE)) == 1);
+        #pragma unroll
+        for (int m = 0; m < size<2>(ts2rrLSE); ++m) {
+            float lse_max = ts2rrLSE(_0{}, _0{}, m);
+            #pragma unroll
+            for (int s = 1; s < size<1>(ts2rrLSE); ++s) { lse_max = max(lse_max, ts2rrLSE(_0{}, s, m)); }
+            MaxOp<float> max_op;
+            lse_max = Allreduce<kSmemThreadsPerColLSEt>::run(lse_max, max_op);
+            int max_valid_idx = -1;
+            #pragma unroll
+            for (int s = 0; s < size<1>(ts2rrLSE); ++s) {
+                if (ts2rrLSE(_0{}, s, m) != -INFINITY) { max_valid_idx = get<0>(ts2rcLSE(_0{}, s, _0{})); }
+            }
+            MaxOp<int> max_int_op;
+            max_valid_split[m] = Allreduce<kSmemThreadsPerColLSEt>::run(max_valid_idx, max_int_op);
+            float lse_max_cur = lse_max == -INFINITY ? 0.0f : lse_max;  // In case all local LSEs are -inf
+            float lse_sum_cur = 0.f;
+            #pragma unroll
+            for (int s = 0; s < size<1>(ts2rrLSE); ++s) {
+                float scale = expf(ts2rrLSE(_0{}, s, m) - lse_max_cur);
+                lse_sum_cur += scale;
+                // if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && thread_idx < 32) { printf("thread_idx = %d, m = %d, s = %d, addr = %p, bank = %d\n", thread_idx, m, s, reinterpret_cast<float *>(&(ts2rsLSE(_0{}, s, m))), reinterpret_cast<int>(&(ts2rsLSE(_0{}, s, m))) / 4 % 32);}
+                // ts2rsLSE(_0{}, m, s) = scale;
+                ts2rrLSE(_0{}, s, m) = scale;
+            }
+            SumOp<float> sum_op;
+            lse_sum_cur = Allreduce<kSmemThreadsPerColLSEt>::run(lse_sum_cur, sum_op);
+            lse_sum(m) = logf(lse_sum_cur) + lse_max;
+            float inv_sum = (lse_sum_cur == 0.f || lse_sum_cur != lse_sum_cur) ? 0.f : 1.f / lse_sum_cur;
+            #pragma unroll
+            for (int s = 0; s < size<1>(ts2rrLSE); ++s) { ts2rrLSE(_0{}, s, m) *= inv_sum; }
+        }
+        // Store the scales exp(lse - lse_logsum) back to smem
+        cute::copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE);
+
+        // Store max_valid_split to smem
+        #pragma unroll
+        for (int m = 0; m < size<2>(ts2rrLSE); ++m) {
+            if (get<0>(ts2rcLSE(_0{}, _0{}, m)) == 0) {  // Only the thread responsible for s=0 writes to smem
+                int mi = int(get<1>(ts2rcLSE(_0{}, _0{}, m)));
+                if (mi < kBlockM) { sMaxValidSplit[mi] = max_valid_split[m]; }
+            }
+        }
+
+        // Step 5: store final LSE back to gmem
+        if (k_block == 0) {
+            auto shape_LSE = select<0, 2, 3>(params.shape_LSE_partial);
+            Tensor mLSE = make_tensor(make_gmem_ptr(params.ptr_LSE + offset * get<0>(params.stride_LSE)), shape_LSE, params.stride_LSE)(_, _, !Varlen ? batch : 0);
+            #pragma unroll
+            for (int m = 0; m < size<2>(ts2rrLSE); ++m) {
+                if (get<0>(ts2rcLSE(_0{}, _0{}, m)) == 0) {  // Only the thread responsible for s=0 writes to gmem
+                    int mi = int(get<1>(ts2rcLSE(_0{}, _0{}, m)));
+                    int idx = m_block * kBlockM + mi;
+                    if (idx < max_idx) {
+                        int m_idx, bidh;
+                        if constexpr (!Varlen) {
+                            bidh = params.seqlen_divmod.divmod(m_idx, idx);
+                        } else {
+                            bidh = seqlen_divmod_dynamic.divmod(m_idx, idx);
+                        }
+                        // printf("thread_idx = %d, m = %d, mi = %d, idx = %d, m_idx = %d, bidh = %d, bidb = %d, lse_sum = %f\n", thread_idx, m, mi, idx, m_idx, bidh, bidb, lse_sum(m));
+                        mLSE(m_idx, bidh) = lse_sum(m);
+                    }
+                }
+            }
+        }
+
+        // Step 6: read O_partial from gmem -> smem -> rmem and accumulate the final O
+        __syncthreads();
+        int thr_max_valid_split = sMaxValidSplit[get<0>(tOcO(_0{}, _0{}, _0{}))];
+        #pragma unroll
+        for (int m = 1; m < size<1>(tOcO); ++m) { thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[get<0>(tOcO(_0{}, m, _0{}))]); }
+        Layout tOrOpartial_layout = gmem_thr_copy_O_partial.partition_S(make_tensor<ElementPartial>(TileShape_MK{})).layout();
+        Tensor tOrOpartial = make_fragment_like<ElementPartial>(tOrOpartial_layout);
+        Tensor tOrO = make_fragment_like<float>(tOrOpartial);
+        clear(tOrO);
+        int stage_load = kStages - 1, stage_compute = 0;
+        #pragma unroll 4 // Already tuned for speed
+        for (int s = 0; s <= thr_max_valid_split; ++s) {
+            Tensor scale = make_tensor<float>(make_shape(size<1>(tOrOpartial)));
+            #pragma unroll
+            for (int m = 0; m < size<1>(tOrOpartial); ++m) { scale(m) = sLSE(s, get<0>(tOcO(_0{}, m, _0{}))); }
+
+            if (s + kStages - 1 <= thr_max_valid_split) { load_O_partial(s + kStages - 1, stage_load); }
+            if constexpr (Has_cp_async) { cute::cp_async_fence(); }
+            stage_load = stage_load < kStages - 1 ? stage_load + 1 : 0;
+            if constexpr (Has_cp_async) { cutlass::arch::cp_async_wait<kStages - 1>(); }
+            // We don't need __syncthreads() because each thread is just reading its own data from smem
+            cute::copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, ElementPartial>{},
+                       tOsOpartial(_, _, _, stage_compute), tOrOpartial);
+            stage_compute = stage_compute < kStages - 1 ? stage_compute + 1 : 0;
+
+            #pragma unroll
+            for (int m = 0; m < size<1>(tOrOpartial); ++m) {
+                if (tObidh(m) >= 0 && scale(m) > 0.f) {
+                    #pragma unroll
+                    for (int k = 0; k < size<2>(tOrOpartial); ++k) {
+                        if (Is_even_K || tOpO(k)) {
+                            Tensor rOpartial = make_tensor_like<float>(tOrOpartial(_, m, k));
+                            flash::convert_type_out(tOrOpartial(_, m, k), rOpartial);
+                            #pragma unroll
+                            for (int i = 0; i < size<0>(tOrOpartial); ++i) {
+                                tOrO(i, m, k) += scale(m) * rOpartial[i];
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        // Step 7: Write the final O to gmem
+        Tensor rO = make_tensor_like<Element>(tOrO);
+        flash::convert_type_out(tOrO, rO);
+        auto shape_O = make_shape(get<0>(params.shape_O_partial), get<1>(params.shape_O_partial) - k_block * kBlockK, get<3>(params.shape_O_partial), get<4>(params.shape_O_partial));
+        Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + offset * get<0>(params.stride_O) + k_block * kBlockK * get<1>(params.stride_O)),
+                                shape_O, params.stride_O)(_, _, _, !Varlen ? batch : 0);
+        Tensor mO_copy = cute::tiled_divide(mO, Shape<_1, Int<kGmemElemsPerLoad>>{});
+        GmemTiledCopy gmem_tiled_copy_O;
+        auto gmem_thr_copy_O = gmem_tiled_copy_O.get_thread_slice(thread_idx);
+
+        #pragma unroll
+        for (int m = 0; m < size<1>(tOcO); ++m) {
+            if (tObidh(m) >= 0)  {
+                #pragma unroll
+                for (int k = 0; k < size<2>(tOcO); ++k) {
+                    int k_idx = get<1>(tOcO(_0{}, _0{}, k)) / kGmemElemsPerLoad;
+                    if (Is_even_K || tOpO(k)) {
+                        cute::copy(gmem_tiled_copy_O, rO(_, m, k), mO_copy(_, tOmidx(m), k_idx, tObidh(m)));
+                    }
+                }
+            }
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_fwd_combine_launch_template.h

@@ -0,0 +1,80 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/arch/arch.h"  // For cutlass::arch::Sm80
+#include "cutlass/device_kernel.h"  // For device_kernel
+#include "cutlass/kernel_launch.h"  // For kernel_launch
+
+#include "static_switch.h"
+#include "flash.h"
+#include "flash_fwd_combine_kernel.h"
+
+using namespace cute;
+
+template <int Arch, int kBlockM, int kBlockK, int kLogMaxSplits, bool IsEvenK, bool Varlen, typename Element, typename ElementPartial>
+void run_flash_fwd_combine(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl) {
+    using ArchTag = std::conditional_t<Arch >= 90, cutlass::arch::Sm90, cutlass::arch::Sm80>;
+    using TileShape_MK = cute::Shape<Int<kBlockM>, Int<kBlockK>>;
+    using CombineKernel = flash::FlashAttnFwdCombine<TileShape_MK, kLogMaxSplits, 256 /*kNThreads*/, 1 /*AlignmentLSE*/,
+                                                     IsEvenK, Varlen, Element, ElementPartial, ArchTag>;
+
+    typename CombineKernel::Arguments args {
+        static_cast<ElementPartial const*>(params.oaccum_ptr),
+        {!Varlen ? params.seqlen_q : params.total_q, params.dv, params.num_splits, params.h, !Varlen ? params.b : 1},  // shape_O_partial
+        {params.oaccum_row_stride, _1{}, params.oaccum_split_stride, params.oaccum_head_stride, !Varlen ? params.oaccum_batch_stride : 0},  // stride_O_partial
+        static_cast<float*>(params.softmax_lseaccum_ptr),
+        {!Varlen ? params.seqlen_q : params.total_q, params.num_splits, params.h, !Varlen ? params.b : 1},  // shape_LSE_partial
+        {_1{}, params.lseaccum_split_stride, params.lseaccum_head_stride, !Varlen ? params.lseaccum_batch_stride : 0},  // stride_LSE_partial
+        static_cast<Element*>(params.o_ptr),
+        {params.o_row_stride, _1{}, params.o_head_stride, !Varlen ? params.o_batch_stride : 0},  // stride_O
+        static_cast<float*>(params.softmax_lse_ptr),
+        {_1{}, !Varlen ? params.seqlen_q : params.total_q, !Varlen ? params.h * params.seqlen_q : 0},  // stride_LSE
+        params.cu_seqlens_q, params.seqused_q, params.num_splits_dynamic_ptr, params.tile_count_semaphore
+    };
+
+    typename CombineKernel::Params kernel_params = CombineKernel::to_underlying_arguments(args);
+    int num_blocks_k = cute::ceil_div(params.dv, kBlockK);
+    int num_blocks_m = cute::ceil_div(params.seqlen_q * params.h, kBlockM);
+    dim3 grid_m(num_blocks_m, num_blocks_k, params.b);
+    auto kernel = cutlass::device_kernel<CombineKernel>;
+    int smem_size = CombineKernel::SharedStorageSize;
+    if (smem_size >= 48 * 1024) {
+        CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    // kernel<<<grid_m, CombineKernel::MaxThreadsPerBlock, smem_size, stream>>>(kernel_params);
+    cutlass::kernel_launch<CombineKernel>(grid_m, CombineKernel::MaxThreadsPerBlock, smem_size, stream, kernel_params, Arch >= 90 && enable_pdl /*launch_with_pdl*/);
+    CHECK_CUDA_KERNEL_LAUNCH();
+}
+
+template<typename T, typename Tpartial, int kBlockK>
+void run_mha_fwd_combine_(Flash_fwd_params &params, cudaStream_t stream, bool enable_pdl) {
+    // We want kBlockM to be as small as possible to maximize parallelism.
+    // E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats).
+    static_assert(kBlockK % 32 == 0, "kBlockK must be a multiple of 32");
+    static constexpr int kBlockM = kBlockK % 128 == 0 ? 8 : (kBlockK % 64 == 0 ? 16 : 32);
+    ARCH_SWITCH(params.arch, Arch, [&] {
+        BOOL_SWITCH(params.cu_seqlens_q || params.seqused_q, Varlen, [&] {
+            if constexpr (kBlockM >= 16) {  // If kBlockM == 8 then the minimum number of splits is 32.
+                if (params.num_splits <= 16) {
+                    run_flash_fwd_combine<Arch, kBlockM, kBlockK, 4, false /*IsEvenK*/, Varlen, T, Tpartial>(params, stream, enable_pdl);
+                    return;
+                }
+            }
+            if (params.num_splits <= 32) {
+                run_flash_fwd_combine<Arch, kBlockM, kBlockK, 5, false /*IsEvenK*/, Varlen, T, Tpartial>(params, stream, enable_pdl);
+            } else if (params.num_splits <= 64) {
+                run_flash_fwd_combine<Arch, kBlockM, kBlockK, 6, false /*IsEvenK*/, Varlen, T, Tpartial>(params, stream, enable_pdl);
+            } else if (params.num_splits <= 128) {
+                run_flash_fwd_combine<Arch, kBlockM, kBlockK, 7, false /*IsEvenK*/, Varlen, T, Tpartial>(params, stream, enable_pdl);
+            } else {
+                run_flash_fwd_combine<Arch, kBlockM, kBlockK, 8, false /*IsEvenK*/, Varlen, T, Tpartial>(params, stream, enable_pdl);
+            }
+        });
+    });
+}

flash-attn/flash_fwd_kernel_sm80.h

@@ -0,0 +1,215 @@
+/******************************************************************************
+ * Copyright (c) 2024, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/kernel_hardware_info.h>
+
+#include "seqlen.h"
+#include "utils.h"
+#include "softmax.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
+class FlashAttnFwdSm80 {
+
+public:
+
+    // Type Aliases
+    using CollectiveMainloop = CollectiveMainloop_;
+    using CollectiveEpilogue = CollectiveEpilogue_;
+    static constexpr bool Is_causal = CollectiveMainloop::Is_causal;
+    static constexpr bool Is_local = CollectiveMainloop::Is_local;
+    static_assert(CollectiveMainloop::Varlen == CollectiveEpilogue::Varlen);
+    static constexpr bool Has_softcap = CollectiveMainloop::Has_softcap;
+    static constexpr bool Varlen = CollectiveMainloop::Varlen;
+    static constexpr bool PagedKV = CollectiveMainloop::PagedKV;
+    static constexpr bool Split = CollectiveMainloop::Split;
+    static constexpr bool Is_FP8 = CollectiveMainloop::Is_FP8;
+    static constexpr bool Transpose_V = CollectiveMainloop::Transpose_V;
+    static constexpr bool AppendKV = CollectiveMainloop::AppendKV;
+    static constexpr bool PackGQA = CollectiveMainloop::PackGQA;
+    static constexpr int NumProducerThreads = CollectiveMainloop::NumProducerThreads;
+    using SeqlenInfo_t = typename CollectiveMainloop::SeqlenInfo_t;
+
+    // Mainloop derived types
+    using TileShape_MNK = typename CollectiveMainloop::TileShape_MNK;
+    using TiledMma = typename CollectiveMainloop::TiledMma;
+    using ArchTag = typename CollectiveMainloop::ArchTag;
+    using MainloopArguments = typename CollectiveMainloop::Arguments;
+    using MainloopParams = typename CollectiveMainloop::Params;
+
+    // Epilogue derived types
+    using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+    using EpilogueParams = typename CollectiveEpilogue::Params;
+
+    static_assert(ArchTag::kMinComputeCapability >= 80);
+
+    using TileScheduler = TileScheduler_;
+    using TileSchedulerArguments = typename flash::TileSchedulerArguments;
+    using TileSchedulerParams = typename TileScheduler::Params;
+
+    static constexpr uint32_t NumThreads = CUTE_STATIC_V(size(TiledMma{}));
+    static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{}));
+    static constexpr uint32_t MinBlocksPerMultiprocessor = NumThreads == 128 ? 2 : 1;
+
+    // Kernel level shared memory storage
+    // We overlap the shared memory for the mainloop and epilogue. However, we only want smem_o to overlap with smem_v + smem_k and not smem_q
+    // and nothing else, so we'll pad in case sizeof(smem_o) > sizeof(smem_v) + sizeof(smem_k).
+    static constexpr int mainloop_smem_padding_ = int(sizeof(typename CollectiveEpilogue::TensorStorage))
+        - int(sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)))
+        - int(sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k)));
+    static constexpr int mainloop_smem_padding = mainloop_smem_padding_ < 0 ? 0 : mainloop_smem_padding_;
+    struct SharedStorage {
+        struct TensorStorage : cute::aligned_struct<128> {
+            union {
+                struct {
+                    cute::array<uint32_t, mainloop_smem_padding / sizeof(uint32_t)> padding_;
+                    typename CollectiveMainloop::TensorStorage mainloop;
+                };
+                // We want smem_o to line up with the start of smem_v
+                typename CollectiveEpilogue::TensorStorage epilogue;
+            };
+        } tensors;
+
+        alignas(16) typename TileScheduler::SharedStorage smem_scheduler;
+
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    // Device side arguments
+    struct Arguments {
+        MainloopArguments mainloop{};
+        EpilogueArguments epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerArguments scheduler{};
+    };
+
+    // Kernel entry point API
+    struct Params {
+        MainloopParams mainloop{};
+        EpilogueParams epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerParams scheduler{};
+    };
+
+    //
+    // Methods
+    //
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        CUTLASS_TRACE_HOST("to_underlying_arguments():");
+
+        // Get SM count if needed, otherwise use user supplied SM count
+        int sm_count = args.hw_info.sm_count;
+        if (sm_count <= 0) {
+            CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+                "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+            sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+        }
+
+        CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
+
+        cutlass::KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
+        return {
+            CollectiveMainloop::to_underlying_arguments(args.mainloop),
+            CollectiveEpilogue::to_underlying_arguments(args.epilogue),
+            hw_info,
+            TileScheduler::to_underlying_arguments(args.scheduler)
+        };
+    }
+
+    // Computes the kernel launch grid shape based on runtime parameters
+    static dim3
+    get_grid_shape(Params const& params) {
+        return TileScheduler::get_grid_shape(params.scheduler, params.hw_info.sm_count * MinBlocksPerMultiprocessor);
+    }
+
+    static dim3
+    get_block_shape() {
+        return dim3(MaxThreadsPerBlock, 1, 1);
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        static constexpr int kBlockM = get<0>(TileShape_MNK{});
+
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+        CollectiveMainloop mainloop;
+        CollectiveEpilogue epilogue;
+
+        TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.smem_scheduler));
+        // Initialize matmul objects.
+        TiledMma tiled_mma;
+
+        scheduler.init_consumer();
+
+        int warp_idx = cutlass::canonical_warp_idx_sync();
+        CUTLASS_PRAGMA_NO_UNROLL
+        for (auto work_tile_info = warp_idx == 0 ? scheduler.template get_initial_work</*IsProducerWarp=*/true>(params.scheduler) : scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+             work_tile_info.is_valid(params.scheduler);
+             work_tile_info = warp_idx == 0 ? scheduler.template get_next_work</*IsProducerWarp=*/true>(params.scheduler, work_tile_info) : scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
+            // Attention output (GEMM-II) accumulator.
+            Tensor tOrO = partition_fragment_C(tiled_mma, select<0, 2>(TileShape_MNK{}));
+            float softmax_scale_log2 = params.mainloop.softmax_scale_log2;
+            // If there's tanh softcap, the scaling will be done before tanh.
+            auto block_coord = work_tile_info.get_block_coord(params.scheduler);
+            int const bidb = get<2>(block_coord);
+            if constexpr (Is_FP8 && !Has_softcap) {
+                int const bidh = get<1>(block_coord);
+                int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh;
+                float const q_descale = params.mainloop.ptr_q_descale == nullptr ? 1.0f : params.mainloop.ptr_q_descale[bidb * get<0>(params.mainloop.stride_q_descale) + bidh_kv * get<1>(params.mainloop.stride_q_descale)];
+                float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)];
+                softmax_scale_log2 *= q_descale * k_descale;
+            }
+            flash::Softmax<2 * (2 * kBlockM / NumThreads), /*Max_offset=*/!Is_FP8 ? 0 : 8> softmax(softmax_scale_log2);
+
+            SeqlenInfo_t seqlen_info{
+                bidb,
+                get<0>(params.mainloop.shape_Q),
+                !PagedKV ? size<0>(params.mainloop.shape_K) : size<0>(params.mainloop.shape_K) * size<1>(params.mainloop.shape_pagetable),
+                get<0>(params.mainloop.shape_K_new),
+                params.mainloop.cu_seqlens_q, params.mainloop.cu_seqlens_k, params.mainloop.cu_seqlens_k_new,
+                params.mainloop.seqused_q, params.mainloop.seqused_k, params.mainloop.leftpad_k,
+                params.mainloop.seqlens_rotary
+            };
+            if constexpr (AppendKV) {
+                bool tile_new_valid = mainloop.store_kv_new(
+                    params.mainloop, threadIdx.x, shared_storage, seqlen_info, block_coord);
+                if (tile_new_valid) { __syncthreads(); }
+            }
+            bool tile_valid = mainloop.mma(
+                params.mainloop, tOrO, softmax, threadIdx.x, seqlen_info, block_coord,
+                shared_storage);
+            scheduler.prefetch_next_work(params.scheduler, work_tile_info);
+            if (tile_valid) {
+                // if (threadIdx.x == 128) { printf("Before epilogue, bid.x = %d, bid.y = %d, bid.z = %d, m_block = %d, bidb = %d, split_idx = %d\n", blockIdx.x, blockIdx.y, blockIdx.z, m_block, bidb, split_idx); }
+                epilogue.store(params.epilogue, tOrO, softmax.row_sum, shared_storage, tiled_mma,
+                               threadIdx.x, block_coord);
+            } else {
+                // Write 0 to gO and -inf to gLSE.
+                epilogue.store_zero(params.epilogue, threadIdx.x, block_coord);
+            }
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_fwd_kernel_sm90.h

@@ -0,0 +1,458 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include <cutlass/cutlass.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/kernel_hardware_info.h>
+#include "cutlass/pipeline/pipeline.hpp"
+
+#include "cutlass/arch/grid_dependency_control.h"
+
+#include "seqlen.h"
+#include "utils.h"
+#include "softmax.h"
+
+namespace flash {
+
+using namespace cute;
+
+template <class CollectiveMainloop_, class CollectiveEpilogue_, class TileScheduler_>
+class FlashAttnFwdSm90 {
+
+public:
+
+    // Type Aliases
+    using CollectiveMainloop = CollectiveMainloop_;
+    using CollectiveEpilogue = CollectiveEpilogue_;
+    static constexpr bool Is_causal = CollectiveMainloop::Is_causal;
+    static constexpr bool Is_local = CollectiveMainloop::Is_local;
+    static_assert(CollectiveMainloop::Varlen == CollectiveEpilogue::Varlen);
+    static constexpr bool Has_softcap = CollectiveMainloop::Has_softcap;
+    static constexpr bool Varlen = CollectiveMainloop::Varlen;
+    static constexpr bool Split = CollectiveMainloop::Split;
+    static constexpr bool Is_FP8 = CollectiveMainloop::Is_FP8;
+    static constexpr bool Transpose_V = CollectiveMainloop::Transpose_V;
+    static constexpr bool AppendKV = CollectiveMainloop::AppendKV;
+    static constexpr bool HasQv = CollectiveMainloop::HasQv;
+    static constexpr bool Use_TMA_Q = CollectiveMainloop::Use_TMA_Q;
+    static constexpr bool Use_TMA_KV = CollectiveMainloop::Use_TMA_KV;
+    static constexpr bool Use_TMA_O = CollectiveEpilogue::Use_TMA_O;
+    static constexpr bool PackGQA = CollectiveMainloop::PackGQA;
+    static constexpr int NumProducerThreads = CollectiveMainloop::NumProducerThreads;
+    static constexpr bool SameHeadDim = CollectiveMainloop::SameHeadDim;
+    static constexpr bool LargeHeadDimV = CollectiveMainloop::LargeHeadDimV;
+    static_assert(CollectiveMainloop::LargeHeadDimV == CollectiveEpilogue::LargeHeadDimV);
+    using SeqlenInfo_t = typename CollectiveMainloop::SeqlenInfo_t;
+
+    // Mainloop derived types
+    using TileShape_MNK_PV = typename CollectiveMainloop::TileShape_MNK_PV;
+    using TiledMmaPV = typename CollectiveMainloop::TiledMmaPV;
+    using ArchTag = typename CollectiveMainloop::ArchTag;
+    using ClusterShape = typename CollectiveMainloop::ClusterShape;
+    using MainloopArguments = typename CollectiveMainloop::Arguments;
+    using MainloopParams = typename CollectiveMainloop::Params;
+    using BarrierQ = std::conditional_t<Use_TMA_Q, cutlass::arch::ClusterTransactionBarrier, cutlass::arch::ClusterBarrier>;
+
+    // Epilogue derived types
+    using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+    using EpilogueParams = typename CollectiveEpilogue::Params;
+
+    static_assert(ArchTag::kMinComputeCapability >= 90);
+
+    using TileScheduler = TileScheduler_;
+    using TileSchedulerArguments = typename flash::TileSchedulerArguments;
+    using TileSchedulerParams = typename TileScheduler::Params;
+
+    static constexpr uint32_t NumLoadWarpGroups = 1;
+    static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMmaPV{})) / cutlass::NumThreadsPerWarpGroup;
+    static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMmaPV{})) + (NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup);
+    static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+    static_assert(NumMmaWarpGroups == 1 || NumMmaWarpGroups == 2 || NumMmaWarpGroups == 3);
+
+    /// Register requirement for Load and Math WGs
+    // If we use cp.async to load K and V, we need more registers for the producer WG.
+    static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 1 ? 56 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV ? 24 : 40) : 32);
+    static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 1 ? 256 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV ? 240 : 232) : 160);
+    // If you want to print from the producer warp, you'd need to increase the number of registers
+    // Otherwise you'll get CUDA error.
+    // static constexpr uint32_t LoadRegisterRequirement = 40;
+    // static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 2 ? 232 : 152;
+
+    // Kernel level shared memory storage
+    // We overlap the shared memory for the mainloop and epilogue. However, we only want smem_o to overlap with smem_v
+    // and nothing else, so we'll pad in case sizeof(smem_o) > sizeof(smem_v).
+    static constexpr int mainloop_smem_padding_ = int(sizeof(typename CollectiveEpilogue::TensorStorage)) - int(sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v)));
+    static constexpr int mainloop_smem_padding = mainloop_smem_padding_ < 0 ? 0 : mainloop_smem_padding_;
+    struct SharedStorage {
+        struct TensorStorage : cute::aligned_struct<128, _1> {
+            union {
+                struct {
+                    cute::array<uint32_t, mainloop_smem_padding / sizeof(uint32_t)> padding_;
+                    typename CollectiveMainloop::TensorStorage mainloop;
+                };
+                // We want smem_o to line up with the start of smem_v
+                typename CollectiveEpilogue::TensorStorage epilogue;
+            };
+        } tensors;
+        struct PipelineStorage : cute::aligned_struct<16, _1> {
+            alignas(16) BarrierQ barrier_Q;
+            alignas(16) BarrierQ barrier_Qv;
+            alignas(16) cutlass::arch::ClusterBarrier barrier_O;
+            alignas(16) typename CollectiveMainloop::MainloopPipelineK::SharedStorage pipeline_k;
+            alignas(16) typename CollectiveMainloop::MainloopPipelineV::SharedStorage pipeline_v;
+            alignas(16) typename CollectiveMainloop::MainloopPipelineVt::SharedStorage pipeline_vt;
+            alignas(16) typename CollectiveMainloop::MainloopPipelineKVNew::SharedStorage pipeline_k_new;
+            alignas(16) typename CollectiveMainloop::MainloopPipelineKVNew::SharedStorage pipeline_v_new;
+            alignas(16) typename TileScheduler::SharedStorage smem_scheduler;
+        } pipelines;
+
+    };
+
+    static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+    // Device side arguments
+    struct Arguments {
+        MainloopArguments mainloop{};
+        EpilogueArguments epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerArguments scheduler{};
+    };
+
+    // Kernel entry point API
+    struct Params {
+        MainloopParams mainloop{};
+        EpilogueParams epilogue{};
+        cutlass::KernelHardwareInfo hw_info{};
+        TileSchedulerParams scheduler{};
+    };
+
+    //
+    // Methods
+    //
+
+    // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+    static
+    Params
+    to_underlying_arguments(Arguments const& args) {
+        CUTLASS_TRACE_HOST("to_underlying_arguments():");
+
+        // Get SM count if needed, otherwise use user supplied SM count
+        int sm_count = args.hw_info.sm_count;
+        if (sm_count <= 0) {
+            CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+                "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+            sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+        }
+
+        CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
+
+        cutlass::KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
+        return {
+            CollectiveMainloop::to_underlying_arguments(args.mainloop),
+            CollectiveEpilogue::to_underlying_arguments(args.epilogue),
+            hw_info,
+            TileScheduler::to_underlying_arguments(args.scheduler)
+        };
+    }
+
+    // Computes the kernel launch grid shape based on runtime parameters
+    static dim3
+    get_grid_shape(Params const& params) {
+        return TileScheduler::get_grid_shape(params.scheduler, params.hw_info.sm_count);
+    }
+
+    static dim3
+    get_block_shape() {
+        return dim3(MaxThreadsPerBlock, 1, 1);
+    }
+
+    CUTLASS_DEVICE
+    void
+    operator()(Params const& params, char* smem_buf) {
+
+        static constexpr int NumMmaThreads = NumMmaWarpGroups * cutlass::NumThreadsPerWarpGroup;
+        static constexpr int MmaThreadOffset = NumLoadWarpGroups * cutlass::NumThreadsPerWarpGroup;
+        static constexpr int kBlockM = get<0>(TileShape_MNK_PV{});
+
+        using MainloopPipelineK = typename CollectiveMainloop::MainloopPipelineK;
+        using MainloopPipelineV = typename CollectiveMainloop::MainloopPipelineV;
+        using MainloopPipelineVt = typename CollectiveMainloop::MainloopPipelineVt;
+        using MainloopPipelineKVNew = typename CollectiveMainloop::MainloopPipelineKVNew;
+        using PipelineState = typename CollectiveMainloop::PipelineState;
+        using PipelineParamsK = typename MainloopPipelineK::Params;
+        using PipelineParamsV = typename MainloopPipelineV::Params;
+        using PipelineParamsVt = typename MainloopPipelineVt::Params;
+        using PipelineParamsKVNew = typename MainloopPipelineKVNew::Params;
+
+        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+        int const lane_predicate = cute::elect_one_sync();
+        int const warp_idx = cutlass::canonical_warp_idx_sync();
+
+        // Issue Tma Descriptor Prefetch from a single thread
+        if (warp_idx == 0 && lane_predicate) {
+            CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
+            CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue);
+        }
+
+        // Obtain warp index
+        int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+        int warp_group_idx = cutlass::canonical_warp_group_idx();
+
+        if (warp_idx == 0 && lane_predicate) {
+            shared_storage.pipelines.barrier_Q.init(Use_TMA_Q ? 1 : NumProducerThreads /*numThreads*/);
+            if constexpr (HasQv) {
+                shared_storage.pipelines.barrier_Qv.init(Use_TMA_Q ? 1 : NumProducerThreads /*numThreads*/);
+            }
+            shared_storage.pipelines.barrier_O.init(size(ClusterShape{}) * (Use_TMA_O ? 1 : NumMmaThreads) /*numThreads*/);
+        }
+
+        // We're counting on pipeline_k to call cutlass::arch::fence_barrier_init();
+        PipelineParamsK pipeline_params_k;
+        pipeline_params_k.role = warp_group_idx == 0
+            ? MainloopPipelineK::ThreadCategory::Producer
+            : MainloopPipelineK::ThreadCategory::Consumer;
+        if constexpr (Use_TMA_KV) {
+            pipeline_params_k.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK;
+            pipeline_params_k.is_leader = warp_group_thread_idx == 0;
+            pipeline_params_k.num_consumers = !LargeHeadDimV ? NumMmaThreads : cutlass::NumThreadsPerWarpGroup;
+        } else {
+            pipeline_params_k.consumer_arv_count = !LargeHeadDimV ? NumMmaThreads : cutlass::NumThreadsPerWarpGroup;
+            pipeline_params_k.producer_arv_count = NumProducerThreads;
+        }
+
+        static_assert(is_same_v<PipelineParamsK, PipelineParamsVt>);
+        PipelineParamsVt pipeline_params_vt = pipeline_params_k;
+        if constexpr (Use_TMA_KV && !SameHeadDim) {
+            pipeline_params_vt.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV;
+            if constexpr (LargeHeadDimV) { pipeline_params_vt.num_consumers = NumMmaThreads; }
+        } else {
+            if constexpr (LargeHeadDimV) { pipeline_params_vt.consumer_arv_count = NumMmaThreads; }
+        }
+
+        MainloopPipelineK pipeline_k = [&] {
+            if constexpr (Use_TMA_KV) {
+                return MainloopPipelineK(shared_storage.pipelines.pipeline_k, pipeline_params_k, ClusterShape{});
+            } else {
+                return MainloopPipelineK(shared_storage.pipelines.pipeline_k, pipeline_params_k);
+            }
+        }();
+        // MainloopPipelineV pipeline_v(shared_storage.pipelines.pipeline_v, pipeline_params_v, ClusterShape{});
+        MainloopPipelineV pipeline_v = [&] {
+            if constexpr (!Transpose_V) {
+                static_assert(is_same_v<PipelineParamsK, PipelineParamsV>);
+                if constexpr (Use_TMA_KV) {
+                    return MainloopPipelineV(shared_storage.pipelines.pipeline_v, pipeline_params_vt, ClusterShape{});
+                } else {
+                    return MainloopPipelineV(shared_storage.pipelines.pipeline_v, pipeline_params_vt);
+                }
+            } else {
+                PipelineParamsV pipeline_params_v;
+                pipeline_params_v.role = warp_group_idx == 0
+                    ? MainloopPipelineV::ThreadCategory::Producer
+                    : MainloopPipelineV::ThreadCategory::Consumer;
+                pipeline_params_v.producer_arv_count = NumProducerThreads;
+                pipeline_params_v.consumer_arv_count = NumMmaThreads;
+                return MainloopPipelineV(shared_storage.pipelines.pipeline_v, pipeline_params_v);
+            }
+        }();
+        // If we need to transpose V (e.g. FP8 and V is row-major), we use pipeline_vt for the TMA, then
+        // the producer WG will read from pipeline_vt and write to pipeline_v.
+        // If we don't need to transpose V, we use pipeline_v for the TMA, and pipeline_vt won't be used.
+        // Technically for pipeline_params_vt, warp0 of WG0 is the producer and all of WG0 are consumers.
+        // However, the thread role isn't used in the pipeline implementation.
+        MainloopPipelineVt pipeline_vt = [&] {
+            if constexpr (Use_TMA_KV) {
+                pipeline_params_vt.num_consumers = NumProducerThreads; // TMA_V is only consumed by the producer WG
+                return MainloopPipelineVt(shared_storage.pipelines.pipeline_vt, pipeline_params_vt, ClusterShape{});
+            } else {
+                pipeline_params_vt.consumer_arv_count = NumProducerThreads; // TMA_V is only consumed by the producer WG
+                return MainloopPipelineVt(shared_storage.pipelines.pipeline_vt, pipeline_params_vt);
+            }
+        }();
+
+        PipelineParamsKVNew pipeline_params_kv_new;
+        pipeline_params_kv_new.role = warp_group_idx == 0
+            ? MainloopPipelineKVNew::ThreadCategory::Producer
+            : MainloopPipelineKVNew::ThreadCategory::Consumer;
+        pipeline_params_kv_new.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK;
+        pipeline_params_kv_new.is_leader = warp_group_thread_idx == 0;
+        pipeline_params_kv_new.num_consumers = NumMmaThreads;
+        auto pipeline_k_new = cute::conditional_return<AppendKV>(MainloopPipelineKVNew(shared_storage.pipelines.pipeline_k_new, pipeline_params_kv_new, ClusterShape{}), nullptr);
+        if constexpr (!SameHeadDim) {
+            pipeline_params_kv_new.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV;
+        }
+        auto pipeline_v_new = cute::conditional_return<AppendKV>(MainloopPipelineKVNew(shared_storage.pipelines.pipeline_v_new, pipeline_params_kv_new, ClusterShape{}), nullptr);
+
+        CollectiveMainloop mainloop;
+        CollectiveEpilogue epilogue;
+
+        // We need this to guarantee that the Pipeline init is visible to all producers and consumer blocks in the Cluster
+        if constexpr (size(ClusterShape{}) > 1) {
+            cute::cluster_arrive_relaxed();
+            cute::cluster_wait();
+        } else {
+            __syncthreads();
+        }
+
+        TileScheduler scheduler(reinterpret_cast<typename TileScheduler::SharedStorage*>(&shared_storage.pipelines.smem_scheduler));
+
+        if (warp_group_idx == 0) {  // Producer
+            cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
+
+            // The pipelines for AppendKV and main attention are different, since e.g. main attention
+            // might use cp.async to load KV (if PagedKVNonTMA) while AppendKV always uses TMA to load
+            // KV_new. Since the pipeline states are different, we have to manually sync to make
+            // sure the two pipelines don't race when accessing smem_k and smem_v.
+            PipelineState smem_pipe_write = cutlass::make_producer_start_state<MainloopPipelineK>();
+            PipelineState smem_pipe_write_new = cutlass::make_producer_start_state<MainloopPipelineKVNew>();
+            int work_idx = 0;
+            int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
+            static constexpr bool SingleProducerWarp = NumProducerThreads == cutlass::NumThreadsPerWarp;
+            if constexpr (SingleProducerWarp) {
+                if (warp_idx_in_warpgroup != 0) { return; }
+            }
+            if (!SingleProducerWarp && warp_idx_in_warpgroup != 0) { scheduler.init_consumer(); }
+
+            cutlass::arch::wait_on_dependent_grids();
+
+            // Load Q, K, V
+            for (auto work_tile_info = SingleProducerWarp || warp_idx_in_warpgroup == 0 ? scheduler.template get_initial_work</*IsProducerWarp=*/true>(params.scheduler) : scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+                 work_tile_info.is_valid(params.scheduler);
+                 work_tile_info = SingleProducerWarp || warp_idx_in_warpgroup == 0 ? scheduler.template get_next_work</*IsProducerWarp=*/true>(params.scheduler, work_tile_info) : scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info)) {
+
+                auto block_coord = work_tile_info.get_block_coord(params.scheduler);
+                SeqlenInfo_t seqlen_info{
+                    get<2>(block_coord) /*bidb*/,
+                    get<0>(params.mainloop.shape_Q),
+                    !params.mainloop.ptr_pagetable ? size<0>(params.mainloop.shape_K) : size<0>(params.mainloop.shape_K) * size<1>(params.mainloop.shape_pagetable),
+                    get<0>(params.mainloop.shape_K_new),
+                    params.mainloop.cu_seqlens_q, params.mainloop.cu_seqlens_k, params.mainloop.cu_seqlens_k_new,
+                    params.mainloop.seqused_q, params.mainloop.seqused_k, params.mainloop.leftpad_k,
+                    params.mainloop.seqlens_rotary
+                };
+                if constexpr (AppendKV) {
+                    bool tile_new_valid = mainloop.load_kv_new(
+                        params.mainloop, pipeline_k_new, pipeline_v_new,
+                        smem_pipe_write_new, shared_storage, seqlen_info, block_coord, work_idx);
+                    if (tile_new_valid) {
+                        // if (threadIdx.x == 0) { printf("Producer: Before sync\n"); }
+                        cutlass::arch::NamedBarrier::sync(NumMmaThreads + NumProducerThreads, static_cast<uint32_t>(FwdNamedBarriers::AppendKV) /*id*/);
+                        // if (threadIdx.x == 0) { printf("Producer: After sync\n"); }
+                    }
+                }
+                auto scheduler_prefetch = [&scheduler, &params, &work_tile_info]() {
+                    scheduler.prefetch_next_work(params.scheduler, work_tile_info);
+                };
+                // pipeline_vt won't be used if we don't need to transpose V.
+                mainloop.load(params.mainloop, pipeline_k, pipeline_v, pipeline_vt, smem_pipe_write,
+                                         shared_storage, scheduler_prefetch, seqlen_info, block_coord, work_idx);
+            }
+            mainloop.load_tail(pipeline_k, pipeline_v, pipeline_vt, smem_pipe_write, shared_storage, work_idx);
+        } else {  // Consumer
+            cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
+
+            // Initialize matmul objects.
+            TiledMmaPV tiled_mma_pv;
+
+            PipelineState smem_pipe_read;
+            PipelineState smem_pipe_read_new;
+            // We don't need separate variables smem_pipe_release_k and smem_pipe_release_v
+            // (like in Cutlass's gemm) because the read and release pipeline states are always the same.
+
+            scheduler.init_consumer();
+            mainloop.mma_init();
+
+            int work_idx = 0;
+            CUTLASS_PRAGMA_NO_UNROLL
+            for (auto work_tile_info = scheduler.template get_initial_work</*IsProducerWarp=*/false>(params.scheduler);
+                 work_tile_info.is_valid(params.scheduler);
+                 // get_next_work will be called before the epilogue
+                 ) {
+                auto block_coord = work_tile_info.get_block_coord(params.scheduler);
+                int const bidb = get<2>(block_coord);
+                SeqlenInfo_t seqlen_info{
+                    bidb,
+                    get<0>(params.mainloop.shape_Q),
+                    !params.mainloop.ptr_pagetable ? size<0>(params.mainloop.shape_K) : size<0>(params.mainloop.shape_K) * size<1>(params.mainloop.shape_pagetable),
+                    get<0>(params.mainloop.shape_K_new),
+                    params.mainloop.cu_seqlens_q, params.mainloop.cu_seqlens_k, params.mainloop.cu_seqlens_k_new,
+                    params.mainloop.seqused_q, params.mainloop.seqused_k, params.mainloop.leftpad_k,
+                    params.mainloop.seqlens_rotary
+                };
+                if constexpr (AppendKV) {
+                    bool tile_new_valid = mainloop.store_kv_new(
+                        params.mainloop, pipeline_k_new, pipeline_v_new, smem_pipe_read_new,
+                        threadIdx.x - MmaThreadOffset, shared_storage, seqlen_info, block_coord);
+                    if (tile_new_valid) {
+                        // if (threadIdx.x == 128) { printf("Consumer: Before sync\n"); }
+                        // We need this sync so that the gmem write from the consumers is visible to the producer
+                        // that might do TMA read after that.
+                        asm volatile ("fence.proxy.async.global;");
+                        cutlass::arch::NamedBarrier::arrive(NumMmaThreads + NumProducerThreads, static_cast<uint32_t>(FwdNamedBarriers::AppendKV) /*id*/);
+                        // arrive is enough, we don't need sync. The producer will sync, which means
+                        // after that sync we're guaranteed that the AppendKV pipeline have finished
+                        // loading and consumer smem_k and smem_v.
+                        // if (threadIdx.x == 128) { printf("Consumer: After sync\n"); }
+                    }
+                }
+                // If there's tanh softcap, the scaling will be done before tanh.
+                float softmax_scale_log2 = params.mainloop.softmax_scale_log2;
+                if constexpr (Is_FP8 && !Has_softcap) {
+                    int const bidh = get<1>(block_coord);
+                    int const bidh_kv = !PackGQA ? params.mainloop.qhead_per_khead_divmod.divide(bidh) : bidh;
+                    float const q_descale = params.mainloop.ptr_q_descale == nullptr ? 1.0f : params.mainloop.ptr_q_descale[bidb * get<0>(params.mainloop.stride_q_descale) + bidh_kv * get<1>(params.mainloop.stride_q_descale)];
+                    float const k_descale = params.mainloop.ptr_k_descale == nullptr ? 1.0f : params.mainloop.ptr_k_descale[bidb * get<0>(params.mainloop.stride_k_descale) + bidh_kv * get<1>(params.mainloop.stride_k_descale)];
+                    softmax_scale_log2 *= q_descale * k_descale;
+                }
+                flash::Softmax<!LargeHeadDimV ? 2 * (2 * kBlockM / NumMmaThreads) : 2, /*Max_offset=*/!Is_FP8 ? 0 : 8> softmax(softmax_scale_log2);
+                // Attention output (GEMM-II) accumulator.
+                Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_MNK_PV{}));
+                bool tile_valid;
+                if constexpr (!LargeHeadDimV) {
+                    tile_valid = mainloop.mma(
+                        params.mainloop, pipeline_k, pipeline_v, smem_pipe_read,
+                        tOrO, softmax, threadIdx.x - MmaThreadOffset, work_idx, seqlen_info, block_coord, shared_storage);
+                } else {  // mma_pv might not compile if !LargeHeadDimV
+                    if (warp_group_idx == 1) {
+                        tile_valid = mainloop.mma(
+                            params.mainloop, pipeline_k, pipeline_v, smem_pipe_read,
+                            tOrO, softmax, threadIdx.x - MmaThreadOffset, work_idx, seqlen_info, block_coord, shared_storage);
+                    } else {
+                        tile_valid = mainloop.mma_pv(
+                            params.mainloop, pipeline_v, smem_pipe_read,
+                            tOrO, softmax, threadIdx.x - MmaThreadOffset, seqlen_info, block_coord, shared_storage);
+                    }
+                }
+                // Do this here before the epilogue so that the next tile is ready to go.
+                work_tile_info = scheduler.template get_next_work</*IsProducerWarp=*/false>(params.scheduler, work_tile_info);
+                if constexpr (Split && Varlen) {
+                    if (!work_tile_info.is_valid(params.scheduler)) {  // Last tile
+                        cutlass::arch::launch_dependent_grids();
+                    }
+                }
+                if (tile_valid) {
+                    // if (threadIdx.x == 128) { printf("Before epilogue, bid.x = %d, bid.y = %d, bid.z = %d, m_block = %d, bidb = %d, split_idx = %d\n", blockIdx.x, blockIdx.y, blockIdx.z, m_block, bidb, split_idx); }
+                    epilogue.store(params.epilogue, tOrO, softmax.row_sum, shared_storage, tiled_mma_pv,
+                                   threadIdx.x - MmaThreadOffset, block_coord);
+                } else {
+                    // Write 0 to gO and -inf to gLSE.
+                    epilogue.store_zero(params.epilogue, threadIdx.x - MmaThreadOffset, block_coord);
+                }
+            }
+            epilogue.store_tail();
+        }
+
+    }
+
+};
+
+} // namespace flash

flash-attn/flash_fwd_launch_template.h

@@ -0,0 +1,223 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/device_kernel.h"  // For device_kernel
+#include <cutlass/kernel_hardware_info.h>
+#include "cutlass/cluster_launch.hpp"
+#include "cutlass/kernel_launch.h"
+
+#include "static_switch.h"
+#include "flash.h"
+#include "tile_size.h"
+#include "tile_scheduler.hpp"
+#include "flash_fwd_kernel_sm90.h"
+#include "flash_fwd_kernel_sm80.h"
+#include "mainloop_fwd_sm90_tma_gmma_ws.hpp"
+#include "mainloop_fwd_sm80.hpp"
+#include "epilogue_fwd.hpp"
+
+using namespace cute;
+
+template <int Arch, int kHeadDim, int kHeadDimV, int ClusterM, typename Element, typename ElementOut,
+          bool Is_causal, bool Is_local, bool Has_softcap, bool Varlen, bool PagedKVNonTMA, bool AppendKV, bool HasQv,
+          bool PackGQA, bool Split, bool V_colmajor>
+void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
+    static_assert(!(Is_causal && Is_local), "Causal and Local cannot be enabled at the same time");
+    static_assert(!(AppendKV && V_colmajor), "AppendKV and V_colmajor cannot be enabled at the same time");
+    static_assert(!(AppendKV && !Varlen), "AppendKV requires Varlen");
+    static constexpr bool Is_FP8 = cute::is_same_v<Element, cutlass::float_e4m3_t> || cute::is_same_v<Element, cutlass::float_e5m2_t>;
+    static constexpr bool FP8_TransposeV = Is_FP8 && !V_colmajor;
+    using ArchTag = std::conditional_t<Arch >= 90, cutlass::arch::Sm90, cutlass::arch::Sm80>;
+
+    // Can't use structured binding since it's not compatible with constexpr
+    static constexpr std::tuple<int, int, bool, bool> kBlockMN_RS_IntraWGOverlap = tile_size_fwd_sm90(kHeadDim, kHeadDimV, Is_causal, Is_local, sizeof(Element) /*element_size*/, V_colmajor, PagedKVNonTMA, Has_softcap);
+    static constexpr std::tuple<int, int, int, int, bool> kBlockMN_kNWarps_Stages_RS = tile_size_fwd_sm8x(Arch == 86 || Arch == 89, kHeadDim, kHeadDimV, Is_causal, Is_local, sizeof(Element) /*element_size*/, PagedKVNonTMA, Varlen && Split, Has_softcap, AppendKV);
+    static constexpr int kBlockM = Arch >= 90 ? std::get<0>(kBlockMN_RS_IntraWGOverlap) : std::get<0>(kBlockMN_kNWarps_Stages_RS);
+    static constexpr int kBlockN = Arch >= 90 ? std::get<1>(kBlockMN_RS_IntraWGOverlap) : std::get<1>(kBlockMN_kNWarps_Stages_RS);
+    static constexpr bool MmaPV_is_RS = std::get<2>(kBlockMN_RS_IntraWGOverlap);
+    static constexpr bool IntraWGOverlap = std::get<3>(kBlockMN_RS_IntraWGOverlap);
+    static constexpr int kNWarps = std::get<2>(kBlockMN_kNWarps_Stages_RS);
+    static constexpr int kStages = Arch >= 90 ? 2 : std::get<3>(kBlockMN_kNWarps_Stages_RS);
+    static constexpr bool Q_in_regs = Arch >= 90 ? false : std::get<4>(kBlockMN_kNWarps_Stages_RS);
+
+    using TileShape_MNK = cute::Shape<Int<kBlockM>, Int<kBlockN>, Int<kHeadDim>>;
+    using TileShape_MNK_PV = cute::Shape<Int<kBlockM>, Int<kHeadDimV>, Int<kBlockN>>;
+    using ClusterShape = cute::Shape<Int<ClusterM>, _1, _1>;
+    using CollectiveMainloop = std::conditional_t<
+        Arch >= 90,
+        flash::CollectiveMainloopFwdSm90<kStages, ClusterShape, TileShape_MNK, kHeadDimV, Element, float, cutlass::arch::Sm90, Is_causal, Is_local, Has_softcap, Varlen, PagedKVNonTMA, AppendKV, HasQv, MmaPV_is_RS, IntraWGOverlap, PackGQA, Split, V_colmajor>,
+        flash::CollectiveMainloopFwdSm80<kNWarps, kStages, Q_in_regs, TileShape_MNK, kHeadDimV, Element, float, cutlass::arch::Sm80, Is_causal, Is_local, Has_softcap, Varlen, PagedKVNonTMA, AppendKV, PackGQA, Split>
+    >;
+    using CollectiveEpilogue = flash::CollectiveEpilogueFwd<TileShape_MNK_PV, ClusterShape, ElementOut, ArchTag, CollectiveMainloop::NumMmaThreads, Varlen, PackGQA, Split, FP8_TransposeV>;
+
+    static constexpr int NumProducerThreads = Arch >= 90 ? CollectiveMainloop::NumProducerThreads : CollectiveMainloop::NumMmaThreads;
+    using SchedulerPersistent = std::conditional_t<Varlen,
+        flash::VarlenDynamicPersistentTileScheduler<kBlockM, CollectiveMainloop::NumMmaThreads, NumProducerThreads, Split, PackGQA, Arch >= 90 /*WarpSpecialized*/>,
+        std::conditional_t<!Is_causal && !Is_local,
+            flash::StaticPersistentTileScheduler<Split>,
+            flash::DynamicPersistentTileScheduler<CollectiveMainloop::NumMmaThreads, NumProducerThreads, Split, PackGQA, Arch >= 90 /*WarpSpecialized*/>
+        >
+    >;
+    using SchedulerSingleTile = flash::SingleTileScheduler<Varlen, Split, PackGQA, kBlockM>;
+    // If Split then we probably don't have enough work for PersistentScheduler to be useful.
+    // However, if Varlen (e.g., during decode where we have max_seqlens), using PersistentScheduler is better
+    // since we'll avoid launching a bunch of thread blocks that immediately exit.
+    // On Sm80, noncausal persistent seems a bit slower.
+    static constexpr bool UsePersistentScheduler = Arch >= 90 ? !(Split && !Varlen) : ((Is_causal && !Varlen) || (Varlen && Split));
+    using Scheduler = std::conditional_t<!UsePersistentScheduler, SchedulerSingleTile, SchedulerPersistent>;
+    using AttnKernel = std::conditional_t<
+        Arch >= 90,
+        flash::enable_sm90_or_later<flash::FlashAttnFwdSm90<CollectiveMainloop, CollectiveEpilogue, Scheduler>>,
+        flash::enable_sm80_to_sm89<flash::FlashAttnFwdSm80<CollectiveMainloop, CollectiveEpilogue, Scheduler>>
+    >;
+
+    bool const is_varlen_q = params.cu_seqlens_q;
+    bool const is_varlen_k = params.cu_seqlens_k;
+    bool const is_varlen_k_new = params.cu_seqlens_knew;
+    int seqlen_q = !is_varlen_q ? params.seqlen_q : params.total_q;
+    int batch_q = !is_varlen_q ? params.b : 1;
+    int batch_k = !is_varlen_k ? (params.kv_batch_idx ? params.b_k : params.b) : 1;
+    typename CollectiveMainloop::StrideV v_strides =
+        cute::conditional_return<!V_colmajor>(
+            make_stride(params.v_row_stride, _1{}, params.v_head_stride, !is_varlen_k ? params.v_batch_stride : 0),
+            make_stride(_1{}, params.v_dim_stride, params.v_head_stride, !is_varlen_k ? params.v_batch_stride : 0));
+    typename CollectiveMainloop::Arguments mainloop_args {
+        static_cast<Element const*>(params.q_ptr),
+        {seqlen_q, params.d, params.h, batch_q},  // shape_Q
+        {params.q_row_stride, _1{}, params.q_head_stride, !is_varlen_q ? params.q_batch_stride : 0},  // stride_Q
+        static_cast<Element*>(params.k_ptr),
+        {!params.page_table ? (!is_varlen_k ? params.seqlen_k : params.total_k) : params.page_size,
+         params.d, params.h_k, !params.page_table ? batch_k : params.num_pages},  // shape_K
+        {params.k_row_stride, _1{}, params.k_head_stride, !is_varlen_k ? params.k_batch_stride : 0},  // stride_K
+        static_cast<Element*>(params.v_ptr),
+        params.dv,  // headdim_v
+        v_strides,  // stride_V
+        static_cast<Element const*>(params.knew_ptr),
+        {!is_varlen_k_new ? params.seqlen_knew : params.total_knew, params.d, params.h_k, !is_varlen_k_new ? params.b : 1},  // shape_K_new
+        {params.knew_row_stride, _1{}, params.knew_head_stride, !is_varlen_k_new ? params.knew_batch_stride : 0},  // stride_K_new
+        static_cast<Element const*>(params.vnew_ptr),
+        {params.vnew_row_stride, _1{}, params.vnew_head_stride, !is_varlen_k_new ? params.vnew_batch_stride : 0}, // stride_V_new
+        static_cast<Element const*>(params.qv_ptr),
+        {params.qv_row_stride, _1{}, params.qv_head_stride, !is_varlen_q ? params.qv_batch_stride : 0},  // stride_Qv
+        static_cast<Element const*>(params.rotary_cos_ptr),
+        {params.seqlen_k, params.rotary_dim / 2},  // shape_rotary, the seqlen shape doesn't matter
+        {params.rotary_dim / 2, _1{}},  // stride_rotary_cos
+        static_cast<Element const*>(params.rotary_sin_ptr),
+        {params.rotary_dim / 2, _1{}},  // stride_rotary_sin
+        params.is_rotary_interleaved,
+        params.page_table,
+        // if page_size is not set, avoid dividing by zero
+        {params.kv_batch_idx ? params.b_k : params.b, !params.page_table ? 0 : params.seqlen_k / params.page_size}, // shape_page_table
+        {params.page_table_batch_stride, _1{}},  // stride_page_table
+        params.scale_softmax,
+        params.q_descale_ptr, params.k_descale_ptr, params.v_descale_ptr,
+        {params.q_descale_batch_stride, params.q_descale_head_stride},
+        {params.k_descale_batch_stride, params.k_descale_head_stride},
+        {params.v_descale_batch_stride, params.v_descale_head_stride},
+        params.window_size_left, params.window_size_right, params.attention_chunk,
+        params.softcap,
+        params.num_splits,
+        params.kv_batch_idx,
+        params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew,
+        params.seqused_q, params.seqused_k,
+        params.leftpad_k, params.seqlens_rotary
+    };
+    typename CollectiveEpilogue::Arguments epilogue_args {
+        static_cast<ElementOut*>(params.o_ptr),
+        {seqlen_q, params.dv, params.h, batch_q, params.num_splits},  // shape_O
+        {params.o_row_stride, _1{}, params.o_head_stride, !is_varlen_q ? params.o_batch_stride : 0, 0}, // stride_O
+        static_cast<float*>(params.oaccum_ptr),
+        {params.oaccum_row_stride, _1{}, params.oaccum_head_stride, !is_varlen_q ? params.oaccum_batch_stride : 0, params.oaccum_split_stride}, // stride_O_partial
+        static_cast<float*>(params.softmax_lse_ptr),
+        {_1{}, seqlen_q, !is_varlen_q ? params.h * seqlen_q : 0, 0},  // stride_LSE
+        static_cast<float*>(params.softmax_lseaccum_ptr),
+        {_1{}, seqlen_q, !is_varlen_q ? params.h * seqlen_q : 0, params.h * seqlen_q * batch_q},  // stride_LSE_partial
+        params.h_k,
+        params.cu_seqlens_q, params.seqused_q
+    };
+
+    int qhead_per_khead = !PackGQA ? 1 : cutlass::ceil_div(params.h, params.h_k);
+    int num_blocks_m = cutlass::ceil_div(params.seqlen_q * qhead_per_khead, get<0>(TileShape_MNK{}));
+    num_blocks_m = cutlass::round_up(num_blocks_m, size<0>(ClusterShape{}));
+    typename flash::TileSchedulerArguments scheduler_args {
+        num_blocks_m, !PackGQA ? params.h : params.h_k, params.b, params.num_splits,
+        params.h / params.h_k,
+        params.seqlen_q,
+        params.seqlen_k, params.d, params.dv, sizeof(Element),
+        params.tile_count_semaphore, params.cu_seqlens_q, params.seqused_q,
+        // params.num_m_blocks_ptr,
+        params.num_splits_dynamic_ptr,
+    };
+
+    if (Varlen && params.num_splits_dynamic_ptr && !params.skip_scheduler_metadata_computation) {
+        prepare_varlen_num_blocks(params, stream, PackGQA, kBlockM, kBlockN, Arch >= 90 /*enable_pdl*/);
+        CHECK_CUDA_KERNEL_LAUNCH();
+    }
+
+    int device;
+    CHECK_CUDA(cudaGetDevice(&device));
+    typename AttnKernel::Params kernel_params = AttnKernel::to_underlying_arguments({
+        mainloop_args, epilogue_args, {device, params.num_sm}, scheduler_args
+    });
+
+    dim3 grid_dims = AttnKernel::get_grid_shape(kernel_params);
+    dim3 block_dims = AttnKernel::get_block_shape();
+    int smem_size = AttnKernel::SharedStorageSize;
+    // int smem_size_q = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_q));
+    // int smem_size_k = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_k));
+    // int smem_size_v = sizeof(decltype((typename CollectiveMainloop::TensorStorage{}).smem_v));
+    // printf("smem_size = %d, q = %d, k = %d, v = %d\n", smem_size, smem_size_q, smem_size_k, smem_size_v);
+    // Get the ptr to kernel function.
+    if constexpr (size(ClusterShape{}) > 1) {
+        void const* kernel = (void const*) cutlass::device_kernel<AttnKernel>;
+        if (smem_size >= 48 * 1024) {
+            CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+        }
+        dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
+        cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, stream};
+        cutlass::launch_kernel_on_cluster(launch_params, kernel, kernel_params);
+    } else {
+        auto kernel = cutlass::device_kernel<AttnKernel>;
+        if (smem_size >= 48 * 1024) {
+            CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+        }
+        // kernel<<<grid_dims, block_dims, smem_size, stream>>>(kernel_params);
+        cutlass::kernel_launch<AttnKernel>(grid_dims, block_dims, smem_size, stream, kernel_params,
+                                           Arch >= 90 && Varlen && params.num_splits_dynamic_ptr && !params.skip_scheduler_metadata_computation /*launch_with_pdl*/);
+    }
+    CHECK_CUDA_KERNEL_LAUNCH();
+}
+
+template<int Arch, typename T, int kHeadDim, int kHeadDimV, bool Split, bool PagedKVNonTMA, bool Has_softcap, bool PackGQA>
+void run_mha_fwd_(Flash_fwd_params &params, cudaStream_t stream) {
+    static_assert(sizeof(T) == 2 || sizeof(T) == 1, "Only 16bit and 8bit are supported");
+    static constexpr bool Is_FP8 = cute::is_same_v<T, cutlass::float_e4m3_t> || cute::is_same_v<T, cutlass::float_e5m2_t>;
+    using T_out = std::conditional_t<!Is_FP8, T, cutlass::bfloat16_t>;
+    CAUSAL_LOCAL_SWITCH(params.is_causal, params.is_local, Is_causal, Is_local, [&] {
+        VCOLMAJOR_SWITCH(params.v_dim_stride != 1, V_colmajor_, [&] {
+            static constexpr bool V_colmajor = V_colmajor_ && sizeof(T) == 1;
+            VARLEN_SWITCH(params.cu_seqlens_q || params.cu_seqlens_k || params.seqused_q || params.seqused_k || params.leftpad_k, Varlen, [&] {
+                // Only needed here to decide if we should use cluster
+                static constexpr int kBlockM = Arch >= 90 ? std::get<0>(tile_size_fwd_sm90(kHeadDim, kHeadDimV, Is_causal, Is_local, sizeof(T) /*element_size*/, V_colmajor, PagedKVNonTMA, Has_softcap)) : 128;
+
+                static constexpr bool Enable_cluster = Arch == 90 && (sizeof(T) == 2 ? (kHeadDim >= 128) : (kHeadDim == 192)) && !Is_causal && !Is_local && !Split && !PagedKVNonTMA && !Varlen;
+                BOOL_SWITCH(params.qv_ptr, HasQV_, [&] {
+                    static constexpr bool HasQv = HasQV_ && Arch == 90 && !Is_FP8 && kHeadDim == 64 && kHeadDimV >= 256;
+                    APPENDKV_SWITCH(params.knew_ptr, AppendKV, [&] {
+                        // Only use Cluster if number of tiles along seqlen_q is even and not varlen
+                        CLUSTER_SWITCH(cutlass::ceil_div(params.seqlen_q * (!PackGQA ? 1 : params.h / params.h_k), kBlockM) % 2 == 0, Use_cluster, [&] {
+                            static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1;
+                            run_flash_fwd<Arch, kHeadDim, kHeadDimV, ClusterM, T, T_out, Is_causal, Is_local, Has_softcap, Varlen, PagedKVNonTMA, AppendKV && Varlen, HasQv, PackGQA, Split, V_colmajor>(params, stream);
+                        });
+                    });
+                });
+            });
+        });
+    });
+}

flash-attn/flash_prepare_scheduler.cu

@@ -0,0 +1,124 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#include "cutlass/fast_math.h"
+#include "cutlass/barrier.h"
+#include "cutlass/arch/barrier.h"
+
+#include "cutlass/arch/grid_dependency_control.h"
+
+#include "flash.h"
+
+namespace flash {
+
+__global__ void prepare_varlen_num_blocks_kernel(
+        int seqlen_q_static, int seqlen_k_static, int seqlen_k_new_static,
+        int const* const cu_seqlens_q, int const* const cu_seqlens_k, int const* const cu_seqlens_k_new,
+        int const* const seqused_q, int const* const seqused_k, int const* const leftpad_k_ptr,
+        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 num_splits_dynamic_ptr,
+        bool enable_pdl) {
+
+    static constexpr int kNumBatchPerWarp = cutlass::NumThreadsPerWarp - 1;
+    static constexpr int kSmemSize = 1;
+    // Assume that there's only one block in the grid
+    __shared__ int total_blocks_smem[kSmemSize];
+
+    // There's only 1 block in the grid, so might as well start launching the main attn kernel
+    if (enable_pdl) { cutlass::arch::launch_dependent_grids(); }
+
+    if (threadIdx.x < kSmemSize) { total_blocks_smem[threadIdx.x] = 0; }
+    __syncthreads();
+
+    if (threadIdx.x == 0 && tile_count_semaphore) { *tile_count_semaphore = 0; }
+
+    int lane = threadIdx.x % cutlass::NumThreadsPerWarp;
+
+    auto get_num_m_blocks = [&](int bidb_start) {
+        int batch_idx = lane + bidb_start;
+        int seqlen;
+        if (seqused_q) {
+            seqlen = batch_idx < num_batch ? seqused_q[batch_idx] : 0;
+        } else if (cu_seqlens_q) {
+            int cur_cu_seqlen = batch_idx <= num_batch ? cu_seqlens_q[batch_idx] : 0;
+            int next_cu_seqlen = __shfl_down_sync(0xffffffff, cur_cu_seqlen, 1);
+            seqlen = next_cu_seqlen - cur_cu_seqlen;
+        } else {
+            seqlen = seqlen_q_static;
+        }
+        seqlen *= qhead_per_khead;
+        return batch_idx < num_batch && lane < kNumBatchPerWarp
+            ? blockm_divmod.div(seqlen + blockm_divmod.divisor - 1) : 0;
+    };
+
+    auto get_num_n_blocks = [&](int bidb_start) {
+        int batch_idx = lane + bidb_start;
+        int leftpad_k = batch_idx < num_batch && leftpad_k_ptr != nullptr ? leftpad_k_ptr[batch_idx] : 0;
+        int seqlen;
+        if (seqused_k) {
+            seqlen = batch_idx < num_batch ? seqused_k[batch_idx] : 0;
+        } else if (cu_seqlens_k) {
+            int cur_cu_seqlen = batch_idx <= num_batch ? cu_seqlens_k[batch_idx] : 0;
+            int next_cu_seqlen = __shfl_down_sync(0xffffffff, cur_cu_seqlen, 1);
+            seqlen = next_cu_seqlen - cur_cu_seqlen;
+        } else {
+            seqlen = seqlen_k_static;
+        }
+        int seqlen_new;
+        if (cu_seqlens_k_new) {
+            int cur_cu_seqlen_new = batch_idx <= num_batch ? cu_seqlens_k_new[batch_idx] : 0;
+            int next_cu_seqlen_new = __shfl_down_sync(0xffffffff, cur_cu_seqlen_new, 1);
+            seqlen_new = next_cu_seqlen_new - cur_cu_seqlen_new;
+        } else {
+            seqlen_new = seqlen_k_new_static;
+        }
+        // if (threadIdx.x == 0) { printf("seqlen = %d, seqlen_new = %d, leftpad_k = %d\n", seqlen, seqlen_new, leftpad_k); }
+        seqlen = seqlen - leftpad_k + seqlen_new;
+        return batch_idx < num_batch && lane < kNumBatchPerWarp
+            ? blockn_divmod.div(seqlen + blockn_divmod.divisor - 1) : 0;
+    };
+
+    int warp_idx = threadIdx.x / cutlass::NumThreadsPerWarp;
+    int bidb_start = kNumBatchPerWarp * warp_idx;
+    int num_m_blocks = get_num_m_blocks(bidb_start);
+    int num_n_blocks = get_num_n_blocks(bidb_start);
+
+    int total_blocks = num_m_blocks * num_n_blocks;
+    // Warp sum
+    #pragma unroll
+    for (int i = cutlass::NumThreadsPerWarp / 2; i >= 1; i /= 2) {
+        total_blocks += __shfl_down_sync(0xffffffff, total_blocks, i);
+    }
+    if (lane == 0) { atomicAdd(total_blocks_smem, total_blocks); }
+    __syncthreads();
+    total_blocks = total_blocks_smem[0];
+    // 10% margin
+    int blocks_per_sm = static_cast<int>(ceilf(float(total_blocks) * 1.1f * float(num_head) / float(num_sm)));
+    // blocks_per_sm = std::max(1, blocks_per_sm);  // 1 is the minimum number of blocks per SM
+    int num_splits_dynamic = std::max(std::min((num_n_blocks + blocks_per_sm - 1) / blocks_per_sm, num_splits_static), 1);
+    if (bidb_start + lane < num_batch && lane < kNumBatchPerWarp) {
+        num_splits_dynamic_ptr[bidb_start + lane] = num_splits_dynamic;
+        // 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);
+    }
+}
+
+} // flash
+
+void prepare_varlen_num_blocks(Flash_fwd_params &params, cudaStream_t stream, bool packgqa,
+                               int blockM, int blockN, bool enable_pdl) {
+    // Only support batch <= 992 (32 warps, each with 31 batches)
+    int qhead_per_khead = !packgqa ? 1 : cutlass::ceil_div(params.h, params.h_k);
+    flash::prepare_varlen_num_blocks_kernel<<<1 /*grid*/, 1024 /*block*/, 0, stream>>>(
+        params.seqlen_q, params.seqlen_k, params.seqlen_knew,
+        params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew,
+        params.seqused_q, params.seqused_k, params.leftpad_k,
+        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.num_splits_dynamic_ptr, enable_pdl);
+}

flash-attn/heuristics.h

@@ -0,0 +1,59 @@
+/******************************************************************************
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <vector>
+
+inline bool should_pack_gqa(bool varlen_q, int seqlen_q, int qhead_per_khead, int blockM) {
+    // If varlen, we don't actually know seqlen_q but only max_seqlen_q.
+    if (varlen_q) return true;
+    // Heuristic: PackGQA is a bit slower but can help if seqlen_q is small or not near a multiple of kBlockM
+    auto round_up = [](int a, int b) { return (a + b - 1) / b * b; };
+    float nopack_gqa_efficiency = float(seqlen_q) / float(round_up(seqlen_q, blockM));
+    float pack_gqa_efficiency = float(seqlen_q * qhead_per_khead) / float(round_up(seqlen_q * qhead_per_khead, blockM));
+    return nopack_gqa_efficiency < 0.9 * pack_gqa_efficiency;
+};
+
+// Find the number of splits that maximizes the occupancy. For example, if we have
+// batch * n_heads = 48 and we have 108 SMs, having 2 splits (efficiency = 0.89) is
+// better than having 3 splits (efficiency = 0.67). However, we also don't want too many
+// splits as that would incur more HBM reads/writes.
+// So we find the best efficiency, then find the smallest number of splits that gets 85%
+// of the best efficiency.
+inline int num_splits_heuristic(int total_mblocks, int num_SMs, int num_n_blocks, int num_m_blocks, int size_one_kv_head, bool is_causal_or_local, int max_splits) {
+    // If we have enough to almost fill the SMs, then just use 1 split
+    // However, in the case of super long seqlen where each head of KV doesn't even fit into
+    // L2 (we assume that L2 size is 50MB), we want to split.
+    if (total_mblocks >= 0.8f * num_SMs) {
+        int const size_l2 = 50 * 1024 * 1024;
+        // Only split if there are enough queries to go over the KV at least twice
+        // Don't split if causal
+        if (size_one_kv_head > size_l2 && num_m_blocks >= num_SMs * 2 && !is_causal_or_local) {
+            return std::min((size_one_kv_head + size_l2 - 1) / size_l2, max_splits);
+        } else {
+            return 1;
+        }
+    }
+    // If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512.
+    if (num_n_blocks <= 4) { return 1; }
+    max_splits = std::min({max_splits, num_SMs, num_n_blocks});
+    float max_efficiency = 0.f;
+    std::vector<float> efficiency;
+    efficiency.reserve(max_splits);
+    for (int num_splits = 1; num_splits <= max_splits; num_splits++) {
+        float n_waves = float(total_mblocks * num_splits) / num_SMs;
+        float eff = n_waves / ceil(n_waves);
+        // printf("num_splits = %d, eff = %f\n", num_splits, eff);
+        if (eff > max_efficiency) { max_efficiency = eff; }
+        efficiency.push_back(eff);
+    }
+    for (int num_splits = 1; num_splits <= max_splits; num_splits++) {
+        if (efficiency[num_splits - 1] >= 0.85 * max_efficiency) {
+            // printf("num_splits chosen = %d\n", num_splits);
+            return num_splits;
+        }
+    }
+    return 1;
+}

flash-attn/instantiations/flash_bwd_hdim128_bf16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<80, cutlass::bfloat16_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<86, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim128_bf16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<90, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim128_bf16_softcap_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<80, cutlass::bfloat16_t, true>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<86, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim128_bf16_softcap_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<90, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim128_bf16_softcapall_sm90.cu

@@ -0,0 +1,6 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_hdim128_bf16_sm90.cu"
+#include "flash_bwd_hdim128_bf16_softcap_sm90.cu"

flash-attn/instantiations/flash_bwd_hdim128_fp16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<80, cutlass::half_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<80, cutlass::half_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::half_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<86, cutlass::half_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim128_fp16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<90, cutlass::half_t, 128, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<90, cutlass::half_t, false>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim128_fp16_softcap_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<80, cutlass::half_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<80, cutlass::half_t, true>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::half_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<86, cutlass::half_t, true>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim128_fp16_softcap_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM128
+template<>
+void run_mha_bwd_<90, cutlass::half_t, 128, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim128<90, cutlass::half_t, true>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim128_fp16_softcapall_sm90.cu

@@ -0,0 +1,6 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_hdim128_fp16_sm90.cu"
+#include "flash_bwd_hdim128_fp16_softcap_sm90.cu"

flash-attn/instantiations/flash_bwd_hdim192_bf16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<80, cutlass::bfloat16_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<86, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim192_bf16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<90, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim192_bf16_softcap_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<80, cutlass::bfloat16_t, true>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<86, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim192_bf16_softcap_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<90, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim192_bf16_softcapall_sm90.cu

@@ -0,0 +1,6 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_hdim192_bf16_sm90.cu"
+#include "flash_bwd_hdim192_bf16_softcap_sm90.cu"

flash-attn/instantiations/flash_bwd_hdim192_fp16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<80, cutlass::half_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<80, cutlass::half_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::half_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<86, cutlass::half_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim192_fp16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<90, cutlass::half_t, 192, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<90, cutlass::half_t, false>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim192_fp16_softcap_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<80, cutlass::half_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<80, cutlass::half_t, true>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::half_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<86, cutlass::half_t, true>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim192_fp16_softcap_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM192
+template<>
+void run_mha_bwd_<90, cutlass::half_t, 192, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim192<90, cutlass::half_t, true>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim192_fp16_softcapall_sm90.cu

@@ -0,0 +1,6 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_hdim192_fp16_sm90.cu"
+#include "flash_bwd_hdim192_fp16_softcap_sm90.cu"

flash-attn/instantiations/flash_bwd_hdim256_bf16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<80, cutlass::bfloat16_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<86, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim256_bf16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<90, cutlass::bfloat16_t, false>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim256_bf16_softcap_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<80, cutlass::bfloat16_t, 256, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<80, cutlass::bfloat16_t, true>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::bfloat16_t, 256, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<86, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim256_bf16_softcap_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<90, cutlass::bfloat16_t, 256, true>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<90, cutlass::bfloat16_t, true>(params, stream);
+}
+#endif

flash-attn/instantiations/flash_bwd_hdim256_bf16_softcapall_sm90.cu

@@ -0,0 +1,6 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_hdim256_bf16_sm90.cu"
+#include "flash_bwd_hdim256_bf16_softcap_sm90.cu"

flash-attn/instantiations/flash_bwd_hdim256_fp16_sm80.cu

@@ -0,0 +1,18 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_SM8x
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<80, cutlass::half_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<80, cutlass::half_t, false>(params, stream);
+}
+template<>
+void run_mha_bwd_<86, cutlass::half_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<86, cutlass::half_t, false>(params, stream);
+}
+#endif
+#endif

flash-attn/instantiations/flash_bwd_hdim256_fp16_sm90.cu

@@ -0,0 +1,12 @@
+// Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
+// Splitting the different template instantiations to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+
+#include "flash_bwd_launch_template.h"
+
+#ifndef FLASHATTENTION_DISABLE_HDIM256
+template<>
+void run_mha_bwd_<90, cutlass::half_t, 256, false>(Flash_bwd_params &params, cudaStream_t stream) {
+    run_mha_bwd_hdim256<90, cutlass::half_t, false>(params, stream);
+}
+#endif