WIP: [GPU] Use micro-sdpa for 1st token calculation of PagedAttention

sshlyapn · sshlyapn · commit e64244fdd268 · 2025-02-10T15:25:02.000+04:00
diff --git a/src/plugins/intel_gpu/src/graph/graph_optimizer/compile_graph.cpp b/src/plugins/intel_gpu/src/graph/graph_optimizer/compile_graph.cpp
@@ -24,10 +24,7 @@ void compile_graph::run(program& p) {
         }
     }
 
-    auto task_executor = p.get_task_executor();
     auto& proc_order = p.get_processing_order();
-    std::vector<ov::threading::Task> tasks;
-    std::exception_ptr exception;
 
     for (size_t idx = 0; idx < proc_order.size(); idx++) {
         auto& node = *(std::next(proc_order.begin(), idx));
@@ -36,37 +33,28 @@ void compile_graph::run(program& p) {
                                !(node->is_type<mutable_data>() && node->get_dependencies().empty());
 
         if (can_select_impl) {
-            tasks.push_back([node, &exception] {
-                try {
-                    const auto& params = node->get_kernel_impl_params();
-                    auto shape_type = ImplementationManager::get_shape_type(*params);
-                    auto selected_impl_manager = node->type()->choose_impl(*node, shape_type);
-                    std::string fail_reason = "";
-                    try {
-                        if (selected_impl_manager) {
-                            node->selected_impl = selected_impl_manager->create(*node, *params);
-                        }
-                    } catch (std::exception& e) {
-                        fail_reason = e.what();
-                    }
-
-                    OPENVINO_ASSERT(shape_type == shape_types::dynamic_shape || node->selected_impl != nullptr,
-                                    "[GPU] Failed to select implementation for"
-                                    "\nname:", node->id(),
-                                    "\ntype: ", node->get_primitive()->type_string(),
-                                    "\noriginal_type: ", node->get_primitive()->origin_op_type_name,
-                                    (!fail_reason.empty() ? fail_reason : ""));
-                } catch(...) {
-                    exception = std::current_exception();
+            // std::cout << "Compiling " << node->id() << "\n";
+            // if (idx + 1 < proc_order.size())
+            //     std::cout << "Compiling next id " << (*(std::next(proc_order.begin(), idx + 1)))->id() << "\n";
+
+            const auto& params = node->get_kernel_impl_params();
+            auto shape_type = ImplementationManager::get_shape_type(*params);
+            auto selected_impl_manager = node->type()->choose_impl(*node, shape_type);
+            std::string fail_reason = "";
+            try {
+                if (selected_impl_manager) {
+                    node->selected_impl = selected_impl_manager->create(*node, *params);
                 }
-            });
+            } catch (std::exception& e) {
+                fail_reason = e.what();
+            }
+
+            OPENVINO_ASSERT(shape_type == shape_types::dynamic_shape || node->selected_impl != nullptr,
+                            "[GPU] Failed to select implementation for"
+                            "\nname:", node->id(),
+                            "\ntype: ", node->get_primitive()->type_string(),
+                            "\noriginal_type: ", node->get_primitive()->origin_op_type_name,
+                            (!fail_reason.empty() ? fail_reason : ""));
         }
     }
-
-    task_executor->run_and_wait(tasks);
-    tasks.clear();
-
-    if (exception) {
-        std::rethrow_exception(exception);
-    }
 }
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
@@ -15,6 +15,7 @@
 #include "sdpa/pa_kv_cache_rotate_kernel_ref.h"
 #include "sdpa/pa_kv_cache_update_kernel_ref.h"
 #include "sdpa/pa_sdpa_kernel_opt.h"
+#include "sdpa/sdpa_kernel_micro.h"
 
 namespace cldnn {
 namespace ocl {
@@ -66,6 +67,33 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         return stage == PagedAttentionStage::MIXED;
     }
 
+    void update_inst_params(primitive_inst& inst) const override {
+        OPENVINO_ASSERT(inst.type() == paged_attention::type_id());
+        OPENVINO_ASSERT(inst.get_impl() == this);
+
+        auto& pa_inst = reinterpret_cast<paged_attention_inst&>(inst);
+        if (is_micro_kernel_used) {
+            auto tile_q_size = get_target_seq_len_block_size(PagedAttentionStage::PREFILL);
+            pa_inst.tile_q_size = tile_q_size;
+            std::cout << "update_inst_params: from micro-sdpa tile_q_size = " << tile_q_size << "\n";
+        } else {
+            pa_inst.tile_q_size = get_target_seq_len_block_size(PagedAttentionStage::PREFILL);
+            std::cout << "update_inst_params: sdpa_opt tile_q_size = " << get_target_seq_len_block_size(PagedAttentionStage::PREFILL) << "\n";
+        }
+    }
+
+    size_t get_target_seq_len_block_size(const PagedAttentionStage& stage) const {
+        if (stage == PagedAttentionStage::PREFILL) {
+            if (is_micro_kernel_used) {
+                return kernel_selector::SDPAKernelMicro::GetTileQSize(_kernels_data[Stage::SDPA]);
+            } else {
+                return 16;
+            }
+        } else {
+            return 16;
+        }
+    }
+
     void load(BinaryInputBuffer& ib) override {
         parent::load(ib);
         ib >> make_data(&has_scores_output, sizeof(bool));
@@ -527,7 +555,8 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
     static kv_cache_update_kernel_params_t get_kv_cache_update_kernel_params(const kernel_impl_params& impl_param,
                                                                              const PagedAttentionStage& stage,
                                                                              const kernel_selector::MultiDataTensor& input_tensors,
-                                                                             bool is_dynamic = false) {
+                                                                             int64_t target_seq_len_block_size,
+                                                                             bool is_dynamic) {
         auto params = get_default_params<kv_cache_update_kernel_params_t>(impl_param, is_dynamic);
 
         const auto& key_tensor = input_tensors[1];
@@ -557,7 +586,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         params.is_prefill = stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED;
 
         if ((stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED) && !is_dynamic)
-            params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage);
+            params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage, target_seq_len_block_size);
 
         const auto& in_offsets_map = impl_param.in_port_to_shape_info_offset;
         std::map<size_t, size_t> in_tensor_to_offset_map = {
@@ -581,13 +610,31 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
     static sdpa_kernel_params_t get_sdpa_kernel_params(const kernel_impl_params& impl_param,
                                                        const PagedAttentionStage& stage,
                                                        const kernel_selector::MultiDataTensor& input_tensors,
-                                                       bool is_dynamic = false) {
+                                                       int64_t target_seq_len_block_size,
+                                                       bool is_dynamic) {
         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];
+        auto get_sdpa_tensor = [&](const layout& input_layout, size_t head_size) {
+            auto new_layout = input_layout;
+            auto orig_shape = new_layout.get_partial_shape();
+            auto new_shape = ov::PartialShape::dynamic(4);
+
+            new_shape[0] = 1;
+            new_shape[1] = orig_shape[0];
+            new_shape[2] = orig_shape[1] / head_size;
+            new_shape[3] = head_size;
+
+            new_layout.set_partial_shape(new_shape);
+
+            std::cout << "Convert layout: " << input_layout.to_short_string() << " -> " << new_layout.to_short_string() << "\n";
+
+            return convert_data_tensor(new_layout);
+        };
+
+        const auto query_tensor = get_sdpa_tensor(impl_param.get_input_layout(0), desc->head_size);
+        const auto key_tensor = get_sdpa_tensor(impl_param.get_input_layout(1), desc->head_size);;
+        const auto value_tensor = get_sdpa_tensor(impl_param.get_input_layout(2), desc->head_size);;
         const auto& subsequence_begins_tensor = input_tensors[6];
         const auto& scale_tensor = input_tensors[9];
         const auto& alibi_tensor = input_tensors[11];
@@ -616,12 +663,17 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         if (has_alibi)
             params.inputs[input_idx++] = alibi_tensor;
 
+        params.outputs[0] = get_sdpa_tensor(impl_param.get_output_layout(0), desc->head_size);;
         if (has_scores_output) {
             params.outputs.resize(2);
             params.outputs[1] = convert_data_tensor(impl_param.get_output_layout(1));
         }
 
         params.conf = get_sdpa_configuration(impl_param, is_dynamic);
+        params.input0_order = {0, 2, 1, 3};
+        params.input1_order = {0, 2, 1, 3};
+        params.input2_order = {0, 2, 1, 3};
+        params.output_order = {0, 1, 2, 3};
 
         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;
@@ -643,7 +695,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
             in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(11)});
 
         if ((stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED) && !is_dynamic)
-            params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage);
+            params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage, target_seq_len_block_size);
 
         if (has_scores_output)
             out_tensor_to_offset_map.insert({1, out_offsets_map.at(1)});
