natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[GPU] SDPA indirect inputs (#24665) 

### Details:
 - Added indirect inputs support for SDPA kernel
 - Added setter for causal flag for ScaledDotProductAttention operation
- Added `ov::intel_gpu::hint::enable_sdpa_optimization` to
`ov::supported_properties` list
- Removed unused `TARGET_SEQ_LEN_BLOCK_SIZE > 1` check from kernel for
single token processing
 - Minor refactoring
- Added `OV_GPU_EnableSDPA` debug option (which allows to force SDPA
kernel for any ScaledDotProductAttention operation _(=1)_ / or
completely disable SDPA kernel _(=0)_, ignoring
`ov::intel_gpu::hint::enable_sdpa_optimization` property)

### Tickets:
 - *CVS-141213*
This commit is contained in:
Sergey Shlyapnikov 2024-05-27 10:11:37 +04:00 committed by GitHub
parent af88a206eb
commit dc523beb3e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
22 changed files with 830 additions and 106 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/core/include/openvino/op/scaled_dot_product_attention.hpp
View File

@ -50,6 +50,10 @@ public:
return m_causal;
}
void set_causal(bool causal) {
m_causal = causal;
}
private:
bool m_causal = false;
};

78
src/plugins/intel_gpu/include/intel_gpu/op/indirect_sdpa.hpp Normal file
View File

@ -0,0 +1,78 @@
// Copyright (C) 2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#pragma once
#include "intel_gpu/op/sdpa.hpp"
#include "openvino/core/node.hpp"
#include "openvino/core/partial_shape.hpp"
#include "openvino/op/op.hpp"
namespace ov {
namespace intel_gpu {
namespace op {
class IndirectSDPA : public ov::intel_gpu::op::SDPA {
public:
OPENVINO_OP("IndirectSDPA", "gpu_opset");
IndirectSDPA() = default;
IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type = ov::element::undefined);
IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& attn_mask,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type = ov::element::undefined);
IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& attn_mask,
const ov::Output<Node>& scale,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type = ov::element::undefined);
bool visit_attributes(ov::AttributeVisitor &visitor) override;
void validate_and_infer_types() override;
std::shared_ptr<Node> clone_with_new_inputs(const ov::OutputVector& new_args) const override;
ov::element::Type get_output_type() const { return m_output_type; }
int64_t get_indirect_axis() const { return m_indirect_axis; }
using ov::intel_gpu::op::SDPA::default_order;
protected:
int64_t m_indirect_axis = -1;
};
} // namespace op
} // namespace intel_gpu
} // namespace ov

1
src/plugins/intel_gpu/include/intel_gpu/plugin/primitives_list.hpp
View File

@ -285,3 +285,4 @@ REGISTER_FACTORY(internal, IndirectGemm);
REGISTER_FACTORY(internal, Convolution);
REGISTER_FACTORY(internal, Placeholder);
REGISTER_FACTORY(internal, SDPA);
REGISTER_FACTORY(internal, IndirectSDPA);

17
src/plugins/intel_gpu/include/intel_gpu/primitives/scaled_dot_product_attention.hpp
View File

@ -19,6 +19,7 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
scaled_dot_product_attention(const primitive_id& id,
const std::vector<cldnn::input_info> inputs,
bool is_causal,
int64_t indirect_axis = -1,
const std::vector<int64_t>& input_q_transpose_order = {},
const std::vector<int64_t>& input_k_transpose_order = {},
const std::vector<int64_t>& input_v_transpose_order = {},
@ -26,17 +27,23 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
const padding& output_padding = padding())
: primitive_base(id, inputs, {output_padding})
, is_causal(is_causal)
, has_attn_mask_input(inputs.size() > 3)
, has_scale_input(inputs.size() > 4)
, indirect_axis(indirect_axis)
, input_q_transpose_order(input_q_transpose_order)
, input_k_transpose_order(input_k_transpose_order)
, input_v_transpose_order(input_v_transpose_order)
, output_transpose_order(output_transpose_order) {}
, output_transpose_order(output_transpose_order) {
auto data_inputs_num = inputs.size();
if (indirect_axis != -1)
data_inputs_num--;
has_attn_mask_input = data_inputs_num > 3;
has_scale_input = data_inputs_num > 4;
}
bool is_causal = false;
bool has_attn_mask_input = false;
bool has_scale_input = false;
int64_t indirect_axis = -1;
std::vector<int64_t> input_q_transpose_order;
std::vector<int64_t> input_k_transpose_order;
@ -48,6 +55,7 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
seed = hash_combine(seed, is_causal);
seed = hash_combine(seed, has_attn_mask_input);
seed = hash_combine(seed, has_scale_input);
seed = hash_combine(seed, indirect_axis);
seed = hash_range(seed, input_q_transpose_order.begin(), input_q_transpose_order.end());
seed = hash_range(seed, input_k_transpose_order.begin(), input_k_transpose_order.end());
seed = hash_range(seed, input_v_transpose_order.begin(), input_v_transpose_order.end());
@ -64,6 +72,7 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
return is_causal == rhs_casted.is_causal &&
has_attn_mask_input == rhs_casted.has_attn_mask_input &&
has_scale_input == rhs_casted.has_scale_input &&
indirect_axis == rhs_casted.indirect_axis &&
input_q_transpose_order == rhs_casted.input_q_transpose_order &&
input_k_transpose_order == rhs_casted.input_k_transpose_order &&
input_v_transpose_order == rhs_casted.input_v_transpose_order &&
@ -75,6 +84,7 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
ob << is_causal;
ob << has_attn_mask_input;
ob << has_scale_input;
ob << indirect_axis;
ob << input_q_transpose_order;
ob << input_k_transpose_order;
ob << input_v_transpose_order;
@ -86,6 +96,7 @@ struct scaled_dot_product_attention : public primitive_base<scaled_dot_product_a
ib >> is_causal;
ib >> has_attn_mask_input;
ib >> has_scale_input;
ib >> indirect_axis;
ib >> input_q_transpose_order;
ib >> input_k_transpose_order;
ib >> input_v_transpose_order;

1
src/plugins/intel_gpu/include/intel_gpu/runtime/debug_configuration.hpp
View File

@ -129,6 +129,7 @@ public:
std::vector<std::string> forced_impl_types; // Force implementation type either ocl or onednn
int max_kernels_per_batch; // Maximum number of kernels in a batch during compiling kernels
int impls_cache_capacity; // The maximum number of entries in the kernel impl cache
int enable_sdpa; // Allows to control SDPA decomposition
int disable_async_compilation; // Disable async compilation
int disable_winograd_conv; // Disable Winograd conv
int disable_dynamic_impl; // Disable dynamic implementation

231
src/plugins/intel_gpu/src/graph/impls/ocl/scaled_dot_product_attention.cpp
View File

