Update fmha_bwd_traits template arguments

ROCm · Dec 28, 2024 · e471acc · e471acc
1 parent 5a93a97
commit e471acc
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/csrc/flash_attn_ck/mha_bwd.cpp b/csrc/flash_attn_ck/mha_bwd.cpp
@@ -25,7 +25,7 @@ fmha_bwd_traits get_ck_fmha_bwd_traits(const mask_info &mask,
                            false, // s_randval
                            deterministic,
                            true, // uses_ext_asm
-                           head_size != 64, // is_v3_atomic_fp32
+                           true, // is_v3_atomic_fp32
                            false, // is_v3_spec
                            1}; // how_v3_bf16_cvt 0:RTNE; 1:RTNA; 2:RTZ
 }
@@ -338,10 +338,6 @@ mha_bwd(const at::Tensor &dout,                   // batch_size x seqlen_q x num
         dv_expanded = dv;
     }
 
-    if (head_size == 64) {
-        dq.zero_();
-    }
-
     auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
         gen_, at::cuda::detail::getDefaultCUDAGenerator());