@@ -756,11 +808,11 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
             (_kernels_data[Stage::KV_CACHE_ROTATE].update_dispatch_data_func)(kv_cache_rotate_kernel_params, _kernels_data[Stage::KV_CACHE_ROTATE]);
         }
 
-        auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+        auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, get_target_seq_len_block_size(stage), impl_param.is_dynamic());
         (_kernels_data[Stage::KV_CACHE_UPDATE].update_dispatch_data_func)(kv_cache_update_kernel_params, _kernels_data[Stage::KV_CACHE_UPDATE]);
 
         if (stage == PagedAttentionStage::PREFILL) {
-            auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+            auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, get_target_seq_len_block_size(stage), impl_param.is_dynamic());
             (_kernels_data[Stage::SDPA].update_dispatch_data_func)(sdpa_kernel_params, _kernels_data[Stage::SDPA]);
         }
 
@@ -779,11 +831,11 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
             input_tensors.emplace_back(convert_data_tensor(input_layout));
 
         const auto& desc = impl_param.typed_desc<paged_attention>();
-        auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+        auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, 0, impl_param.is_dynamic());
         auto& kv_cache_update_kernel_selector = kv_cache_update_kernel_selector_t::Instance();
         kernels_data.push_back(kv_cache_update_kernel_selector.get_best_kernel(kv_cache_update_kernel_params));
 
