refactor rope part and optimize mask memory for Phi3 (#2828)

Author: blzheng

csrc/cpu/aten/MaskedMultiHeadAttention.cpp

@@ -6,6 +6,7 @@ namespace torch_ipex {
 namespace cpu {
 
 IPEX_DEFINE_DISPATCH(masked_multihead_self_attention_kernel_stub);
+IPEX_DEFINE_DISPATCH(prepare_4d_causal_attention_mask_kernel_stub);
 
 /*
  *Caculate the masked multihead attention for decoder layer in decoder only
@@ -54,6 +55,21 @@ masked_multihead_self_attention_forward_cpu(
       add_casual_mask);
 }
 
+at::Tensor prepare_4d_causal_attention_mask_forward_cpu(
+    at::Tensor& attention_mask,
+    at::Tensor& inputs_embeds,
+    at::Tensor& past_kv_len,
+    at::Tensor& finfo_min,
+    int64_t sliding_window) {
+  return prepare_4d_causal_attention_mask_kernel_stub(
+      kCPU,
+      attention_mask,
+      inputs_embeds,
+      past_kv_len,
+      finfo_min,
+      sliding_window);
+}
+
 } // namespace cpu
 } // namespace torch_ipex
 
@@ -69,4 +85,13 @@ TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
       c10::DispatchKey::CPU,
       torch_ipex::cpu::masked_multihead_self_attention_forward_cpu);
 }
+
+TORCH_LIBRARY_FRAGMENT(torch_ipex, m) {
+  m.def(
+      "prepare_4d_causal_attention_mask(Tensor attention_mask, Tensor inputs_embeds, Tensor past_kv_len, Tensor finfo_min, int sliding_window)-> (Tensor)");
+  m.impl(
+      "prepare_4d_causal_attention_mask",
+      c10::DispatchKey::CPU,
+      torch_ipex::cpu::prepare_4d_causal_attention_mask_forward_cpu);
+}
 } // namespace

csrc/cpu/aten/MaskedMultiHeadAttention.h

@@ -22,7 +22,14 @@ masked_multihead_self_attention(
     const c10::optional<at::Tensor>& head_mask /* optional */,
     const c10::optional<at::Tensor>& attention_mask /* optional */,
     c10::optional<bool> add_casual_mask /* optional */);
-}
+
+at::Tensor prepare_4d_causal_attention_mask_forward_cpu(
+    at::Tensor& attention_mask,
+    at::Tensor& inputs_embeds,
+    at::Tensor& past_kv_len,
+    at::Tensor& finfo_min,
+    int64_t sliding_window);
+} // namespace
 
 using masked_multihead_self_attention_kernel_fn =
     std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor> (*)(
@@ -42,6 +49,16 @@ using masked_multihead_self_attention_kernel_fn =
 IPEX_DECLARE_DISPATCH(
     masked_multihead_self_attention_kernel_fn,
     masked_multihead_self_attention_kernel_stub);
+using prepare_4d_causal_attention_mask_kernel_fn = at::Tensor (*)(
+    at::Tensor& attention_mask,
+    at::Tensor& inputs_embeds,
+    at::Tensor& past_kv_len,
+    at::Tensor& finfo_min,
+    int64_t sliding_window);
+
+IPEX_DECLARE_DISPATCH(
+    prepare_4d_causal_attention_mask_kernel_fn,
+    prepare_4d_causal_attention_mask_kernel_stub);
 
 } // namespace cpu
 } // namespace torch_ipex