Skip to content

Commit 8638537

Browse files
sshlyapnsshlyapn
committed
[GPU] Use sdpa-micro kernel for prefill processing in PagedAttention
1 parent ec9dfae commit 8638537

File tree

43 files changed

+436
-205
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

43 files changed

+436
-205
lines changed

‎src/plugins/intel_gpu/src/graph/impls/ocl/multi_stage_primitive.hpp

+5-5
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ struct multi_stage_primitive : public typed_primitive_impl<PType> {
7676
ob << _kernels_data.size();
7777
for (auto& kd : _kernels_data) {
7878
ob << make_data(&kd.internalBufferDataType, sizeof(kernel_selector::Datatype));
79-
ob << kd.internalBufferSizes;
79+
ob << kd.internalBuffers;
8080
ob << kd.kernels;
8181
ob << kd.kernelName;
8282
}
@@ -90,7 +90,7 @@ struct multi_stage_primitive : public typed_primitive_impl<PType> {
9090
for (size_t i = 0; i < kernels_size; i++) {
9191
kernel_selector::kernel_data kd;
9292
ib >> make_data(&kd.internalBufferDataType, sizeof(kernel_selector::Datatype));
93-
ib >> kd.internalBufferSizes;
93+
ib >> kd.internalBuffers;
9494
ib >> kd.kernels;
9595
ib >> kd.kernelName;
9696
_kernels_data[i] = kd;
@@ -160,14 +160,14 @@ struct multi_stage_primitive : public typed_primitive_impl<PType> {
160160
std::vector<layout> get_internal_buffer_layouts_impl() const override {
161161
std::vector<layout> layouts;
162162
for (auto& kd : _kernels_data) {
163-
if (kd.internalBufferSizes.empty())
163+
if (kd.internalBuffers.empty())
164164
continue;
165165

166166
auto dtype = from_data_type(kd.internalBufferDataType);
167167
const auto bpp = data_type_traits::size_of(dtype);
168-
for (auto size : kd.internalBufferSizes) {
168+
for (const auto& buffer : kd.internalBuffers) {
169169
layout inbuf_layout = {dtype, format::bfyx, // simple linear format (flattern to x channel)
170-
{1, 1, 1, (tensor::value_type)(size / bpp)}};
170+
{1, 1, 1, (tensor::value_type)(buffer.byte_count / bpp)}};
171171
layouts.push_back(inbuf_layout);
172172
}
173173
}

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

+81-29
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "sdpa/pa_kv_cache_rotate_kernel_ref.h"
1616
#include "sdpa/pa_kv_cache_update_kernel_ref.h"
1717
#include "sdpa/pa_sdpa_kernel_opt.h"
18+
#include "sdpa/sdpa_kernel_micro.h"
1819

1920
namespace cldnn {
2021
namespace ocl {
@@ -66,10 +67,34 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
6667
return stage == PagedAttentionStage::MIXED;
6768
}
6869

70+
void update_inst_params(primitive_inst& inst) const override {
71+
OPENVINO_ASSERT(inst.type() == paged_attention::type_id());
72+
OPENVINO_ASSERT(inst.get_impl() == this);
73+
74+
auto& pa_inst = reinterpret_cast<paged_attention_inst&>(inst);
75+
pa_inst.query_block_size = get_query_block_size(PagedAttentionStage::PREFILL);
76+
pa_inst.use_micro_sdpa = use_micro_sdpa;
77+
}
78+
79+
size_t get_query_block_size(const PagedAttentionStage& stage) const {
80+
const auto default_block_size = 16;
81+
82+
if (stage == PagedAttentionStage::PREFILL) {
83+
#ifdef ENABLE_ONEDNN_FOR_GPU
84+
if (use_micro_sdpa)
85+
return kernel_selector::SDPAKernelMicro::GetTileQSize(_kernels_data[Stage::SDPA]);
86+
#endif
87+
return default_block_size;
88+
} else {
89+
return default_block_size;
90+
}
91+
}
92+
6993
void load(BinaryInputBuffer& ib) override {
7094
parent::load(ib);
7195
ib >> make_data(&has_scores_output, sizeof(bool));
7296
ib >> make_data(&has_rotated_blocks, sizeof(bool));
97+
ib >> make_data(&use_micro_sdpa, sizeof(bool));
7398
if (is_dynamic()) {
7499
auto& kv_cache_update_kernel_selector = kv_cache_update_kernel_selector_t::Instance();
75100
auto kv_cache_update_kernel_impl = kv_cache_update_kernel_selector.GetImplementation(_kernels_data[Stage::KV_CACHE_UPDATE].kernelName);
@@ -95,9 +120,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
95120
parent::save(ob);
96121
ob << make_data(&has_scores_output, sizeof(bool));
97122
ob << make_data(&has_rotated_blocks, sizeof(bool));
123+
ob << make_data(&use_micro_sdpa, sizeof(bool));
98124
}
99125

100-
std::vector<layout> get_internal_buffer_layouts_impl() const override {
126+
std::vector<kernel_selector::InternalBuffer> get_internal_buffers_desc() const {
101127
/*
102128
* Internal buffers allocation owners and users:
103129
* +--------------------------------------+--------------------+--------------------+
@@ -117,6 +143,8 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
117143
* +--------------------------------------+--------------------+--------------------+
118144
* | PA_SDPA (mixed mode) + scores output | [3, 4, 5, 6, 7, 8] | |
119145
* +--------------------------------------+--------------------+--------------------+
146+
* | SDPA (1st token, micro-kernel) | [last (8/9)] | |
147+
* +--------------------------------------+--------------------+--------------------+
120148
*
121149
* Description:
122150
* 0, 1, 2 - Buffers used for proper blocks distribution for kv_cache_update and
@@ -129,24 +157,32 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
129157
* Filled in PA/SDPA kernels.
130158
* 8 - Optional buffer used for mixed PA execution mode, mapping gws idx to subsequence id.
131159
* Filled in paged_attention_inst::on_execute() call.
160+
* last - Used for defining query block index for the currently processing subsequence and mapping
161+
* gws index to subsequence idx. Values stored in pairs like:
162+
* [block_idx0, subsequence_idx0, block_idx1, subsequence_idx0, ..., block_idx0, subsequence_idx1].
163+
* Filled in paged_attention_inst::on_execute() call for sdpa-micro kernel only.
132164
*/
133165

134-
auto add_internal_buffers = [](std::vector<layout>& layouts, const kernel_selector::KernelData& kd) {
135-
if (kd.internalBufferSizes.empty())
136-
return;
137-
138-
auto dtype = from_data_type(kd.internalBufferDataType);
139-
const auto bpp = data_type_traits::size_of(dtype);
140-
for (auto size : kd.internalBufferSizes) {
141-
layout inbuf_layout = {dtype, format::bfyx, // simple linear format (flattern to x channel)
142-
{1, 1, 1, (tensor::value_type)(size / bpp)}};
143-
layouts.push_back(inbuf_layout);
144-
}
166+
auto add_internal_buffers = [](std::vector<kernel_selector::InternalBuffer>& internal_buffers,
167+
const kernel_selector::KernelData& kd) {
168+
internal_buffers.insert(internal_buffers.end(), kd.internalBuffers.begin(), kd.internalBuffers.end());
145169
};
146170

171+
std::vector<kernel_selector::InternalBuffer> internal_buffers;
172+
add_internal_buffers(internal_buffers, _kernels_data[Stage::KV_CACHE_UPDATE]);
173+
add_internal_buffers(internal_buffers, _kernels_data[Stage::PA_SDPA]);
174+
175+
if (use_micro_sdpa)
176+
add_internal_buffers(internal_buffers, _kernels_data[Stage::SDPA]);
177+
178+
return internal_buffers;
179+
}
180+
181+
std::vector<layout> get_internal_buffer_layouts_impl() const override {
147182
std::vector<layout> layouts;
148-
add_internal_buffers(layouts, _kernels_data[Stage::KV_CACHE_UPDATE]);
149-
add_internal_buffers(layouts, _kernels_data[Stage::PA_SDPA]);
183+
184+
for (const auto& buffer : get_internal_buffers_desc())
185+
layouts.emplace_back(ov::PartialShape{static_cast<int64_t>(buffer.byte_count)}, ov::element::u8, format::bfyx);
150186

151187
return layouts;
152188
}
@@ -245,12 +281,13 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
245281
}
246282

247283
std::set<size_t> get_lockable_internal_buffers() const override {
248-
size_t mixed_mode_buffer = has_scores_output ? 8 : 6;
249-
250-
std::set<size_t> lockable_ids = { 0, 1, 2, /* SDPA and KV_CACHE_UPDATE indexes configuration */
251-
mixed_mode_buffer /* PA_SDPA multiple tokens mode */ };
252-
if (has_scores_output)
253-
lockable_ids.insert(4 /* Precalculated accumulated sequence length offsets for each subsequence */);
284+
std::set<size_t> lockable_ids;
285+
const auto& internal_buffers = get_internal_buffers_desc();
286+
for (size_t i = 0; i < internal_buffers.size(); i++) {
287+
if (internal_buffers[i].lockable) {
288+
lockable_ids.insert(i);
289+
}
290+
}
254291

255292
return lockable_ids;
256293
};
@@ -271,12 +308,12 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
271308
size_t internal_buffers_offset = 0;
272309
size_t internal_buffers_count = 0;
273310
if (stage == Stage::PA_SDPA) {
274-
internal_buffers_offset = _kernels_data[Stage::KV_CACHE_UPDATE].internalBufferSizes.size();
275-
internal_buffers_count = _kernels_data[Stage::PA_SDPA].internalBufferSizes.size();
311+
internal_buffers_offset = _kernels_data[Stage::KV_CACHE_UPDATE].internalBuffers.size();
312+
internal_buffers_count = _kernels_data[Stage::PA_SDPA].internalBuffers.size();
276313
} else if (stage == Stage::KV_CACHE_UPDATE) {
277-
internal_buffers_count = _kernels_data[Stage::KV_CACHE_UPDATE].internalBufferSizes.size();
314+
internal_buffers_count = _kernels_data[Stage::KV_CACHE_UPDATE].internalBuffers.size();
278315
} else if (stage == Stage::SDPA) {
279-
internal_buffers_count = _kernels_data[Stage::KV_CACHE_UPDATE].internalBufferSizes.size();
316+
internal_buffers_count = _kernels_data[Stage::KV_CACHE_UPDATE].internalBuffers.size();
280317

281318
const auto desc = instance.get_node().as<paged_attention>().get_primitive();
282319
if (desc->has_scores_output()) {
@@ -304,6 +341,10 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
304341
intermediate_memories.begin() + internal_buffers_offset,
305342
intermediate_memories.begin() + internal_buffers_offset + internal_buffers_count);
306343

344+
if (use_micro_sdpa && stage == Stage::SDPA) {
345+
args.intermediates.push_back(intermediate_memories.back());
346+
}
347+
307348
GPU_DEBUG_TRACE_DETAIL << "Execute stage=" << stage << " kernel=" << kd_idx << " " << _kernels_data[stage].kernelName << " start_offset="
308349
<< internal_buffers_offset << " count=" << internal_buffers_count << "\n";
309350

@@ -581,7 +622,8 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
581622
static sdpa_kernel_params_t get_sdpa_kernel_params(const kernel_impl_params& impl_param,
582623
const PagedAttentionStage& stage,
583624
const kernel_selector::MultiDataTensor& input_tensors,
584-
bool is_dynamic = false) {
625+
int64_t query_block_size,
626+
bool is_dynamic) {
585627
const auto desc = impl_param.typed_desc<paged_attention>();
586628
auto params = get_default_params<sdpa_kernel_params_t>(impl_param, is_dynamic);
587629

@@ -623,6 +665,12 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
623665

624666
params.conf = get_sdpa_configuration(impl_param, is_dynamic);
625667

668+
const std::vector<int64_t> default_order = {0, 1, 2, 3};
669+
params.input0_order = default_order;
670+
params.input1_order = default_order;
671+
params.input2_order = default_order;
672+
params.output_order = default_order;
673+
626674
const auto& in_offsets_map = impl_param.in_port_to_shape_info_offset;
627675
const auto& out_offsets_map = impl_param.out_port_to_shape_info_offset;
628676
std::map<size_t, size_t> in_tensor_to_offset_map = {
@@ -643,7 +691,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
643691
in_tensor_to_offset_map.insert({input_idx++, in_offsets_map.at(11)});
644692

645693
if ((stage == PagedAttentionStage::PREFILL || stage == PagedAttentionStage::MIXED) && !is_dynamic)
646-
params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage);
694+
params.conf.paged_attention_aligned_seq_len = get_aligned_seq_len(impl_param, stage, query_block_size);
647695

648696
if (has_scores_output)
649697
out_tensor_to_offset_map.insert({1, out_offsets_map.at(1)});
@@ -760,7 +808,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
760808
(_kernels_data[Stage::KV_CACHE_UPDATE].update_dispatch_data_func)(kv_cache_update_kernel_params, _kernels_data[Stage::KV_CACHE_UPDATE]);
761809

762810
if (stage == PagedAttentionStage::PREFILL) {
763-
auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
811+
auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, get_query_block_size(stage), impl_param.is_dynamic());
764812
(_kernels_data[Stage::SDPA].update_dispatch_data_func)(sdpa_kernel_params, _kernels_data[Stage::SDPA]);
765813
}
766814

@@ -782,8 +830,7 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
782830
auto kv_cache_update_kernel_params = get_kv_cache_update_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
783831
auto& kv_cache_update_kernel_selector = kv_cache_update_kernel_selector_t::Instance();
784832
kernels_data.push_back(kv_cache_update_kernel_selector.get_best_kernel(kv_cache_update_kernel_params));
785-
786-
auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, impl_param.is_dynamic());
833+
auto sdpa_kernel_params = get_sdpa_kernel_params(impl_param, stage, input_tensors, 0, impl_param.is_dynamic());
787834
auto& sdpa_kernel_selector = sdpa_kernel_selector_t::Instance();
788835
kernels_data.push_back(sdpa_kernel_selector.get_best_kernel(sdpa_kernel_params));
789836

@@ -801,12 +848,17 @@ struct paged_attention_impl : multi_stage_primitive<paged_attention> {
801848
impl->has_scores_output = desc->has_scores_output();
802849
impl->has_rotated_blocks = desc->has_rotated_blocks;
803850

851+
if (!kernels_data[Stage::SDPA].kernels[0].micro_kernels.empty()) {
852+
impl->use_micro_sdpa = true;
853+
}
854+
804855
return impl;
805856
}
806857

807858
private:
808859
bool has_scores_output = false;
809860
bool has_rotated_blocks = false;
861+
bool use_micro_sdpa = false;
810862
};
811863

812864
namespace detail {

‎src/plugins/intel_gpu/src/graph/impls/ocl/primitive_base.hpp

+5-5
Original file line numberDiff line numberDiff line change
@@ -74,15 +74,15 @@ struct typed_primitive_impl_ocl : public typed_primitive_impl<PType> {
7474
void save(BinaryOutputBuffer& ob) const override {
7575
primitive_impl::save(ob);
7676
ob << make_data(&_kernel_data.internalBufferDataType, sizeof(kernel_selector::Datatype));
77-
ob << _kernel_data.internalBufferSizes;
77+
ob << _kernel_data.internalBuffers;
7878
ob << _kernel_data.kernels;
7979
ob << _kernel_data.kernelName;
8080
}
8181

8282
void load(BinaryInputBuffer& ib) override {
8383
primitive_impl::load(ib);
8484
ib >> make_data(&_kernel_data.internalBufferDataType, sizeof(kernel_selector::Datatype));
85-
ib >> _kernel_data.internalBufferSizes;
85+
ib >> _kernel_data.internalBuffers;
8686
ib >> _kernel_data.kernels;
8787
ib >> _kernel_data.kernelName;
8888
}
@@ -185,15 +185,15 @@ struct typed_primitive_impl_ocl : public typed_primitive_impl<PType> {
185185
}
186186

187187
std::vector<layout> get_internal_buffer_layouts_impl() const override {
188-
if (_kernel_data.internalBufferSizes.empty())
188+
if (_kernel_data.internalBuffers.empty())
189189
return {};
190190

191191
std::vector<layout> layouts;
192192
auto dtype = from_data_type(_kernel_data.internalBufferDataType);
193193
const auto bpp = data_type_traits::size_of(dtype);
194-
for (auto size : _kernel_data.internalBufferSizes) {
194+
for (const auto& buffer : _kernel_data.internalBuffers) {
195195
layout inbuf_layout = {dtype, format::bfyx, // simple linear format (flattern to x channel)
196-
{1, 1, 1, (tensor::value_type)(size / bpp)}};
196+
{1, 1, 1, (tensor::value_type)(buffer.byte_count / bpp)}};
197197
layouts.push_back(inbuf_layout);
198198
}
199199
return layouts;

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

+4-4
Original file line numberDiff line numberDiff line change
@@ -71,19 +71,19 @@ struct scaled_dot_product_attention_impl : multi_stage_primitive<scaled_dot_prod
7171
// _kernels_data[1] - sdpa_opt (indirect)
7272
// => use internal buffers from [1] kernel
7373
size_t kernel_idx = _kernels_data.size();
74-
if (_kernels_data.size() >= 1 && !_kernels_data[0].internalBufferSizes.empty()) {
74+
if (_kernels_data.size() >= 1 && !_kernels_data[0].internalBuffers.empty()) {
7575
kernel_idx = 0;
76-
} else if (_kernels_data.size() >= 2 && !_kernels_data[1].internalBufferSizes.empty()) {
76+
} else if (_kernels_data.size() >= 2 && !_kernels_data[1].internalBuffers.empty()) {
7777
kernel_idx = 1;
7878
}
7979

8080
std::vector<layout> layouts;
8181
if (kernel_idx < _kernels_data.size()) {
8282
auto dtype = from_data_type(_kernels_data[kernel_idx].internalBufferDataType);
8383
const auto bpp = data_type_traits::size_of(dtype);
84-
for (auto size : _kernels_data[kernel_idx].internalBufferSizes) {
84+
for (const auto& buffer : _kernels_data[kernel_idx].internalBuffers) {
8585
layout inbuf_layout = {dtype, format::bfyx, // simple linear format (flattern to x channel)
86-
{1, 1, 1, (tensor::value_type)(size / bpp)}};
86+
{1, 1, 1, (tensor::value_type)(buffer.byte_count / bpp)}};
8787
layouts.push_back(inbuf_layout);
8888
}
8989
}

‎src/plugins/intel_gpu/src/graph/include/paged_attention_inst.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,8 @@ class typed_primitive_inst<paged_attention> : public typed_primitive_inst_base<p
6262
memory::ptr rotation_deltas_memory_ptr() const { return input_memory_ptr(14); }
6363
memory::ptr rotation_trig_lut_memory_ptr() const { return input_memory_ptr(15); }
6464

65-
std::shared_ptr<network> prefill_network;
65+
bool use_micro_sdpa = false;
66+
size_t query_block_size = 0;
6667

6768
protected:
6869
void on_execute() override;

‎src/plugins/intel_gpu/src/graph/include/primitive_inst.h

+2-1
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ struct primitive_impl {
5959
virtual std::set<size_t> get_lockable_internal_buffers() const { return {}; }
6060
virtual void set_node_params(const program_node&) {}
6161
virtual const std::string& get_type_info() const = 0;
62+
virtual void update_inst_params(primitive_inst& instance) const {}
6263
virtual void set_arguments(primitive_inst& instance) = 0;
6364
virtual void set_arguments(primitive_inst& instance, kernel_arguments_data& args) = 0;
6465
virtual event::ptr execute(const std::vector<event::ptr>& events, primitive_inst& instance) = 0;
@@ -413,7 +414,7 @@ class primitive_inst {
413414
std::vector<memory::ptr> allocate_outputs(kernel_impl_params* updated_params = nullptr,
414415
bool reset_mem = true,
415416
bool runtime_alloc = false);
416-
memory::ptr allocate_internal_buffer(size_t idx, bool reset = true);
417+
memory::ptr allocate_internal_buffer(const layout& layout, size_t idx, bool reset = true, bool lockable = false);
417418
void allocate_shape_info_memory();
418419
static std::vector<primitive_inst*> build_exec_deps(
419420
std::vector<std::pair<primitive_inst*, int32_t>> const& mem_deps);

0 commit comments

Comments
 (0)