fix beam search perf regression (#4952) (#4959)

Author: guanbaoy

intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/transformer_modules/CacheUtils.py

@@ -166,19 +166,26 @@ def get_kv_slice_for_decoding(
             layer_idx: int,
             key: torch.Tensor,
         ) -> Tuple[torch.Tensor, torch.Tensor]:
-            # return the key and value cache for decoding in shape of BNFH
+            # return the key and value cache for decoding
+            prompt_len = (
+                0 if len(self.key_prompt) == 0 else self.key_prompt[layer_idx].size(2)
+            )
             seqlen = self.update_or_get_seq_cnt(layer_idx) + key.size(2)
             if self.cache_format == CacheFormat.FBNH:
-                key = self.key_cache[layer_idx][:seqlen, :, :, :].permute(1, 2, 0, 3)
-                value = self.value_cache[layer_idx][:seqlen, :, :, :].permute(
+                key = self.key_cache[layer_idx][prompt_len:seqlen, :, :, :].permute(
+                    1, 2, 0, 3
+                )
+                value = self.value_cache[layer_idx][prompt_len:seqlen, :, :, :].permute(
                     1, 2, 0, 3
                 )
             elif self.cache_format == CacheFormat.BNFH:
-                key = self.key_cache[layer_idx][:, :, :seqlen, :]
-                value = self.value_cache[layer_idx][:, :, :seqlen, :]
+                key = self.key_cache[layer_idx][:, :, prompt_len:seqlen, :]
+                value = self.value_cache[layer_idx][:, :, prompt_len:seqlen, :]
             elif self.cache_format == CacheFormat.BFNH:
-                key = self.key_cache[layer_idx][:, :seqlen, :, :].permute(0, 2, 1, 3)
-                value = self.value_cache[layer_idx][:, :seqlen, :, :].permute(
+                key = self.key_cache[layer_idx][:, prompt_len:seqlen, :, :].permute(
+                    0, 2, 1, 3
+                )
+                value = self.value_cache[layer_idx][:, prompt_len:seqlen, :, :].permute(
                     0, 2, 1, 3
                 )
             return key, value

intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/transformer_modules/XPUAttentionfp16.py

@@ -239,12 +239,7 @@ def sdp(self, query, key, value, past_key_value, attention_mask, head_mask, alib
             key_prompt, value_prompt = past_key_value.get_prompt_for_beam_search(
                 self.layer_idx
             )
-            prompt_length = key_prompt.size(2)
             curr_len = key.size(2)
-            # TODO: remove this after ifmha support combined kv cache with both prompt
-            # and decode in [bs, curr_len, num_head, head_dim] layout
-            key = key[:, :, prompt_length:, :]
-            value = value[:, :, prompt_length:, :]
             # TODO: remove this after ifmha support [bs, curr_len, num_head, head_dim] layout
             if (
                 isinstance(past_key_value, IPEXStaticCache)
@@ -260,7 +255,6 @@ def sdp(self, query, key, value, past_key_value, attention_mask, head_mask, alib
                         0,
                     )
                 )
-
             attention_output = torch.xpu.IpexSDP_Index(
                 query,
                 key_prompt,
@@ -401,7 +395,6 @@ def forward(
         value = value.view(
             [value.shape[0], value.shape[1], self.num_kv_heads, self.head_dim]
         )
-
         # apply rope to qk
         query, key, value = self.rotary_embedding(
             query, key, value, past_key_value, position_ids, self.layer_idx, curr_len