Skip to content

Commit 7b51189

Browse files
authored
[Cherry-pick] Fix alibi blocked shape issue (#5007)(#5009)
1 parent b06ff2c commit 7b51189

File tree

1 file changed

+6
-1
lines changed
  • intel_extension_for_pytorch/transformers/models/xpu/optimize_transformers/modules/transformer_modules

1 file changed

+6
-1
lines changed

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,12 @@ def sdp(self, query, key, value, past_key_value, attention_mask, head_mask, alib
229229
# if attention_mask is not None:
230230
# attention_mask = self.get_blocked_attn_mask(attention_mask)
231231
if alibi is not None:
232-
alibi = self.get_blocked_alibi(alibi, key.size(2))
232+
if isinstance(past_key_value, IPEXStaticCache):
233+
alibi = self.get_blocked_alibi(
234+
alibi, past_key_value.get_seq_length() + key.size(2)
235+
)
236+
else:
237+
alibi = self.get_blocked_alibi(alibi, key.size(2))
233238
if (
234239
self.beam_idx is not None
235240
and query.size(-2) == 1

0 commit comments

Comments
 (0)