[GPU] Fix sdpa_micro kernel query input remainder check and kernel se…

…lection logic in case of scalar-value attention mask
sshlyapn · Nov 20, 2024 · 041078f · 041078f
1 parent 1319c64
commit 041078f
Showing 1 changed file with 5 additions and 1 deletion.
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
@@ -347,6 +347,10 @@ bool SDPAKernelMicro::Validate(const Params& p) const {
     if (params.conf.is_kv_compressed)
         return false;
 
+    // Do not use sdpa_micro kernel with a scalar-value mask
+    if (params.inputs.size() > 3 && !params.inputs[3].is_dynamic() && params.inputs[3].LogicalSize() == 1)
+        return false;
+
     return true;
 }
 
@@ -391,7 +395,7 @@ JitConstants SDPAKernelMicro::GetJitConstants(const sdpa_params& params, const m
     bool d_full = (head_size == d_max);
     bool v_full = (head_size == tile_v);
     bool k_full = !n_keys.is_dynamic && (n_keys.v % tile_k) == 0;
-    bool q_full = !n_queries.is_dynamic && (n_queries.v % tile_q) != 0;
+    bool q_full = !n_queries.is_dynamic && (n_queries.v % tile_q) == 0;
 
     auto Q_num_heads_dim = get_num_heads(Q, params.input0_order);
     auto K_num_heads_dim = get_num_heads(K, params.input1_order);