Repository: openvinotoolkit/openvino
Pull request #30394: [GPU] Add PA scores window calculation support
Author: login: sshlyapn
Labels: name: category: GPU
description: OpenVINO GPU plugin
name: Code Freeze
Created: 2025-04-30T15:58:30Z | Merged: 2025-05-20T08:05:19Z

### Details:
 - Add scores window calculation support

### Tickets:
 - *ticket-id*



Diff:
diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
index 68754ea5a15cbb..dcbb0b04de30ba 100644
--- a/src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
+++ b/src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
@@ -13,6 +13,26 @@ namespace cldnn {
 struct paged_attention : public primitive_base<paged_attention> {
     CLDNN_DECLARE_PRIMITIVE(paged_attention)
 
+    enum PagedAttentionInputIdx {
+        QUERY = 0,
+        KEY = 1,
+        VALUE = 2,
+        KEY_CACHE = 3,
+        VALUE_CACHE = 4,
+        PAST_LENS = 5,
+        SUBSEQUENCE_BEGINS = 6,
+        BLOCK_INDICES = 7,
+        BLOCK_INDICES_BEGINS = 8,
+        SCALE = 9,
+        SLIDING_WINDOW = 10,
+        ALIBI = 11,
+        MAX_CONTEXT_LEN = 12,
+        SCORE_AGGREGATION = 13,
+        ROTATED_BLOCK_INDICES = 14,
+        ROTATION_DELTAS = 15,
+        ROTATION_TRIG_LUT = 16,
+    };
+
     static constexpr size_t block_size = 16;
 
     paged_attention() : primitive_base("", {}) {}
@@ -56,6 +76,7 @@ struct paged_attention : public primitive_base<paged_attention> {
         ob << has_score_aggregation;
         ob << has_rotated_blocks;
         ob << sliding_window;
+        ob << has_score_aggregation;
 
         if (scale_val.has_value()) {
             ob << true;
@@ -75,6 +96,7 @@ struct paged_attention : public primitive_base<paged_attention> {
         ib >> has_score_aggregation;
         ib >> has_rotated_blocks;
         ib >> sliding_window;
+        ib >> has_score_aggregation;
 
         bool has_scale;
         ib >> has_scale;
@@ -94,7 +116,7 @@ struct paged_attention : public primitive_base<paged_attention> {
     size_t kv_heads_num = 0;
     size_t sliding_window = 0;
     bool has_alibi = false;
-    bool has_score_aggregation = false;
     bool has_rotated_blocks = false;
+    bool has_score_aggregation = false;
 };
 }  // namespace cldnn
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
index eee7dd6dd68ad4..17fb6af7ea1d44 100644
--- a/src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
+++ b/src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
@@ -60,17 +60,16 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
     };
 
     PagedAttentionStage get_paged_attention_stage(const kernel_impl_params& impl_param) const {
-        const auto& query_shape = impl_param.get_input_layout(0).get_partial_shape();
-        const auto& past_lens_shape = impl_param.get_input_layout(5).get_partial_shape();
+        const auto& query_shape = impl_param.get_input_layout(PagedAttentionInputIdx::QUERY).get_partial_shape();
+        const auto& past_lens_shape = impl_param.get_input_layout(PagedAttentionInputIdx::PAST_LENS).get_partial_shape();
 
         if (query_shape.is_static() && past_lens_shape.is_static()) {
             if (query_shape[0].get_length() == past_lens_shape[0].get_length()) {
                 return PagedAttentionStage::GENERATE;
             }
 
-            const auto past_lens_idx = 5;
             const auto& memory_deps = impl_param.memory_deps;
-            const auto past_lens_mem = memory_deps.at(past_lens_idx);
+            const auto past_lens_mem = memory_deps.at(PagedAttentionInputIdx::PAST_LENS);
             mem_lock<int32_t, mem_lock_type::read> past_lens_mem_lock(past_lens_mem, *impl_param.strm);
 
             const auto past_lens_size = past_lens_mem_lock.size();
@@ -146,26 +145,33 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
     std::vector<BufferDescriptor> get_internal_buffer_descs(const kernel_impl_params&) const override {
         /*
-        * Internal buffers allocation owners and users:
-        * +--------------------------------------+--------------------+--------------------+
-        * | Stage                                | Allocates & uses   | Reuses             |
-        * +--------------------------------------+--------------------+--------------------+
-        * | KV_CACHE_UPDATE                      | [0, 1, 2]          |                    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | SDPA (1st token)                     |                    | [0, 1, 2]          |
-        * +--------------------------------------+--------------------+--------------------+
-        * | PA_SDPA (2nd+ token)                 | [5, 6, 7]          |                    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | PA_SDPA (mixed mode)                 | [5, 6, 7, 8]       |                    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | SDPA (1st token) + scores output     |                    | [0, 1, 2, 3, 4]    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | PA_SDPA (2nd+ token) + scores output | [3, 4, 5, 6, 7]    |                    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | PA_SDPA (mixed mode) + scores output | [3, 4, 5, 6, 7, 8] |                    |
-        * +--------------------------------------+--------------------+--------------------+
-        * | SDPA (1st token, micro-kernel)       | [last (8/9)]       |                    |
-        * +--------------------------------------+--------------------+--------------------+
+        * Internal buffers allocation owners and users (numbers represent unique buffers names,
+        * not the real indexes in _intermediates_memory structure):
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | Stage                                            | Allocates & uses      | Reuses             |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | KV_CACHE_UPDATE                                  | [0, 1, 2]             |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | SDPA (1st token)                                 |                       | [0, 1, 2]          |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (2nd+ token)                             | [5, 6, 7]             |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (mixed mode)                             | [5, 6, 7, 8]          |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | SDPA (1st token) + scores output                 |                       | [0, 1, 2, 3, 4]    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (2nd+ token) + scores output             | [3, 4, 5, 6, 7]       |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (mixed mode) + scores output             | [3, 4, 5, 6, 7, 8]    |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | SDPA (1st token) + scores output aggregation     |                       | [0, 1, 2, 3, 4, 5] |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (2nd+ token) + scores output aggregation | [3, 4, 5, 6, 7, 8]    |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | PA_SDPA (mixed mode) + scores output aggregation | [3, 4, 5, 6, 7, 8, 9] |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
+        * | SDPA (1st token, micro-kernel)                   | [last (8/9/10)]       |                    |
+        * +--------------------------------------------------+-----------------------+--------------------+
         *
         * Description:
         * 0, 1, 2 - Buffers used for proper blocks distribution for kv_cache_update and
@@ -225,7 +231,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
             args.outputs = { instance.key_cache_memory_ptr(),
                              instance.value_cache_memory_ptr() };
         } else if (stage == Stage::SDPA) {
-            args.inputs = {  instance.input_memory_ptr(0),
+            args.inputs = {  instance.query_memory_ptr(),
                              instance.key_memory_ptr(),
                              instance.value_memory_ptr(),
                              instance.subsequence_begins_memory_ptr() };
@@ -244,7 +250,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
                 // 2nd+ token calculation or mixed stage tokens calculation
                 args.shape_info = instance.shape_info_memory_ptr();
 
-                args.inputs = { instance.input_memory_ptr(0),
+                args.inputs = { instance.query_memory_ptr(),
                                 instance.key_cache_memory_ptr(),
                                 instance.value_cache_memory_ptr(),
                                 instance.past_lens_memory_ptr(),
@@ -297,6 +303,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
     void prepare_internal_buffers(paged_attention_inst& instance, const PagedAttentionStage& stage) {
         const auto& desc = instance.get_impl_params()->typed_desc<paged_attention>();
         const bool has_scores_output = desc->has_scores_output();
+        const bool has_score_aggregation = desc->has_score_aggregation;
 
         if ((stage == PagedAttentionStage::UNKNOWN) ||
             (stage == PagedAttentionStage::GENERATE && !has_scores_output))
@@ -309,6 +316,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         mem_lock<int32_t, mem_lock_type::read> past_lens_mem_lock(past_lens_mem, stream);
         mem_lock<int32_t, mem_lock_type::read> subsequence_begins_mem_lock(subsequence_begins_mem, stream);
         std::unique_ptr<mem_lock<int32_t, mem_lock_type::write>> subsequence_offsets_lock = nullptr;
+        std::unique_ptr<mem_lock<int32_t, mem_lock_type::write>> cumulative_window_size_sum_lock = nullptr;
 
         if (has_scores_output) {
             const size_t subsequence_offsets_idx = 4;
@@ -318,6 +326,28 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
             auto subsequence_offsets_mem = intermediates_memories[subsequence_offsets_idx];
             subsequence_offsets_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(subsequence_offsets_mem, stream));
+
+            if (has_score_aggregation) {
+                const size_t cumulative_window_size_sum_idx = 5;
+                OPENVINO_ASSERT(intermediates_memories.size() > cumulative_window_size_sum_idx,
+                                "[GPU] Unexpected number of intermediates buffers for Paged Attention for scores aggregation");
+
+                auto cumulative_window_size_sum_mem = intermediates_memories[cumulative_window_size_sum_idx];
+                cumulative_window_size_sum_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(cumulative_window_size_sum_mem, stream));
+
+                mem_lock<int32_t, mem_lock_type::read> score_aggregation_mem_lock(instance.score_aggregation_memory_ptr(), stream);
+
+                // Transform window sizes to cumulative buffer for offsets precalculation
+                // For example:
+                // Original score aggregation buffer content: {4, 2, 1, 1}
+                // Cumulative window size sum buffer content: {0, 4, 6, 7, 8}
+                size_t cumulative_sum = 0;
+                cumulative_window_size_sum_lock->operator[](0) = static_cast<int32_t>(cumulative_sum);
+                for (size_t i = 0; i < score_aggregation_mem_lock.size(); i++) {
+                    cumulative_sum += score_aggregation_mem_lock[i];
+                    cumulative_window_size_sum_lock->operator[](i + 1) = static_cast<int32_t>(cumulative_sum);
+                }
+            }
         }
 
         if (stage == PagedAttentionStage::GENERATE) {
@@ -358,7 +388,12 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         std::unique_ptr<mem_lock<int32_t, mem_lock_type::write>> micro_sdpa_block_starts_and_gws_mapping_lock = nullptr;
 
         if (stage == PagedAttentionStage::MIXED) {
-            const size_t sequential_gws_subseq_mapping_idx = has_scores_output ? 8 : 6;
+            size_t sequential_gws_subseq_mapping_idx = 6;
+            if (has_score_aggregation) {
+                sequential_gws_subseq_mapping_idx = 9;
+            } else if (has_scores_output) {
+                sequential_gws_subseq_mapping_idx = 8;
+            }
 
             OPENVINO_ASSERT(intermediates_memories.size() > sequential_gws_subseq_mapping_idx,
                             "[GPU] Unexpected number of intermediates buffers for Paged Attention for mixed stage");
@@ -463,6 +498,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
                 // Add intermediate buffers for PagedAttention scores calculation:
                 // softmax_results, subsequence_offsets, exp_sums, max_logits, tmp_out
                 internal_buffers_count += 5;
+
+                // Add cumulative_score_aggregation_sum buffer
+                if (desc->has_score_aggregation)
+                   internal_buffers_count += 1;
             }
         }
 
@@ -571,14 +610,12 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         auto calculate_aligned_seq_len = [&]() {
             const auto& input_mem = impl_param.memory_deps;
-            const auto subsequence_begins_input_idx = 6;
-            const auto subsequence_begins_mem = input_mem.at(subsequence_begins_input_idx);
+            const auto subsequence_begins_mem = input_mem.at(PagedAttentionInputIdx::SUBSEQUENCE_BEGINS);
             mem_lock<int32_t, mem_lock_type::read> subsequence_begins_mem_lock(subsequence_begins_mem, *impl_param.strm);
 
             auto aligned_seq_len = 0;
             if (stage == PagedAttentionStage::MIXED) {
-                const auto past_lens_idx = 5;
-                const auto past_lens_mem = input_mem.at(past_lens_idx);
+                const auto past_lens_mem = input_mem.at(PagedAttentionInputIdx::PAST_LENS);
                 mem_lock<int32_t, mem_lock_type::read> past_lens_mem_lock(past_lens_mem, *impl_param.strm);
 
                 for (size_t i = 0; i < subsequence_begins_mem_lock.size() - 1; i++) {
@@ -624,8 +661,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         if (stage == PagedAttentionStage::PREFILL) {
             const auto desc = impl_param.typed_desc<paged_attention>();
             if (static_cast<int64_t>(paged_attention::block_size) == target_seq_len_block_size) {
-                const auto block_indices_input_idx = 7;
-                const auto& block_indices_ps = impl_param.get_input_layout(block_indices_input_idx).get_partial_shape();
+                const auto& block_indices_ps = impl_param.get_input_layout(PagedAttentionInputIdx::BLOCK_INDICES).get_partial_shape();
 
                 aligned_seq_len = block_indices_ps[0].get_length() * target_seq_len_block_size;
             } else {
@@ -651,6 +687,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         config.is_paged_attention = true;
         config.paged_attention_block_size = static_cast<int64_t>(paged_attention::block_size);
         config.paged_attention_sliding_window = desc->sliding_window;
+        config.has_score_aggregation = desc->has_score_aggregation;
 
         if (desc->scale_val.has_value()) {
             config.has_const_scale_val = true;
@@ -669,12 +706,23 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         if (desc->has_scores_output() && !is_dynamic) {
             const auto& input_mem = impl_param.memory_deps;
-            const auto max_context_len = input_mem.at(12);
+            const auto max_context_len = input_mem.at(PagedAttentionInputIdx::MAX_CONTEXT_LEN);
             mem_lock<int32_t, mem_lock_type::read> max_context_len_mem_lock(max_context_len, *impl_param.strm);
             config.paged_attention_max_len = max_context_len_mem_lock[0];
+
+            if (desc->has_score_aggregation) {
+                const auto score_aggregation = input_mem.at(PagedAttentionInputIdx::SCORE_AGGREGATION);
+                mem_lock<int32_t, mem_lock_type::read> score_aggregation_mem_lock(score_aggregation, *impl_param.strm);
+
+                auto total_tokens_num = 0;
+                for (size_t i = 0; i < score_aggregation_mem_lock.size(); i++) {
+                    total_tokens_num += score_aggregation_mem_lock[i];
+                }
+                config.paged_attention_snap_kv_tokens = total_tokens_num;
+            }
         }
 
-        if (data_type_traits::is_i8_u8(impl_param.get_input_layout(3).data_type)) {
+        if (data_type_traits::is_i8_u8(impl_param.get_input_layout(PagedAttentionInputIdx::KEY_CACHE).data_type)) {
             config.is_kv_compressed = true;
             config.use_asymmetric_quantization = true;
         }
@@ -687,10 +735,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
                                                                              bool is_dynamic = false) {
         auto params = get_default_params<kv_cache_rotate_kernel_params_t>(impl_param, is_dynamic);
 
-        const auto& key_cache_tensor = input_tensors[3];
-        const auto& rotated_block_indices_tensor = input_tensors[14];
-        const auto& rotation_deltas_tensor = input_tensors[15];
-        const auto& rotation_trig_lut_tensor = input_tensors[16];
+        const auto& key_cache_tensor = input_tensors[PagedAttentionInputIdx::KEY_CACHE];
+        const auto& rotated_block_indices_tensor = input_tensors[PagedAttentionInputIdx::ROTATED_BLOCK_INDICES];
+        const auto& rotation_deltas_tensor = input_tensors[PagedAttentionInputIdx::ROTATION_DELTAS];
+        const auto& rotation_trig_lut_tensor = input_tensors[PagedAttentionInputIdx::ROTATION_TRIG_LUT];
 
         const auto inputs_number = 3;
         const auto outputs_number = 1;
@@ -701,17 +749,17 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         params.inputs[2] = rotation_trig_lut_tensor;
         params.outputs[0] = key_cache_tensor;
 
-        params.original_cache_dt = to_data_type(impl_param.get_input_layout(1).data_type);
+        params.original_cache_dt = to_data_type(impl_param.get_input_layout(PagedAttentionInputIdx::KEY).data_type);
         params.conf = get_sdpa_configuration(impl_param, is_dynamic);
 
         const auto& in_offsets_map = impl_param.in_port_to_shape_info_offset;
         std::map<size_t, size_t> in_tensor_to_offset_map = {
-            {0, in_offsets_map.at(14)},
-            {1, in_offsets_map.at(15)},
-            {2, in_offsets_map.at(16)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::ROTATED_BLOCK_INDICES)},
+            {1, in_offsets_map.at(PagedAttentionInputIdx::ROTATION_DELTAS)},
+            {2, in_offsets_map.at(PagedAttentionInputIdx::ROTATION_TRIG_LUT)},
         };
         std::map<size_t, size_t> out_tensor_to_offset_map = {
-            {0, in_offsets_map.at(3)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::KEY_CACHE)},
         };
 
         params.set_dynamic_shape_offsets(in_tensor_to_offset_map, out_tensor_to_offset_map);
@@ -725,14 +773,14 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
                                                                              bool is_dynamic = false) {
         auto params = get_default_params<kv_cache_update_kernel_params_t>(impl_param, is_dynamic);
 
-        const auto& key_tensor = input_tensors[1];
-        const auto& value_tensor = input_tensors[2];
-        const auto& key_cache_tensor = input_tensors[3];
-        const auto& value_cache_tensor = input_tensors[4];
-        const auto& past_lens_tensor = input_tensors[5];
-        const auto& block_indices_tensor = input_tensors[7];
-        const auto& block_indices_begins_tensor = input_tensors[8];
-        const auto& subsequence_begins_tensor = input_tensors[6];
+        const auto& key_tensor = input_tensors[PagedAttentionInputIdx::KEY];
+        const auto& value_tensor = input_tensors[PagedAttentionInputIdx::VALUE];
+        const auto& key_cache_tensor = input_tensors[PagedAttentionInputIdx::KEY_CACHE];
+        const auto& value_cache_tensor = input_tensors[PagedAttentionInputIdx::VALUE_CACHE];
+        const auto& past_lens_tensor = input_tensors[PagedAttentionInputIdx::PAST_LENS];
+        const auto& block_indices_tensor = input_tensors[PagedAttentionInputIdx::BLOCK_INDICES];
+        const auto& block_indices_begins_tensor = input_tensors[PagedAttentionInputIdx::BLOCK_INDICES_BEGINS];
+        const auto& subsequence_begins_tensor = input_tensors[PagedAttentionInputIdx::SUBSEQUENCE_BEGINS];
 
         const auto inputs_number = 6;
         const auto outputs_number = 2;
@@ -756,16 +804,16 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         const auto& in_offsets_map = impl_param.in_port_to_shape_info_offset;
         std::map<size_t, size_t> in_tensor_to_offset_map = {
-            {0, in_offsets_map.at(1)},
-            {1, in_offsets_map.at(2)},
-            {2, in_offsets_map.at(5)},
-            {3, in_offsets_map.at(7)},
-            {4, in_offsets_map.at(8)},
-            {5, in_offsets_map.at(6)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::KEY)},
+            {1, in_offsets_map.at(PagedAttentionInputIdx::VALUE)},
+            {2, in_offsets_map.at(PagedAttentionInputIdx::PAST_LENS)},
+            {3, in_offsets_map.at(PagedAttentionInputIdx::BLOCK_INDICES)},
+            {4, in_offsets_map.at(PagedAttentionInputIdx::BLOCK_INDICES_BEGINS)},
+            {5, in_offsets_map.at(PagedAttentionInputIdx::SUBSEQUENCE_BEGINS)},
         };
         std::map<size_t, size_t> out_tensor_to_offset_map = {
-            {0, in_offsets_map.at(3)},
-            {1, in_offsets_map.at(4)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::KEY_CACHE)},
+            {1, in_offsets_map.at(PagedAttentionInputIdx::VALUE_CACHE)},
         };
 
         params.set_dynamic_shape_offsets(in_tensor_to_offset_map, out_tensor_to_offset_map);
@@ -781,14 +829,14 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         const auto desc = impl_param.typed_desc<paged_attention>();
         auto params = get_default_params<sdpa_kernel_params_t>(impl_param, is_dynamic);
 
-        const auto& query_tensor = input_tensors[0];
-        const auto& key_tensor = input_tensors[1];
-        const auto& value_tensor = input_tensors[2];
-        const auto& subsequence_begins_tensor = input_tensors[6];
-        const auto& scale_tensor = input_tensors[9];
-        const auto& alibi_tensor = input_tensors[11];
+        const auto& query_tensor = input_tensors[PagedAttentionInputIdx::QUERY];
+        const auto& key_tensor = input_tensors[PagedAttentionInputIdx::KEY];
+        const auto& value_tensor = input_tensors[PagedAttentionInputIdx::VALUE];
+        const auto& subsequence_begins_tensor = input_tensors[PagedAttentionInputIdx::SUBSEQUENCE_BEGINS];
+        const auto& scale_tensor = input_tensors[PagedAttentionInputIdx::SCALE];
+        const auto& alibi_tensor = input_tensors[PagedAttentionInputIdx::ALIBI];
 
-        const auto has_alibi = impl_param.get_input_layout(11).count() > 0;
+        const auto has_alibi = impl_param.get_input_layout(PagedAttentionInputIdx::ALIBI).count() > 0;
         const auto has_scale_input = !desc->scale_val.has_value();
         const auto has_scores_output = desc->has_scores_output();
 
@@ -833,10 +881,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         const auto& in_offsets_map = impl_param.in_port_to_shape_info_offset;
         const auto& out_offsets_map = impl_param.out_port_to_shape_info_offset;
         std::map<size_t, size_t> in_tensor_to_offset_map = {
-            {0, in_offsets_map.at(0)},
-            {1, in_offsets_map.at(1)},
-            {2, in_offsets_map.at(2)},
-            {3, in_offsets_map.at(6)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::QUERY)},
+            {1, in_offsets_map.at(PagedAttentionInputIdx::KEY)},
+            {2, in_offsets_map.at(PagedAttentionInputIdx::VALUE)},
+            {3, in_offsets_map.at(PagedAttentionInputIdx::SUBSEQUENCE_BEGINS)},
         };
         std::map<size_t, size_t> out_tensor_to_offset_map = {
             {0, out_offsets_map.at(0)},
@@ -844,10 +892,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         input_idx = 4;
         if (has_scale_input)
-            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(9)});
+            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(PagedAttentionInputIdx::SCALE)});
 
         if (has_alibi)
-            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(11)});
+            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(PagedAttentionInputIdx::ALIBI)});
 
         if ((stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED) && !is_dynamic)
             params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage, query_block_size);
@@ -869,17 +917,17 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         const auto desc = impl_param.typed_desc<paged_attention>();
         auto params = get_default_params<pa_sdpa_kernel_params_t>(impl_param, is_dynamic);
 
-        const auto& query_tensor = input_tensors[0];
-        const auto& key_cache_tensor = input_tensors[3];
-        const auto& value_cache_tensor = input_tensors[4];
-        const auto& past_lens_tensor = input_tensors[5];
-        const auto& block_indices_tensor = input_tensors[7];
-        const auto& block_indices_begins_tensor = input_tensors[8];
-        const auto& subsequence_begins_tensor = input_tensors[6];
-        const auto& scale_tensor = input_tensors[9];
-        const auto& alibi_tensor = input_tensors[11];
-
-        const auto has_alibi = impl_param.get_input_layout(11).count() > 0;
+        const auto& query_tensor = input_tensors[PagedAttentionInputIdx::QUERY];
+        const auto& key_cache_tensor = input_tensors[PagedAttentionInputIdx::KEY_CACHE];
+        const auto& value_cache_tensor = input_tensors[PagedAttentionInputIdx::VALUE_CACHE];
+        const auto& past_lens_tensor = input_tensors[PagedAttentionInputIdx::PAST_LENS];
+        const auto& block_indices_tensor = input_tensors[PagedAttentionInputIdx::BLOCK_INDICES];
+        const auto& block_indices_begins_tensor = input_tensors[PagedAttentionInputIdx::BLOCK_INDICES_BEGINS];
+        const auto& subsequence_begins_tensor = input_tensors[PagedAttentionInputIdx::SUBSEQUENCE_BEGINS];
+        const auto& scale_tensor = input_tensors[PagedAttentionInputIdx::SCALE];
+        const auto& alibi_tensor = input_tensors[PagedAttentionInputIdx::ALIBI];
+
+        const auto has_alibi = impl_param.get_input_layout(PagedAttentionInputIdx::ALIBI).count() > 0;
         const auto has_scale_input = !desc->scale_val.has_value();
         const auto has_scores_output = desc->has_scores_output();
 
@@ -926,13 +974,13 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         const auto& out_offsets_map = impl_param.out_port_to_shape_info_offset;
 
         std::map<size_t, size_t> in_tensor_to_offset_map = {
-            {0, in_offsets_map.at(0)},
-            {1, in_offsets_map.at(3)},
-            {2, in_offsets_map.at(4)},
-            {3, in_offsets_map.at(5)},
-            {4, in_offsets_map.at(7)},
-            {5, in_offsets_map.at(8)},
-            {6, in_offsets_map.at(6)},
+            {0, in_offsets_map.at(PagedAttentionInputIdx::QUERY)},
+            {1, in_offsets_map.at(PagedAttentionInputIdx::KEY_CACHE)},
+            {2, in_offsets_map.at(PagedAttentionInputIdx::VALUE_CACHE)},
+            {3, in_offsets_map.at(PagedAttentionInputIdx::PAST_LENS)},
+            {4, in_offsets_map.at(PagedAttentionInputIdx::BLOCK_INDICES)},
+            {5, in_offsets_map.at(PagedAttentionInputIdx::BLOCK_INDICES_BEGINS)},
+            {6, in_offsets_map.at(PagedAttentionInputIdx::SUBSEQUENCE_BEGINS)},
         };
         std::map<size_t, size_t> out_tensor_to_offset_map = {
             {0, out_offsets_map.at(0)},
@@ -940,10 +988,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
         input_idx = 7;
         if (has_scale_input)
-            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(9)});
+            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(PagedAttentionInputIdx::SCALE)});
 
         if (has_alibi)
-            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(11)});
+            in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(PagedAttentionInputIdx::ALIBI)});
 
         if (has_scores_output)
             out_tensor_to_offset_map.insert({1, out_offsets_map.at(1)});
diff --git a/src/plugins/intel_gpu/src/graph/include/paged_attention_inst.h b/src/plugins/intel_gpu/src/graph/include/paged_attention_inst.h
index ec86156c1b5825..af3b1da301cb57 100644
--- a/src/plugins/intel_gpu/src/graph/include/paged_attention_inst.h
+++ b/src/plugins/intel_gpu/src/graph/include/paged_attention_inst.h
@@ -9,6 +9,8 @@
 
 namespace cldnn {
 
+using PagedAttentionInputIdx = cldnn::paged_attention::PagedAttentionInputIdx;
+
 template <>
 struct typed_program_node<paged_attention> : public typed_program_node_base<paged_attention> {
 private:
@@ -18,11 +20,21 @@ struct typed_program_node<paged_attention> : public typed_program_node_base<page
     using parent::parent;
 
     std::set<size_t> get_lockable_input_ids() const override {
-        return { 5 /* past_lens */, 6 /* subsequence_begins */, 12 /* max_context_len */ };
+        std::set<size_t> input_ports = { PagedAttentionInputIdx::PAST_LENS,
+                                         PagedAttentionInputIdx::SUBSEQUENCE_BEGINS,
+                                         PagedAttentionInputIdx::MAX_CONTEXT_LEN };
+
+        if (typed_desc()->has_score_aggregation)
+            input_ports.insert(PagedAttentionInputIdx::SCORE_AGGREGATION);
+
+        return input_ports;
     }
 
     std::vector<size_t> get_shape_infer_dependencies() const override {
-        return { 5 /* past_lens */, 6 /* subsequence_begins */, 12 /* max_context_len */ };
+        const auto lockable_input_ids = get_lockable_input_ids();
+        std::vector<size_t> input_ports(lockable_input_ids.begin(), lockable_input_ids.end());
+
+        return input_ports;
     }
 };
 
@@ -42,20 +54,20 @@ class typed_primitive_inst<paged_attention> : public typed_primitive_inst_base<p
     typed_primitive_inst(network& network, const paged_attention_node& desc);
     typed_primitive_inst(network& network) : parent(network) {}
 
-    memory::ptr query_memory_ptr() const { return input_memory_ptr(0); }
-    memory::ptr key_memory_ptr() const { return input_memory_ptr(1); }
-    memory::ptr value_memory_ptr() const { return input_memory_ptr(2); }
-    memory::ptr key_cache_memory_ptr() const { return input_memory_ptr(3); }
-    memory::ptr value_cache_memory_ptr() const { return input_memory_ptr(4); }
-    memory::ptr past_lens_memory_ptr() const { return input_memory_ptr(5); }
-    memory::ptr subsequence_begins_memory_ptr() const { return input_memory_ptr(6); }
-    memory::ptr block_indices_memory_ptr() const { return input_memory_ptr(7); }
-    memory::ptr block_indices_begins_memory_ptr() const { return input_memory_ptr(8); }
-    memory::ptr alibi_memory_ptr() const { return input_memory_ptr(11); }
-    memory::ptr score_aggregation_window_ptr() const { return input_memory_ptr(13); }
-    memory::ptr rotated_block_indices_memory_ptr() const { return input_memory_ptr(14); }
-    memory::ptr rotation_deltas_memory_ptr() const { return input_memory_ptr(15); }
-    memory::ptr rotation_trig_lut_memory_ptr() const { return input_memory_ptr(16); }
+    memory::ptr query_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::QUERY); }
+    memory::ptr key_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::KEY); }
+    memory::ptr value_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::VALUE); }
+    memory::ptr key_cache_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::KEY_CACHE); }
+    memory::ptr value_cache_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::VALUE_CACHE); }
+    memory::ptr past_lens_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::PAST_LENS); }
+    memory::ptr subsequence_begins_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::SUBSEQUENCE_BEGINS); }
+    memory::ptr block_indices_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::BLOCK_INDICES); }
+    memory::ptr block_indices_begins_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::BLOCK_INDICES_BEGINS); }
+    memory::ptr alibi_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::ALIBI); }
+    memory::ptr score_aggregation_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::SCORE_AGGREGATION); }
+    memory::ptr rotated_block_indices_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::ROTATED_BLOCK_INDICES); }
+    memory::ptr rotation_deltas_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::ROTATION_DELTAS); }
+    memory::ptr rotation_trig_lut_memory_ptr() const { return input_memory_ptr(PagedAttentionInputIdx::ROTATION_TRIG_LUT); }
 };
 
 using paged_attention_inst = typed_primitive_inst<paged_attention>;