@ -2,34 +2,188 @@
// SPDX-License-Identifier: Apache-2.0
//
#include "primitive_base.hpp"
#include "multi_stage_primitive.hpp"
#include "scaled_dot_product_attention_inst.h"
#include "kv_cache_inst.h"
#include "sdpa/sdpa_kernel_selector.h"
#include "sdpa/sdpa_kernel_base.h"
namespace cldnn {
namespace ocl {
struct scaled_dot_product_attention_impl : typed_primitive_impl_ocl<scaled_dot_product_attention> {
using parent = typed_primitive_impl_ocl<scaled_dot_product_attention>;
// SDPA impl may create 2 versions of the kernel internally
// 1. Default SDPA kernels
// 2. SDPA kernels with indirect access to one of the inputs
// This feature is used to avoid perf drop when we create single kernel which checks batch size in runtime
// Can be reverted once performance of the kernel is improved
struct scaled_dot_product_attention_impl : multi_stage_primitive<scaled_dot_product_attention> {
using parent = multi_stage_primitive<scaled_dot_product_attention>;
using parent::parent;
using kernel_selector_t = kernel_selector::sdpa_kernel_selector;
using kernel_params_t = kernel_selector::sdpa_params;
DECLARE_OBJECT_TYPE_SERIALIZATION(cldnn::ocl::scaled_dot_product_attention_impl)
const uint32_t default_sdpa = 0;
const uint32_t indirect_sdpa = 1;
std::unique_ptr<primitive_impl> clone() const override {
return make_unique<scaled_dot_product_attention_impl>(*this);
}
scaled_dot_product_attention_impl() = default;
scaled_dot_product_attention_impl(const std::vector<kernel_selector::kernel_data>& kd) : parent(kd) {
this->can_reuse_memory = true;
}
void load(BinaryInputBuffer& ib) override {
parent::load(ib);
if (is_dynamic()) {
auto& kernel_selector = kernel_selector_t::Instance();
auto kernel_impl = kernel_selector.GetImplementation(_kernel_data.kernelName);
kernel_impl->GetUpdateDispatchDataFunc(_kernel_data);
auto kernel_impl = kernel_selector.GetImplementation(_kernels_data[default_sdpa].kernelName);
kernel_impl->GetUpdateDispatchDataFunc(_kernels_data[default_sdpa]);
if (_kernels_data.size() == 2) {
auto bt_kernel_impl = kernel_selector.GetImplementation(_kernels_data[indirect_sdpa].kernelName);
bt_kernel_impl->GetUpdateDispatchDataFunc(_kernels_data[indirect_sdpa]);
}
}
}
protected:
std::vector<layout> get_internal_buffer_layouts_impl() const override {
// TODO: current implementation is supposed to have the same kernel version for both indirect/default paths,
// considering this, we may assume that both indirect/default kernels have absolutely the same intermediate
// buffers number and its' sizes (since update_dispatch_data is called for both kernels too), and
// do not double memory allocations during reallocate_if_needed() function call
std::vector<layout> layouts;
if (_kernels_data.size() > 0) {
auto dtype = from_data_type(_kernels_data[0].internalBufferDataType);
const auto bpp = data_type_traits::size_of(dtype);
for (auto size : _kernels_data[0].internalBufferSizes) {
layout inbuf_layout = {dtype, format::bfyx, // simple linear format (flattern to x channel)
{1, 1, 1, (tensor::value_type)(size / bpp)}};
layouts.push_back(inbuf_layout);
}
}
return layouts;
}
static size_t get_beam_table_id(std::shared_ptr<const scaled_dot_product_attention> primitive) {
GPU_DEBUG_TRACE << "get_beam_table_id " << primitive->input_size() - 1 << "\n";
return primitive->input_size() - 1;
}
static bool has_indirect_inputs(const kernel_impl_params& impl_param) {
const auto& desc = impl_param.typed_desc<scaled_dot_product_attention>();
return desc->indirect_axis != -1;
}
kernel_arguments_data get_arguments(const scaled_dot_product_attention_inst& instance, size_t stage) const override {
kernel_arguments_data args;
auto inputs_num = instance.inputs_memory_count();
if (instance.has_indirect_inputs() && stage == default_sdpa)
inputs_num--;
for (size_t i = 0; i < inputs_num; i++) {
args.inputs.push_back(instance.input_memory_ptr(i));
}
if (instance.has_fused_primitives()) {
size_t count = instance.get_fused_mem_count();
for (size_t i = 0; i < count; i++) {
args.fused_op_inputs.push_back(instance.fused_memory(i));
}
}
for (size_t i = 0; i < instance.outputs_memory_count(); i++) {
args.outputs.push_back(instance.output_memory_ptr(i));
}
args.shape_info = instance.shape_info_memory_ptr();
return args;
}
void set_arguments_impl(scaled_dot_product_attention_inst& instance) override {}
event::ptr execute_stage(const std::vector<event::ptr>& events, scaled_dot_product_attention_inst& instance, size_t stage) {
stream& stream = instance.get_network().get_stream();
std::vector<event::ptr> tmp_events(events);
std::vector<event::ptr> all_events;
size_t kernel_offset = 0;
for (size_t s = 0; s < stage; s++) {
kernel_offset += _kernels_data[s].kernels.size();
}
for (size_t kd_idx = 0; kd_idx < _kernels_data[stage].kernels.size(); ++kd_idx) {
if (_kernels_data[stage].kernels[kd_idx].skip_execution)
continue;
size_t idx_final = kernel_offset + kd_idx;
// If any user of the desc's users is CPU implementation or network's output, set desc as a output event (event won't be nullptr)
bool needs_completion_event = instance.needs_completion_event();
auto& params = _kernels_data[stage].kernels[kd_idx].params;
auto args = get_arguments(instance, stage);
args.scalars = &params.scalars;
for (size_t i = 0; i < instance.get_intermediates_memories().size(); i++)
args.intermediates.push_back(instance.get_intermediates_memories()[i]);
stream.set_arguments(*_kernels[idx_final], _kernels_data[stage].kernels[kd_idx].params, args);
const auto& gws = params.workGroups.global;
const auto& lws = params.workGroups.local;
GPU_DEBUG_TRACE_DETAIL << "Enqueue stage " << stage << " kernel " << idx_final << ": gws=[" << gws[0] << ", " << gws[1] << ", " << gws[2] << "] "
<< "lws=[" << lws[0] << ", " << lws[1] << ", " << lws[2] << "]"
<< (needs_completion_event ? " has_completion_event=true" : "") << std::endl;
auto ev = stream.enqueue_kernel(*_kernels[idx_final], params, args, tmp_events, needs_completion_event);
if (_kernels_data[stage].needs_sub_kernels_sync) {
tmp_events = {ev};
}
all_events.push_back(ev);
}
return aggregate_events(all_events, stream, all_events.size() > 1);
}
bool need_indirect_load(const scaled_dot_product_attention_inst& instance) const {
auto desc = instance.get_typed_desc<scaled_dot_product_attention>();
if (!instance.has_indirect_inputs())
return false;
const auto& params = *instance.get_impl_params();
const auto indirect_axis = desc->indirect_axis;
if (params.input_layouts[get_beam_table_id(desc)].get_partial_shape()[indirect_axis].get_length() == 1)
return false;
const auto& deps = instance.dependencies();
const auto indirect_dep_idx = 1;
const auto& indirect_dep = deps[indirect_dep_idx].first;
if (dynamic_cast<const kv_cache_inst*>(indirect_dep) == nullptr) {
return true;
}
auto state_layout = indirect_dep->get_impl_params()->get_input_layout(0);
bool is_prefill = state_layout.count() == 0;
return !is_prefill;
}
event::ptr execute_impl(const std::vector<event::ptr>& events, scaled_dot_product_attention_inst& instance) override {
if (need_indirect_load(instance))
return execute_stage(events, instance, indirect_sdpa);
else
return execute_stage(events, instance, default_sdpa);
}
static kernel_selector::sdpa_configuration get_sdpa_configuration(const kernel_impl_params& impl_param) {
kernel_selector::sdpa_configuration config;
@ -44,16 +198,16 @@ struct scaled_dot_product_attention_impl : typed_primitive_impl_ocl<scaled_dot_p
return transposed_pshape;
};
const auto& prim = impl_param.typed_desc<scaled_dot_product_attention>();
const auto query_shape = transpose_pshape(impl_param.get_input_layout(0).get_partial_shape(), prim->input_q_transpose_order);
const auto key_shape = transpose_pshape(impl_param.get_input_layout(1).get_partial_shape(), prim->input_k_transpose_order);
const auto value_shape = transpose_pshape(impl_param.get_input_layout(2).get_partial_shape(), prim->input_v_transpose_order);
const auto& desc = impl_param.typed_desc<scaled_dot_product_attention>();
const auto query_shape = transpose_pshape(impl_param.get_input_layout(0).get_partial_shape(), desc->input_q_transpose_order);
const auto key_shape = transpose_pshape(impl_param.get_input_layout(1).get_partial_shape(), desc->input_k_transpose_order);
const auto value_shape = transpose_pshape(impl_param.get_input_layout(2).get_partial_shape(), desc->input_v_transpose_order);
OPENVINO_ASSERT(key_shape == value_shape, "[GPU] The shapes of key and value inputs are expected to be equal");
for (size_t i = 0; i < query_shape.size(); ++i) {
if (query_shape[i].is_static() && key_shape[i].is_static() && value_shape[i].is_static()) {
if (query_shape[i].get_length() > key_shape[i].get_length()) {
config.broadcast_axis = prim->input_k_transpose_order[i];
config.broadcast_axis = desc->input_k_transpose_order[i];
config.group_size = query_shape[i].get_length() / key_shape[i].get_length();
}
}
@ -62,44 +216,73 @@ struct scaled_dot_product_attention_impl : typed_primitive_impl_ocl<scaled_dot_p
if (query_shape[query_shape.size() - 1].is_static())
config.head_size = query_shape[query_shape.size() - 1].get_length();
config.is_causal = prim->is_causal;
config.is_causal = desc->is_causal;
return config;
}
static kernel_params_t get_kernel_params(const kernel_impl_params& impl_param, bool is_dynamic) {
public:
static kernel_params_t get_kernel_params(const kernel_impl_params& impl_param, bool is_dynamic, bool indirect = false) {
const auto& desc = impl_param.typed_desc<scaled_dot_product_attention>();
auto params = get_default_params<kernel_selector::sdpa_params>(impl_param, is_dynamic);
const auto inputs_num = impl_param.input_layouts.size();
params.inputs.resize(inputs_num);
for (size_t i = 0; i < inputs_num; i++) {
auto data_inputs_num = impl_param.input_layouts.size();
if (has_indirect_inputs(impl_param))
data_inputs_num--;
params.inputs.resize(data_inputs_num);
for (size_t i = 0; i < data_inputs_num; i++) {
params.inputs[i] = convert_data_tensor(impl_param.get_input_layout(i));
}
params.conf = get_sdpa_configuration(impl_param);
const auto& prim = impl_param.typed_desc<scaled_dot_product_attention>();
params.input0_order = prim->input_q_transpose_order;
params.input1_order = prim->input_k_transpose_order;
params.input2_order = prim->input_v_transpose_order;
params.output_order = prim->output_transpose_order;
params.input0_order = desc->input_q_transpose_order;
params.input1_order = desc->input_k_transpose_order;
params.input2_order = desc->input_v_transpose_order;
params.output_order = desc->output_transpose_order;
if (indirect && has_indirect_inputs(impl_param)) {
params.beam_table = convert_data_tensor(impl_param.get_input_layout(get_beam_table_id(desc)));
params.indirect_axis = desc->indirect_axis;
}
params.set_dynamic_shape_offsets();
// Need to adjust sdpa kernel offset to consider beam table input
if (has_indirect_inputs(impl_param)) {
auto out_offset = params.outputs[0].get_dynamic_shape_offset();
if (indirect)
params.beam_table.SetDynamicShapeOffset(out_offset);
params.outputs[0].SetDynamicShapeOffset(out_offset + kernel_selector::DataTensor::max_rank());
}
return params;
}
static std::unique_ptr<primitive_impl> create(const typed_program_node<scaled_dot_product_attention>& arg, const kernel_impl_params& impl_param) {
std::vector<kernel_selector::kernel_data> kernels_data;
auto sdpa_kernel_params = get_kernel_params(impl_param, impl_param.is_dynamic());
auto& sdpa_kernel_selector = kernel_selector_t::Instance();
auto kd = sdpa_kernel_selector.get_best_kernel(sdpa_kernel_params);
auto& kernel_selector = kernel_selector_t::Instance();
kernels_data.push_back(kernel_selector.get_best_kernel(sdpa_kernel_params));
return cldnn::make_unique<scaled_dot_product_attention_impl>(kd);
if (has_indirect_inputs(impl_param)) {
auto indirect_kernel_params = get_kernel_params(impl_param, impl_param.is_dynamic(), true);
kernels_data.push_back(kernel_selector.get_best_kernel(indirect_kernel_params));
}
return cldnn::make_unique<scaled_dot_product_attention_impl>(kernels_data);
}
void update_dispatch_data(const kernel_impl_params& impl_param) override {
auto kernel_params = get_kernel_params(impl_param, true);
(_kernel_data.update_dispatch_data_func)(kernel_params, _kernel_data);
(_kernels_data[default_sdpa].update_dispatch_data_func)(kernel_params, _kernels_data[default_sdpa]);
if (_kernels_data.size() == 2) {
auto kernel_params = get_kernel_params(impl_param, true);
(_kernels_data[indirect_sdpa].update_dispatch_data_func)(kernel_params, _kernels_data[indirect_sdpa]);
}
}
};

3
src/plugins/intel_gpu/src/graph/include/scaled_dot_product_attention_inst.h
View File

@ -32,6 +32,9 @@ public:
static std::vector<layout> calc_output_layouts(scaled_dot_product_attention_node const& /*node*/, const kernel_impl_params& impl_param);
static layout calc_output_layout(scaled_dot_product_attention_node const& node, kernel_impl_params const& impl_param);
static std::string to_string(scaled_dot_product_attention_node const& node);
bool has_indirect_inputs() const {
return get_typed_desc<scaled_dot_product_attention>()->indirect_axis != -1;
}
typed_primitive_inst(network& network, scaled_dot_product_attention_node const& desc);
};

131
src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_opt.cl
View File

@ -96,6 +96,24 @@ inline uint FUNC(get_input2_index)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint
#endif
}
#ifdef BEAM_TABLE_TYPE
inline uint FUNC(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
#if BEAM_TABLE_SIMPLE
return GET_DATA_INDEX_6D_SAFE(BEAM_TABLE, b, f, w, z, y, x);
#else
# error sdpa_ref.cl : Unsupported beam table format
#endif
}
inline uint FUNC(get_bt_index_key)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
return FUNC_CALL(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_TENSOR INPUT1_DIMS_ORDER);
}
inline uint FUNC(get_bt_index_value)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
return FUNC_CALL(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_TENSOR INPUT2_DIMS_ORDER);
}
#endif
#define VALUE_BLOCK_READ(ptr, offset) BLOCK_READN(INPUT2_TYPE, 1, ptr, offset)
#define SUBGROUPS_PER_WG (HEAD_SIZE / SUBGROUP_SIZE)
@ -117,6 +135,9 @@ KERNEL(sdpa_opt)(
const __global INPUT4_TYPE* scale,
#endif
__global OUTPUT_TYPE* output,
#ifdef BEAM_TABLE_TYPE
const __global BEAM_TABLE_TYPE* beam_table,
#endif
__global SOFTMAX_ACCUMULATOR_TYPE* exp_sums,
__global SOFTMAX_ACCUMULATOR_TYPE* max_logits,
__global OUTPUT_TYPE* tmp_out
@ -125,12 +146,7 @@ KERNEL(sdpa_opt)(
const uint batch_idx = get_global_id(0);
const uint b0_idx = batch_idx / NUM_HEADS; /* BATCH dim */
const uint b1_idx = batch_idx % NUM_HEADS; /* HEADS_NUM dim */
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint target_seq_idx = (uint)get_global_id(1) * TARGET_SEQ_LEN_BLOCK_SIZE;
#else
const uint target_seq_idx = get_global_id(1);
#endif
const uint lid = get_local_id(2);
const uint head_size_idx = lid;
@ -173,12 +189,7 @@ KERNEL(sdpa_opt)(
// Query input loading to SLM
#define QUERY_STEP_LOCAL SUBGROUP_SIZE * SUBGROUPS_PER_WG
uint query_local_offset = sgid * SUBGROUP_SIZE + sglid;
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint seq_idx_end = min(TARGET_SEQ_LEN - target_seq_idx, (uint)TARGET_SEQ_LEN_BLOCK_SIZE);
#else
const uint seq_idx_end = 1;
#endif
#ifdef INPUT0_DIMS_ORDER
uint query_offset = FUNC_CALL(get_input0_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, target_seq_idx, (sgid * SUBGROUP_SIZE));
uint query_offset_next_seq = FUNC_CALL(get_input0_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, target_seq_idx + 1, (sgid * SUBGROUP_SIZE));
@ -207,9 +218,14 @@ KERNEL(sdpa_opt)(
// HEAD_SIZE / SUBGROUPS_PER_WG times in the loop and saves the result to the qk_local SLM buffer
for (uint seq_len = sgid; seq_len < partition_seq_len; seq_len += (HEAD_SIZE / SUBGROUP_SIZE)) {
#ifdef INPUT1_DIMS_ORDER
uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len, 0);
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_key)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len, 0)];
#else
uint key_offset = INPUT1_GET_INDEX(b0_idx, b1_idx, start_partition_idx + seq_len, 0);
const uint b_idx = b0_idx;
#endif
const uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + seq_len, 0);
#else
const uint key_offset = INPUT1_GET_INDEX(b0_idx, b1_idx, start_partition_idx + seq_len, 0);
#endif
INPUT0_TYPE acc[TARGET_SEQ_LEN_BLOCK_SIZE] = {INPUT0_VAL_ZERO};
@ -316,11 +332,7 @@ KERNEL(sdpa_opt)(
barrier(CLK_LOCAL_MEM_FENCE);
INPUT0_TYPE qk_val[TARGET_SEQ_LEN_BLOCK_SIZE];
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint seq_idx_end = min(TARGET_SEQ_LEN - target_seq_idx, (uint)TARGET_SEQ_LEN_BLOCK_SIZE);
#else
const uint seq_idx_end = 1;
#endif
for (uint seq_idx = 0; seq_idx < seq_idx_end; seq_idx++) {
// Iterate over all values QK values in SLM and apply scale and attention mask
for (uint seq_len = sgid * SUBGROUP_SIZE + sglid; seq_len < partition_seq_len; seq_len += (HEAD_SIZE)) {
@ -349,11 +361,7 @@ KERNEL(sdpa_opt)(
{
// SoftMax calculation
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint seq_idx_end = min(TARGET_SEQ_LEN - target_seq_idx, (uint)TARGET_SEQ_LEN_BLOCK_SIZE);
#else
const uint seq_idx_end = 1;
#endif
// Find the maximum value of qk in the subgroup
for (uint seq_idx = 0; seq_idx < seq_idx_end; seq_idx++) {
qk_max[seq_idx] = sub_group_reduce_max(qk_max[seq_idx]);
@ -446,20 +454,26 @@ KERNEL(sdpa_opt)(
{
// Gemm2 calculation
OUTPUT_TYPE acc[TARGET_SEQ_LEN_BLOCK_SIZE] = {OUTPUT_VAL_ZERO};
#ifndef BEAM_TABLE_TYPE
#ifdef INPUT2_DIMS_ORDER
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, 0, 0);
uint value_offset_next_seq = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, 1, 0);
const uint value_pitch = value_offset_next_seq - value_offset;
#else
const uint value_pitch = HEAD_SIZE;
#endif
#endif
for (uint seq_len = 0; seq_len < partition_seq_len / SUBGROUP_SIZE; seq_len++) {
#ifdef BEAM_TABLE_TYPE
uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, (head_size_idx / SUBGROUP_SIZE) * SUBGROUP_SIZE)];
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, (head_size_idx / SUBGROUP_SIZE) * SUBGROUP_SIZE);
#else
#ifdef INPUT2_DIMS_ORDER
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#else
uint value_offset = INPUT2_GET_INDEX(b0_idx, b1_idx, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#endif
#endif
OUTPUT_TYPE qk_val[TARGET_SEQ_LEN_BLOCK_SIZE];
@ -468,19 +482,30 @@ KERNEL(sdpa_opt)(
}
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
#ifdef BEAM_TABLE_TYPE
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, sub_group_broadcast(value_offset, i));
#else
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, value_offset);
#endif
unroll_for (uint seq_idx = 0; seq_idx < TARGET_SEQ_LEN_BLOCK_SIZE; seq_idx++) {
acc[seq_idx] = mad(sub_group_broadcast(qk_val[seq_idx], i), value_val, acc[seq_idx]);
}
#ifndef BEAM_TABLE_TYPE
value_offset += value_pitch;
#endif
}
}
const uint seq_len_leftovers_start = (partition_seq_len / SUBGROUP_SIZE) * SUBGROUP_SIZE;
for (uint seq_len = seq_len_leftovers_start; seq_len < partition_seq_len; seq_len++) {
#ifdef INPUT2_DIMS_ORDER
const uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len, head_size_idx);
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len, head_size_idx)];
#else
const uint b_idx = b0_idx;
#endif
const uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + seq_len, head_size_idx);
#else
const uint value_offset = INPUT2_GET_INDEX(b0_idx, b1_idx, start_partition_idx + seq_len, head_size_idx);
#endif
@ -500,11 +525,7 @@ KERNEL(sdpa_opt)(
// If the number of partitions is greater than 1, save results to the temporary buffer;
// otherwise, save results directly to the main output.
if (num_of_partitions > 1) {
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint seq_idx_end = min(TARGET_SEQ_LEN - target_seq_idx, (uint)TARGET_SEQ_LEN_BLOCK_SIZE);
#else
const uint seq_idx_end = 1;
#endif
for (uint seq_idx = 0; seq_idx < seq_idx_end; seq_idx++) {
// Data layout of tmp_output buf: [batch, heads_num, q_len, partition_idx, head_size]
const uint tmp_out_offset = b0_idx * (NUM_HEADS * TARGET_SEQ_LEN * num_of_partitions * HEAD_SIZE) +
@ -515,15 +536,11 @@ KERNEL(sdpa_opt)(
tmp_out[tmp_out_offset] = acc[seq_idx];
}
} else {
#if TARGET_SEQ_LEN_BLOCK_SIZE > 1
const uint seq_idx_end = min(TARGET_SEQ_LEN - target_seq_idx, (uint)TARGET_SEQ_LEN_BLOCK_SIZE);
#else
const uint seq_idx_end = 1;
#endif
for (uint seq_idx = 0; seq_idx < seq_idx_end; seq_idx++) {
const uint output_offset = OUTPUT_GET_INDEX(b0_idx, b1_idx, target_seq_idx + seq_idx, head_size_idx);
const uint output_offset = OUTPUT_GET_INDEX(b0_idx, b1_idx, target_seq_idx + seq_idx, head_size_idx);
output[output_offset] = acc[seq_idx];
output[output_offset] = acc[seq_idx];
}
}
} // Gemm2 calculation end
@ -545,6 +562,9 @@ KERNEL(sdpa_opt)(
const __global INPUT4_TYPE* scale,
#endif
__global OUTPUT_TYPE* output,
#ifdef BEAM_TABLE_TYPE
const __global BEAM_TABLE_TYPE* beam_table,
#endif
__global SOFTMAX_ACCUMULATOR_TYPE* exp_sums,
__global SOFTMAX_ACCUMULATOR_TYPE* max_logits,
__global OUTPUT_TYPE* tmp_out
@ -637,6 +657,10 @@ KERNEL(sdpa_opt)(
// Main Gemm1 calculation loop
uint seq_len = sgid * TARGET_SEQ_LEN_BLOCK_SIZE;
for (; seq_len < partition_seq_len; seq_len += SUBGROUPS_PER_WG * SUBGROUP_SIZE) {
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_key)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len + sglid, 0)];
const uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + seq_len + sglid, 0);
#else
#ifdef INPUT1_DIMS_ORDER
uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len, 0);
uint key_offset_next_seq = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len + 1, 0);
@ -644,6 +668,7 @@ KERNEL(sdpa_opt)(
#else
uint key_offset = INPUT1_GET_INDEX(b0_idx, b1_idx, start_partition_idx + seq_len, 0);
const uint key_pitch = HEAD_SIZE;
#endif
#endif
INPUT0_TYPE acc[TARGET_SEQ_LEN_BLOCK_SIZE] = {INPUT0_VAL_ZERO};
@ -660,7 +685,11 @@ KERNEL(sdpa_opt)(
}
unroll_for (uint key_row_idx = 0; key_row_idx < TARGET_SEQ_LEN_BLOCK_SIZE; key_row_idx++) {
#ifdef BEAM_TABLE_TYPE
INPUT1_TYPE key_vals = KEY_BLOCK_READ(key_input, sub_group_broadcast(key_offset, key_row_idx) + head_idx_index);
#else
INPUT1_TYPE key_vals = KEY_BLOCK_READ(key_input, key_offset + key_row_idx * key_pitch + head_idx_index);
#endif
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
acc[key_row_idx] = mad(sub_group_broadcast(key_vals, i), queries_vec[i], acc[key_row_idx]);
@ -801,10 +830,15 @@ KERNEL(sdpa_opt)(
#endif
for (uint seq_len = 0; seq_len < partition_seq_len / SUBGROUP_SIZE; seq_len++) {
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE)];
const uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE);
#else
#ifdef INPUT2_DIMS_ORDER
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#else
uint value_offset = INPUT2_GET_INDEX(b0_idx, b1_idx, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#endif
#endif
OUTPUT_TYPE qk_val[TARGET_SEQ_LEN_BLOCK_SIZE];
@ -813,12 +847,18 @@ KERNEL(sdpa_opt)(
}
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
#ifdef BEAM_TABLE_TYPE
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, sub_group_broadcast(value_offset, i));
#else
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, value_offset);
#endif
unroll_for (uint seq_idx = 0; seq_idx < TARGET_SEQ_LEN_BLOCK_SIZE; seq_idx++) {
acc[seq_idx] = mad(sub_group_broadcast(qk_val[seq_idx], i), value_val, acc[seq_idx]);
}
#ifndef BEAM_TABLE_TYPE
value_offset += value_pitch;
#endif
}
}
@ -833,17 +873,31 @@ KERNEL(sdpa_opt)(
qk_val[seq_idx] = qk_local[qk_offset];
qk_offset += SEQ_LEN_PARTITION_SIZE;
}
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len_leftovers_start + sglid, sgid * SUBGROUP_SIZE)];
const uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + seq_len_leftovers_start + sglid, sgid * SUBGROUP_SIZE);
#else
#ifdef INPUT2_DIMS_ORDER
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + seq_len_leftovers_start, head_size_idx);
#else
uint value_offset = INPUT2_GET_INDEX(b0_idx, b1_idx, start_partition_idx + seq_len_leftovers_start, head_size_idx);
#endif
#endif
for (uint seq_len_idx = 0; seq_len_idx < partition_seq_len - seq_len_leftovers_start; seq_len_idx++) {
#ifdef BEAM_TABLE_TYPE
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, sub_group_broadcast(value_offset, seq_len_idx));
#else
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, value_offset);
#endif
for (uint seq_idx = 0; seq_idx < TARGET_SEQ_LEN_BLOCK_SIZE; seq_idx++) {
acc[seq_idx] = mad(sub_group_broadcast(qk_val[seq_idx], seq_len_idx), value_val, acc[seq_idx]);
}
#ifndef BEAM_TABLE_TYPE
value_offset += value_pitch;
#endif
}
}
@ -890,10 +944,15 @@ KERNEL(sdpa_opt)(
#endif
for (uint seq_len = 0; seq_len < partition_seq_len / SUBGROUP_SIZE; seq_len++) {
#ifdef BEAM_TABLE_TYPE
const uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE)];
const uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE) + sglid, sgid * SUBGROUP_SIZE);
#else
#ifdef INPUT2_DIMS_ORDER
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0_idx, b1_idx, 0, 0, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#else
uint value_offset = INPUT2_GET_INDEX(b0_idx, b1_idx, start_partition_idx + (seq_len * SUBGROUP_SIZE), head_size_idx);
#endif
#endif
OUTPUT_TYPE qk_val[TARGET_SEQ_LEN_BLOCK_SIZE];
@ -902,12 +961,18 @@ KERNEL(sdpa_opt)(
}
unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) {
#ifdef BEAM_TABLE_TYPE
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, sub_group_broadcast(value_offset, i));
#else
INPUT2_TYPE value_val = VALUE_BLOCK_READ(value_input, value_offset);
#endif
unroll_for (uint seq_idx = 0; seq_idx < TARGET_SEQ_LEN_BLOCK_SIZE; seq_idx++) {
acc[seq_idx] = mad(sub_group_broadcast(qk_val[seq_idx], i), value_val, acc[seq_idx]);
}
#ifndef BEAM_TABLE_TYPE
value_offset += value_pitch;
#endif
}
}

