WIP: [GPU] PagedAttention OCL kernel kv-cache compression

Author: sshlyapn

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

@@ -553,6 +553,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         params.outputs[1] = value_cache_tensor;
 
         params.conf = get_sdpa_configuration(impl_param, is_dynamic);
+        if (ov::element::Type(impl_param.get_input_layout(3).data_type).size() == 1) {
+            params.conf.is_kv_compressed = true;
+            params.conf.use_asymmetric_quantization = true;
+        }
 
         params.is_prefill = stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED;
 
@@ -692,6 +696,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         params.inputs[input_idx++] = subsequence_begins_tensor;
 
         params.conf = get_sdpa_configuration(impl_param, is_dynamic);
+        if (ov::element::Type(impl_param.get_input_layout(3).data_type).size() == 1) {
+            params.conf.is_kv_compressed = true;
+            params.conf.use_asymmetric_quantization = true;
+        }
 
         if (has_scale_input)
             params.inputs[input_idx++] = scale_tensor;
@@ -779,28 +787,50 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
             input_tensors.emplace_back(convert_data_tensor(input_layout));
 
         const auto& desc = impl_param.typed_desc<paged_attention>();
-        auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
-        auto& kv_cache_update_kernel_selector = kv_cache_update_kernel_selector_t::Instance();
-        kernels_data.push_back(kv_cache_update_kernel_selector.get_best_kernel(kv_cache_update_kernel_params));
+        try {
+            auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+            auto& kv_cache_update_kernel_selector = kv_cache_update_kernel_selector_t::Instance();
+            kernels_data.push_back(kv_cache_update_kernel_selector.get_best_kernel(kv_cache_update_kernel_params));
+        } catch (std::exception& e) {
+            std::cout << "PagedAttention1 error: " << e.what() << "\n";
+            std::rethrow_exception(std::current_exception());
+        }
 
-        auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
-        auto& sdpa_kernel_selector = sdpa_kernel_selector_t::Instance();
-        kernels_data.push_back(sdpa_kernel_selector.get_best_kernel(sdpa_kernel_params));
+        try {
+            auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+            auto& sdpa_kernel_selector = sdpa_kernel_selector_t::Instance();
+            kernels_data.push_back(sdpa_kernel_selector.get_best_kernel(sdpa_kernel_params));
+        } catch (std::exception& e) {
+            std::cout << "PagedAttention2 error: " << e.what() << "\n";
+            std::rethrow_exception(std::current_exception());
+        }
 
-        auto pa_sdpa_kernel_params = get_pa_sdpa_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
-        auto& pa_sdpa_kernel_selector = pa_sdpa_kernel_selector_t::Instance();
-        kernels_data.push_back(pa_sdpa_kernel_selector.get_best_kernel(pa_sdpa_kernel_params));
+        try {
+            auto pa_sdpa_kernel_params = get_pa_sdpa_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+            auto& pa_sdpa_kernel_selector = pa_sdpa_kernel_selector_t::Instance();
+            kernels_data.push_back(pa_sdpa_kernel_selector.get_best_kernel(pa_sdpa_kernel_params));
+        } catch (std::exception& e) {
+            std::cout << "PagedAttention3 error: " << e.what() << "\n";
+            std::rethrow_exception(std::current_exception());
+        }
 
-        if (desc->has_rotated_blocks) {
-            auto kv_cache_rotate_kernel_params = get_kv_cache_rotate_kernel_params(impl_param, input_tensors, impl_param.is_dynamic());
-            auto& kv_cache_rotate_kernel_selector = kv_cache_rotate_kernel_selector_t::Instance();
-            kernels_data.push_back(kv_cache_rotate_kernel_selector.get_best_kernel(kv_cache_rotate_kernel_params));
+        try {
+            if (desc->has_rotated_blocks) {
+                auto kv_cache_rotate_kernel_params = get_kv_cache_rotate_kernel_params(impl_param, input_tensors, impl_param.is_dynamic());
+                auto& kv_cache_rotate_kernel_selector = kv_cache_rotate_kernel_selector_t::Instance();
+                kernels_data.push_back(kv_cache_rotate_kernel_selector.get_best_kernel(kv_cache_rotate_kernel_params));
+            }
+        } catch (std::exception& e) {
+            std::cout << "PagedAttention4 error: " << e.what() << "\n";
+            std::rethrow_exception(std::current_exception());
         }
 
+
         auto impl = std::make_unique<paged_attention_impl>(kernels_data);
         impl->has_scores_output = desc->has_scores_output();
         impl->has_rotated_blocks = desc->has_rotated_blocks;
 
+
         return impl;
     }
 

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

@@ -1947,6 +1947,9 @@ void primitive_inst::prepare_primitive() {
 void primitive_inst::execute() {
     GPU_DEBUG_PROFILED_STAGE(instrumentation::pipeline_stage::inference);
     if (get_flag(ExecutionFlags::SKIP)) {
+        if (_node->is_type<read_value>())
+            get_network().get_stream().finish();
+
         set_out_event(get_network().get_stream().aggregate_events(_impl_params->dep_events));
         return;
     }