diff --git a/src/plugins/intel_gpu/src/graph/paged_attention.cpp b/src/plugins/intel_gpu/src/graph/paged_attention.cpp
index 8963e579d7ab20..94ea3023aefec9 100644
--- a/src/plugins/intel_gpu/src/graph/paged_attention.cpp
+++ b/src/plugins/intel_gpu/src/graph/paged_attention.cpp
@@ -85,6 +85,8 @@ std::string paged_attention_inst::to_string(const paged_attention_node& node) {
     paged_attention_info.add("has_rotated_blocks", desc->has_rotated_blocks);
     paged_attention_info.add("key_cache_dt", node.get_input_layout(3).data_type);
     paged_attention_info.add("value_cache_dt", node.get_input_layout(4).data_type);
+    paged_attention_info.add("score_output", desc->has_scores_output());
+    paged_attention_info.add("score_aggregation", desc->has_score_aggregation);
     node_info->add("paged_attention primitive info", paged_attention_info);
     node_info->dump(primitive_description);
 
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
index c7ccb424ba0ea3..331287f5b63fe1 100644
--- 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
@@ -67,6 +67,9 @@ KERNEL(pa_sdpa_opt)(
 #if PAGED_ATTENTION_SCORES_OUTPUT
     __global SOFTMAX_ACCUMULATOR_TYPE* softmax_results,
     const __global int* subsequence_offsets,
+#if HAS_SCORE_AGGREGATION
+    const __global int* cumulative_score_aggregation_sum,
+#endif
 #endif
     __global SOFTMAX_ACCUMULATOR_TYPE* exp_sums,
     __global SOFTMAX_ACCUMULATOR_TYPE* max_logits,
@@ -404,9 +407,21 @@ KERNEL(pa_sdpa_opt)(
             const uint subsequence_idx = gws_subseq_mapping[seq_idx];
             const uint subsequence_start_pos = subsequence_begins[subsequence_idx];
             const uint subsequence_end_pos = subsequence_begins[subsequence_idx + 1];
+
+#if HAS_SCORE_AGGREGATION
+            const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+            const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+            const int global_first_row_idx = max((int)subsequence_end_pos - window_size, 0);
+            const int local_row_idx = seq_idx - global_first_row_idx;
+            const bool save_softmax_results = local_row_idx >= 0;
+#else
+            const int scores_offset = subsequence_idx;
+            const int local_row_idx = 0;
             const bool save_softmax_results = seq_idx == subsequence_end_pos - 1;
+#endif
 #else
-            const uint subsequence_idx = seq_idx;
+            const uint scores_offset = seq_idx;
+            const int local_row_idx = 0;
             const bool save_softmax_results = true;
 #endif // MULTI_TOKENS_PROCESSING
             // PagedAttention is supposed to save only last "row" of the QK matrix multiplication,
@@ -417,11 +432,17 @@ KERNEL(pa_sdpa_opt)(
                     if (q_idx >= iter_heads_num)
                         break;
 #endif
-                    const uint output_offset = subsequence_idx * HEADS_NUM * total_partitions_num * SEQ_LEN_PARTITION_SIZE +
+                    const uint output_offset = (scores_offset + local_row_idx) * HEADS_NUM * total_partitions_num * SEQ_LEN_PARTITION_SIZE +
                                                (head_num_idx + q_idx) * total_partitions_num * SEQ_LEN_PARTITION_SIZE +
                                                partition_idx * SEQ_LEN_PARTITION_SIZE;
                     for (uint i = sgid * SUBGROUP_SIZE + sglid; i < SEQ_LEN_PARTITION_SIZE; i += SUBGROUPS_PER_WG * SUBGROUP_SIZE) {
+#if HAS_SCORE_AGGREGATION && MULTI_TOKENS_PROCESSING
+                        // In case of multi tokens processing we have to fill the buffer with zeros to guarantee correct result
+                        // for reduction operation
+                        softmax_results[output_offset + i] = partition_idx * SEQ_LEN_PARTITION_SIZE + i >= seq_len ? SOFTMAX_ACCUMULATOR_VAL_ZERO : slm_qk_vals[i];
+#else
                         softmax_results[output_offset + i] = slm_qk_vals[i];
+#endif
                     }
                 }
             }
@@ -617,6 +638,9 @@ KERNEL(pa_sdpa_finalization_stage)(
 #if PAGED_ATTENTION_SCORES_OUTPUT
     __global SOFTMAX_ACCUMULATOR_TYPE* softmax_results,
     const __global int* subsequence_offsets,
+#if HAS_SCORE_AGGREGATION
+    const __global int* cumulative_score_aggregation_sum,
+#endif
 #endif
     const __global SOFTMAX_ACCUMULATOR_TYPE* exp_sums,
     const __global SOFTMAX_ACCUMULATOR_TYPE* max_logits,
@@ -759,6 +783,9 @@ KERNEL(pa_sdpa_scores_calculation)(
     __global OUTPUT1_TYPE* scores_output,
     const __global SOFTMAX_ACCUMULATOR_TYPE* softmax_output,
     const __global int* subsequence_offsets,
+#if HAS_SCORE_AGGREGATION
+    const __global int* cumulative_score_aggregation_sum,
+#endif
     const __global SOFTMAX_ACCUMULATOR_TYPE* exp_sums,
     const __global SOFTMAX_ACCUMULATOR_TYPE* max_logits,
     const __global OUTPUT_TYPE* tmp_out,
@@ -788,107 +815,139 @@ KERNEL(pa_sdpa_scores_calculation)(
 
     SOFTMAX_ACCUMULATOR_TYPE total_score = SOFTMAX_ACCUMULATOR_VAL_ZERO;
     if (seq_len <= partition_size) {
-        // If seq_len is less than the partition size, just reduce the results over the heads
-        for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
-            const uint input_offset = subsequence_idx * HEADS_NUM * max_seq_len + head_idx * max_seq_len + partition_global_idx;
-            SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
-            total_score += softmax_value;
+#if HAS_SCORE_AGGREGATION
+        const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+        const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+#else
+        const int scores_offset = subsequence_idx;
+        const int window_size = 1;
+#endif
+        // If seq_len is less than the partition size, just reduce the results over the heads and window size
+        for (uint score_row_idx = 0; score_row_idx < window_size; score_row_idx++) {
+            for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
+                const uint input_offset = (scores_offset + score_row_idx) * HEADS_NUM * max_seq_len +
+                                          head_idx * max_seq_len + partition_global_idx;
+                SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
+                total_score += softmax_value;
+            }
         }
     } else if (seq_len <= partition_size * MAX_PARTITIONS_NUM) {
         // Optimized version for longer prompts (up to partition_size * MAX_PARTITIONS_NUM, ~64K tokens)
 
-        // Depending on the previous kernel exp_sums and max_logits might have different structure:
-        // For ordinary 1st and 2nd token kernels, there is only a single entry per subsequence.
-        // However, for mixed mode execution, exp_sums and max_logits include information for all
-        // tokens of each subsequence, but only the last one is needed for score calculation.
-        const uint subsequence_pos = is_mixed_mode ? subsequence_end - 1 : subsequence_idx;
+#if HAS_SCORE_AGGREGATION
+        const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+        const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+#else
+        const int scores_offset = subsequence_idx;
+        const int window_size = 1;
+#endif
+        for (uint score_row_idx = 0; score_row_idx < window_size; score_row_idx++) {
+            // Depending on the previous kernel exp_sums and max_logits might have different structure:
+            // For ordinary 1st and 2nd token kernels, there is only a single entry per subsequence.
+            // However, for mixed mode execution, exp_sums and max_logits include information for all
+            // tokens of each subsequence, but only the last one is needed for score calculation.
+            const uint subsequence_pos = is_mixed_mode ? subsequence_end - window_size + score_row_idx : scores_offset + score_row_idx;
+
+            for (uint head_idx = sgid; head_idx < HEADS_NUM; head_idx += sgid_num) {
+                SOFTMAX_ACCUMULATOR_TYPE max_logit[MAX_PARTITIONS_NUM / SUBGROUP_SIZE];
+                SOFTMAX_ACCUMULATOR_TYPE exp_sum[MAX_PARTITIONS_NUM / SUBGROUP_SIZE];
+
+                const uint exp_sums_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
+                for (int i = 0; i < num_of_partitions / SUBGROUP_SIZE; i++) {
+                    max_logit[i] = max_logits[exp_sums_offset + i * SUBGROUP_SIZE + sglid];
+                    exp_sum[i] = exp_sums[exp_sums_offset + i * SUBGROUP_SIZE + sglid];
+                }
 
-        for (uint head_idx = sgid; head_idx < HEADS_NUM; head_idx += sgid_num) {
-            SOFTMAX_ACCUMULATOR_TYPE max_logit[MAX_PARTITIONS_NUM / SUBGROUP_SIZE];
-            SOFTMAX_ACCUMULATOR_TYPE exp_sum[MAX_PARTITIONS_NUM / SUBGROUP_SIZE];
+                const uint partitions_leftovers = num_of_partitions % SUBGROUP_SIZE;
+                if (partitions_leftovers != 0) {
+                    const uint idx = num_of_partitions / SUBGROUP_SIZE;
+                    max_logit[idx] = sglid >= partitions_leftovers ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[exp_sums_offset + idx * SUBGROUP_SIZE + sglid];
+                    exp_sum[idx] = sglid >= partitions_leftovers ? SOFTMAX_ACCUMULATOR_VAL_ZERO : exp_sums[exp_sums_offset + idx * SUBGROUP_SIZE + sglid];
+                }
 
-            const uint exp_sums_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
-            for (int i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
-                max_logit[i] = max_logits[exp_sums_offset + i * SUBGROUP_SIZE + sglid];
-                exp_sum[i] = exp_sums[exp_sums_offset + i * SUBGROUP_SIZE + sglid];
-            }
+                SOFTMAX_ACCUMULATOR_TYPE global_max_logit = max_logit[0];
+                for (uint i = 1; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
+                    global_max_logit = SOFTMAX_ACCUMULATOR_MAX_FUNC(global_max_logit, max_logit[i]);
+                }
 
-            const uint partitions_leftovers = num_of_partitions % SUBGROUP_SIZE;
-            if (partitions_leftovers != 0) {
-                const uint idx = num_of_partitions / SUBGROUP_SIZE;
-                max_logit[idx] = sglid >= partitions_leftovers ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[exp_sums_offset + idx * SUBGROUP_SIZE + sglid];
-                exp_sum[idx] = sglid >= partitions_leftovers ? SOFTMAX_ACCUMULATOR_VAL_ZERO : exp_sums[exp_sums_offset + idx * SUBGROUP_SIZE + sglid];
-            }
+                global_max_logit = sub_group_reduce_max(global_max_logit);
 
-            SOFTMAX_ACCUMULATOR_TYPE global_max_logit = max_logit[0];
-            for (uint i = 1; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
-                global_max_logit = SOFTMAX_ACCUMULATOR_MAX_FUNC(global_max_logit, max_logit[i]);
-            }
+                SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
+                for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
+                    SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = exp_sum[i] * native_exp(max_logit[i] - global_max_logit);
+                    if (i * SUBGROUP_SIZE + sglid == partition_idx)
+                        slm_exp_sums[head_idx] = adjusted_exp_sum;
+                    global_exp_sum += adjusted_exp_sum;
+                }
 
-            global_max_logit = sub_group_reduce_max(global_max_logit);
+                global_exp_sum = sub_group_reduce_add(global_exp_sum);
 
-            SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
-            for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
-                SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = exp_sum[i] * native_exp(max_logit[i] - global_max_logit);
-                // slm_exp_sums[head_idx][i * SUBGROUP_SIZE + sglid] = adjusted_exp_sum;
-                if (i * SUBGROUP_SIZE + sglid == partition_idx)
-                    slm_exp_sums[head_idx] = adjusted_exp_sum;
-                global_exp_sum += adjusted_exp_sum;
+                slm_global_exp_sum[head_idx] = global_exp_sum;
             }
 
-            global_exp_sum = sub_group_reduce_add(global_exp_sum);
-
-            slm_global_exp_sum[head_idx] = global_exp_sum;
-        }
-
-        barrier(CLK_LOCAL_MEM_FENCE);
+            barrier(CLK_LOCAL_MEM_FENCE);
 
-        for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
-            SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = slm_exp_sums[head_idx];
-            SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = slm_global_exp_sum[head_idx];
+            for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
+                SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = slm_exp_sums[head_idx];
+                SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = slm_global_exp_sum[head_idx];
 
-            const uint input_offset = subsequence_idx * HEADS_NUM * max_seq_len + head_idx * max_seq_len + partition_global_idx;
-            SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
+                const uint input_offset = (scores_offset + score_row_idx) * HEADS_NUM * max_seq_len +
+                                          head_idx * max_seq_len +
+                                          partition_global_idx;
+                SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
 
-            softmax_value = softmax_value * adjusted_exp_sum / global_exp_sum;
-            total_score += softmax_value;
+                softmax_value = softmax_value * adjusted_exp_sum / global_exp_sum;
+                total_score += softmax_value;
+            }
         }
     } else {
         // Non optimized fallback version
-        const uint subsequence_pos = is_mixed_mode ? subsequence_end - 1 : subsequence_idx;
-        for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
-            SOFTMAX_ACCUMULATOR_TYPE global_max_logit = SOFTMAX_ACCUMULATOR_VAL_MIN;
-            const uint max_logits_base_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
-            for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
-                const uint partition_offset = i * SUBGROUP_SIZE + sglid;
-                SOFTMAX_ACCUMULATOR_TYPE max_logit = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[max_logits_base_offset + partition_offset];
-                global_max_logit = SOFTMAX_ACCUMULATOR_MAX_FUNC(global_max_logit, max_logit);
-            }
+#if HAS_SCORE_AGGREGATION
+        const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+        const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+#else
+        const int scores_offset = subsequence_idx;
+        const int window_size = 1;
+#endif
+        for (uint score_row_idx = 0; score_row_idx < window_size; score_row_idx++) {
+            const uint subsequence_pos = is_mixed_mode ? subsequence_end - window_size + score_row_idx : subsequence_idx;
+            for (uint head_idx = 0; head_idx < HEADS_NUM; head_idx++) {
+                SOFTMAX_ACCUMULATOR_TYPE global_max_logit = SOFTMAX_ACCUMULATOR_VAL_MIN;
+                const uint max_logits_base_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
+                for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
+                    const uint partition_offset = i * SUBGROUP_SIZE + sglid;
+                    SOFTMAX_ACCUMULATOR_TYPE max_logit = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[max_logits_base_offset + partition_offset];
+                    global_max_logit = SOFTMAX_ACCUMULATOR_MAX_FUNC(global_max_logit, max_logit);
+                }
 
-            global_max_logit = sub_group_reduce_max(global_max_logit);
-
-            SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
-            SOFTMAX_ACCUMULATOR_TYPE partition_adjusted_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
-            const uint exp_sums_base_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
-            for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
-                const uint partition_offset = i * SUBGROUP_SIZE + sglid;
-                SOFTMAX_ACCUMULATOR_TYPE exp_sum = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_ZERO : exp_sums[exp_sums_base_offset + partition_offset];
-                SOFTMAX_ACCUMULATOR_TYPE max_logit = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[max_logits_base_offset + partition_offset];
-                SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = exp_sum * native_exp(max_logit - global_max_logit);
-                global_exp_sum += adjusted_exp_sum;
-
-                // Save and broadcast the adjusted exp_sum for the currently being processed partition
-                if (i == partition_idx / SUBGROUP_SIZE)
-                    partition_adjusted_exp_sum = sub_group_broadcast(adjusted_exp_sum, partition_idx % SUBGROUP_SIZE);
-            }
+                global_max_logit = sub_group_reduce_max(global_max_logit);
+
+                SOFTMAX_ACCUMULATOR_TYPE global_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
+                SOFTMAX_ACCUMULATOR_TYPE partition_adjusted_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
+                const uint exp_sums_base_offset = subsequence_pos * HEADS_NUM * max_partitions_num + head_idx * max_partitions_num;
+                for (uint i = 0; i < CEIL_DIV(num_of_partitions, SUBGROUP_SIZE); i++) {
+                    const uint partition_offset = i * SUBGROUP_SIZE + sglid;
+                    SOFTMAX_ACCUMULATOR_TYPE exp_sum = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_ZERO : exp_sums[exp_sums_base_offset + partition_offset];
+                    SOFTMAX_ACCUMULATOR_TYPE max_logit = partition_offset >= num_of_partitions ? SOFTMAX_ACCUMULATOR_VAL_MIN : max_logits[max_logits_base_offset + partition_offset];
+                    SOFTMAX_ACCUMULATOR_TYPE adjusted_exp_sum = exp_sum * native_exp(max_logit - global_max_logit);
+                    global_exp_sum += adjusted_exp_sum;
+
+                    // Save and broadcast the adjusted exp_sum for the currently being processed partition
+                    if (i == partition_idx / SUBGROUP_SIZE) {
+                        partition_adjusted_exp_sum = sub_group_broadcast(adjusted_exp_sum, partition_idx % SUBGROUP_SIZE);
+                    }
+                }
 
-            global_exp_sum = sub_group_reduce_add(global_exp_sum);
+                global_exp_sum = sub_group_reduce_add(global_exp_sum);
 
-            const uint input_offset = subsequence_idx * HEADS_NUM * max_seq_len + head_idx * max_seq_len + partition_global_idx;
-            SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
+                const uint input_offset = (scores_offset + score_row_idx) * HEADS_NUM * max_seq_len +
+                                          head_idx * max_seq_len +
+                                          partition_global_idx;
+                SOFTMAX_ACCUMULATOR_TYPE softmax_value = softmax_output[input_offset];
 
-            softmax_value = softmax_value * partition_adjusted_exp_sum / global_exp_sum;
-            total_score += softmax_value;
+                softmax_value = softmax_value * partition_adjusted_exp_sum / global_exp_sum;
+                total_score += softmax_value;
+            }
         }
     }
 
diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl
index 7782c613508908..7b140ae01c71d9 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl
+++ b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl
@@ -859,6 +859,9 @@ KERNEL(sdpa_opt)(
 #if PAGED_ATTENTION_SCORES_OUTPUT
     , __global SOFTMAX_ACCUMULATOR_TYPE* softmax_results
     , const __global int* subsequence_offsets
+#if IS_PAGED_ATTENTION && HAS_SCORE_AGGREGATION
+    , const __global int* cumulative_score_aggregation_sum
+#endif
     , __global SOFTMAX_ACCUMULATOR_TYPE* exp_sums
     , __global SOFTMAX_ACCUMULATOR_TYPE* max_logits
     , __global OUTPUT_TYPE* tmp_out
@@ -1024,7 +1027,7 @@ KERNEL(sdpa_opt)(
                     query_offset += query_pitch;
                     query_local_offset++;
                 }
-            #else 
+            #else
             // Load query_input to slm_query. E.g., if TARGET_SEQ_LEN_BLOCK_SIZE = 16 && K_HEAD_SIZE == 72 && # subgroups = 5
             //     16 (TARGET_SEQ_LEN_BLOCK_SIZE)
             // ------------
@@ -1388,30 +1391,47 @@ KERNEL(sdpa_opt)(
                     const uint subsequence_idx = gws_seq_indexes_correspondence[target_seq_dim];
                     const uint subsequence_end_pos = subsequence_begins[subsequence_idx + 1];
 
-                    // PagedAttention is supposed to save only last "row" of the QK matrix multiplication,
-                    // so save SEQ_LEN_PARTITION_SIZE elements for each partition
+                    // PagedAttention is supposed to save last N "rows" of the QK matrix multiplication,
+                    // so save SEQ_LEN_PARTITION_SIZE * N elements for each partition
+#if IS_PAGED_ATTENTION && HAS_SCORE_AGGREGATION
+                    const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+                    const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+                    const int global_first_row_idx = max((int)subsequence_end_pos - window_size, 0);
+                    if (subsequence_end_pos > global_first_row_idx) {
+                        const bool save_row = m >= max(global_first_row_idx - (int)block_start_pos, 0);
+                        const int local_row_idx = block_start_pos + m - global_first_row_idx;
+#else
+                    const int scores_offset = subsequence_idx;
                     if (subsequence_end_pos == block_end_pos) {
-                        const uint last_row_idx = block_end_pos - block_start_pos - 1;
-                        if (m == last_row_idx) {
+                        const bool save_row = m == (block_end_pos - block_start_pos - 1);
+                        const int local_row_idx = 0;
+#endif
+                        if (save_row) {
                             const uint partition_idx = start_partition_idx / SEQ_LEN_PARTITION_SIZE;
 
                             SOFTMAX_ACCUMULATOR_TYPE correction_factor = native_exp(qk_max_new - qk_max_cur);
 
                             if (sglid == 0) {
                                 const uint max_partitions_num = aligned_max_context_len / SEQ_LEN_PARTITION_SIZE;
-                                const uint exp_sums_output_offset = subsequence_idx * NUM_HEADS * max_partitions_num +
+                                const uint exp_sums_output_offset = (scores_offset + local_row_idx) * NUM_HEADS * max_partitions_num +
                                                                     num_heads_dim * max_partitions_num +
                                                                     partition_idx;
                                 exp_sums[exp_sums_output_offset] = exp_sum_new * correction_factor;
                                 max_logits[exp_sums_output_offset] = qk_max_cur;
                             }
 
-                            const uint output_offset = subsequence_idx * NUM_HEADS * aligned_max_context_len +
-                                                    num_heads_dim * aligned_max_context_len +
-                                                    partition_idx * SEQ_LEN_PARTITION_SIZE;
+                            const uint output_offset = (scores_offset + local_row_idx) * NUM_HEADS * aligned_max_context_len +
+                                                       num_heads_dim * aligned_max_context_len +
+                                                       partition_idx * SEQ_LEN_PARTITION_SIZE;
                             for (uint i = sglid; i < partition_seq_len; i += SUBGROUP_SIZE) {
                                 softmax_results[output_offset + i] = TO_SOFTMAX_ACCUMULATOR_TYPE(slm_qk_vals[m][i]) / exp_sum_new;
                             }
+#if HAS_SCORE_AGGREGATION
+                            const uint full_partition_seq_len = min((uint)SOURCE_SEQ_LEN - start_partition_idx, (uint)SEQ_LEN_PARTITION_SIZE);
+                            for (uint i = partition_seq_len + sglid; i < full_partition_seq_len; i += SUBGROUP_SIZE) {
+                                softmax_results[output_offset + i] = SOFTMAX_ACCUMULATOR_VAL_ZERO;
+                            }
+#endif
                         }
                     }
                 }
@@ -1492,34 +1512,49 @@ KERNEL(sdpa_opt)(
 
                 // PagedAttention is supposed to save only last "row" of the QK matrix multiplication,
                 // so save SEQ_LEN_PARTITION_SIZE elements for each partition
+#if IS_PAGED_ATTENTION && HAS_SCORE_AGGREGATION
+                const int scores_offset = cumulative_score_aggregation_sum[subsequence_idx];
+                const int window_size = cumulative_score_aggregation_sum[subsequence_idx + 1] - scores_offset;
+                const int global_first_row_idx = max((int)subsequence_end_pos - window_size, 0);
+                if (subsequence_end_pos > global_first_row_idx) {
+                    const int local_row_idx = block_start_pos + sglid - global_first_row_idx;
+                    const bool save_row = sglid >= max(global_first_row_idx - (int)block_start_pos, 0) && local_row_idx < window_size;
+#else
+                const int scores_offset = subsequence_idx;
                 if (subsequence_end_pos == block_end_pos) {
+                    const int local_row_idx = 0;
+                    const bool save_row = sglid == block_end_pos - block_start_pos - 1;
+#endif
                     const uint last_row_idx = block_end_pos - block_start_pos - 1;
-                    if (sglid == last_row_idx) {
+                    if (save_row) {
                         const uint partition_idx = start_partition_idx / SEQ_LEN_PARTITION_SIZE;
 
                         if (sgid == 0) {
                             const uint max_partitions_num = aligned_max_context_len / SEQ_LEN_PARTITION_SIZE;
-                            const uint exp_sums_output_offset = subsequence_idx * NUM_HEADS * max_partitions_num +
+                            const uint exp_sums_output_offset = (scores_offset + local_row_idx) * NUM_HEADS * max_partitions_num +
                                                                 num_heads_dim * max_partitions_num +
                                                                 partition_idx;
                             exp_sums[exp_sums_output_offset] = exp_sum_new;
                             max_logits[exp_sums_output_offset] = qk_max_new;
                         }
 
-                        const uint output_offset = subsequence_idx * NUM_HEADS * aligned_max_context_len +
-                                                num_heads_dim * aligned_max_context_len +
-                                                partition_idx * SEQ_LEN_PARTITION_SIZE + sgid * TARGET_SEQ_LEN_BLOCK_SIZE;
+                        const uint output_offset = (scores_offset + local_row_idx) * NUM_HEADS * aligned_max_context_len +
+                                                   num_heads_dim * aligned_max_context_len +
+                                                   partition_idx * SEQ_LEN_PARTITION_SIZE + sgid * TARGET_SEQ_LEN_BLOCK_SIZE;
                         for (uint i = 0; i < TARGET_SEQ_LEN_BLOCK_SIZE; i++) {
+#if HAS_SCORE_AGGREGATION
+                            softmax_results[output_offset + i] = sgid * TARGET_SEQ_LEN_BLOCK_SIZE + i >= partition_seq_len ? SOFTMAX_ACCUMULATOR_VAL_ZERO : qk_acc[i];
+#else
                             softmax_results[output_offset + i] = qk_acc[i];
+#endif
                         }
-
                     }
                 }
             }
 #endif  /*end of PAGED_ATTENTION_SCORES_OUTPUT*/
         }
 #endif /*end of softmax calc*/
-        
+
         barrier(CLK_LOCAL_MEM_FENCE);
 
         // QK*V calculation
@@ -1576,7 +1611,7 @@ KERNEL(sdpa_opt)(
 #endif
 #endif
                     #ifdef V_HEAD_SIZE_LEFTOVER
-                    // splitting to two cases for supressing reg spill : block_read & eltwise read 
+                    // splitting to two cases for supressing reg spill : block_read & eltwise read
                     if (sgid < SUBGROUPS_PER_WG - 1) {
                     // block_read values
                     #endif
@@ -1671,7 +1706,7 @@ KERNEL(sdpa_opt)(
                         qk_val[seq_idx] = slm_qk_vals[seq_idx][seq_len * SUBGROUP_SIZE + sglid];
                     }
                 #ifdef V_HEAD_SIZE_LEFTOVER
-                    // splitting to two cases for supressing reg spill : block_read & eltwise read 
+                    // splitting to two cases for supressing reg spill : block_read & eltwise read
                     if (sgid < SUBGROUPS_PER_WG - 1) {
                     // block_read values
                 #endif
@@ -1777,7 +1812,7 @@ KERNEL(sdpa_opt)(
                         else
                             value_packed = (sglid < V_HEAD_SIZE_LEFTOVER) ? value_input[value_offset] : INPUT2_VAL_ZERO;
                         #endif // BEAM_TABLE_TYPE
- 
+
                     #else // !V_HEAD_SIZE_LEFTOVER
                         #ifdef BEAM_TABLE_TYPE
                         const INPUT2_TYPE value_packed = VALUE_BLOCK_READ(value_input, sub_group_broadcast(value_offset, seq_len_idx));
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
index 3cb159a80436fe..3145ed5fafd757 100644
--- 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
@@ -89,6 +89,8 @@ KernelsData PagedAttentionSDPAKernelOpt::GetKernelsData(const Params& p) const {
                                                KernelsTypes::FINALIZATION_MULTI_TOKENS };
 
     const auto has_scores_output = params.outputs.size() > 1;
+    const auto has_score_aggregation = params.conf.has_score_aggregation;
+
     if (has_scores_output) {
         kernels_type.push_back(KernelsTypes::SCORES_CALCULATION);
     }
@@ -149,6 +151,11 @@ KernelsData PagedAttentionSDPAKernelOpt::GetKernelsData(const Params& p) const {
             // Intermediate softmax results for scores output calculation and precalculated accumulated
             // sequence length offsets for each subsequence
             internal_buffers_num += 2;
+
+            if (has_score_aggregation) {
+               // Precalculated cumulative sum of aggregation score windows
+               internal_buffers_num += 1;
+            }
         }
 
         // Softmax's exp_sums, max_logits and intermediate output
@@ -296,6 +303,9 @@ JitConstants PagedAttentionSDPAKernelOpt::GetJitConstants(const pa_sdpa_params&
 
     if (params.outputs.size() > 1) {
         jit.AddConstant(MakeJitConstant("PAGED_ATTENTION_SCORES_OUTPUT", 1));
+        if (params.conf.has_score_aggregation) {
+            jit.AddConstant(MakeJitConstant("HAS_SCORE_AGGREGATION", 1));
+        }
     }
 
     const size_t score_aggregation_idx = 13;
@@ -465,7 +475,8 @@ void PagedAttentionSDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) cons
             auto subsequences_number = past_lens.Batch().v;
             auto softmax_buf_dt_size = BytesPerElement(softmax_acc_dt);
 
-            auto softmax_buf_elements_count = subsequences_number * prim_params.conf.heads_num * num_of_partitions * partition_size;
+            auto tokens_number = prim_params.conf.has_score_aggregation ? prim_params.conf.paged_attention_snap_kv_tokens : subsequences_number;
+            auto softmax_buf_elements_count = tokens_number * prim_params.conf.heads_num * num_of_partitions * partition_size;
             auto softmax_buf_size = softmax_buf_elements_count * softmax_buf_dt_size;
 
             // Softmax intermediate output
@@ -473,9 +484,14 @@ void PagedAttentionSDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) cons
             // Precalculated accumulated sequence length offsets for each subsequence
             kd.internalBuffers.emplace_back(subsequences_number * BytesPerElement(Datatype::INT32), lockable);
 
+            if (prim_params.conf.has_score_aggregation) {
+                // Cumulative window size sum buffer
+                kd.internalBuffers.emplace_back((subsequences_number + 1) * BytesPerElement(Datatype::INT32), lockable);
+            }
+
             if (prim_params.stage == PagedAttentionStage::PREFILL) {
                 // Recalculate buf_size as in case of PREFILL stage it's not needed to allocate buffer per each input token
-                buf_elements_count = subsequences_number * prim_params.conf.heads_num * num_of_partitions;
+                buf_elements_count = tokens_number * prim_params.conf.heads_num * num_of_partitions;
                 buf_size = buf_elements_count * buf_dt_size;
 
                 // Intermediate tmp output buffer is not used for PREFILL stage
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
index 06a486f044ab7f..c0408d5ee65129 100644
--- 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
@@ -97,9 +97,13 @@ struct sdpa_configuration {
     // Paged Attention configuration
     bool is_paged_attention = false;
     size_t paged_attention_sliding_window = 0;
-    int64_t paged_attention_aligned_seq_len = -1;
     int64_t paged_attention_block_size = 0;
+
+    // Runtime Paged Attention params
+    int64_t paged_attention_aligned_seq_len = -1;
     int64_t paged_attention_max_len = 0;
+    int64_t paged_attention_snap_kv_tokens = 0;
+
     bool has_const_scale_val = false;
     float scale_val = 0.f;
     bool has_const_attn_mask_val = false;
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
index 477b9b5d95e7d1..6790383c0a8305 100644
--- 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
@@ -499,7 +499,7 @@ bool SDPAKernelMicro::Validate(const Params& p) const {
     }
 
     // Scores output is not supported
-    if (params.conf.is_paged_attention && params.outputs.size() > 1)
+    if (params.conf.is_paged_attention && (params.outputs.size() > 1 || params.conf.has_score_aggregation))
         return false;
 
     if (params.conf.is_paged_attention && params.conf.paged_attention_sliding_window != 0) {
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp
index b4c1da4d4a4b6a..c761ef949388e6 100644
--- a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp
+++ b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp
@@ -214,6 +214,9 @@ JitConstants SDPAKernelOpt::GetJitConstants(const sdpa_params& params, size_t ke
 
         if (params.outputs.size() > 1) {
             jit.AddConstant(MakeJitConstant("PAGED_ATTENTION_SCORES_OUTPUT", 1));
+            if (params.conf.has_score_aggregation) {
+                jit.AddConstant(MakeJitConstant("HAS_SCORE_AGGREGATION", 1));
+            }
         }
     } else {
         size_t scale_idx = params.conf.has_const_attn_mask_val ? 3 : 4;
@@ -376,13 +379,17 @@ KernelsData SDPAKernelOpt::GetKernelsData(const Params& params) const {
 
         if (prim_params.conf.is_paged_attention && prim_params.outputs.size() > 1) {
             // Intermediate buffers for PagedAttention scores calculation:
-            // softmax_results, subsequence_offsets, exp_sums, max_logits, tmp_out
+            // softmax_results, subsequence_offsets, (cumulative_score_aggregation_sum), exp_sums, max_logits, tmp_out
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 3});
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 4});
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 5});
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 6});
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 7});
 
+            if (prim_params.conf.has_score_aggregation) {
+                kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 8});
+            }
+
             // Scalar used for proper offset calculation of intermediate data buffers
             kernel.params.arguments.push_back({ArgumentDescriptor::Types::SCALAR, 0});
         }
