Clear dq before launch kernel

ROCm · Dec 27, 2024 · e3dd0e2 · e3dd0e2
1 parent 12cd298
commit e3dd0e2
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/csrc/flash_attn_ck/mha_bwd.cpp b/csrc/flash_attn_ck/mha_bwd.cpp
@@ -338,6 +338,10 @@ 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());