Optimizes flash attention backward kernel configurations

Adjusts kernel parameters across different head dimensions to improve memory usage and performance on various GPU architectures.

Updates shared memory requirements and CTA counts for better utilization on sm86, sm89, A100, and H100 GPUs.

Enables double buffering and adjusts block sizes to reduce memory footprint while maintaining or improving performance across different hardware configurations.

Author: LoserCheems

csrc/src/flash_bwd_launch_template.h

@@ -138,12 +138,11 @@ void run_mha_bwd_hdim32(Flash_bwd_params &params, cudaStream_t stream) {
       C10_CUDA_CHECK(status_);
     }
     if (max_smem_per_block >= 104 * 1024) {             // H100 and A100
-        // 104KB
+        // 104KB, 1 CTAs in A100, 2 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, false, T>, Is_causal>(params, stream);
     } else {                                            // sm86 and sm89
-        // 96KB
-        // We need to adjust no_double_buffer to save some smem, because is_v_in_regs=true will still allocate smem that may overflow
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, true, T>, Is_causal>(params, stream);
+        // 96KB, 2 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, true, false, T>, Is_causal>(params, stream);
     }
 }
 
@@ -158,17 +157,17 @@ void run_mha_bwd_hdim64(Flash_bwd_params &params, cudaStream_t stream) {
     if (status_ != cudaSuccess) {
       C10_CUDA_CHECK(status_);
     }
-    // printf("max_smem_per_block = %d\n", max_smem_per_block);
-    // Changing AtomLayoutMdQ from 2 to 4 takes the same time
-    // This is slightly faster. We want to split M more so we need fewer registers to store LSE.
     if (max_smem_per_block >= 144 * 1024) {             // H100 and A100
-        // 144KB
+        // In fwd, multi-CTA configurations are faster, but in bwd, their speeds are very close.
+        // 56KB, 1 CTAs in sm86 and sm 89, 2 CTAs in A100, 4 CTAs in H100.
+        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_causal>(params, stream);
+        // 72KB, 1 CTAs in sm86 and sm 89, 2 CTAs in A100, 3 CTAs in H100.
+        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_causal>(params, stream);
+        // 144KB, N/A CTAs in sm86 and sm 89, 1 CTAs in A100, 1 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, false, T>, Is_causal>(params, stream);
-        // This has a lot of register spilling
-        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, true, false, T>>(params, stream);
     } else {                                            // sm86 and sm89
-        // 88KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>, Is_causal>(params, stream);
+        // 72KB, 1 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_causal>(params, stream);
     }
     // M=128, N=64 is quite slow, I think because we need to read/write dQaccum twice as many times
 }
@@ -186,11 +185,11 @@ void run_mha_bwd_hdim96(Flash_bwd_params &params, cudaStream_t stream) {
     }
     // printf("max_smem_per_block = %d\n", max_smem_per_block);
     if (max_smem_per_block >= 116 * 1024) {             // H100 and A100
-        // 116KB
+        // 116KB, 1 CTAs in A100, 1 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>, Is_causal>(params, stream);
     } else {                                            // sm86 and sm89
-        // 80KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 2, 4, 4, false, false, T>, Is_causal>(params, stream);
+        // 92KB, 1 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_causal>(params, stream);
     }
 }
 
@@ -205,20 +204,12 @@ void run_mha_bwd_hdim128(Flash_bwd_params &params, cudaStream_t stream) {
     if (status_ != cudaSuccess) {
       C10_CUDA_CHECK(status_);
     }
-    // printf("max_smem_per_block = %d\n", max_smem_per_block);
-    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 64, 8, 2, 2, 2, false, false, T>>(params, stream);
-    // This is faster, in the case of sequence-parallel bwd (where we need fewer registers).
-    // Out of these three, the 2nd one is slightly faster (2% faster than the first). Idk why.
-    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 2, 2, 2, false, false, T>>(params, stream);
-    if (max_smem_per_block >= 224 * 1024) {             // H100
-        // 224KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 2, 4, 2, false, false, T>, Is_causal>(params, stream);
-    } else if (max_smem_per_block >= 144 * 1024) {      // A100
-        // 144KB
+    if (max_smem_per_block >= 144 * 1024) {             // H100 and A100
+        // 144KB, 1 CTAs in A100, 1 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 2, false, false, T>, Is_causal>(params, stream);
     } else {                                            // sm86 and sm89
-        // 88KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, true, T>, Is_causal>(params, stream);
+        // 88KB, 1 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, true, false, T>, Is_causal>(params, stream);
     }
 }
 
@@ -233,15 +224,12 @@ void run_mha_bwd_hdim192(Flash_bwd_params &params, cudaStream_t stream) {
     if (status_ != cudaSuccess) {
       C10_CUDA_CHECK(status_);
     }
-    if (max_smem_per_block >= 208 * 1024) {             // H100
-        // 208KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 4, 2, 2, false, false, T>, Is_causal>(params, stream);
-    } else if (max_smem_per_block >= 152 * 1024) {      // A100
-        // 152KB
+    if (max_smem_per_block >= 136 * 1024) {             // H100 and A100
+        // 136KB, 1 CTAs in A100, 1 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_causal>(params, stream);
     } else {                                            // sm86 and sm89
-        // 88KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 64, 8, 4, 2, 2, false, true, T>, Is_causal>(params, stream);
+        // 96KB, 1 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, true, true, T>, Is_causal>(params, stream);
     }
 }
 
@@ -256,15 +244,15 @@ void run_mha_bwd_hdim256(Flash_bwd_params &params, cudaStream_t stream) {
     if (status_ != cudaSuccess) {
       C10_CUDA_CHECK(status_);
     }
-    if (max_smem_per_block >= 200 * 1024) {             // H100
-        // 200KB
+    if (max_smem_per_block >= 176 * 1024) {             // H100
+        // 176KB, 1 CTAs in H100.
         run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_causal>(params, stream);
-    } else if (max_smem_per_block >= 132 * 1024) {      // A100
-        // 132KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 64, 8, 4, 2, 2, false, false, T>, Is_causal>(params, stream);
+    } else if (max_smem_per_block >= 144 * 1024) {      // A100
+        // 144KB, 1 CTAs in A100.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, true, T>, Is_causal>(params, stream);
     } else {                                            // sm86 and sm89
-        // 82KB
-        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 32, 8, 4, 1, 2, true, false, T>, Is_causal>(params, stream);
+        // 96KB, 1 CTAs in sm86 and sm 89.
+        run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 32, 8, 4, 1, 2, true, true, T>, Is_causal>(params, stream);
     }
 }