Enable multi head size support for KV cache

In continue batching, head size for key and value will be different.
Add support for it for sdpa.

Author: clee30

src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp

@@ -35,7 +35,8 @@ struct paged_attention : public primitive_base<paged_attention> {
 
         auto rhs_casted = downcast<const paged_attention>(rhs);
 
-        return head_size == rhs_casted.head_size &&
+        return k_head_size == rhs_casted.k_head_size &&
+               v_head_size == rhs_casted.v_head_size &&
                heads_num == rhs_casted.heads_num &&
                kv_heads_num == rhs_casted.kv_heads_num &&
                sliding_window == rhs_casted.sliding_window &&
@@ -46,7 +47,8 @@ struct paged_attention : public primitive_base<paged_attention> {
 
     void save(BinaryOutputBuffer& ob) const override {
         primitive_base<paged_attention>::save(ob);
-        ob << head_size;
+        ob << k_head_size;
+        ob << v_head_size;
         ob << heads_num;
         ob << kv_heads_num;
         ob << has_alibi;
@@ -63,7 +65,8 @@ struct paged_attention : public primitive_base<paged_attention> {
 
     void load(BinaryInputBuffer& ib) override {
         primitive_base<paged_attention>::load(ib);
-        ib >> head_size;
+        ib >> k_head_size;
+        ib >> v_head_size;
         ib >> heads_num;
         ib >> kv_heads_num;
         ib >> has_alibi;
@@ -82,7 +85,8 @@ struct paged_attention : public primitive_base<paged_attention> {
     }
 
     std::optional<float> scale_val{};
-    size_t head_size = 0;
+    size_t k_head_size = 0;
+    size_t v_head_size = 0;
     size_t heads_num = 0;
     size_t kv_heads_num = 0;
     size_t sliding_window = 0;

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

@@ -640,7 +640,8 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         kernel_selector::sdpa_configuration config;
 
         const auto desc = impl_param.typed_desc<paged_attention>();
-        config.head_size = desc->head_size;
+        config.k_head_size = desc->k_head_size;
+        config.v_head_size = desc->v_head_size;
         config.heads_num = desc->heads_num;
         config.kv_heads_num = desc->kv_heads_num;
         config.has_alibi_input = desc->has_alibi;

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

@@ -266,7 +266,7 @@ struct scaled_dot_product_attention_impl : multi_stage_primitive<scaled_dot_prod
         }
 
         if (query_shape[query_shape.size() - 1].is_static())
-            config.head_size = query_shape[query_shape.size() - 1].get_length();
+            config.k_head_size = query_shape[query_shape.size() - 1].get_length();
 
         config.is_causal = desc->is_causal;
 

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

@@ -21,14 +21,21 @@ layout paged_attention_inst::calc_output_layout(const paged_attention_node& /*no
 
 template<typename ShapeType>
 std::vector<layout> paged_attention_inst::calc_output_layouts(paged_attention_node const& /*node*/, kernel_impl_params const& impl_param) {
-    auto data_layout = impl_param.get_input_layout(0);
+    auto q_layout = impl_param.get_input_layout(0);
+    auto v_layout = impl_param.get_input_layout(2);
+    auto data_layout = q_layout;
+
+    if (v_layout.is_static()) {
+        ShapeType v_shape = v_layout.get_shape();
+        data_layout = layout{v_shape, data_layout.data_type, data_layout.format};
+    }
+
     data_layout.data_padding = padding();
 
     const auto& key_cache_ps = impl_param.get_input_layout(3).get_partial_shape();
     bool valid_block_size = key_cache_ps[3].is_dynamic() || key_cache_ps[3].get_length() == paged_attention::block_size;
     OPENVINO_ASSERT(valid_block_size, "[GPU] Incorrect block size for Paged Attention operation. "
                                       "Expected ", paged_attention::block_size, ", but got ", key_cache_ps[3].get_length());
-
     std::vector<layout> output_layouts{ data_layout };
 
     const auto& desc = impl_param.typed_desc<paged_attention>();
@@ -67,7 +74,8 @@ std::string paged_attention_inst::to_string(const paged_attention_node& node) {
 
     json_composite paged_attention_info;
     paged_attention_info.add("paged_attention_block_size", desc->block_size);
-    paged_attention_info.add("head_size", desc->head_size);
+    paged_attention_info.add("k_head_size", desc->k_head_size);
+    paged_attention_info.add("v_head_size", desc->v_head_size);
     paged_attention_info.add("heads_num", desc->heads_num);
     paged_attention_info.add("kv_heads_num", desc->kv_heads_num);
     paged_attention_info.add("scale", desc->scale_val.value_or(1.0f));
@@ -85,7 +93,8 @@ paged_attention_inst::typed_primitive_inst(network& network, const paged_attenti
     : parent(network, node) {
     const auto desc = node.get_primitive();
 
-    const auto head_size = desc->head_size;
+    const auto k_head_size = desc->k_head_size;
+    const auto v_head_size = desc->v_head_size;
     const auto heads_num = desc->heads_num;
     const auto kv_heads_num = desc->kv_heads_num;
     const auto pa_block_size = desc->block_size;
@@ -97,6 +106,7 @@ paged_attention_inst::typed_primitive_inst(network& network, const paged_attenti
     }
 
     OPENVINO_ASSERT(heads_num % kv_heads_num == 0);
-    OPENVINO_ASSERT(head_size % pa_block_size == 0);
+    OPENVINO_ASSERT(k_head_size % pa_block_size == 0);
+    OPENVINO_ASSERT(v_head_size % pa_block_size == 0);
 }
 }  // namespace cldnn

src/plugins/intel_gpu/src/graph/primitive_inst.cpp

@@ -404,7 +404,7 @@ void primitive_inst::update_shape() {
         input_shape_changed = true;
     }
 
-    if (!_node->is_type<kv_cache>() && !input_shape_changed && _impl_params->get_output_layout().is_static())
+    if (!_node->is_type<kv_cache>() && !_node->is_type<strided_slice>() && !input_shape_changed && _impl_params->get_output_layout().is_static())
         return;
 
     std::vector<event::ptr> dependencies_events;
@@ -456,7 +456,6 @@ void primitive_inst::update_shape() {
 
     _impl_params->memory_deps = memory_deps;
 
-
     auto new_layouts = _node->type()->calc_output_layouts(*_node, *_impl_params);
     for (size_t idx = 0; idx != new_layouts.size(); ++idx) {
         auto& new_layout = new_layouts[idx];