Remove sources

README.md

@@ -1,13 +1,14 @@
 ---
 license: apache-2.0
 tags:
-- kernel
+  - kernel
 ---
-  
+
 # vllm-flash-attn3
 
 This is an implementation of Flash Attention 3 CUDA kernels with support for attention sinks. The attention sinks implementation was contributed to Flash Attention by the [vLLM team](https://huggingface.co/vllm-project). The [transformers team](https://huggingface.co/transformers-community) packaged the implementation and pre-built it for use with the [kernels library](https://github.com/huggingface/kernels).
 
+Kernel source: https://github.com/huggingface/kernels-community/tree/main/vllm-flash-attn3
 
 ## Quickstart
 
@@ -43,7 +44,7 @@ torch.cuda.manual_seed(42)
 # Parameters
 batch_size = 2
 seqlen_q = 128  # Query sequence length
-seqlen_k = 256  # Key sequence length  
+seqlen_k = 256  # Key sequence length
 nheads = 8      # Number of attention heads
 d = 64          # Head dimension
 
@@ -65,7 +66,6 @@ print(f"\nAttention computation successful!")
 print(f"Output tensor stats - Mean: {output.mean().item():.4f}, Std: {output.std().item():.4f}")
 ```
 
-
 ## How to Use
 
 When loading your model with transformers, provide this repository id as the source of the attention implementation:
@@ -91,4 +91,5 @@ This will automatically resolve and download the appropriate code for your archi
 
 - [Tri Dao](https://huggingface.co/tridao) and team for Flash Attention and [Flash Attention 3](https://tridao.me/blog/2024/flash3/).
 - The [vLLM team](https://huggingface.co/vllm-project) for their implementation and their contribution of attention sinks.
-- The [transformers team](https://huggingface.co/transformers-community) for packaging, testing, building and making it available for use with the [kernels library](https://github.com/huggingface/kernels).
+- The [transformers team](https://huggingface.co/transformers-community) for packaging, testing, building and making it available for use with the [kernels library](https://github.com/huggingface/kernels).
+

build.toml

@@ -1,593 +0,0 @@
-[general]
-name = "vllm_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",
-]
-cxx-flags = ["-DFLASHATTENTION_DISABLE_PYBIND"]
-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

@@ -1,168 +0,0 @@
-{
-  "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": 1751014803,
-        "narHash": "sha256-9Xfq2k3uPfB602NwQF+zAY2GQZiKUN1G7Q6XiDCUR8Y=",
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "rev": "bbc4e712ff2046e217818e97de2201e2b996756e",
-        "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

@@ -1,51 +0,0 @@
-{
-  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;
-      # Building with CDUA later than 12.4 fails with:
-      #
-      # error: 'ptxas' died due to signal 11 (Invalid memory reference)
-      #
-      # So, build for 12.4 only and copy to all the other build variants
-      # by hand (which works fine thanks to backward compat).
-      #
-      # Still need to check if upstream FA3 has the same issue.
-      torchVersions = [
-        {
-          torchVersion = "2.6";
-          cudaVersion = "12.4";
-          cxx11Abi = false;
-          systems = [ "x86_64-linux" ];
-          upstreamVariant = true;
-        }
-        {
-          torchVersion = "2.6";
-          cudaVersion = "12.4";
-          cxx11Abi = true;
-          systems = [ "x86_64-linux" ];
-          upstreamVariant = true;
-        }
-        {
-          torchVersion = "2.7";
-          cudaVersion = "12.4";
-          cxx11Abi = true;
-          systems = [
-            "x86_64-linux"
-            "aarch64-linux"
-          ];
-          upstreamVariant = true;
-        }
-      ];
-    };
-}

flash-attn/block.h

@@ -1,94 +0,0 @@
-/******************************************************************************
- * 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& 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 ; }
-            n_block_max = std::min(n_block_max,
-                                   cute::ceil_div(m_idx_max + seqlen_k - seqlen_q + window_size_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); }
-            n_block_min = std::max(int(0), (m_idx_min + seqlen_k - seqlen_q - window_size_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& 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, 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) {
-
-        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};
-    }
-
-};
-
-} // namespace flash

flash-attn/copy_sm90_bulk_reduce.hpp

@@ -1,49 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,19 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,523 +0,0 @@
-/******************************************************************************
- * 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;
-        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;
-        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_dK, 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.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_dK);
-            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_dK, 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 tdKVpdKV = make_tensor<bool>(make_shape(size<2>(tdKVgdV)));
-            #pragma unroll
-            for (int k = 0; k < size(tdKVpdKV); ++k) { tdKVpdKV(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, tdKVpdKV, kBlockN);
-            flash::copy</*Is_even_MN=*/EvenN, /*Is_even_K=*/true, /*Clear_OOB_MN=*/false>(
-                gmem_tiled_copy_dKV, tdKVsdK, tdKVrdK, tdKVcdKV, tdKVpdKV, 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, tdKVpdKV, 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, tdKVpdKV, 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_dK, 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 tdKVpdKV = make_tensor<bool>(make_shape(size<2>(tdKVgdK)));
-        #pragma unroll
-        for (int k = 0; k < size(tdKVpdKV); ++k) { tdKVpdKV(k) = get<1>(tdKVcdKV(_0{}, _0{}, k)) < get<1>(params.shape_dK); }
-        // 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, tdKVpdKV, 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, tdKVpdKV, 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;
-        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;
-        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.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_dKaccum, 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

@@ -1,484 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,220 +0,0 @@
-/******************************************************************************
- * 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;
-
-    // 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;
-
-    // The S extra matrix, (num_heads)
-    void *__restrict__ s_aux_ptr;
-};
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-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

@@ -1,1623 +0,0 @@
-/******************************************************************************
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
- ******************************************************************************/
-
-// Include these 2 headers instead of torch/extension.h since we don't need all of the torch headers.
-#include <torch/nn/functional.h>
-#include <torch/version.h>  // For TORCH_VERSION* macros
-#include <ATen/cuda/CUDAContext.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"
-
-// Copied from https://github.com/pytorch/pytorch/commit/7931eee5c5ebcdf468bff4d308510b03355cd909
-// This is so that we can pass in torch.dtype as a parameter to the function.
-#if TORCH_VERSION_MAJOR < 2 || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 4)
-
-#include <pybind11/pybind11.h>
-#include <pybind11/stl.h>
-
-namespace pybind11::detail {
-
-    template <>
-    struct type_caster<at::ScalarType> {
-    public:
-        // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes)
-        PYBIND11_TYPE_CASTER(at::ScalarType, _("torch.dtype"));
-        // PYBIND11_TYPE_CASTER defines a member field called value. at::ScalarType
-        // cannot be default-initialized, we provide this constructor to explicitly
-        // initialize that field. The value doesn't matter as it will be overwritten
-        // after a successful call to load.
-        type_caster() : value(at::kFloat) {}
-        bool load(handle src, bool) {
-            PyObject* obj = src.ptr();
-            if (THPDtype_Check(obj)) {
-                value = reinterpret_cast<THPDtype*>(obj)->scalar_type;
-                return true;
-            }
-            return false;
-        }
-        static handle cast(
-                           const at::ScalarType& src,
-                           return_value_policy /* policy */,
-                           handle /* parent */) {
-            return Py_NewRef(torch::getTHPDtype(src));
-        }
-    };
-
-} // namespace pybind11::detail
-
-#endif
-
-#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,
-                      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;
-    params.is_local = (window_size_left >= 0 || window_size_right >= 0) && !params.is_causal;
-
-    // TODO: check this
-    if (window_size_left < 0 && window_size_right >= 0) { window_size_left = seqlen_k - 1; }
-    if (window_size_left >= 0 && window_size_right < 0) { window_size_right = seqlen_q - 1; }
-    params.window_size_left = window_size_left;
-    params.window_size_right = window_size_right;
-
-    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,
-                      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,
-                     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;
-}
-
-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, [&] {
-                        if (!params.is_e4m3) {
-                            if (params.is_bf16) {
-                                #ifndef FLASHATTENTION_DISABLE_HDIM64
-                                if (params.d <= 64) {
-                                    if (params.dv > 256 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 64, 512, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else if (params.dv > 64 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 64, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else {
-                                        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 (params.dv <= 128 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::bfloat16_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else {
-                                        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 (params.dv > 256 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::half_t, 64, 512, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else if (params.dv > 64 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::half_t, 64, 256, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else {
-                                        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 (params.dv <= 128 && Arch == 90) {
-                                        return run_mha_fwd_<Arch, cutlass::half_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                    } else {
-                                        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 (params.dv <= 128 && Arch == 90) {
-                                    return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream);
-                                } else {
-                                    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_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, use_one_mma_wg(params));
-    // 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(
-        int batch_size,
-        int max_seqlen_q,
-        int max_seqlen_k,
-        int num_heads,
-        int num_heads_k,
-        int headdim,
-        int headdim_v,
-        at::ScalarType qkv_dtype,
-        const at::Tensor &seqused_k, // b
-        std::optional<const at::Tensor> &cu_seqlens_q_,  // b+1
-        std::optional<const at::Tensor> &cu_seqlens_k_,  // b+1
-        std::optional<const at::Tensor> &cu_seqlens_k_new_,  // b+1
-        std::optional<const at::Tensor> &seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
-        std::optional<const at::Tensor> &leftpad_k_, // b
-        std::optional<int> page_size,
-        int max_seqlen_k_new,  // 0 means we're not appending new KV
-        bool is_causal,
-        int window_size_left,
-        int window_size_right,
-        bool has_softcap,
-        int num_splits,
-        std::optional<bool> pack_gqa_,
-        int const 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) {
-        // 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;
-    params.is_local = (window_size_left >= 0 || window_size_right >= 0) && !params.is_causal;
-    if (window_size_left < 0 && window_size_right >= 0) { window_size_left = max_seqlen_k - 1; }
-    if (window_size_left >= 0 && window_size_right < 0) { window_size_right = max_seqlen_q - 1; }
-    params.window_size_left = window_size_left;
-    params.window_size_right = window_size_right;
-    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);
-    // Determine if we should pack GQA before num_splits since it impacts use_one_mma_wg (in get_num_splits)
-    params.pack_gqa = pack_gqa_.has_value() ? pack_gqa_.value() : get_pack_gqa(params);
-    params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits;
-    // Always enable PackGQA for Split
-    params.pack_gqa = params.num_splits > 1;
-
-    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, use_one_mma_wg(params));
-        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::vector<at::Tensor>
-mha_fwd(at::Tensor &q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
-        const 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.
-        const 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<const 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<const 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<const 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<const at::Tensor> &cu_seqlens_q_,  // b+1
-        std::optional<const at::Tensor> &cu_seqlens_k_,  // b+1
-        std::optional<const at::Tensor> &cu_seqlens_k_new_,  // b+1
-        std::optional<const at::Tensor> &seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
-        std::optional<const at::Tensor> &seqused_k_, // b. If given, only this many elements of each batch element's keys are used.
-        std::optional<int> max_seqlen_q_,
-        // TODO: check if we need max_seqlen_k
-        std::optional<int> max_seqlen_k_,
-        std::optional<const at::Tensor> &page_table_, // (b_k, max_num_pages_per_seq)
-        std::optional<const at::Tensor> &kv_batch_idx_, // b. indices to index into the KV cache
-        std::optional<const at::Tensor> &leftpad_k_, // b
-        std::optional<const at::Tensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
-        std::optional<const at::Tensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
-        std::optional<const 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)
-        float const softmax_scale,
-        bool is_causal,
-        int window_size_left,
-        int window_size_right,
-        float const softcap,
-        bool const is_rotary_interleaved,   // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
-        std::optional<at::Tensor> &scheduler_metadata_,  // (b + 1)
-        int num_splits,
-        std::optional<bool> pack_gqa_,
-        int const sm_margin,
-        std::optional<const at::Tensor> &s_aux_ // (h)
-        ) {
-
-    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 = !max_seqlen_k_.has_value() ? (!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);
-    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) {
-        // 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; }
-    // There's a case where is_causal=false, window_size=(-1, 0). Then set_params_fprop will set params.is_causal=true.
-    // If we don't have is_causal here matching params.is_causal, we might get the wrong kBlockM.
-    is_causal = window_size_left < 0 && 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,
-                     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);
-    // Determine if we should pack GQA before num_splits since it impacts use_one_mma_wg (in get_num_splits)
-    params.pack_gqa = pack_gqa_.has_value() ? pack_gqa_.value() : get_pack_gqa(params);
-    params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits;
-    // Always enable PackGQA for Split
-    params.pack_gqa = params.num_splits > 1;
-
-    // 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;
-        }
-    }
-
-    if(s_aux_.has_value()) {
-        auto s_aux = s_aux_.value();
-        TORCH_CHECK(s_aux.scalar_type() == at::ScalarType::BFloat16,
-            "We only support bf16 dtype for S extra.");
-        CHECK_DEVICE(s_aux);
-        CHECK_SHAPE(s_aux, num_heads);
-        CHECK_CONTIGUOUS(s_aux);
-        params.s_aux_ptr = s_aux.data_ptr();
-    } else {
-        params.s_aux_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};
-}
-
-void run_mha_bwd(Flash_bwd_params &params, cudaStream_t stream) {
-    #ifndef FLASHATTENTION_DISABLE_BACKWARD
-        // 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, [&] {
-            if (!params.is_bf16) {
-                #ifndef FLASHATTENTION_DISABLE_FP16
-                #ifndef FLASHATTENTION_DISABLE_HDIM64
-                if (params.d <= 64) { return run_mha_bwd_<Arch, cutlass::half_t, 64, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM96
-                if (params.d <= 96) { return run_mha_bwd_<Arch, cutlass::half_t, 96, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM128
-                if (params.d <= 128) { return run_mha_bwd_<Arch, cutlass::half_t, 128, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM192
-                if (params.d <= 192) { return run_mha_bwd_<Arch, cutlass::half_t, 192, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM256
-                if (params.d <= 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 <= 64) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 64, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM96
-                if (params.d <= 96) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 96, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM128
-                if (params.d <= 128) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 128, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM192
-                if (params.d <= 192) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 192, Has_softcap>(params, stream); }
-                #endif
-                #ifndef FLASHATTENTION_DISABLE_HDIM256
-                if (params.d <= 256) { return run_mha_bwd_<Arch, cutlass::bfloat16_t, 256, Has_softcap>(params, stream); }
-                #endif
-            }
-        });
-    });
-    #endif
-}
-
-
-// b: batch_size
-// s_q: seqlen_q
-// s_k: seqlen_k
-// h: num_heads
-// h_k: num_heads_k
-// d: head_size
-std::vector<at::Tensor> mha_bwd(
-    const at::Tensor &dout,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
-    const at::Tensor &q,     // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
-    const at::Tensor &k,     // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
-    const at::Tensor &v,     // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
-    const at::Tensor &out,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
-    const 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, d) or (total_k, h_k, d) if there is cu_seqlens_k
-    std::optional<const at::Tensor> &cu_seqlens_q_,   // b+1
-    std::optional<const at::Tensor> &cu_seqlens_k_,   // b+1
-    std::optional<const at::Tensor> &seqused_q_, // b. If given, only this many elements of each batch element's queries and outputs are used.
-    std::optional<const at::Tensor> &seqused_k_, // b. If given, only this many elements of each batch element's keys are used.
-    std::optional<int> max_seqlen_q_,
-    std::optional<int> max_seqlen_k_,
-    float const softmax_scale,
-    bool is_causal,
-    int window_size_left,
-    int window_size_right,
-    float const softcap,
-    bool const deterministic,
-    int const 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 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");
-    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");
-
-    // 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(head_size);
-    // 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);
-        CHECK_SHAPE(dout, batch_size, seqlen_q, num_heads, head_size);
-    } else {
-        CHECK_SHAPE(q, total_q, num_heads, head_size);
-        CHECK_SHAPE(out, total_q, num_heads, head_size);
-        CHECK_SHAPE(dout, total_q, num_heads, head_size);
-        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);
-    } else {
-        CHECK_SHAPE(k, total_k, num_heads_k, head_size);
-        CHECK_SHAPE(v, total_k, num_heads_k, head_size);
-        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);
-        } else {
-            CHECK_SHAPE(dv, total_k, num_heads_k, head_size);
-        }
-    } 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_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_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,
-                     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;  // We don't support hdim_v being different from hdim_qk for now
-
-    // 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::vector<at::Tensor>
-mha_combine(const at::Tensor &out_partial,         // num_splits x batch_size x seqlen x num_heads x head_size
-            const 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};
-}
-
-#ifndef FLASHATTENTION_DISABLE_PYBIND
-
-#include <torch/python.h>
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "FlashAttention";
-    m.def("fwd", &mha_fwd, "Forward pass");
-    m.def("bwd", &mha_bwd, "Backward pass");
-    m.def("fwd_combine", &mha_combine, "Combine partial attention outputs");
-    m.def("get_scheduler_metadata", &mha_fwd_get_scheduler_metadata, "Get scheduler metadata for varlen forward pass");
-}
-
-#endif

flash-attn/flash_bwd_kernel_sm80.h

@@ -1,173 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,282 +0,0 @@
-
-/******************************************************************************
- * 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

@@ -1,377 +0,0 @@
-/******************************************************************************
- * 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.d, 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 = 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),
-        {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),
-        {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,
-        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::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.d_rounded * seqlen_k_rounded, !is_varlen_k ? params.h_k * params.d_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.d_rounded, params.h_k, batch_k},  // shape_dVaccum
-            {_1{}, seqlen_k_rounded * params.d_rounded, !is_varlen_k ? params.d_rounded * params.seqlen_k_rounded * params.h_k : 0},  // stride_dVaccum
-            static_cast<Element*>(params.dv_ptr),
-            {seqlen_k, params.d, 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

@@ -1,256 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,252 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,13 +0,0 @@
-// 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

@@ -1,702 +0,0 @@
-/******************************************************************************
- * 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 BlockCoord {
-        int block_m;
-        int block_k;
-        int bidb;
-    };
-
-    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;
-        BlockCoord block_coord;
-    };
-
-    static constexpr int SharedStorageSize = sizeof(SharedStorage);
-
-    // Device side arguments
-    struct Arguments {
-        int b;
-        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 CollectiveParams {
-        int b;
-        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
-    CollectiveParams
-    to_underlying_arguments(Arguments const& args) {
-        assert(get<1>(args.shape_LSE_partial) <= kMaxSplits);
-        return {
-            args.b,
-            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
-        };
-    }
-
-    struct SchedulerArguments {
-        int b;
-        int seqlen_q;
-        int total_q;
-        int num_heads;
-        int dv;
-        int const* cu_seqlens_q;
-        int const* seqused_q;
-    };
-
-    struct StaticTileScheduler {
-        struct Params {};
-        static Params to_underlying_arguments(SchedulerArguments const& args) { return {}; }
-
-        SharedStorage& shared_storage;
-        CUTE_DEVICE StaticTileScheduler(SharedStorage& shared_storage): shared_storage(shared_storage) {}
-
-        static dim3 get_grid_shape(SchedulerArguments const& args) {
-            unsigned int num_blocks_k = cute::ceil_div(args.dv, kBlockK);
-            unsigned int num_blocks_m = cute::ceil_div(args.seqlen_q * args.num_heads, kBlockM);
-            return {num_blocks_m, num_blocks_k, static_cast<unsigned int>(args.b)};
-        }
-
-        CUTE_DEVICE BlockCoord get_block_coord(Params const& params) {
-            int block_m = blockIdx.x;
-            int block_k = blockIdx.y;
-            int bidb = blockIdx.z;
-            return {block_m, block_k, bidb};
-        }
-    };
-
-    struct StaticVarlenTileScheduler {
-        //
-        // For varlen we have two Scheduling algos:
-        //  1) STANDARD, same as StaticTileScheduler
-        //  2) LINEARIZE_M_AND_BATCH, this flattens the tiled M dimension and
-        //     batch dimension into a linear tile index. The grid is then a
-        //     2D grid of (tile_id, k_block). We then map the linear tile id
-        //     to (m_block, bidb) in the get_block_coord function. This mapping
-        //     is non-trivial since each batch element can have a different
-        //     number of m_blocks. This has overhead when computing the block
-        //     coordinates, but it is more efficient when prefills and decodes
-        //     are mixed since in that case the STANDARD scheduling algo will
-        //     have a lot of empty (no work) blocks in the grid.
-        //
-
-        enum SchedulingAlgo {
-            STANDARD,           // Same as StaticTileScheduler
-            LINEARIZE_M_AND_BATCH,  // Linearize the M and batch dimensions into a single tile index
-        };
-
-        struct Params {
-            int b;
-            int num_heads;
-            int const* const cu_seqlens_q;
-            int const* const seqused_q;
-            SchedulingAlgo algo;
-        };
-
-        SharedStorage& shared_storage;
-        CUTE_DEVICE StaticVarlenTileScheduler(SharedStorage& shared_storage): shared_storage(shared_storage) {}
-
-        static SchedulingAlgo choose_scheduling_algo(SchedulerArguments const& args) {
-            // Choose the scheduling algorithm based on how dense the grid of tiles that
-            // do actual work is. If the grid is more then 50% sparse, we linearize the M
-            // and batch. If the grid is more than 50% dense, we use the standard scheduling
-            // algorithm since its more efficient at calculating the block coordinates.
-            // NOTE: in varlen case args.seqlen_q is the max seqlen_q across all batches
-            // use lower bound to estimate when the density is more than 50%
-            int lower_bound_on_non_empty_tiles = cute::ceil_div(args.total_q, kBlockM);
-            int grid_size = args.b * cute::ceil_div(args.seqlen_q, kBlockM);
-            return 2 * lower_bound_on_non_empty_tiles >= grid_size ? 
-                SchedulingAlgo::STANDARD : 
-                SchedulingAlgo::LINEARIZE_M_AND_BATCH;
-        }
-
-        static Params to_underlying_arguments(SchedulerArguments const& args) { 
-            return {
-                args.b,
-                args.num_heads,
-                args.cu_seqlens_q,
-                args.seqused_q,
-                choose_scheduling_algo(args)
-            }; 
-        }
-
-        static dim3 get_grid_shape(SchedulerArguments const& args) {
-            unsigned int num_blocks_k = cute::ceil_div(args.dv, kBlockK);
-
-            switch (choose_scheduling_algo(args)) {
-            case SchedulingAlgo::STANDARD: {
-                unsigned int num_blocks_k = cute::ceil_div(args.dv, kBlockK);
-                unsigned int num_blocks_m = cute::ceil_div(args.seqlen_q * args.num_heads, kBlockM);
-                return {num_blocks_m, num_blocks_k, static_cast<unsigned int>(args.b)};
-            }
-            case SchedulingAlgo::LINEARIZE_M_AND_BATCH: {
-                // rough worst case upper bound on the number of blocks required 
-                //  (assuming each batch has an additional partial block)
-                unsigned int num_blocks_m = cute::ceil_div(args.total_q * args.num_heads, kBlockM) + args.b;
-                return {num_blocks_m, num_blocks_k, 1};
-            }}
-
-            // rough worst case upper bound on the number of blocks required 
-            //  (assuming each batch has an additional partial block)
-            unsigned int num_blocks_m = cute::ceil_div(args.total_q * args.num_heads, kBlockM) + args.b;
-            return {num_blocks_m, num_blocks_k, 1};
-        }
-
-        CUTE_DEVICE BlockCoord get_block_coord_linearized_m_and_batch(Params const& params) {
-            int num_heads = params.num_heads;
-            int curr_tile_id = blockIdx.x;
-
-            // Scan through the batches find the batch that contains the current
-            // tile_id. Compute using only the first warp of the block.
-            if (threadIdx.x < 32) {
-                // We compute linearized tile index start and ends for each batch
-                // in groups of 32 in parallel
-                int group_start_bidb = -(cutlass::NumThreadsPerWarp);
-                int group_end_bidb = 0;
-                int group_end_tile_id = 0;
-                int group_start_tile_id = 0;
-                int group_total_num_tiles = 0;
-
-                int local_num_m_blocks = 0;
-                int local_num_m_blocks_cumulative = 0;
-
-                do {
-                    group_start_bidb += cutlass::NumThreadsPerWarp;
-                    group_end_bidb += cutlass::NumThreadsPerWarp;
-
-                    auto get_num_m_blocks = [&](int bidb) {
-                        if (bidb >= params.b) return 0;
-                        flash::SeqlenInfo<Varlen, kBlockM> seqlen_info{bidb, 0, params.cu_seqlens_q, params.seqused_q};
-                        return cute::ceil_div(seqlen_info.seqlen * num_heads, Int<kBlockM>{}());
-                    };
-
-                    // Cumulative number of blocks for the next 31 batches
-                    local_num_m_blocks = get_num_m_blocks(group_start_bidb + threadIdx.x);
-                    local_num_m_blocks_cumulative = warp_prefix_sum(local_num_m_blocks);
-                    // Total number of blocks for the next 32 batches
-                    group_total_num_tiles = warp_shfl_get_last(local_num_m_blocks_cumulative);
-                    
-                    group_start_tile_id = group_end_tile_id;
-                    group_end_tile_id += group_total_num_tiles;
-                } while (curr_tile_id >= group_end_tile_id && group_end_bidb < params.b);
-                
-                int local_batch_end_tile_id = group_start_tile_id + local_num_m_blocks_cumulative;
-                // Find the last batch idx in the group where `local_batch_end_tile_id <= curr_tile_id`
-                // these values below are now common to all threads in the warp
-                int batch_idx_in_group = warp_last_true_laneid(local_batch_end_tile_id <= curr_tile_id);
-                int batch_num_m_blocks = warp_shfl_get(local_num_m_blocks, batch_idx_in_group);
-                int batch_m_start_tile_id = group_start_tile_id + (batch_idx_in_group > 0 ? 
-                    warp_shfl_get(local_num_m_blocks_cumulative, batch_idx_in_group - 1) : 0);
-                
-                int bidb = group_start_bidb + batch_idx_in_group;
-                int block_m = curr_tile_id - batch_m_start_tile_id;
-                // NOTE(lucas): not sure why this causes a block_k unused warning
-                //  just inlined `blockIdx.y` to suppress the warning
-                // int block_k = blockIdx.y;
-                // shared_storage.block_coord = {block_m, block_k, bidb};
-                BlockCoord block_coord{block_m, static_cast<int>(blockIdx.y), bidb};
-                if (threadIdx.x == 0) { shared_storage.block_coord = block_coord; }
-            }
-
-            __syncthreads();
-            return shared_storage.block_coord;
-        }
-
-
-        CUTE_DEVICE BlockCoord get_block_coord_standard(Params const& params) {
-            int block_m = blockIdx.x;
-            int block_k = blockIdx.y;
-            int bidb = blockIdx.z;
-            return {block_m, block_k, bidb};
-        }
-
-        CUTE_DEVICE BlockCoord get_block_coord(Params const& params) {
-            switch (params.algo) {
-                case SchedulingAlgo::STANDARD:
-                    return get_block_coord_standard(params);
-                case SchedulingAlgo::LINEARIZE_M_AND_BATCH:
-                    return get_block_coord_linearized_m_and_batch(params);
-            }
-            return {0, 0, 0};  // Should never reach here
-        }
-    };
-
-    using TileScheduler = std::conditional_t<
-        Varlen,
-        StaticVarlenTileScheduler,
-        StaticTileScheduler
-    >;
-
-    using SchedulerParams = typename TileScheduler::Params;
-
-    struct Params {
-        CollectiveParams params;
-        SchedulerParams scheduler_params;
-    };
-
-    CUTLASS_DEVICE
-    void
-    operator()(Params const& kernel_params, char* smem_buf) {
-        CollectiveParams const& params = kernel_params.params;
-
-        SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
-        TileScheduler tile_scheduler{shared_storage};
-
-        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;
-
-        BlockCoord block_coord = tile_scheduler.get_block_coord(kernel_params.scheduler_params);
-
-        int const m_block = block_coord.block_m;
-        int const k_block = block_coord.block_k;
-        int const batch = block_coord.bidb;
-
-        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;
-        }
-
-        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);
-
-        bool block_coord_valid = 
-            block_coord.block_m < cute::ceil_div(max_idx, Int<kBlockM>{}) &&
-            block_coord.bidb < params.b;
-        if (!block_coord_valid) { return; }
-
-        int const num_splits = params.num_splits_dynamic_ptr ? params.num_splits_dynamic_ptr[batch] : get<1>(params.shape_LSE_partial);
-        if (num_splits <= 1) { 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

@@ -1,88 +0,0 @@
-/******************************************************************************
- * 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 {
-        params.b,
-        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::SchedulerArguments scheduler_args  {
-        params.b, params.seqlen_q, params.total_q, params.h, params.dv,
-        params.cu_seqlens_q, params.seqused_q
-    };
-
-    typename CombineKernel::Params kernel_params = {
-        CombineKernel::to_underlying_arguments(args),
-        CombineKernel::TileScheduler::to_underlying_arguments(scheduler_args)
-    };
-
-    dim3 grid_m = CombineKernel::TileScheduler::get_grid_shape(scheduler_args);
-    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

@@ -1,215 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,468 +0,0 @@
-/******************************************************************************
- * 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;
-
-    using SmemLayoutSAux = typename CollectiveMainloop::SmemLayoutSAux;
-
-    // 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;
-
-        const int num_heads = get<2>(params.mainloop.shape_Q);
-        Tensor gS_aux = make_tensor(make_gmem_ptr(params.mainloop.ptr_S_aux), make_shape(num_heads));
-        Tensor sS_aux = make_tensor(make_smem_ptr(shared_storage.tensors.mainloop.smem_s_aux.data()), SmemLayoutSAux{});
-
-        if(params.mainloop.ptr_S_aux && threadIdx.x < num_heads) {
-            sS_aux(threadIdx.x) = gS_aux(threadIdx.x);
-        }
-
-        // 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

@@ -1,231 +0,0 @@
-/******************************************************************************
- * 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"
-#include "heuristics.h"
-
-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, bool Use_one_mma_wg>
-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>;
-    using ElementS = cutlass::bfloat16_t;
-
-    // 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, Use_one_mma_wg);
-    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, ElementS>,
-        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, ElementS>
-    >;
-    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.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,
-        static_cast<ElementS const*>(params.s_aux_ptr)
-    };
-    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, [&] {
-                BOOL_SWITCH(use_one_mma_wg(params), Use_one_mma_wg_, [&] {
-                    // Avoid over compiliation by making sure this only get set if it is actually used, i.e. we currently only support one mma wg for 128 head dim and hopper
-                    static constexpr bool Use_one_mma_wg = Use_one_mma_wg_ && Arch >= 90 && kHeadDim == 128;
-
-                    // 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, Use_one_mma_wg)) : 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, Use_one_mma_wg>(params, stream);
-                            });
-                        });
-                    });
-                });
-            });
-        });
-    });
-}

flash-attn/flash_prepare_scheduler.cu

@@ -1,124 +0,0 @@
-/******************************************************************************
- * 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

@@ -1,65 +0,0 @@
-/******************************************************************************
- * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
- ******************************************************************************/
-
-#pragma once
-
-#include <vector>
-#include "flash.h"
-
-inline bool use_one_mma_wg(Flash_fwd_params const& params) {
-    return params.arch >= 90 && params.d == 128 && 
-        params.seqlen_q * (!params.pack_gqa ? 1 : params.h / params.h_k) <= 64;
-};
-
-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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,6 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,6 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,6 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,6 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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

@@ -1,12 +0,0 @@
-// 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

@@ -1,6 +0,0 @@
-// 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

@@ -1,18 +0,0 @@
-// 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