diff --git a/src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp b/src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp
index 83efa03fe8be5d..8940ff352291df 100644
--- a/src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp
+++ b/src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp
@@ -57,11 +57,12 @@ static void CreatePagedAttentionExtensionOp(ProgramBuilder& p, const std::shared
     prim.kv_heads_num = kv_heads_num;
     prim.heads_num = heads_num;
 
-    const size_t scale_idx = 9;
-    const size_t sliding_window_idx = 10;
-    const size_t alibi_idx = 11;
+    const size_t scale_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCALE;
+    const size_t sliding_window_idx = cldnn::paged_attention::PagedAttentionInputIdx::SLIDING_WINDOW;
+    const size_t alibi_idx = cldnn::paged_attention::PagedAttentionInputIdx::ALIBI;
+    const size_t score_aggregation_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCORE_AGGREGATION;
 
-    std::shared_ptr<ov::op::v0::Constant> scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
+    auto scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
     if (scale_const) {
         OPENVINO_ASSERT(ov::shape_size(scale_const->get_output_shape(0)) == 1);
         prim.scale_val = scale_const->cast_vector<float>()[0];
@@ -69,28 +70,22 @@ static void CreatePagedAttentionExtensionOp(ProgramBuilder& p, const std::shared
         prim.scale_val = std::optional<float>();
     }
 
-    std::shared_ptr<ov::op::v0::Constant> sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
+    auto sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
     if (sliding_windows_const) {
         OPENVINO_ASSERT(ov::shape_size(sliding_windows_const->get_output_shape(0)) == 1);
         prim.sliding_window = sliding_windows_const->cast_vector<size_t>()[0];
     }
 
-    std::shared_ptr<ov::op::v0::Constant> alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
+    auto alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
     OPENVINO_ASSERT(alibi_const != nullptr);
     prim.has_alibi = ov::shape_size(alibi_const->get_output_shape(0)) > 0;
 
-
-    const size_t score_aggregation_idx = 13;
-
     auto score_aggregation_input = ov::as_type_ptr<ov::op::v0::Parameter>(op->get_input_node_shared_ptr(score_aggregation_idx));
-
-    if (!score_aggregation_input) {
-        prim.has_score_aggregation = false;
-    } else {
-        prim.has_score_aggregation = !score_aggregation_input->get_output_partial_shape(0).is_dynamic();
+    if (score_aggregation_input && score_aggregation_input->get_output_partial_shape(0).is_dynamic() && op->get_output_size() > 1) {
+        prim.has_score_aggregation = true;
     }
 
-    prim.has_rotated_blocks = op->get_input_size() == 17;
+    prim.has_rotated_blocks = op->get_input_size() > cldnn::paged_attention::PagedAttentionInputIdx::ROTATION_TRIG_LUT;
 
     prim.num_outputs = 1;
 
diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
index 2194935d60392e..56d1d6e2c42f35 100644
--- a/src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
+++ b/src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
@@ -53,6 +53,13 @@ using namespace ::tests;
 * shape: [max_num_batched_tokens / BLOCK_SIZE, head_size]​ || [max_num_batched_tokens, head_size], type: f16
 */
 
+
+enum class ScoresMode {
+    DISABLED = 0,
+    LAST_TOKEN,
+    SNAPKV
+};
+
 struct SubsequenceDescriptor {
     int num_tokens;
     int past_len;
@@ -72,6 +79,7 @@ struct PagedAttentionManager {
     int v_head_size;
     int block_size;
     bool kv_cache_compression;
+    bool has_score_aggregation;
     CacheRotationDescriptor rotation_config;
     std::vector<SubsequenceDescriptor> subsequence_descs;
 
@@ -88,6 +96,9 @@ struct PagedAttentionManager {
     std::vector<int> max_context_len;
     std::vector<int> score_aggregation_window;
 
+    // score aggregation related inputs
+    std::vector<int> score_aggregation;
+
     // rotation related inputs
     std::vector<int> rotated_block_indices;
     std::vector<int> rotation_deltas;
@@ -106,12 +117,14 @@ struct PagedAttentionManager {
                           int v_head_size,
                           int block_size,
                           bool kv_cache_compression,
+                          bool has_score_aggregation,
                           CacheRotationDescriptor rotation_config)
         : num_heads(num_heads)
         , k_head_size(k_head_size)
         , v_head_size(v_head_size)
         , block_size(block_size)
         , kv_cache_compression(kv_cache_compression)
+        , has_score_aggregation(has_score_aggregation)
         , rotation_config(rotation_config)
         , subsequence_descs(subsequence_descs)
         , test_engine(engine)
@@ -172,6 +185,15 @@ struct PagedAttentionManager {
                 rotation_trig_lut = generate_rotation_trig_lut_data(rg, max_context_len[0], k_head_size);
             }
         }
+
+        if (has_score_aggregation) {
+            for (const auto& subsequence_desc : subsequence_descs) {
+                const auto max_tokens = 10;
+                auto max_window_size = std::min(subsequence_desc.num_tokens, max_tokens);
+                auto window_size = rg.generate_random_val<int>(1, max_window_size);
+                score_aggregation.push_back(window_size);
+            }
+        }
     }
 
     memory::ptr get_query_memory() {
@@ -351,8 +373,8 @@ struct PagedAttentionManager {
         return get_memory_from_vec(max_context_len);
     }
 
-    memory::ptr get_score_aggregation_window_memory() {
-        return get_memory_from_vec(score_aggregation_window);
+    memory::ptr get_score_aggregation() {
+        return get_memory_from_vec(score_aggregation);
     }
 
     memory::ptr get_rotated_block_indices_memory() {
@@ -559,6 +581,8 @@ struct PagedAttentionReference {
                 }
             }
 
+            auto window_size = pam.has_score_aggregation ? pam.score_aggregation[i] : 1;
+
             auto subsequence_ref_results = run_reference(pam.query_data[i],
                                                          key_data,
                                                          pam.value_data[i],
@@ -567,6 +591,7 @@ struct PagedAttentionReference {
                                                          pam.num_heads,
                                                          pam.k_head_size,
                                                          pam.v_head_size,
+                                                         window_size,
                                                          pam.get_default_scale());
 
             // concatenate all subsequences into one vector
@@ -591,6 +616,7 @@ struct PagedAttentionReference {
                       int num_heads,
                       int k_head_size,
                       int v_head_size,
+                      int window_size,
                       float scale) {
         auto query_shape = ov::PartialShape{1, num_queries, num_heads, k_head_size};
         auto key_shape = ov::PartialShape{1, num_keys, num_heads, k_head_size};
@@ -645,10 +671,11 @@ struct PagedAttentionReference {
         auto output_scores_mem = outputs.at("scores_data").get_memory();
 
         return { get_output_data_vec(output_data_mem, num_queries, v_head_size, num_heads),
-                 get_output_scores_vec(output_scores_mem, num_queries, num_keys, num_heads) };
+                 get_output_scores_vec(output_scores_mem, window_size, num_queries, num_keys, num_heads) };
     }
 
     std::vector<ov::float16> get_output_scores_vec(memory::ptr scores_output,
+                                                   int window_size,
                                                    int num_queries,
                                                    int num_keys,
                                                    int num_heads) {
@@ -656,11 +683,14 @@ struct PagedAttentionReference {
 
         std::vector<ov::float16> output_scores(num_keys, 0);
         mem_lock<ov::float16, mem_lock_type::read> mem_ptr(scores_output, test_stream);
-        for (int head_idx = 0; head_idx < num_heads; head_idx++) {
-            for (int score_idx = 0; score_idx < num_keys; score_idx++) {
-                output_scores[score_idx] += mem_ptr[head_idx * num_queries * num_keys +
-                                                     (num_queries - 1) * num_keys +
-                                                     score_idx];
+        for (int row_idx = 0; row_idx < window_size; row_idx++) {
+            for (int head_idx = 0; head_idx < num_heads; head_idx++) {
+                for (int score_idx = 0; score_idx < num_keys; score_idx++) {
+                    auto scores_offset = head_idx * num_queries * num_keys +
+                                         (num_queries - window_size + row_idx) * num_keys +
+                                         score_idx;
+                    output_scores[score_idx] += mem_ptr[scores_offset];
+                }
             }
         }
 
@@ -792,6 +822,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
                                   p.v_head_size,
                                   p.block_size,
                                   p.kv_cache_compression,
+                                  p.scores_mode == ScoresMode::SNAPKV,
                                   p.rotation_config);
 
         if (p.kv_cache_compression)
@@ -802,7 +833,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         auto value_mem = pam.get_value_memory();
 
         auto key_cache_mem = pam.get_key_cache_memory();
-       auto value_cache_mem = pam.get_value_cache_memory();
+        auto value_cache_mem = pam.get_value_cache_memory();
 
         auto past_lens_mem = pam.get_past_lens_memory();
         auto subsequence_begins_mem = pam.get_subsequence_begins_memory();
@@ -814,7 +845,8 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         auto alibi_mem = pam.get_alibi_memory();
         auto max_context_len_mem = pam.get_max_context_len_memory();
 
-        auto score_aggregation_window_mem = pam.get_score_aggregation_window_memory();
+        // scores calculation related memory buffers
+        auto score_aggregation_mem = pam.get_score_aggregation();
 
         // cache rotation related memory buffers
         auto rotated_block_indices_mem = pam.get_rotated_block_indices_memory();
@@ -834,7 +866,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         auto sliding_window_layout = sliding_window_mem->get_layout();
         auto alibi_layout = alibi_mem->get_layout();
         auto max_context_len_layout = max_context_len_mem->get_layout();
-        auto score_aggregation_window_layout = score_aggregation_window_mem->get_layout();
+        auto score_aggregation_window_layout = score_aggregation_mem->get_layout();
         auto rotated_block_indices_layout = rotated_block_indices_mem->get_layout();
         auto rotation_deltas_layout = rotation_deltas_mem->get_layout();
         auto rotation_trig_lut_layout = rotation_trig_lut_mem->get_layout();
@@ -849,6 +881,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         subsequence_begins_layout.set_partial_shape(ov::PartialShape{ -1 });
         block_indices_layout.set_partial_shape(ov::PartialShape{ -1 });
         block_indices_begins_layout.set_partial_shape(ov::PartialShape{ -1 });
+        score_aggregation_window_layout.set_partial_shape(ov::PartialShape{ -1 });
         rotated_block_indices_layout.set_partial_shape(ov::PartialShape{ -1 });
         rotation_deltas_layout.set_partial_shape(ov::PartialShape{ -1, -1 });
         rotation_trig_lut_layout.set_partial_shape(ov::PartialShape{ -1, p.k_head_size });
@@ -914,8 +947,9 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         pa_prim.heads_num = p.num_heads;
         pa_prim.scale_val = pam.get_default_scale();
         pa_prim.has_alibi = false;
-        pa_prim.num_outputs = p.scores_output ? 2 : 1;
+        pa_prim.num_outputs = p.scores_mode == ScoresMode::DISABLED ? 1 : 2;
         pa_prim.has_rotated_blocks = p.rotation_config.apply_rotation;
+        pa_prim.has_score_aggregation = p.scores_mode == ScoresMode::SNAPKV;
 
         topology topology;
 
@@ -938,7 +972,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
             reorder("output_data", input_info("paged_attention", 0), format::bfyx, data_types::f16)
         );
 
-        if (p.scores_output) {
+        if (p.scores_mode != ScoresMode::DISABLED) {
             topology.add(reorder("output_scores", input_info("paged_attention", 1), format::bfyx, data_types::f16));
         }
 
@@ -971,7 +1005,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         network->set_input_data("sliding_window", sliding_window_mem);
         network->set_input_data("alibi", alibi_mem);
         network->set_input_data("max_context_len", max_context_len_mem);
-        network->set_input_data("score_aggregation_window", score_aggregation_window_mem);
+        network->set_input_data("score_aggregation_window", score_aggregation_mem);
 
         if (p.rotation_config.apply_rotation) {
             network->set_input_data("rotated_block_indices", rotated_block_indices_mem);
@@ -985,7 +1019,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
         cldnn::memory::ptr output_scores_mem = nullptr;
 
         output_data_mem = outputs.at("output_data").get_memory();
-        if (p.scores_output) {
+        if (p.scores_mode != ScoresMode::DISABLED) {
             output_scores_mem = outputs.at("output_scores").get_memory();
         }
 
@@ -998,7 +1032,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
             ASSERT_EQ(data_output_mem->count(), ref_data.first.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(data_output_mem, get_test_stream());
             for (size_t i = 0; i < data_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance);
+                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance) << " at index=" << i;
             }
         }
 
@@ -1006,7 +1040,7 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
             ASSERT_EQ(scores_output_mem->count(), ref_data.second.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(scores_output_mem, get_test_stream());
             for (size_t i = 0; i < scores_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.second[i], tolerance);
+                ASSERT_NEAR(mem_ptr[i], ref_data.second[i], tolerance) << " at index=" << i;
             }
         }
     }
@@ -1020,7 +1054,7 @@ struct paged_attention_test_params {
     int block_size;
     bool kv_cache_compression;
     bool dynamic_paddings;
-    bool scores_output;
+    ScoresMode scores_mode;
     CacheRotationDescriptor rotation_config;
     bool disable_flashattn_v2;
 };
@@ -1034,8 +1068,9 @@ TEST_P(paged_attention_test, basic) {
 
 const auto ENABLE_CACHE_COMPRESSION = true;
 const auto DISABLE_CACHE_COMPRESSION = false;
-const auto ENABLE_SCORES = true;
-const auto DISABLE_SCORES = false;
+const auto DISABLE_SCORES = ScoresMode::DISABLED;
+const auto ENABLE_SCORES = ScoresMode::LAST_TOKEN;
+const auto ENABLE_SCORES_SNAPKV = ScoresMode::SNAPKV;
 const auto PER_BLOCK_ROTATION = CacheRotationDescriptor{ true, true };
 const auto PER_TOKEN_ROTATION = CacheRotationDescriptor{ true, false };
 const auto DISABLE_ROTATION = CacheRotationDescriptor{ false, false };
@@ -1045,6 +1080,16 @@ const auto ENABLE_FA_V2 = false;
 const auto DISABLE_FA_V2 = true;
 
 INSTANTIATE_TEST_SUITE_P(smoke_paged_attention, paged_attention_test, ::testing::ValuesIn(std::vector<paged_attention_test_params>{
+    /* with scores output, use SnapKV */
+    paged_attention_test_params{ {{10, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, ENABLE_FA_V2 }, // 1st token
+    paged_attention_test_params{ {{36, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // 1st token
+    paged_attention_test_params{ {{1024, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, ENABLE_FA_V2 }, // 1st token long
+    paged_attention_test_params{ {{10, 0}, {30, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // 1st token + 1st token
+    paged_attention_test_params{ {{128, 0}, {256, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, ENABLE_FA_V2 }, // 1st token + 1st token
+    paged_attention_test_params{ {{128, 0}, {256, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // 1st token + 1st token
+    paged_attention_test_params{ {{1, 10}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // 2nd token
+    paged_attention_test_params{ {{1, 34}, {1, 515}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // 2nd token + 2nd token
+    paged_attention_test_params{ {{1, 34}, {25, 0}, {10, 34}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES_SNAPKV, DISABLE_ROTATION, DISABLE_FA_V2 }, // mixed: 2nd token + 1st token + part of 1st token
     /* with scores output */
     paged_attention_test_params{ {{10, 0}}, 2, 64, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES, DISABLE_ROTATION, ENABLE_FA_V2 }, // 1st token
     paged_attention_test_params{ {{10, 0}}, 2, 16, 64, 16, DISABLE_CACHE_COMPRESSION, STATIC_INPUT_PAD, ENABLE_SCORES, DISABLE_ROTATION, ENABLE_FA_V2 }, // 1st token


Reviews:
user: login: yeonbok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-30T23:13:43Z
user: login: yeonbok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-01T04:53:28Z
user: login: sshlyapn
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-01T06:40:42Z
user: login: sshlyapn
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-01T06:46:35Z
user: login: isanghao
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T07:14:37Z
user: login: yeonbok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T20:07:54Z
user: login: yeonbok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T23:04:07Z
user: login: yeonbok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T23:04:34Z
user: login: sshlyapn
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T14:35:44Z
user: login: sshlyapn
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T14:36:04Z
user: login: Lyamin-Roman
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T15:03:47Z
user: login: sshlyapn
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T15:34:03Z
user: login: yeonbok
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T17:32:21Z

Review comments:
diff_hunk: @@ -316,6 +324,30 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
             auto subsequence_offsets_mem = intermediates_memories[subsequence_offsets_idx];
             subsequence_offsets_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(subsequence_offsets_mem, stream));
+
+            if (has_score_aggregation) {
+                const size_t cumulative_window_size_sum_idx = 5;
+                OPENVINO_ASSERT(intermediates_memories.size() > subsequence_offsets_idx,
+                                "[GPU] Unexpected number of intermediates buffers for Paged Attention for scores aggregation");
+
+                auto cumulative_window_size_sum_mem = intermediates_memories[cumulative_window_size_sum_idx];
+                cumulative_window_size_sum_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(cumulative_window_size_sum_mem, stream));
+
+                mem_lock<int32_t, mem_lock_type::read> score_aggregation_mem_lock(instance.score_aggregation_memory_ptr(), stream);
+
+                // Transform window sizes to cumulative buffer for offsets precalculation
+                // For example:
+                // Original score aggregation buffer content: {4, 2, 1, 1}
+                // Cumulative window size sum buffer content: {0, 4, 6, 7, 8}
+                size_t cumulative_sum = 0;
+                cumulative_window_size_sum_lock->operator[](0) = cumulative_sum;
+                std::cout << "cumulative_window_size_sum_lock[0]=" << cumulative_sum << "\n";
path: src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
user: login: yeonbok
body: Please remove debug print.
created_at: 2025-04-30T23:13:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,14 +13,34 @@ namespace cldnn {
 struct paged_attention : public primitive_base<paged_attention> {
     CLDNN_DECLARE_PRIMITIVE(paged_attention)
 
+    enum PagedAttentionInputIdx {
+        QUERY = 0,
+        KEY = 1,
+        VALUE = 2,
+        KEY_CACHE = 3,
+        VALUE_CACHE = 4,
+        PAST_LENS = 5,
+        SUBSEQUENCE_BEGINS = 6,
+        BLOCK_INDICES = 7,
+        BLOCK_INDICES_BEGINS = 8,
+        SCALE = 9,
+        SLIDING_WINDOW = 10,
+        ALIBI = 11,
+        MAX_CONTEXT_LEN = 12,
+        SCORE_AGGREGATION = 13,
+        ROTATED_BLOCK_INDICES = 14,
+        ROTATION_DELTAS = 15,
+        ROTATION_TRIG_LUT = 16,
+    };
+
     static constexpr size_t block_size = 16;
 
     paged_attention() : primitive_base("", {}) {}
 
     paged_attention(const primitive_id& id,
                     const std::vector<input_info>& inputs)
         : primitive_base(id, inputs) {
-        OPENVINO_ASSERT((inputs.size() == 13) || (inputs.size() == 16),
+        OPENVINO_ASSERT((inputs.size() == 14) || (inputs.size() == 17),
path: src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
user: login: yeonbok
body: score_aggregation is optional but do we need 14 at least? 
created_at: 2025-05-01T04:53:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -316,6 +324,30 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
             auto subsequence_offsets_mem = intermediates_memories[subsequence_offsets_idx];
             subsequence_offsets_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(subsequence_offsets_mem, stream));
+
+            if (has_score_aggregation) {
+                const size_t cumulative_window_size_sum_idx = 5;
+                OPENVINO_ASSERT(intermediates_memories.size() > subsequence_offsets_idx,
+                                "[GPU] Unexpected number of intermediates buffers for Paged Attention for scores aggregation");
+
+                auto cumulative_window_size_sum_mem = intermediates_memories[cumulative_window_size_sum_idx];
+                cumulative_window_size_sum_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(cumulative_window_size_sum_mem, stream));
+
+                mem_lock<int32_t, mem_lock_type::read> score_aggregation_mem_lock(instance.score_aggregation_memory_ptr(), stream);
+
+                // Transform window sizes to cumulative buffer for offsets precalculation
+                // For example:
+                // Original score aggregation buffer content: {4, 2, 1, 1}
+                // Cumulative window size sum buffer content: {0, 4, 6, 7, 8}
+                size_t cumulative_sum = 0;
+                cumulative_window_size_sum_lock->operator[](0) = cumulative_sum;
+                std::cout << "cumulative_window_size_sum_lock[0]=" << cumulative_sum << "\n";
path: src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
user: login: sshlyapn
body: Removed, thanks!
created_at: 2025-05-01T06:40:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,14 +13,34 @@ namespace cldnn {
 struct paged_attention : public primitive_base<paged_attention> {
     CLDNN_DECLARE_PRIMITIVE(paged_attention)
 
+    enum PagedAttentionInputIdx {
+        QUERY = 0,
+        KEY = 1,
+        VALUE = 2,
+        KEY_CACHE = 3,
+        VALUE_CACHE = 4,
+        PAST_LENS = 5,
+        SUBSEQUENCE_BEGINS = 6,
+        BLOCK_INDICES = 7,
+        BLOCK_INDICES_BEGINS = 8,
+        SCALE = 9,
+        SLIDING_WINDOW = 10,
+        ALIBI = 11,
+        MAX_CONTEXT_LEN = 12,
+        SCORE_AGGREGATION = 13,
+        ROTATED_BLOCK_INDICES = 14,
+        ROTATION_DELTAS = 15,
+        ROTATION_TRIG_LUT = 16,
+    };
+
     static constexpr size_t block_size = 16;
 
     paged_attention() : primitive_base("", {}) {}
 
     paged_attention(const primitive_id& id,
                     const std::vector<input_info>& inputs)
         : primitive_base(id, inputs) {
-        OPENVINO_ASSERT((inputs.size() == 13) || (inputs.size() == 16),
+        OPENVINO_ASSERT((inputs.size() == 14) || (inputs.size() == 17),
path: src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
user: login: sshlyapn
body: Score_aggregation has become required input, but at ov::Model level it can either be Constant (scores disabled) or Parameter (containing window sizes)
created_at: 2025-05-01T06:46:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -51,28 +51,35 @@ static void CreatePagedAttentionExtensionOp(ProgramBuilder& p, const std::shared
     prim.kv_heads_num = kv_heads_num;
     prim.heads_num = heads_num;
 
-    const size_t scale_idx = 9;
-    const size_t sliding_window_idx = 10;
-    const size_t alibi_idx = 11;
+    const size_t scale_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCALE;
+    const size_t sliding_window_idx = cldnn::paged_attention::PagedAttentionInputIdx::SLIDING_WINDOW;
+    const size_t alibi_idx = cldnn::paged_attention::PagedAttentionInputIdx::ALIBI;
+    const size_t score_aggregation_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCORE_AGGREGATION;
 
-    std::shared_ptr<ov::op::v0::Constant> scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
+    auto scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
     if (scale_const) {
         OPENVINO_ASSERT(ov::shape_size(scale_const->get_output_shape(0)) == 1);
         prim.scale_val = scale_const->cast_vector<float>()[0];
     } else {
         prim.scale_val = std::optional<float>();
     }
 
-    std::shared_ptr<ov::op::v0::Constant> sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
+    auto sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
     if (sliding_windows_const) {
         OPENVINO_ASSERT(ov::shape_size(sliding_windows_const->get_output_shape(0)) == 1);
         prim.sliding_window = sliding_windows_const->cast_vector<size_t>()[0];
     }
 
-    std::shared_ptr<ov::op::v0::Constant> alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
+    auto alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
     OPENVINO_ASSERT(alibi_const != nullptr);
     prim.has_alibi = ov::shape_size(alibi_const->get_output_shape(0)) > 0;
-    prim.has_rotated_blocks = op->get_input_size() == 16;
+
+    auto score_aggregation_input = ov::as_type_ptr<ov::op::v0::Parameter>(op->get_input_node_shared_ptr(score_aggregation_idx));
+    if (score_aggregation_input && score_aggregation_input->get_output_partial_shape(0).is_dynamic() && op->get_output_size() > 1) {
+        prim.has_score_aggregation = true;
+    }
+
+    prim.has_rotated_blocks = op->get_input_size() == 17;
path: src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp
user: login: isanghao
body: nit:  instead of number `17`, what about: `op->get_input_size() > ROTATION_TRIG_LUT`?
created_at: 2025-05-02T07:08:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -316,6 +324,28 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
 
             auto subsequence_offsets_mem = intermediates_memories[subsequence_offsets_idx];
             subsequence_offsets_lock.reset(new mem_lock<int32_t, mem_lock_type::write>(subsequence_offsets_mem, stream));
+
+            if (has_score_aggregation) {
+                const size_t cumulative_window_size_sum_idx = 5;
+                OPENVINO_ASSERT(intermediates_memories.size() > subsequence_offsets_idx,
path: src/plugins/intel_gpu/src/graph/impls/ocl/paged_attention.cpp
user: login: yeonbok
body: cumulative_windoew_size_sum_idx
created_at: 2025-05-02T20:07:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -979,15 +1026,16 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
             ASSERT_EQ(data_output_mem->count(), ref_data.first.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(data_output_mem, get_test_stream());
             for (size_t i = 0; i < data_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance);
+                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance) << " at index=" << i;
             }
         }
 
         if (scores_output_mem) {
             ASSERT_EQ(scores_output_mem->count(), ref_data.second.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(scores_output_mem, get_test_stream());
             for (size_t i = 0; i < scores_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.second[i], tolerance);
+                std::cout << i << ". Compare " << mem_ptr[i] << " vs " << ref_data.second[i] << "\n";
path: src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
user: login: yeonbok
body: Please remove debug print
created_at: 2025-05-02T23:04:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -167,6 +182,16 @@ struct PagedAttentionManager {
                 rotation_trig_lut = generate_rotation_trig_lut_data(rg, max_context_len[0], head_size);
             }
         }
+
+        if (has_score_aggregation) {
+            for (const auto& subsequence_desc : subsequence_descs) {
+                const auto max_tokens = 10;
+                auto max_window_size = std::min(subsequence_desc.num_tokens, max_tokens);
+                auto window_size = rg.generate_random_val<int>(1, max_window_size);
+                score_aggregation.push_back(window_size);
+                std::cout << "Set scores window " << window_size << "\n";
path: src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
user: login: yeonbok
body: Please remove if this is for debug
created_at: 2025-05-02T23:04:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -51,28 +51,35 @@ static void CreatePagedAttentionExtensionOp(ProgramBuilder& p, const std::shared
     prim.kv_heads_num = kv_heads_num;
     prim.heads_num = heads_num;
 
-    const size_t scale_idx = 9;
-    const size_t sliding_window_idx = 10;
-    const size_t alibi_idx = 11;
+    const size_t scale_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCALE;
+    const size_t sliding_window_idx = cldnn::paged_attention::PagedAttentionInputIdx::SLIDING_WINDOW;
+    const size_t alibi_idx = cldnn::paged_attention::PagedAttentionInputIdx::ALIBI;
+    const size_t score_aggregation_idx = cldnn::paged_attention::PagedAttentionInputIdx::SCORE_AGGREGATION;
 
-    std::shared_ptr<ov::op::v0::Constant> scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
+    auto scale_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(scale_idx));
     if (scale_const) {
         OPENVINO_ASSERT(ov::shape_size(scale_const->get_output_shape(0)) == 1);
         prim.scale_val = scale_const->cast_vector<float>()[0];
     } else {
         prim.scale_val = std::optional<float>();
     }
 
-    std::shared_ptr<ov::op::v0::Constant> sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
+    auto sliding_windows_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(sliding_window_idx));
     if (sliding_windows_const) {
         OPENVINO_ASSERT(ov::shape_size(sliding_windows_const->get_output_shape(0)) == 1);
         prim.sliding_window = sliding_windows_const->cast_vector<size_t>()[0];
     }
 
-    std::shared_ptr<ov::op::v0::Constant> alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
+    auto alibi_const = ov::as_type_ptr<ov::op::v0::Constant>(op->get_input_node_shared_ptr(alibi_idx));
     OPENVINO_ASSERT(alibi_const != nullptr);
     prim.has_alibi = ov::shape_size(alibi_const->get_output_shape(0)) > 0;
-    prim.has_rotated_blocks = op->get_input_size() == 16;
+
+    auto score_aggregation_input = ov::as_type_ptr<ov::op::v0::Parameter>(op->get_input_node_shared_ptr(score_aggregation_idx));
+    if (score_aggregation_input && score_aggregation_input->get_output_partial_shape(0).is_dynamic() && op->get_output_size() > 1) {
+        prim.has_score_aggregation = true;
+    }
+
+    prim.has_rotated_blocks = op->get_input_size() == 17;
path: src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp
user: login: sshlyapn
body: Applied, thanks!
created_at: 2025-05-19T14:35:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -979,15 +1026,16 @@ struct PagedAttentionTest : public ::testing::TestWithParam<T> {
             ASSERT_EQ(data_output_mem->count(), ref_data.first.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(data_output_mem, get_test_stream());
             for (size_t i = 0; i < data_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance);
+                ASSERT_NEAR(mem_ptr[i], ref_data.first[i], tolerance) << " at index=" << i;
             }
         }
 
         if (scores_output_mem) {
             ASSERT_EQ(scores_output_mem->count(), ref_data.second.size());
             mem_lock<ov::float16, mem_lock_type::read> mem_ptr(scores_output_mem, get_test_stream());
             for (size_t i = 0; i < scores_output_mem->count(); i++) {
-                ASSERT_NEAR(mem_ptr[i], ref_data.second[i], tolerance);
+                std::cout << i << ". Compare " << mem_ptr[i] << " vs " << ref_data.second[i] << "\n";
path: src/plugins/intel_gpu/tests/unit/test_cases/paged_attention_gpu_test.cpp
user: login: sshlyapn
body: Removed, thanks!
created_at: 2025-05-19T14:36:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -56,6 +76,7 @@ struct paged_attention : public primitive_base<paged_attention> {
         ob << has_score_aggregation;
         ob << has_rotated_blocks;
         ob << sliding_window;
+        ob << has_score_aggregation;
path: src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
user: login: Lyamin-Roman
body: nit: Duplicate, the same in `load(..)`
```suggestion
```
Overall LGTM
created_at: 2025-05-19T14:43:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -56,6 +76,7 @@ struct paged_attention : public primitive_base<paged_attention> {
         ob << has_score_aggregation;
         ob << has_rotated_blocks;
         ob << sliding_window;
+        ob << has_score_aggregation;
path: src/plugins/intel_gpu/include/intel_gpu/primitives/paged_attention.hpp
user: login: sshlyapn
body: Thanks, I’ll fix it in follow-up PR
created_at: 2025-05-19T15:34:03Z
author_association: CONTRIBUTOR

PR comments:
user: login: yeonbok
created_at: 2025-05-01T05:22:11Z
body: It seems that this new feature is supported by sdpa_opt only shouldn't we prevent it from sdpa_micro? (Sorry if I missed something)
author_association: CONTRIBUTOR
user: login: sshlyapn
created_at: 2025-05-01T06:49:03Z
body: > It seems that this new feature is supported by sdpa_opt only shouldn't we prevent it from sdpa_micro? (Sorry if I missed something)

You are right. The thing is scores aggregation can only be enabled with scores output, and micro_sdpa is not allowed in such a case. But I updated sdpa_micro with has_score_aggregation check for clarity
author_association: CONTRIBUTOR
user: login: yeonbok
created_at: 2025-05-02T23:05:12Z
body: LGTM except some minor comments.
author_association: CONTRIBUTOR