Fix bug in mem_eff kernel with attention mask and MQA (#952)

facebookresearch · Feb 23, 2024 · 4f7ffaa · 4f7ffaa
1 parent 292be24
commit 4f7ffaa
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/xformers/csrc/attention/cuda/fmha/kernel_forward.h b/xformers/csrc/attention/cuda/fmha/kernel_forward.h
@@ -327,6 +327,7 @@ struct AttentionKernel {
           return false;
         }
         q_strideM = q_strideH;
+        bias_strideM = bias_strideH;
         num_queries = num_heads;
         num_heads = 1; // unused but here for intent
         // remove causal since n_query = 1