Parallelize splitkv alignment templated kernels, remove flag (#2683)

Author: janeyx99

csrc/flash_attn/src/flash_fwd_launch_template.h

@@ -160,25 +160,34 @@ void run_flash_splitkv_fwd(Flash_fwd_params &params, cudaStream_t stream) {
     }
 }
 
+// The num_splits==1 blocksize-aligned splitkv template. If a user specifies
+// num_splits=1, we assume they want bitwise identical numerics across the split
+// KV and standard kernels so we align kBlockN to match.
+// We technically can combine this into one dispatch under run_flash_splitkv_fwd
+// but that pathologically slowed down build time by doubling the number of kernels
+// in a single file, which made build go from minutes to hours. Thus, it is pulled
+// into its own function so it can be explicitly instantiated in a separate file
+// (flash_fwd_split_align_*.cu) and compiled in parallel instead of serializing in
+// one ptxas invocation.
+template<typename T, int Headdim, bool Is_causal>
+void run_mha_fwd_splitkv_align(Flash_fwd_params &params, cudaStream_t stream) {
+    constexpr static int kBlockM = 64;
+    constexpr static int kBlockN_standard = Headdim <= 64 ? 128 : 64;
+    run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN_standard, 4, false, false, T>, Is_causal>(params, stream);
+}
+
 template<typename T, int Headdim, bool Is_causal>
 void run_mha_fwd_splitkv_dispatch(Flash_fwd_params &params, cudaStream_t stream) {
     constexpr static int kBlockM = 64;
     // TD [2023-08-28]: nvcc segfaults for headdim 96 with block size 64 x 256,
     // and for headdim 192 with block size 64 x 128.
     constexpr static int kBlockN = Headdim <= 64 ? 256 : (Headdim <= 128 ? 128 : 64);
-#ifndef FLASHATTENTION_DISABLE_SPLIT_ALIGNMENT
-    // If a user specifies num_splits=1, we assume they want bitwise identical
-    // numerics across the split KV and standard kernels so we align kBLockN to
-    // match.
-    // This compiles a second splitkv kernel tree (a different kBlockN), which
-    // could push build time to hours+. Define FLASHATTENTION_DISABLE_SPLIT_ALIGNMENT
-    // to skip.
     if (params.num_splits == 1) {
-        constexpr static int kBlockN_standard = Headdim <= 64 ? 128 : 64;
-        run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN_standard, 4, false, false, T>, Is_causal>(params, stream);
+        // Defined in flash_fwd_split_align_*.cu; declared extern in the main
+        // flash_fwd_split_*.cu so this call does not re-instantiate the tree here.
+        run_mha_fwd_splitkv_align<T, Headdim, Is_causal>(params, stream);
         return;
     }
-#endif
     run_flash_splitkv_fwd<Flash_fwd_kernel_traits<Headdim, kBlockM, kBlockN, 4, false, false, T>, Is_causal>(params, stream);
 }
 

csrc/flash_attn/src/flash_fwd_split_align_hdim128_bf16_causal_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::bfloat16_t, 128, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim128_bf16_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::bfloat16_t, 128, false>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim128_fp16_causal_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::half_t, 128, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim128_fp16_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::half_t, 128, false>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim192_bf16_causal_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::bfloat16_t, 192, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim192_bf16_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::bfloat16_t, 192, false>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim192_fp16_causal_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::half_t, 192, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim192_fp16_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::half_t, 192, false>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE

csrc/flash_attn/src/flash_fwd_split_align_hdim256_bf16_causal_sm80.cu

@@ -0,0 +1,11 @@
+// Copyright (c) 2024, Tri Dao.
+// Splitting the different head dimensions to different files to speed up compilation.
+// This file is auto-generated. See "generate_kernels.py"
+#include "namespace_config.h"
+#include "flash_fwd_launch_template.h"
+
+namespace FLASH_NAMESPACE {
+
+template void run_mha_fwd_splitkv_align<cutlass::bfloat16_t, 256, true>(Flash_fwd_params &params, cudaStream_t stream);
+
+} // namespace FLASH_NAMESPACE