[GPU] GQA optimization of PagedAttention OCL kernel for long sequences

sshlyapn · sshlyapn · commit 66cd717107ac · 2025-03-10T20:14:41.000+04:00
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
@@ -660,7 +660,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         if (desc->heads_num != desc->kv_heads_num) {
             config.broadcast_axis = 1;
-            config.group_size = desc->heads_num / desc->kv_heads_num;
+            config.kv_group_size = desc->heads_num / desc->kv_heads_num;
         }
 
         if (desc->has_scores_output() && !is_dynamic) {
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp
@@ -252,7 +252,7 @@ struct scaled_dot_product_attention_impl : multi_stage_primitive<scaled_dot_prod
         if (query_shape[num_heads_dim].is_static() && key_shape[num_heads_dim].is_static() && value_shape[num_heads_dim].is_static()) {
             if (query_shape[num_heads_dim].get_length() > key_shape[num_heads_dim].get_length()) {
                 config.broadcast_axis = desc->input_k_transpose_order[num_heads_dim];
-                config.group_size = query_shape[num_heads_dim].get_length() / key_shape[num_heads_dim].get_length();
+                config.kv_group_size = query_shape[num_heads_dim].get_length() / key_shape[num_heads_dim].get_length();
             }
         }
 
diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/pa_sdpa_opt.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/pa_sdpa_opt.cl
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/pa_sdpa_kernel_opt.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/pa_sdpa_kernel_opt.cpp
@@ -13,13 +13,28 @@ namespace kernel_selector {
 namespace {
 enum KernelsTypes {
     SINGLE_TOKEN = 0,
+    SINGLE_TOKEN_GQA,
     MULTI_TOKENS,
     FINALIZATION,
     FINALIZATION_MULTI_TOKENS,
     SCORES_CALCULATION,
     TOTAL_KERNELS_NUM
 };
 
+static size_t get_heads_per_wi(const pa_sdpa_params& params) {
+    if (params.conf.kv_group_size > 1) {
+        std::vector<size_t> preferable_head_nums = {4, 3, 2};
+        for (const auto& heads_num : preferable_head_nums) {
+            const auto leftovers = params.conf.kv_group_size % heads_num;
+            if (leftovers == 0 || heads_num - leftovers <= 1) {
+                return heads_num;
+            }
+        }
+    }
+
+    return 1;
+}
+
 constexpr size_t subgroup_size = 16;
 constexpr size_t seq_len_partition_size = 256;
 constexpr size_t paged_attention_block_size = 16;
@@ -29,6 +44,7 @@ size_t get_sg_number_scale_factor(const pa_sdpa_params& params, size_t head_size
     if (params.conf.is_kv_compressed) {
         const size_t optimal_scale_factor = 2;
         if (kernel_type == KernelsTypes::SINGLE_TOKEN ||
+            kernel_type == KernelsTypes::SINGLE_TOKEN_GQA ||
             kernel_type == KernelsTypes::MULTI_TOKENS) {
             if (head_size * optimal_scale_factor <= params.engineInfo.maxWorkGroupSize) {
                 return optimal_scale_factor;
@@ -45,6 +61,8 @@ static std::string GetKernelName(std::string base_name, KernelsTypes type) {
 
     if (type == KernelsTypes::SINGLE_TOKEN) {
         kernel_name += "_single_token";
+    } else if (type == KernelsTypes::SINGLE_TOKEN_GQA) {
+        kernel_name += "_single_token_gqa";
     } else if (type == KernelsTypes::MULTI_TOKENS) {
         kernel_name += "_multi_tokens_seq";
     } else if (type == KernelsTypes::FINALIZATION) {
@@ -65,6 +83,7 @@ KernelsData PagedAttentionSDPAKernelOpt::GetKernelsData(const Params& p) const {
 
     const auto& params = static_cast<const pa_sdpa_params&>(p);
     std::vector<KernelsTypes> kernels_type = { KernelsTypes::SINGLE_TOKEN,
+                                               KernelsTypes::SINGLE_TOKEN_GQA,
                                                KernelsTypes::MULTI_TOKENS,
                                                KernelsTypes::FINALIZATION,
                                                KernelsTypes::FINALIZATION_MULTI_TOKENS };
@@ -90,7 +109,7 @@ KernelsData PagedAttentionSDPAKernelOpt::GetKernelsData(const Params& p) const {
 
         int inputs_num = static_cast<int>(params.inputs.size());
         int outputs_num = 1;
-        if (kernel_type == KernelsTypes::SINGLE_TOKEN) {
+        if (kernel_type == KernelsTypes::SINGLE_TOKEN || kernel_type == KernelsTypes::SINGLE_TOKEN_GQA) {
             // SINGLE_TOKEN kernel doesn't use the subsequence_begins input
             inputs_num -= 1;
         } else if (kernel_type == KernelsTypes::FINALIZATION) {
@@ -221,6 +240,7 @@ JitConstants PagedAttentionSDPAKernelOpt::GetJitConstants(const pa_sdpa_params&
     jit.AddConstant(MakeJitConstant("HEAD_SIZE", config.head_size));
     jit.AddConstant(MakeJitConstant("HEADS_NUM", config.heads_num));
     jit.AddConstant(MakeJitConstant("KV_HEADS_NUM", config.kv_heads_num));
+    jit.AddConstant(MakeJitConstant("KV_HEADS_GROUP_SIZE", config.kv_group_size));
     jit.AddConstant(MakeJitConstant("SEQ_LEN_PARTITION_SIZE", seq_len_partition_size));
     jit.AddConstant(MakeJitConstant("PAGED_ATTENTION_BLOCK_SIZE", paged_attention_block_size));
     jit.AddConstant(MakeJitConstant("SUBGROUP_SIZE", subgroup_size));
@@ -236,8 +256,13 @@ JitConstants PagedAttentionSDPAKernelOpt::GetJitConstants(const pa_sdpa_params&
         jit.AddConstant(MakeJitConstant("ADJUSTED_HEAD_SIZE", params.conf.head_size));
     }
 
-    if (config.broadcast_axis != -1) {
-        jit.AddConstant(MakeJitConstant("BROADCAST_GROUP_SIZE", config.group_size));
+    if (kernel_idx == KernelsTypes::SINGLE_TOKEN_GQA) {
+        auto heads_per_wi = get_heads_per_wi(params);
+        jit.AddConstant(MakeJitConstant("HEADS_PER_WI", heads_per_wi));
+        jit.AddConstant(MakeJitConstant("ITERATIONS_PER_KV_HEADS_GROUP", CeilDiv(config.kv_group_size, heads_per_wi)));
+        jit.AddConstant(MakeJitConstant("HEADS_LEFTOVERS_NUM", config.kv_group_size % heads_per_wi));
+    } else {
+        jit.AddConstant(MakeJitConstant("HEADS_PER_WI", 1));
     }
 
     auto sdpa_stage = 0;
@@ -293,6 +318,16 @@ CommonDispatchData PagedAttentionSDPAKernelOpt::SetDefault(const pa_sdpa_params&
                                   heads_num,
                                   head_size * num_of_partitions * sg_scale };
             dispatch_data.lws = { 1, 1, head_size * sg_scale };
+        } else if (kernel_idx == KernelsTypes::SINGLE_TOKEN_GQA) {
+            auto sg_scale = get_sg_number_scale_factor(params, head_size, kernel_idx);
+
+            auto kv_groups = heads_num / params.conf.kv_group_size;
+            auto gqa_heads_num = kv_groups * CeilDiv(params.conf.kv_group_size, get_heads_per_wi(params));
+
+            dispatch_data.gws = { total_tokens,
+                                  gqa_heads_num,
+                                  head_size * num_of_partitions * sg_scale };
+            dispatch_data.lws = { 1, 1, head_size * sg_scale };
         } else if (kernel_idx == KernelsTypes::SCORES_CALCULATION) {
             const auto& past_lens = params.inputs[3];
             const auto subsequences_number = past_lens.Batch().v;
@@ -334,13 +369,30 @@ void PagedAttentionSDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) cons
         const auto scores_calc_only = prim_params.stage == PagedAttentionStage::PREFILL && has_scores_output;
         const auto multi_tokens_mode = prim_params.stage == PagedAttentionStage::MIXED;
 
-        auto dispatch_data1 = SetDefault(prim_params, KernelsTypes::SINGLE_TOKEN);
-        kd.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.global = dispatch_data1.gws;
-        kd.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.local = dispatch_data1.lws;
-        kd.kernels[KernelsTypes::SINGLE_TOKEN].skip_execution = multi_tokens_mode || scores_calc_only;
-
-        kd.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.global = dispatch_data1.gws;
-        kd.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.local = dispatch_data1.lws;
+        // Apply GQA optimization starting from a certain sequence length value
+        const auto min_gqa_sequence_len = 8 * seq_len_partition_size;
+        // Apply GQA only if there is a single subsequence in the request,
+        // as multiple subsequences might have significantly different lengths
+        const auto max_subsequences_num = 1;
+        const auto subsequences_num = prim_params.inputs[0].Batch().v;
+        const auto can_use_gqa_kernel = prim_params.conf.paged_attention_max_len >= static_cast<int64_t>(min_gqa_sequence_len) &&
+                                        subsequences_num <= max_subsequences_num &&
+                                        prim_params.conf.kv_group_size > 1 &&
+                                        !multi_tokens_mode &&
+                                        !scores_calc_only;
+
+        auto dispatch_data = SetDefault(prim_params, KernelsTypes::SINGLE_TOKEN_GQA);
+        kd.kernels[KernelsTypes::SINGLE_TOKEN_GQA].params.workGroups.global = dispatch_data.gws;
+        kd.kernels[KernelsTypes::SINGLE_TOKEN_GQA].params.workGroups.local = dispatch_data.lws;
+        kd.kernels[KernelsTypes::SINGLE_TOKEN_GQA].skip_execution = multi_tokens_mode || scores_calc_only || !can_use_gqa_kernel;
+
+        dispatch_data = SetDefault(prim_params, KernelsTypes::SINGLE_TOKEN);
+        kd.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.global = dispatch_data.gws;
+        kd.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.local = dispatch_data.lws;
+        kd.kernels[KernelsTypes::SINGLE_TOKEN].skip_execution = multi_tokens_mode || scores_calc_only || can_use_gqa_kernel;
+
+        kd.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.global = dispatch_data.gws;
+        kd.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.local = dispatch_data.lws;
         kd.kernels[KernelsTypes::MULTI_TOKENS].skip_execution = !multi_tokens_mode || scores_calc_only;
 
         size_t partition_size = 0;
@@ -351,13 +403,13 @@ void PagedAttentionSDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) cons
         }
         const size_t num_of_partitions = CeilDiv(prim_params.conf.paged_attention_max_len, partition_size);
 
-        auto dispatch_data2 = SetDefault(prim_params, KernelsTypes::FINALIZATION);
-        kd.kernels[KernelsTypes::FINALIZATION].params.workGroups.global = dispatch_data2.gws;
-        kd.kernels[KernelsTypes::FINALIZATION].params.workGroups.local = dispatch_data2.lws;
+        dispatch_data = SetDefault(prim_params, KernelsTypes::FINALIZATION);
+        kd.kernels[KernelsTypes::FINALIZATION].params.workGroups.global = dispatch_data.gws;
+        kd.kernels[KernelsTypes::FINALIZATION].params.workGroups.local = dispatch_data.lws;
         kd.kernels[KernelsTypes::FINALIZATION].skip_execution = num_of_partitions == 1 || multi_tokens_mode || scores_calc_only;
 
-        kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].params.workGroups.global = dispatch_data2.gws;
-        kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].params.workGroups.local = dispatch_data2.lws;
+        kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].params.workGroups.global = dispatch_data.gws;
+        kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].params.workGroups.local = dispatch_data.lws;
         kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].skip_execution = num_of_partitions == 1 || !multi_tokens_mode || scores_calc_only;
 
         ScalarDescriptor num_of_partitions_scalar;
@@ -369,7 +421,7 @@ void PagedAttentionSDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) cons
         kd.kernels[KernelsTypes::FINALIZATION_MULTI_TOKENS].params.scalars[0] = num_of_partitions_scalar;
 
         if (has_scores_output) {
-            auto dispatch_data = SetDefault(prim_params, KernelsTypes::SCORES_CALCULATION);
+            dispatch_data = SetDefault(prim_params, KernelsTypes::SCORES_CALCULATION);
             kd.kernels[KernelsTypes::SCORES_CALCULATION].params.workGroups.global = dispatch_data.gws;
             kd.kernels[KernelsTypes::SCORES_CALCULATION].params.workGroups.local = dispatch_data.lws;
             kd.kernels[KernelsTypes::SCORES_CALCULATION].skip_execution = false;
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.cpp
@@ -70,10 +70,10 @@ JitConstants SDPAKernelBase::GetJitConstants(const sdpa_params& params) const {
     auto jit = MakeBaseParamsJitConstants(params);
 
     if (params.conf.broadcast_axis != -1) {
-        jit.AddConstant(MakeJitConstant("BROADCAST_GROUP_SIZE", params.conf.group_size));
+        jit.AddConstant(MakeJitConstant("BROADCAST_GROUP_SIZE", params.conf.kv_group_size));
         jit.AddConstant(MakeJitConstant("DO_BROADCAST_KEY_VALUE", GetBroadcastInputStr(params.inputs[0].GetDims().size(),
                                                                                        params.conf.broadcast_axis,
-                                                                                       params.conf.group_size)));
+                                                                                       params.conf.kv_group_size)));
     } else {
         jit.AddConstant(MakeJitConstant("BROADCAST_GROUP_SIZE", 1));
     }
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.h b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.h
@@ -83,7 +83,7 @@ struct sdpa_configuration {
     int64_t kv_heads_num = -1;
 
     // GQA configuration
-    int64_t group_size = -1;
+    int64_t kv_group_size = 1;
     int64_t broadcast_axis = -1;
 
     bool is_causal = false;
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_selector.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_selector.cpp
@@ -17,7 +17,18 @@ sdpa_kernel_selector::sdpa_kernel_selector() {
     Attach<SDPAKernelOpt>();
     Attach<SDPAKernelRef>();
 #ifdef ENABLE_ONEDNN_FOR_GPU
-    Attach<SDPAKernelMicro>();
+    int DISABLE_MICRO = 0;
+    if (const auto env_var = std::getenv("DISABLE_MICRO")) {
+        std::istringstream ss(env_var);
+        ss >> DISABLE_MICRO;
+        static bool printed = false;
+        if (!printed) {
+            std::cout << "Set DISABLE_MICRO=" << DISABLE_MICRO << "\n";
+            printed = true;
+        }
+    }
+    if (!DISABLE_MICRO)
+        Attach<SDPAKernelMicro>();
 #endif
 }
 

Original file line number	Diff line number	Diff line change
`@@ -660,7 +660,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {`
`660`	`660`
`661`	`661`	`if (desc->heads_num != desc->kv_heads_num) {`
`662`	`662`	`config.broadcast_axis = 1;`
`663`		`- config.group_size = desc->heads_num / desc->kv_heads_num;`
	`663`	`+ config.kv_group_size = desc->heads_num / desc->kv_heads_num;`
`664`	`664`	`}`
`665`	`665`
`666`	`666`	`if (desc->has_scores_output() && !is_dynamic) {`
Original file line number	Diff line number	Diff line change
`@@ -252,7 +252,7 @@ struct scaled_dot_product_attention_impl : multi_stage_primitive<scaled_dot_prod`
`252`	`252`	`if (query_shape[num_heads_dim].is_static() && key_shape[num_heads_dim].is_static() && value_shape[num_heads_dim].is_static()) {`
`253`	`253`	`if (query_shape[num_heads_dim].get_length() > key_shape[num_heads_dim].get_length()) {`
`254`	`254`	`config.broadcast_axis = desc->input_k_transpose_order[num_heads_dim];`
`255`		`- config.group_size = query_shape[num_heads_dim].get_length() / key_shape[num_heads_dim].get_length();`
	`255`	`+ config.kv_group_size = query_shape[num_heads_dim].get_length() / key_shape[num_heads_dim].get_length();`
`256`	`256`	`}`
`257`	`257`	`}`
`258`	`258`