Optimize Flash Dynamic Mask Attention Kernel Configurations

LoserCheems · web-flow · commit a32c7302b659 · 2025-09-03T22:48:56.000+08:00
diff --git a/csrc/src/flash_bwd_launch_template.h b/csrc/src/flash_bwd_launch_template.h
@@ -183,7 +183,6 @@ void run_mha_bwd_hdim96(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);
     if (max_smem_per_block >= 116 * 1024) {             // H100 and A100
         // 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);

Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,6 @@ void run_mha_bwd_hdim96(Flash_bwd_params &params, cudaStream_t stream) {`
`183`	`183`	`if (status_ != cudaSuccess) {`
`184`	`184`	`C10_CUDA_CHECK(status_);`
`185`	`185`	`}`
`186`		`- // printf("max_smem_per_block = %d\n", max_smem_per_block);`
`187`	`186`	`if (max_smem_per_block >= 116 * 1024) { // H100 and A100`
`188`	`187`	`// 116KB, 1 CTAs in A100, 1 CTAs in H100.`
`189`	`188`	`run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>, Is_causal>(params, stream);`