1. fix nan in alibi: kernel assume alibi is same dtype as key/value, pass f32 alibi but interprete as f16 cause nan (#5068) (#5080)

2. always use 2d load for bias, for boundary check

Co-authored-by: WeiZhu <[email protected]>

Author: rogerxfeng8

csrc/gpu/aten/operators/xetla/kernels/SDP/fmha_forward.hpp

@@ -396,15 +396,16 @@ class fmha_forward_t {
       // B, N, 1, T
       // gid * T + startT
       if constexpr (kUseAlibi && !kVarlen) {
-        int32_t batch_start = gid * args.uAT;
-        int32_t start_x = batch_start + startT;
-        uint32_t end_x = startT + kBc;
+        int32_t start_x = startT;
+        uint32_t end_x = start_x + kBc;
         uint32_t boundary_x = args.uT;
         end_x = end_x > boundary_x ? boundary_x : end_x;
-        end_x += batch_start;
+
+        int32_t start_y = gid;
+        uint32_t end_y = start_y + 1;
 
         mem_desc_Ai.init(
-            args.A_ptr, {end_x, 1, args.uAT * args.uN * args.uB}, {start_x, 0});
+            args.A_ptr, {end_x, end_y, args.uAT}, {start_x, start_y});
       }
 
       // B, N or N

csrc/gpu/aten/operators/xetla/kernels/include/subgroup/tile/impl/tile_op_functor.hpp

@@ -609,7 +609,7 @@ struct bias_add_op_t<
     using bias_payload_t = mem_payload_t<
         mem_desc_bias_t,
         bias_tile_desc_t,
-        msg_type_v<bias_tile_desc_t, mem_desc_bias_t>,
+        msg_type::block_2d,
         arch_tag>;
     coord_t bias_coord(coord.x, 0);
     mem_desc_bias_t mem_desc_bias(args.base, args.shape, bias_coord);

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

@@ -177,7 +177,7 @@ def all_reduce_if_necessary(self, reduce_target):
             dist.all_reduce(reduce_target, group=self.tp_group)
         return
 
-    def get_blocked_alibi(self, alibi, seq_len):
+    def get_blocked_alibi(self, alibi, seq_len, dtype):
         if self.layer_idx == 0:
             cache_len = (
                 self.max_position
@@ -190,7 +190,7 @@ def get_blocked_alibi(self, alibi, seq_len):
                 cache_len,
             ]  # [beam*num_head, q_len, kv_len]
             IPEXAttention.blocked_alibi = torch.empty(
-                shape, device=alibi.device, dtype=alibi.dtype
+                shape, device=alibi.device, dtype=dtype
             )
             kv_len = alibi.shape[2]
             IPEXAttention.blocked_alibi[:, :, 0:kv_len] = alibi
@@ -228,13 +228,14 @@ def sdp(self, query, key, value, past_key_value, attention_mask, head_mask, alib
 
         # if attention_mask is not None:
         #     attention_mask = self.get_blocked_attn_mask(attention_mask)
+        # use key/value's data type as alibi's data type
         if alibi is not None:
             if isinstance(past_key_value, IPEXStaticCache):
                 alibi = self.get_blocked_alibi(
-                    alibi, past_key_value.get_seq_length() + key.size(2)
+                    alibi, past_key_value.get_seq_length() + key.size(2), key.dtype
                 )
             else:
-                alibi = self.get_blocked_alibi(alibi, key.size(2))
+                alibi = self.get_blocked_alibi(alibi, key.size(2), key.dtype)
         if (
             self.beam_idx is not None
             and query.size(-2) == 1
@@ -304,7 +305,6 @@ def sdp(self, query, key, value, past_key_value, attention_mask, head_mask, alib
                 if not self.is_beam_search() and query.size(-2) == 1:
                     key = key.permute(2, 0, 1, 3).contiguous().permute(1, 2, 0, 3)
                     value = value.permute(2, 0, 1, 3).contiguous().permute(1, 2, 0, 3)
-
                 attention_output = torch.xpu.IpexSDP(
                     query,
                     key,