36
src/plugins/intel_gpu/src/kernel_selector/cl_kernels/sdpa_ref.cl
View File

@ -93,6 +93,24 @@ inline uint FUNC(get_input2_index)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint
#endif
}
#ifdef BEAM_TABLE_TYPE
inline uint FUNC(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
#if BEAM_TABLE_SIMPLE
return GET_DATA_INDEX_6D_SAFE(BEAM_TABLE, b, f, w, z, y, x);
#else
# error sdpa_ref.cl : Unsupported beam table format
#endif
}
inline uint FUNC(get_bt_index_key)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
return FUNC_CALL(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_TENSOR INPUT1_DIMS_ORDER);
}
inline uint FUNC(get_bt_index_value)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
return FUNC_CALL(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_TENSOR INPUT2_DIMS_ORDER);
}
#endif
#define APPLY_SCALE_TO_QUERY 1
KERNEL(sdpa_ref)(
@ -107,6 +125,9 @@ KERNEL(sdpa_ref)(
const __global INPUT4_TYPE* scale,
#endif
__global OUTPUT_TYPE* output,
#ifdef BEAM_TABLE_TYPE
const __global BEAM_TABLE_TYPE* beam_table,
#endif
__global OUTPUT_TYPE* tmp_buf
)
{
@ -129,7 +150,12 @@ KERNEL(sdpa_ref)(
OUTPUT_TYPE acc = 0;
for (uint h = 0; h < HEAD_SIZE /* head_size */; h++) {
uint query_offset = FUNC_CALL(get_input0_index)(OPTIONAL_SHAPE_INFO_TENSOR b0, b1, 0, 0, target_seq_idx, h);
uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b0, b1, 0, 0, s, h);
#ifdef BEAM_TABLE_TYPE
uint b_idx = beam_table[FUNC_CALL(get_bt_index_key)(OPTIONAL_SHAPE_INFO_TENSOR b0, b1, 0, 0, s, h)];
#else
uint b_idx = b0;
#endif
uint key_offset = FUNC_CALL(get_input1_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1, 0, 0, s, h);
#if APPLY_SCALE_TO_QUERY
INPUT0_TYPE q_val = query_input[query_offset] * scale_val;
@ -202,7 +228,13 @@ KERNEL(sdpa_ref)(
uint tmp_buf_offset = b0 * (NUM_HEADS * TARGET_SEQ_LEN * SOURCE_SEQ_LEN) +
b1 * (TARGET_SEQ_LEN * SOURCE_SEQ_LEN) +
target_seq_idx * (SOURCE_SEQ_LEN) + s;
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b0, b1, 0, 0, s, head_size_idx);
#ifdef BEAM_TABLE_TYPE
uint b_idx = beam_table[FUNC_CALL(get_bt_index_value)(OPTIONAL_SHAPE_INFO_TENSOR b0, b1, 0, 0, s, head_size_idx)];
#else
uint b_idx = b0;
#endif
uint value_offset = FUNC_CALL(get_input2_index)(OPTIONAL_SHAPE_INFO_TENSOR b_idx, b1, 0, 0, s, head_size_idx);
acc += tmp_buf[tmp_buf_offset] * value_input[value_offset];
}

27
src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.cpp
View File

@ -85,15 +85,30 @@ JitConstants SDPAKernelBase::GetJitConstants(const sdpa_params& params) const {
return true;
};
if ((!params.input0_order.empty() && !is_default_order(params.input0_order)) || params.conf.broadcast_axis != -1) {
auto use_index_calc_func = [&](const std::vector<int64_t> order, bool is_query = false) {
if (!params.input0_order.empty() && !is_default_order(params.input0_order))
return true;
if (params.conf.broadcast_axis != -1)
return true;
if (params.indirect_axis != -1 && !is_query)
return true;
return false;
};
if (params.indirect_axis != -1)
jit.AddConstant(MakeJitConstant("BEAM_TABLE", params.beam_table));
if (use_index_calc_func(params.input0_order, true))
jit.AddConstant(MakeJitConstant("INPUT0_DIMS_ORDER", GetDimsOrder(params.input0_order)));
}
if ((!params.input1_order.empty() && !is_default_order(params.input1_order)) || params.conf.broadcast_axis != -1) {
if (use_index_calc_func(params.input1_order))
jit.AddConstant(MakeJitConstant("INPUT1_DIMS_ORDER", GetDimsOrder(params.input1_order)));
}
if ((!params.input2_order.empty() && !is_default_order(params.input2_order)) || params.conf.broadcast_axis != -1) {
if (use_index_calc_func(params.input2_order))
jit.AddConstant(MakeJitConstant("INPUT2_DIMS_ORDER", GetDimsOrder(params.input2_order)));
}
TransposedDimensionAccessHelperJit dims_q(params.inputs[0], params.input0_order);
jit.AddConstant(MakeJitConstant("TARGET_SEQ_LEN", dims_q.y()));

3
src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_base.h
View File

@ -99,6 +99,9 @@ struct sdpa_params : public base_params {
std::vector<int64_t> input1_order;
std::vector<int64_t> input2_order;
std::vector<int64_t> output_order;
int64_t indirect_axis = -1;
DataTensor beam_table;
sdpa_configuration conf;
};

24
src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_opt.cpp
View File

@ -29,6 +29,22 @@ static size_t get_seq_len_partition_size() {
return seq_len;
}
static std::string GetKernelName(std::string base_name, KernelsTypes type, bool is_indirect) {
auto kernel_name = base_name;
if (is_indirect)
kernel_name += "_ind";
if (type == KernelsTypes::SINGLE_TOKEN) {
kernel_name += "_single_token";
} else if (type == KernelsTypes::MULTI_TOKENS) {
kernel_name += "_multi_tokens";
} else if (type == KernelsTypes::FINALIZATION) {
kernel_name += "_finalization";
}
return kernel_name;
}
ParamsKey SDPAKernelOpt::GetSupportedKey() const {
ParamsKey k;
k.EnableInputDataType(Datatype::F16);
@ -104,7 +120,7 @@ CommonDispatchData SDPAKernelOpt::SetDefault(const sdpa_params& params, size_t k
CeilDiv(target_seq_len, target_seq_len_block_size),
head_size * num_of_partitions };
dispatch_data.lws = { 1, 1, head_size };
} else if (kernel_idx == 2) {
} else if (kernel_idx == KernelsTypes::FINALIZATION) {
dispatch_data.gws = { batch_size * heads_num,
target_seq_len,
16 };
@ -134,8 +150,7 @@ KernelsData SDPAKernelOpt::GetKernelsData(const Params& params) const {
const auto& prim_params = dynamic_cast<const sdpa_params&>(params);
for (size_t kernel_idx = 0; kernel_idx < kernels_num; kernel_idx++) {
auto dispatch_data = SetDefault(prim_params, kernel_idx);
auto kernel_name = kernel_idx == 0 ? kernelName + "_single_token" :
kernel_idx == 1 ? kernelName + "_multi_tokens" : kernelName + "_finalization";
auto kernel_name = GetKernelName(kernelName, static_cast<KernelsTypes>(kernel_idx), prim_params.indirect_axis != -1);
auto entry_point = GetEntryPoint(kernel_name, prim_params.layerID, params);
auto jit_constants = GetJitConstants(prim_params, kernel_idx);
auto jit = CreateJit(kernel_name, jit_constants, entry_point);
@ -171,6 +186,9 @@ KernelsData SDPAKernelOpt::GetKernelsData(const Params& params) const {
auto tmp_out_elements_count = (num_of_partitions == 1) ? 1 : output.LogicalSize() * num_of_partitions;
auto tmp_out_size = tmp_out_elements_count * tmp_out_dt_size;
if (prim_params.indirect_axis != -1 && kernel_idx != KernelsTypes::FINALIZATION)
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INPUT, static_cast<uint32_t>(prim_params.inputs.size())});
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 0});
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 1});
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 2});

