[GPU] Add support of per-head mask for sdpa_micro kernel (#28828)

### Details: - Add support of per-head attention mask for sdpa_micro kernel, as some models have different mask configurations for each head ### Tickets: - CVS-161777 - Closes #28445
openvinotoolkit · Feb 5, 2025 · 7d18f00 · 7d18f00
1 parent c27f796
commit 7d18f00
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_micro.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_micro.cl
@@ -210,6 +210,9 @@ KERNEL(micro_sdpa)(OPTIONAL_SHAPE_INFO_ARG
     Q += (QRY_OFF(b1, b0, 0, 0) + INPUT0_OFFSET);
     V += (VAL_OFF(b1, b0_kv, 0, 0) + INPUT2_OFFSET) / VAL_ELEMENTS_PER_BYTE;
     A += DST_OFF(b1, b0, 0, 0, 0);
+#if WITH_ATTN_MASK
+    msk += MSK_OFF(b1 % MSK_D0, b0 % MSK_D1, 0, 0);
+#endif
 
 #if KEY_SCALES
     K_scales += KEY_COMP_OFF(b1, b0_kv, 0, 0);

diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_micro.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_micro.cpp
@@ -604,22 +604,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;
@@ -635,6 +650,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"));
+    }
+
     if (params.conf.is_kv_compressed) {
         jit.AddConstant(MakeJitConstant("KEY_SCALE", params.key_cache_comp_scale));
         jit.AddConstant(MakeJitConstant("VAL_SCALE", params.value_cache_comp_scale));