WA: Add PA explicit configuration via env variable

Author: sshlyapn

src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp

@@ -163,6 +163,33 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         config.x_block_size = desc->x_block_size;
         config.max_context_len = 1;
 
+        if (!impl_param.is_dynamic()) {
+            auto query_shape = impl_param.get_input_layout(0).get_shape();
+            auto key_cache_shape = impl_param.get_input_layout(3).get_shape();
+            auto value_cache_shape = impl_param.get_input_layout(4).get_shape();
+
+            auto actual_head_size = value_cache_shape[2];
+            auto actual_heads_num = query_shape[2] / actual_head_size;
+            auto actual_kv_heads_num = value_cache_shape[1];
+            auto actual_block_size = value_cache_shape[3];
+            auto actual_x_block_size = key_cache_shape[4];
+
+            bool valid_params = config.head_size == actual_head_size &&
+                                config.heads_num == actual_heads_num &&
+                                config.kv_heads_num == actual_kv_heads_num &&
+                                config.block_size == actual_block_size &&
+                                config.x_block_size == actual_x_block_size;
+
+            OPENVINO_ASSERT(valid_params, "[GPU] Got unexpected parameters for PA operation. ",
+                            "Currently they need to be specified explicitly (this should be fixed soon by PA model conversion improvement). ",
+                            "Please use the following environment variables for proper PA configuration: ",
+                            "PA_HEAD_SIZE=", actual_head_size, " ",
+                            "PA_HEADS_NUM=", actual_heads_num, " ",
+                            "PA_KV_HEADS_NUM=", actual_kv_heads_num, " ",
+                            "PA_BLOCK_SIZE=", actual_block_size, " ",
+                            "PA_X_BLOCK_SIZE=", actual_x_block_size);
+        }
+
         const size_t simd_size = 16;
         OPENVINO_ASSERT(config.head_size % simd_size == 0, "[GPU] Head size is expected to be divisible by 16");
 

src/plugins/intel_gpu/src/plugin/ops/custom.cpp

@@ -101,6 +101,19 @@ class CustomLayerAttributeVisitor : public ov::AttributeVisitor {
     std::map<std::string, std::string> m_values;
 };
 
+template <typename T>
+T convert_to(const std::string &str) {
+    std::istringstream ss(str);
+    T res;
+    ss >> res;
+    return res;
+}
+
+template <>
+std::string convert_to(const std::string &str) {
+    return str;
+}
+
 void CreatePagedAttention(ProgramBuilder& p, const std::shared_ptr<ov::Node>& op) {
     validate_inputs_count(op, {13});
     auto inputs = p.GetInputInfo(op);
@@ -126,6 +139,26 @@ void CreatePagedAttention(ProgramBuilder& p, const std::shared_ptr<ov::Node>& op
     prim.block_size = 16;
     prim.x_block_size = 8;
 
+    if (const auto env_var = std::getenv("PA_HEAD_SIZE")) {
+        prim.head_size = convert_to<size_t>(env_var);
+    }
+
+    if (const auto env_var = std::getenv("PA_HEADS_NUM")) {
+        prim.heads_num = convert_to<size_t>(env_var);
+    }
+
+    if (const auto env_var = std::getenv("PA_KV_HEADS_NUM")) {
+        prim.kv_heads_num = convert_to<size_t>(env_var);
+    }
+
+    if (const auto env_var = std::getenv("PA_BLOCK_SIZE")) {
+        prim.block_size = convert_to<size_t>(env_var);
+    }
+
+    if (const auto env_var = std::getenv("PA_X_BLOCK_SIZE")) {
+        prim.x_block_size = convert_to<size_t>(env_var);
+    }
+
     prim.num_outputs = op->get_output_size();
     prim.output_data_types = get_output_data_types(op);
     prim.output_paddings = get_output_paddings(op);