5
src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_ref.cpp
View File

@ -13,6 +13,8 @@ ParamsKey SDPAKernelRef::GetSupportedKey() const {
ParamsKey k;
k.EnableInputDataType(Datatype::F16);
k.EnableInputDataType(Datatype::F32);
// beam table input
k.EnableInputDataType(Datatype::INT32);
k.EnableOutputDataType(Datatype::F16);
k.EnableOutputDataType(Datatype::F32);
@ -72,6 +74,9 @@ KernelsData SDPAKernelRef::GetKernelsData(const Params& params) const {
"", false, false, static_cast<int>(prim_params.inputs.size()),
GetFusedPrimitiveInputsCount(params), 1, prim_params.is_shape_agnostic);
if (prim_params.indirect_axis != -1)
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INPUT, static_cast<uint32_t>(prim_params.inputs.size())});
kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 0});
kd.internalBufferSizes.clear();

24
src/plugins/intel_gpu/src/plugin/ops/scaled_dot_product_attention.cpp
View File

@ -6,6 +6,7 @@
#include "intel_gpu/plugin/common_utils.hpp"
#include "intel_gpu/op/sdpa.hpp"
#include "intel_gpu/op/indirect_sdpa.hpp"
#include "openvino/op/scaled_dot_product_attention.hpp"
@ -15,6 +16,7 @@ namespace ov {
namespace op {
namespace internal {
using SDPA = ov::intel_gpu::op::SDPA;
using IndirectSDPA = ov::intel_gpu::op::IndirectSDPA;
} // namespace internal
} // namespace op
} // namespace ov
@ -41,9 +43,30 @@ static void CreateSDPAOp(ProgramBuilder& p, const std::shared_ptr<ov::op::intern
auto layerName = layer_type_name_ID(op);
bool is_causal = op->get_causal();
int64_t indirect_axis = -1;
auto sdpa_prim = cldnn::scaled_dot_product_attention(layerName,
inputs,
is_causal,
indirect_axis,
op->get_input0_transpose_order(),
op->get_input1_transpose_order(),
op->get_input2_transpose_order(),
op->get_output_transpose_order());
p.add_primitive(*op, sdpa_prim);
}
static void CreateIndirectSDPAOp(ProgramBuilder& p, const std::shared_ptr<ov::op::internal::IndirectSDPA>& op) {
validate_inputs_count(op, {4, 5, 6});
auto inputs = p.GetInputInfo(op);
auto layerName = layer_type_name_ID(op);
bool is_causal = op->get_causal();
int64_t indirect_axis = op->get_indirect_axis();
auto sdpa_prim = cldnn::scaled_dot_product_attention(layerName,
inputs,
is_causal,
indirect_axis,
op->get_input0_transpose_order(),
op->get_input1_transpose_order(),
op->get_input2_transpose_order(),
@ -53,6 +76,7 @@ static void CreateSDPAOp(ProgramBuilder& p, const std::shared_ptr<ov::op::intern
}
REGISTER_FACTORY_IMPL(internal, SDPA);
REGISTER_FACTORY_IMPL(internal, IndirectSDPA);
REGISTER_FACTORY_IMPL(v13, ScaledDotProductAttention);
} // namespace intel_gpu

1
src/plugins/intel_gpu/src/plugin/plugin.cpp
View File

@ -542,6 +542,7 @@ std::vector<ov::PropertyName> Plugin::get_supported_properties() const {
ov::PropertyName{ov::intel_gpu::hint::host_task_priority.name(), PropertyMutability::RW},
ov::PropertyName{ov::intel_gpu::hint::queue_priority.name(), PropertyMutability::RW},
ov::PropertyName{ov::intel_gpu::hint::queue_throttle.name(), PropertyMutability::RW},
ov::PropertyName{ov::intel_gpu::hint::enable_sdpa_optimization.name(), PropertyMutability::RW},
ov::PropertyName{ov::intel_gpu::enable_loop_unrolling.name(), PropertyMutability::RW},
ov::PropertyName{ov::intel_gpu::disable_winograd_convolution.name(), PropertyMutability::RW},
ov::PropertyName{ov::cache_dir.name(), PropertyMutability::RW},

138
src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.cpp
View File

@ -6,7 +6,9 @@
#include <memory>
#include "intel_gpu/op/gemm.hpp"
#include "intel_gpu/op/sdpa.hpp"
#include "intel_gpu/op/indirect_gemm.hpp"
#include "intel_gpu/op/indirect_sdpa.hpp"
#include "intel_gpu/op/kv_cache.hpp"
#include "intel_gpu/op/read_value.hpp"
#include "intel_gpu/plugin/common_utils.hpp"
@ -42,7 +44,7 @@ void replace_node_unsafe(const std::shared_ptr<ov::Node>& target, const std::sha
namespace ov {
namespace intel_gpu {
IndirectKVCache::IndirectKVCache() {
IndirectGemmOpt::IndirectGemmOpt() {
using namespace ov::pass::pattern;
auto beam_idx = wrap_type<ov::op::v0::Parameter>();
@ -108,9 +110,141 @@ IndirectKVCache::IndirectKVCache() {
return true;
};
auto m = std::make_shared<ov::pass::pattern::Matcher>(matmul, "IndirectKVCache");
auto m = std::make_shared<ov::pass::pattern::Matcher>(matmul, "IndirectGemmOpt");
this->register_matcher(m, callback);
}
IndirectSDPAOpt::IndirectSDPAOpt() {
using namespace ov::pass::pattern;
using ov::pass::pattern::op::Or;
auto beam_idx = wrap_type<ov::op::v0::Parameter>();
auto gather_input_0 = wrap_type<ov::intel_gpu::op::ReadValue>();
auto gather_input_1 = wrap_type<ov::intel_gpu::op::ReadValue>();
auto axis_const = wrap_type<ov::op::v0::Constant>(
ov::op::util::constant_predicate<int64_t>([](const std::vector<int64_t>& value) -> bool {
return value.size() == 1 && (value[0] == 0 || value[0] == 1);
}));
auto gather_past_0 = wrap_type<ov::op::v8::Gather>({gather_input_0, beam_idx, axis_const});
auto gather_past_1 = wrap_type<ov::op::v8::Gather>({gather_input_1, beam_idx, axis_const});
auto kv_cache_0 = wrap_type<ov::intel_gpu::op::KVCache>({gather_past_0, any_input()});
auto kv_cache_1 = wrap_type<ov::intel_gpu::op::KVCache>({gather_past_1, any_input()});
auto input_attn_mask = any_input();
auto input_scale = any_input();
auto sdpa_without_attn_mask_m = wrap_type<ov::intel_gpu::op::SDPA>({ any_input(), kv_cache_0, kv_cache_1 });
auto sdpa_with_attn_mask_m = wrap_type<ov::intel_gpu::op::SDPA>({ any_input(), kv_cache_0, kv_cache_1, input_attn_mask });
auto sdpa_with_attn_mask_and_scale_m =
wrap_type<ov::intel_gpu::op::SDPA>({ any_input(), kv_cache_0, kv_cache_1, input_attn_mask, input_scale });
auto sdpa_m = std::make_shared<Or>(OutputVector{sdpa_without_attn_mask_m, sdpa_with_attn_mask_m, sdpa_with_attn_mask_and_scale_m});
ov::matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) {
if (transformation_callback(m.get_match_root())) {
return false;
}
const auto& pattern_map = m.get_pattern_value_map();
auto kv_cache_node_0 = std::dynamic_pointer_cast<ov::intel_gpu::op::KVCache>(pattern_map.at(kv_cache_0).get_node_shared_ptr());
auto kv_cache_node_1 = std::dynamic_pointer_cast<ov::intel_gpu::op::KVCache>(pattern_map.at(kv_cache_1).get_node_shared_ptr());
auto beam_idx_node = pattern_map.at(beam_idx).get_node_shared_ptr();
auto gather_input_node_0 = pattern_map.at(gather_input_0).get_node_shared_ptr();
auto gather_input_node_1 = pattern_map.at(gather_input_1).get_node_shared_ptr();
auto gather_node_0 = std::dynamic_pointer_cast<ov::op::v8::Gather>(pattern_map.at(gather_past_0).get_node_shared_ptr());
auto gather_node_1 = std::dynamic_pointer_cast<ov::op::v8::Gather>(pattern_map.at(gather_past_1).get_node_shared_ptr());
auto gather_axis_0 = gather_node_0->get_axis();
auto gather_axis_1 = gather_node_1->get_axis();
OPENVINO_ASSERT(gather_axis_0 == gather_axis_1);
ov::replace_node(gather_node_0, gather_input_node_0);
ov::replace_node(gather_node_1, gather_input_node_1);
auto indirect_kv_cache_0 = std::make_shared<op::KVCache>(gather_input_node_0,
kv_cache_node_0->get_input_node_shared_ptr(1),
beam_idx_node,
kv_cache_node_0->get_variable(),
kv_cache_node_0->get_concat_axis(),
gather_axis_0,
kv_cache_node_0->get_output_element_type(0));
auto indirect_kv_cache_1 = std::make_shared<op::KVCache>(gather_input_node_1,
kv_cache_node_1->get_input_node_shared_ptr(1),
beam_idx_node,
kv_cache_node_1->get_variable(),
kv_cache_node_1->get_concat_axis(),
gather_axis_1,
kv_cache_node_1->get_output_element_type(0));
indirect_kv_cache_0->set_friendly_name(kv_cache_node_0->get_friendly_name());
indirect_kv_cache_1->set_friendly_name(kv_cache_node_1->get_friendly_name());
ov::copy_runtime_info(kv_cache_node_0, indirect_kv_cache_0);
ov::copy_runtime_info(kv_cache_node_1, indirect_kv_cache_1);
replace_node_unsafe(kv_cache_node_0, indirect_kv_cache_0);
replace_node_unsafe(kv_cache_node_1, indirect_kv_cache_1);
auto sdpa = std::dynamic_pointer_cast<op::SDPA>(m.get_match_root());
auto order_in0 = sdpa->get_input0_transpose_order();
auto order_in1 = sdpa->get_input1_transpose_order();
auto order_in2 = sdpa->get_input2_transpose_order();
auto order_out = sdpa->get_output_transpose_order();
auto is_causal = sdpa->get_causal();
std::shared_ptr<op::SDPA> indirect_sdpa;
if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) {
indirect_sdpa = std::make_shared<ov::intel_gpu::op::IndirectSDPA>(sdpa->get_input_node_shared_ptr(0),
sdpa->get_input_node_shared_ptr(1),
sdpa->get_input_node_shared_ptr(2),
indirect_kv_cache_0->output(1), // beam table
is_causal,
gather_axis_1,
order_in0,
order_in1,
order_in2,
order_out);
} else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
indirect_sdpa = std::make_shared<ov::intel_gpu::op::IndirectSDPA>(sdpa->get_input_node_shared_ptr(0),
sdpa->get_input_node_shared_ptr(1),
sdpa->get_input_node_shared_ptr(2),
sdpa->get_input_node_shared_ptr(3),
indirect_kv_cache_0->output(1), // beam table
is_causal,
gather_axis_1,
order_in0,
order_in1,
order_in2,
order_out);
} else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) {
indirect_sdpa = std::make_shared<ov::intel_gpu::op::IndirectSDPA>(sdpa->get_input_node_shared_ptr(0),
sdpa->get_input_node_shared_ptr(1),
sdpa->get_input_node_shared_ptr(2),
sdpa->get_input_node_shared_ptr(3),
sdpa->get_input_node_shared_ptr(4),
indirect_kv_cache_0->output(1), // beam table
is_causal,
gather_axis_1,
order_in0,
order_in1,
order_in2,
order_out);
}
OPENVINO_ASSERT(indirect_sdpa != nullptr);
indirect_sdpa->set_friendly_name(sdpa->get_friendly_name());
ov::copy_runtime_info(sdpa, indirect_sdpa);
ov::replace_node(sdpa, indirect_sdpa);
return true;
};
auto m = std::make_shared<ov::pass::pattern::Matcher>(sdpa_m, "IndirectSDPAOpt");
this->register_matcher(m, callback);
}
IndirectKVCache::IndirectKVCache() {
add_matcher<IndirectGemmOpt>();
add_matcher<IndirectSDPAOpt>();
}
} // namespace intel_gpu
} // namespace ov