-        auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
+        auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, 0, impl_param.is_dynamic());
         auto& sdpa_kernel_selector = sdpa_kernel_selector_t::Instance();
         kernels_data.push_back(sdpa_kernel_selector.get_best_kernel(sdpa_kernel_params));
 
@@ -801,12 +853,19 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
         impl->has_scores_output = desc->has_scores_output();
         impl->has_rotated_blocks = desc->has_rotated_blocks;
 
+        if (!kernels_data[Stage::SDPA].kernels[0].micro_kernels.empty()) {
+            std::cout << "Micro SDPA is choosen!\n";
+            std::cout << "tile_q_size = " << kernel_selector::SDPAKernelMicro::GetTileQSize(kernels_data[Stage::SDPA]) << "\n";
+            impl->is_micro_kernel_used = true;
+        }
+
         return impl;
     }
 
 private:
     bool has_scores_output = false;
     bool has_rotated_blocks = false;
+    bool is_micro_kernel_used = false;
 };
 
 namespace detail {
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
@@ -62,7 +62,7 @@ class typed_primitive_inst<paged_attention> : public typed_primitive_inst_base<p
     memory::ptr rotation_deltas_memory_ptr() const { return input_memory_ptr(14); }
     memory::ptr rotation_trig_lut_memory_ptr() const { return input_memory_ptr(15); }
 
-    std::shared_ptr<network> prefill_network;
+    size_t tile_q_size = 0;
 
 protected:
     void on_execute() override;
diff --git a/src/plugins/intel_gpu/src/graph/include/primitive_inst.h b/src/plugins/intel_gpu/src/graph/include/primitive_inst.h
@@ -59,6 +59,7 @@ struct primitive_impl {
     virtual std::set<size_t> get_lockable_internal_buffers() const { return {}; }
     virtual void set_node_params(const program_node&) {}
     virtual const std::string& get_type_info() const = 0;
+    virtual void update_inst_params(primitive_inst& instance) const {}
     virtual void set_arguments(primitive_inst& instance) = 0;
     virtual void set_arguments(primitive_inst& instance, kernel_arguments_data& args) = 0;
     virtual event::ptr execute(const std::vector<event::ptr>& events, primitive_inst& instance) = 0;
diff --git a/src/plugins/intel_gpu/src/graph/paged_attention.cpp b/src/plugins/intel_gpu/src/graph/paged_attention.cpp
@@ -114,6 +114,9 @@ void paged_attention_inst::on_execute() {
         (stage == PagedAttentionStage::GENERATE && !has_scores_output))
         return;
 
+    OPENVINO_ASSERT(_impl != nullptr, "[GPU] impl shouldn't be nullptr");
+    _impl->update_inst_params(*this);
+
     auto& stream = get_network().get_stream();
     const auto past_lens_mem = past_lens_memory_ptr();
     const auto subsequence_begins_mem = subsequence_begins_memory_ptr();
@@ -179,7 +182,7 @@ void paged_attention_inst::on_execute() {
 
     size_t index = 0;
     size_t subsequence_offsets_acc = 0;
-    const auto target_seq_len_block_size = 16; // TODO: Get block size from the impl
+    const auto target_seq_len_block_size = static_cast<int>(tile_q_size);
     for (size_t i = 0; i < subsequence_begins_mem_lock.size() - 1; i++) {
         const auto past_len = past_lens_mem_lock[i];
         const auto seq_start = subsequence_begins_mem_lock[i];
diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_micro.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_micro.cl
@@ -143,26 +143,51 @@ KERNEL(micro_sdpa)(OPTIONAL_SHAPE_INFO_ARG
         const global QRY_DATA_T *Q,
         const global VAL_DATA_T *V,
         global half *A,
+#if IS_PAGED_ATTENTION
+    const __global INPUT3_TYPE* subsequence_begins,
+#endif
 #if WITH_ATTN_MASK
         const global half *msk,
 #endif
 #if WITH_SCALE
         global SCALE_DATA_T *scale_ptr,
 #endif
-        int d, int k, int q
+        int d,
+#if IS_PAGED_ATTENTION
+        const __global int* blocked_indexes_start,
+        const __global int* blocked_indexes_end,
+        const __global int* gws_seq_indexes_correspondence
+#else
+        int k,
+        int q
+#endif
 #ifdef KV_COMPRESSED
         , const global KEY_ATTR_SCALES_DATA_T *K_scales
         , const global KEY_ATTR_ZP_DATA_T *K_zp
         , const global VAL_ATTR_SCALES_DATA_T *V_scales
         , const global VAL_ATTR_ZP_DATA_T *V_zp
 #endif
         ) {
+#if IS_PAGED_ATTENTION
+    const uint q_tile_idx = get_group_id(0);
+    const uint block_start_pos = blocked_indexes_start[q_tile_idx];
+    const uint block_end_pos = blocked_indexes_end[q_tile_idx];
+    const uint subsequence_q_tile_idx = block_start_pos - subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
+    // const uint sequence_idx_end = block_end_pos - block_start_pos;
+    const uint subsequence_begin = subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
+    const int k = subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx] + 1] - subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
+    const int q = k;
+#endif
     uint sg_ij = sub_group_broadcast(get_local_id(1), 0);
     uint b0 = get_group_id(1);
     uint b1 = get_group_id(2);
     uint b0_kv = b0 / KV_GROUP_SIZE;
 
+#if IS_PAGED_ATTENTION
+    uint wg_j0 = subsequence_q_tile_idx;
+#else
     uint wg_j0 = get_group_id(0) * ugemm_kq_wg_tile_n;
+#endif
 
     /* Leading dimension for matrices */
     uint ldk = TRANSPOSE_K ? KEY_S3 : KEY_S2;
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
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_micro.h b/src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_micro.h
