Rebase with master

Author: ashwins990

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/softmax_kernel.hpp

@@ -578,14 +578,15 @@ inline void scale_add2_reduce_max(ov::float16* a,
 
         if (has_alibi) {
             svfloat16_t v_lookup = svld1_f16(pg_f16, reinterpret_cast<const float16_t*>(alibi_lookup + i));
-            v_a = svmla_f16_x(pg_f16, v_a, v_lookup, v_alibi_slope);
+            v_a = svmla_f16_z(pg_f16, v_a, v_lookup, v_alibi_slope);
         }
 
         if (has_attn_mask) {
             static_assert(std::is_same_v<T, float> || std::is_same_v<T, ov::float16>,
                           "attn_mask must be float or float16 type.");
             if constexpr (std::is_same_v<T, float>) {
                 svfloat16_t zero = svdup_n_f16(0.0f);
+                size_t inc_low = (vec_len + 1) / 2;
                 size_t inc_high = vec_len / 2;
                 svbool_t pg_f32_low = svwhilelt_b32(0, static_cast<int>(inc_low));
                 svbool_t pg_f32_high = svwhilelt_b32(0, static_cast<int>(inc_high));
@@ -607,9 +608,13 @@ inline void scale_add2_reduce_max(ov::float16* a,
                 v_a = svuzp1(low_f16_out, high_f16_out);
             } else if constexpr (std::is_same_v<T, ov::float16>) {
                 svfloat16_t v_mask = svld1_f16(pg_f16, reinterpret_cast<const float16_t*>(attn_mask + i));
-                v_a = svadd_f16_z(pg_f16, v_a, v_mask);
+                v_a = svadd_f16_x(pg_f16, v_a, v_mask);
             }
         }
+
+        if (has_causal_mask) {
+            svuint8_t v_maski8 = svld1_u8(pg_u8, causal_mask + i);
+            svuint16_t v_maski16 = svtrn1_u16(svreinterpret_u16_u8(v_maski8), svdup_n_u16(0));
             svbool_t kmask = svcmpeq_u16(pg_u16, v_maski16, v_zeroi16);
             kmask = sveor_z(pg_u16, kmask, mask_xor);
             v_a = svsel_f16(kmask, v_nfltmax, v_a);
@@ -1443,12 +1448,16 @@ inline void attn_softmax_kernel<ov::float16>(ov::float16* a,
     }
 
     ov::float16 sum = 0.0f;
+    ov::float16 clamped_sink_value = 0.0f;
     if (sink != nullptr) {
-        max = std::max(max, static_cast<const ov::float16>(*sink));
+        clamped_sink_value = static_cast<const ov::float16>(*sink);
+        clamped_sink_value =
+            std::isinf(clamped_sink_value) ? std::numeric_limits<ov::float16>::max() : clamped_sink_value;
+        max = std::max(max, clamped_sink_value);
     }
     exp_reduce_sum_f32(a, max, len, sum);
     if (sink != nullptr) {
-        sum += std::exp(*sink - max);
+        sum += std::exp(clamped_sink_value - max);
     }
     if (dst_precision == ov::element::f32) {
         ov::float16 scalar = 1.0f / sum;