13
src/plugins/intel_gpu/src/plugin/transformations/indirect_kv_cache.hpp
View File

@ -36,11 +36,22 @@ namespace intel_gpu {
/// ┌────┴──────┐ ┌────┴──────┴───┐
/// │ Gemm │ | IndirectGemm |
/// └───────────┘ └───────────────┘
class IndirectKVCache : public ov::pass::MatcherPass {
class IndirectKVCache : public ov::pass::GraphRewrite {
public:
OPENVINO_RTTI("IndirectKVCache", "0");
IndirectKVCache();
};
class IndirectGemmOpt : public ov::pass::MatcherPass {
public:
OPENVINO_RTTI("IndirectGemmOpt", "0");
IndirectGemmOpt();
};
class IndirectSDPAOpt : public ov::pass::MatcherPass {
public:
OPENVINO_RTTI("IndirectSDPAOpt", "0");
IndirectSDPAOpt();
};
} // namespace intel_gpu
} // namespace ov

113
src/plugins/intel_gpu/src/plugin/transformations/op/indirect_sdpa.cpp Normal file
View File

@ -0,0 +1,113 @@
// Copyright (C) 2024 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "intel_gpu/op/indirect_sdpa.hpp"
#include "openvino/core/partial_shape.hpp"
namespace ov {
namespace intel_gpu {
namespace op {
IndirectSDPA::IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type)
: ov::intel_gpu::op::SDPA(Q, K, V, order_q, order_k, order_v, order_out, is_causal, output_type)
, m_indirect_axis(indirect_axis) {
set_argument(3, beam_table);
validate_and_infer_types();
}
IndirectSDPA::IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& attn_mask,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type)
: ov::intel_gpu::op::SDPA(Q, K, V, attn_mask, order_q, order_k, order_v, order_out, is_causal, output_type)
, m_indirect_axis(indirect_axis) {
set_argument(4, beam_table);
validate_and_infer_types();
}
IndirectSDPA::IndirectSDPA(const ov::Output<Node>& Q,
const ov::Output<Node>& K,
const ov::Output<Node>& V,
const ov::Output<Node>& attn_mask,
const ov::Output<Node>& scale,
const ov::Output<Node>& beam_table,
const bool is_causal,
const int64_t indirect_axis,
const std::vector<int64_t>& order_q,
const std::vector<int64_t>& order_k,
const std::vector<int64_t>& order_v,
const std::vector<int64_t>& order_out,
const ov::element::Type output_type)
: ov::intel_gpu::op::SDPA(Q, K, V, attn_mask, scale, order_q, order_k, order_v, order_out, is_causal, output_type)
, m_indirect_axis(indirect_axis) {
set_argument(5, beam_table);
validate_and_infer_types();
}
std::shared_ptr<ov::Node> IndirectSDPA::clone_with_new_inputs(const ov::OutputVector& new_args) const {
check_new_args_count(this, new_args);
if (new_args.size() == 4) {
return std::make_shared<IndirectSDPA>(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3),
m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type);
} else if (new_args.size() == 5) {
return std::make_shared<IndirectSDPA>(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4),
m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type);
} else {
return std::make_shared<IndirectSDPA>(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4), new_args.at(5),
m_is_causal, m_indirect_axis, m_order_q, m_order_k, m_order_v, m_order_out, m_output_type);
}
}
void IndirectSDPA::validate_and_infer_types() {
const auto input_size = get_input_size();
NODE_VALIDATION_CHECK(this,
input_size == 4 || input_size == 5 || input_size == 6,
"Number of inputs is incorrect. Current value is: ",
input_size,
", expected 4, 5 or 6.");
std::vector<ov::PartialShape> input_shapes;
for (size_t i = 0; i < input_size - 1; i++) {
input_shapes.push_back(get_input_partial_shape(i));
}
auto out_shapes = shape_infer(this,
input_shapes,
m_order_q,
m_order_k,
m_order_v,
m_order_out);
auto output_type = m_output_type == ov::element::undefined ? get_input_element_type(0) : m_output_type;
set_output_type(0, output_type, out_shapes[0]);
}
bool IndirectSDPA::visit_attributes(ov::AttributeVisitor &visitor) {
SDPA::visit_attributes(visitor);
visitor.on_attribute("indirect_axis", m_indirect_axis);
return true;
}
} // namespace op
} // namespace intel_gpu
} // namespace ov

