[Paddle Inference] Fix mmha when src_mask is not equal to zero (#57936)

xiaoxiaohehe001 · web-flow · commit d78ec37590e1 · 2023-10-11T10:30:50.000+08:00
* fix_mmha_scrmask

* fix_mmha_scrmask

* remove_mask_to_qk_smem_act_time_step

* remove_add_!

* mask_bhi
diff --git a/paddle/phi/kernels/fusion/gpu/masked_multihead_attention.cu b/paddle/phi/kernels/fusion/gpu/masked_multihead_attention.cu
@@ -609,6 +609,11 @@ __global__ void masked_multihead_attention_kernel(
     //    bi * (params.timestep + 1) + params.timestep];
     // qk += static_cast<float>(mask);
     qk *= params.inv_sqrt_dh;
+    if (params.attn_mask) {
+      auto mask_bhi = params.mask_broadcast_num_heads ? bi : bhi;
+      T mask = params.attn_mask[mask_bhi * params.mask_length + act_time_step];
+      qk += static_cast<float>(mask);
+    }
     qk_max = qk;
     qk_smem[act_time_step] = qk;
   }
