[GPU] Add support of per-head mask for sdpa_micro kernel

Author: sshlyapn

src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_micro.cl

@@ -187,6 +187,9 @@ KERNEL(micro_sdpa)(OPTIONAL_SHAPE_INFO_ARG
     Q += QRY_OFF(b1, b0, 0, 0) + INPUT0_OFFSET;
     V += VAL_OFF(b1, (b0 / KV_GROUP_SIZE), 0, 0) + INPUT2_OFFSET;
     A += DST_OFF(b1, b0, 0, 0, 0);
+#if WITH_ATTN_MASK
+    msk += MSK_OFF(b1 % MSK_D0, b0 % MSK_D1, 0, 0);
+#endif
 
     __builtin_assume_aligned(K, K_ALIGN);
     __builtin_assume_aligned(Q, Q_ALIGN);

src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_micro.cpp

@@ -439,22 +439,37 @@ JitConstants SDPAKernelMicro::GetJitConstants(const sdpa_params& params, const m
     auto convert_strides = [](std::string target_prefix, std::string source_prefix, const std::vector<int64_t> order) {
         JitConstants definitions({});
 
-        std::vector<std::string> target_definitions = {
+        std::vector<std::string> target_stride_definitions = {
             target_prefix + "_S0",
             target_prefix + "_S1",
             target_prefix + "_S2",
             target_prefix + "_S3",
         };
 
-        std::vector<std::string> source_definitions = {
+        std::vector<std::string> source_stride_definitions = {
             source_prefix + "_BATCH_PITCH",
             source_prefix + "_FEATURE_PITCH",
             source_prefix + "_Y_PITCH",
             source_prefix + "_X_PITCH",
         };
 
-        for (size_t i = 0; i < target_definitions.size(); i++) {
-            definitions.AddConstant(MakeJitConstant(target_definitions[i], source_definitions[order[i]]));
+        std::vector<std::string> target_size_definitions = {
+            target_prefix + "_D0",
+            target_prefix + "_D1",
+            target_prefix + "_D2",
+            target_prefix + "_D3",
+        };
+
+        std::vector<std::string> source_size_definitions = {
+            source_prefix + "_BATCH_NUM",
+            source_prefix + "_FEATURE_NUM",
+            source_prefix + "_SIZE_Y",
+            source_prefix + "_SIZE_X",
+        };
+
+        for (size_t i = 0; i < target_stride_definitions.size(); i++) {
+            definitions.AddConstant(MakeJitConstant(target_stride_definitions[i], source_stride_definitions[order[i]]));
+            definitions.AddConstant(MakeJitConstant(target_size_definitions[i], source_size_definitions[order[i]]));
         }
 
         return definitions;
@@ -470,6 +485,11 @@ JitConstants SDPAKernelMicro::GetJitConstants(const sdpa_params& params, const m
     jit.Merge(unit_parameters("VAL"));
     jit.Merge(unit_parameters("DST"));
 
+    if (params.inputs.size() > 3) {
+        jit.Merge(convert_strides("MSK", "INPUT3", {0, 1, 2, 3}));
+        jit.Merge(unit_parameters("MSK"));
+    }
+
     return jit;
 }