14
src/plugins/intel_gpu/src/plugin/transformations/op/sdpa.cpp
View File

@ -30,6 +30,7 @@ SDPA::SDPA(const ov::Output<Node>& Q,
, m_is_causal(is_causal)
, m_output_type(output_type) {
set_arguments({Q, K, V});
set_causal(is_causal);
validate_and_infer_types();
}
@ -50,6 +51,7 @@ SDPA::SDPA(const ov::Output<Node>& Q,
, m_is_causal(is_causal)
, m_output_type(output_type) {
set_arguments({Q, K, V, attn_mask});
set_causal(is_causal);
validate_and_infer_types();
}
@ -71,13 +73,23 @@ SDPA::SDPA(const ov::Output<Node>& Q,
, m_is_causal(is_causal)
, m_output_type(output_type) {
set_arguments({Q, K, V, attn_mask, scale});
set_causal(is_causal);
validate_and_infer_types();
}
std::shared_ptr<ov::Node> SDPA::clone_with_new_inputs(const ov::OutputVector& new_args) const {
check_new_args_count(this, new_args);
return std::make_shared<SDPA>(new_args.at(0), new_args.at(1), new_args.at(2), m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type);
if (new_args.size() == 3) {
return std::make_shared<SDPA>(new_args.at(0), new_args.at(1), new_args.at(2),
m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type);
} else if (new_args.size() == 4) {
return std::make_shared<SDPA>(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3),
m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type);
} else {
return std::make_shared<SDPA>(new_args.at(0), new_args.at(1), new_args.at(2), new_args.at(3), new_args.at(4),
m_order_q, m_order_k, m_order_v, m_order_out, m_is_causal, m_output_type);
}
}
void SDPA::validate_and_infer_types() {

63
src/plugins/intel_gpu/src/plugin/transformations/transpose_fusion.cpp
View File

@ -80,14 +80,7 @@ TransposeSDPAMatcher::TransposeSDPAMatcher() {
ov::matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](Matcher& m) {
const auto& pattern_map = m.get_pattern_value_map();
std::shared_ptr<ov::op::v13::ScaledDotProductAttention> sdpa;
if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) {
sdpa = std::dynamic_pointer_cast<ov::op::v13::ScaledDotProductAttention>(pattern_map.at(sdpa_without_attn_mask_m).get_node_shared_ptr());
} else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
sdpa = std::dynamic_pointer_cast<ov::op::v13::ScaledDotProductAttention>(pattern_map.at(sdpa_with_attn_mask_m).get_node_shared_ptr());
} else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) {
sdpa = std::dynamic_pointer_cast<ov::op::v13::ScaledDotProductAttention>(pattern_map.at(sdpa_with_attn_mask_and_scale_m).get_node_shared_ptr());
}
auto sdpa = std::dynamic_pointer_cast<ov::op::v13::ScaledDotProductAttention>(m.get_match_root());
if (!sdpa || transformation_callback(sdpa)) {
return false;
@ -101,33 +94,41 @@ TransposeSDPAMatcher::TransposeSDPAMatcher() {
size_t input_k_output_idx = sdpa->get_input_source_output(1).get_index();
size_t input_v_output_idx = sdpa->get_input_source_output(2).get_index();
if (pattern_map.count(transpose_q_m) > 0) {
auto tranpose_a_order = std::dynamic_pointer_cast<ov::op::v0::Constant>(pattern_map.at(transpose_q_order_m).get_node_shared_ptr());
order_q = tranpose_a_order->cast_vector<int64_t>();
if (order_q.back() != static_cast<int64_t>(order_q.size() - 1)) // Allow any transposes without head_size dim position change
auto process_transpose = [](const std::shared_ptr<Node>& transpose_node,
const std::shared_ptr<Node>& transpose_order_const_node,
std::vector<int64_t>& order,
size_t& output_idx) {
auto transpose_order_const = std::dynamic_pointer_cast<ov::op::v0::Constant>(transpose_order_const_node);
order = transpose_order_const->cast_vector<int64_t>();
// Allow any transposes without head_size dim position change
if (order.back() != static_cast<int64_t>(order.size() - 1))
return false;
auto tranpose_a = std::dynamic_pointer_cast<ov::op::v1::Transpose>(pattern_map.at(transpose_q_m).get_node_shared_ptr());
input_q_output_idx = tranpose_a->get_input_source_output(0).get_index();
}
if (pattern_map.count(transpose_k_m) > 0) {
auto tranpose_b_order = std::dynamic_pointer_cast<ov::op::v0::Constant>(pattern_map.at(transpose_k_order_m).get_node_shared_ptr());
order_k = tranpose_b_order->cast_vector<int64_t>();
if (order_k.back() != static_cast<int64_t>(order_k.size() - 1)) // Allow any transposes without head_size dim position change
return false;
auto transpose = std::dynamic_pointer_cast<ov::op::v1::Transpose>(transpose_node);
output_idx = transpose->get_input_source_output(0).get_index();
auto tranpose_b = std::dynamic_pointer_cast<ov::op::v1::Transpose>(pattern_map.at(transpose_k_m).get_node_shared_ptr());
input_k_output_idx = tranpose_b->get_input_source_output(0).get_index();
}
if (pattern_map.count(transpose_v_m) > 0) {
auto tranpose_c_order = std::dynamic_pointer_cast<ov::op::v0::Constant>(pattern_map.at(transpose_v_order_m).get_node_shared_ptr());
order_v = tranpose_c_order->cast_vector<int64_t>();
if (order_v.back() != static_cast<int64_t>(order_v.size() - 1)) // Allow any transposes without head_size dim position change
return false;
return true;
};
auto tranpose_c = std::dynamic_pointer_cast<ov::op::v1::Transpose>(pattern_map.at(transpose_k_m).get_node_shared_ptr());
input_v_output_idx = tranpose_c->get_input_source_output(0).get_index();
}
bool can_fuse_transposes = true;
if (pattern_map.count(transpose_q_m) > 0)
can_fuse_transposes &= process_transpose(pattern_map.at(transpose_q_m).get_node_shared_ptr(),
pattern_map.at(transpose_q_order_m).get_node_shared_ptr(),
order_q, input_q_output_idx);
if (pattern_map.count(transpose_k_m) > 0)
can_fuse_transposes &= process_transpose(pattern_map.at(transpose_k_m).get_node_shared_ptr(),
pattern_map.at(transpose_k_order_m).get_node_shared_ptr(),
order_k, input_k_output_idx);
if (pattern_map.count(transpose_v_m) > 0)
can_fuse_transposes &= process_transpose(pattern_map.at(transpose_v_m).get_node_shared_ptr(),
pattern_map.at(transpose_v_order_m).get_node_shared_ptr(),
order_v, input_v_output_idx);
if (!can_fuse_transposes)
return false;
auto input_q = ov::Output<Node>(pattern_map.at(input_q_m).get_node_shared_ptr(), input_q_output_idx);
auto input_k = ov::Output<Node>(pattern_map.at(input_k_m).get_node_shared_ptr(), input_k_output_idx);

5
src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp
View File

@ -13,6 +13,7 @@
#include <vector>
#include "intel_gpu/plugin/transformations_pipeline.hpp"
#include "intel_gpu/runtime/debug_configuration.hpp"
#include "intel_gpu/runtime/itt.hpp"
#include "low_precision/convolution.hpp"
#include "low_precision/convolution_backprop_data.hpp"
@ -307,6 +308,10 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
manager.register_pass<ov::pass::CommonOptimizations>();
pass_config->set_callback<ov::pass::ScaledDotProductAttentionDecomposition>([&](const std::shared_ptr<const ov::Node> node){
GPU_DEBUG_IF(cldnn::debug_configuration::get_instance()->enable_sdpa != -1) {
GPU_DEBUG_CODE(return cldnn::debug_configuration::get_instance()->enable_sdpa == 1);
}
if (!config.get_property(ov::intel_gpu::hint::enable_sdpa_optimization))
return false;

4
src/plugins/intel_gpu/src/runtime/debug_configuration.cpp
View File

@ -171,6 +171,8 @@ static void print_help_messages() {
message_list.emplace_back("OV_GPU_DisableDynamicImpl", "Disable dynamic implementation");
message_list.emplace_back("OV_GPU_DisableRuntimeBufferFusing", "Disable runtime buffer fusing");
message_list.emplace_back("OV_GPU_DisableMemoryReuse", "Disable memory reuse");
message_list.emplace_back("OV_GPU_EnableSDPA", "This allows the enforcement of SDPA decomposition logic: 0 completely disables SDPA kernel usage, "
"and 1 enables it for all the cases.");
message_list.emplace_back("OV_GPU_DumpMemoryPool", "Dump memory pool contents of each iteration");
message_list.emplace_back("OV_GPU_DumpMemoryPoolIters", "List of iterations to dump memory pool status, separated by space.");
message_list.emplace_back("OV_GPU_DumpMemoryPoolPath", "Enable dumping memory pool status to csv file and set the dest path");
@ -232,6 +234,7 @@ debug_configuration::debug_configuration()
, serialize_compile(0)
, max_kernels_per_batch(0)
, impls_cache_capacity(-1)
, enable_sdpa(-1)
, disable_async_compilation(0)
, disable_winograd_conv(0)
, disable_dynamic_impl(0)
@ -280,6 +283,7 @@ debug_configuration::debug_configuration()
get_gpu_debug_env_var("ForceImplTypes", forced_impl_types_str);
get_gpu_debug_env_var("MaxKernelsPerBatch", max_kernels_per_batch);
get_gpu_debug_env_var("ImplsCacheCapacity", impls_cache_capacity);
get_gpu_debug_env_var("EnableSDPA", enable_sdpa);
get_gpu_debug_env_var("DisableAsyncCompilation", disable_async_compilation);
get_gpu_debug_env_var("DisableWinogradConv", disable_winograd_conv);
get_gpu_debug_env_var("DisableDynamicImpl", disable_dynamic_impl);