[GPU] SDPA minor improvements (#24770)
### Details: - Change the default softmax accumulator type from FP16 to FP32 (this doesn't affect performance, as well as accuracy at first glance; however, let's stick to higher precision here) - Applied review comments from PR https://github.com/openvinotoolkit/openvino/pull/24466 - Minor internal operations refactoring (SDPA and IndirectSDPA) - Add new tests for GQA optimization and for indirect inputs
This commit is contained in:
Sergey Shlyapnikov
GitHub
parent
cea8e1f5e4
commit
ba8d6c5de5
|
|
@ -19,36 +19,7 @@ public:
|
|||
|
||||
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,
|
||||
IndirectSDPA(const OutputVector& data_inputs,
|
||||
const ov::Output<Node>& beam_table,
|
||||
const bool is_causal,
|
||||
const int64_t indirect_axis,
|
||||
|
|
|
|||
|
|
@ -19,37 +19,12 @@ public:
|
|||
|
||||
SDPA() = default;
|
||||
|
||||
SDPA(const ov::Output<Node>& Q,
|
||||
const ov::Output<Node>& K,
|
||||
const ov::Output<Node>& V,
|
||||
SDPA(const OutputVector& inputs,
|
||||
const bool is_causal,
|
||||
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 bool is_causal,
|
||||
const ov::element::Type output_type = ov::element::undefined);
|
||||
|
||||
SDPA(const ov::Output<Node>& Q,
|
||||
const ov::Output<Node>& K,
|
||||
const ov::Output<Node>& V,
|
||||
const ov::Output<Node>& attn_mask,
|
||||
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 bool is_causal,
|
||||
const ov::element::Type output_type = ov::element::undefined);
|
||||
|
||||
SDPA(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 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 bool is_causal,
|
||||
const ov::element::Type output_type = ov::element::undefined);
|
||||
|
||||
bool visit_attributes(ov::AttributeVisitor &visitor) override;
|
||||
|
|
@ -73,11 +48,11 @@ public:
|
|||
}
|
||||
|
||||
protected:
|
||||
bool m_is_causal;
|
||||
std::vector<int64_t> m_order_q;
|
||||
std::vector<int64_t> m_order_k;
|
||||
std::vector<int64_t> m_order_v;
|
||||
std::vector<int64_t> m_order_out;
|
||||
bool m_is_causal;
|
||||
ov::element::Type m_output_type;
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -72,7 +72,6 @@ protected:
|
|||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -20,7 +20,27 @@ layout scaled_dot_product_attention_inst::calc_output_layout(scaled_dot_product_
|
|||
kernel_impl_params const& impl_param) {
|
||||
auto desc = impl_param.typed_desc<scaled_dot_product_attention>();
|
||||
|
||||
return impl_param.get_input_layout(0);
|
||||
auto transpose_shape = [&desc](const ov::PartialShape& shape, const std::vector<int64_t>& order) {
|
||||
if (desc->input_q_transpose_order.empty())
|
||||
return shape;
|
||||
|
||||
auto shape_transposed = ov::PartialShape(shape);
|
||||
auto rank_diff = shape.size() - order.size();
|
||||
for (size_t i = 0; i < order.size(); i++) {
|
||||
size_t idx = static_cast<size_t>(order[i]);
|
||||
shape_transposed[i + rank_diff] = shape[idx + rank_diff];
|
||||
}
|
||||
|
||||
return shape_transposed;
|
||||
};
|
||||
|
||||
auto input0_layout = impl_param.get_input_layout(0);
|
||||
auto default_out_dt = data_type_traits::is_floating_point(input0_layout.data_type) ? input0_layout.data_type : data_types::f32;
|
||||
auto output_type = desc->output_data_types[0].value_or(default_out_dt);
|
||||
auto output_format = input0_layout.format;
|
||||
auto output_shape = transpose_shape(input0_layout.get_partial_shape(), desc->input_q_transpose_order); // output shape matches with Q input shape
|
||||
|
||||
return { layout{output_shape, output_type, output_format, desc->output_paddings[0]} };
|
||||
}
|
||||
|
||||
template<typename ShapeType>
|
||||
|
|
|
|||
|
|
@ -20,16 +20,16 @@
|
|||
|
||||
inline uint FUNC(get_input0_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w, uint z, uint y, uint x) {
|
||||
#if INPUT0_SIMPLE
|
||||
return GET_DATA_INDEX_6D_SAFE(INPUT0, b, f, w, z, y, x);
|
||||
return GET_DATA_INDEX_6D(INPUT0, b, f, w, z, y, x);
|
||||
#else
|
||||
#if INPUT0_DIMS == 4
|
||||
return INPUT0_GET_INDEX_SAFE(b, f, y, x);
|
||||
return INPUT0_GET_INDEX(b, f, y, x);
|
||||
#elif INPUT0_DIMS == 5
|
||||
return INPUT0_GET_INDEX_SAFE(b, f, z, y, x);
|
||||
return INPUT0_GET_INDEX(b, f, z, y, x);
|
||||
#elif INPUT0_DIMS == 6
|
||||
return INPUT0_GET_INDEX_SAFE(b, f, w, z, y, x);
|
||||
return INPUT0_GET_INDEX(b, f, w, z, y, x);
|
||||
#else
|
||||
# error sdpa_ref.cl : Unsupported input 0 format
|
||||
# error sdpa_opt.cl : Unsupported input 0 format
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
|
|
@ -47,16 +47,16 @@ inline uint FUNC(get_input1_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, ui
|
|||
DO_BROADCAST_KEY_VALUE;
|
||||
#endif
|
||||
#if INPUT1_SIMPLE
|
||||
return GET_DATA_INDEX_6D_SAFE(INPUT1, b, f, w, z, y, x);
|
||||
return GET_DATA_INDEX_6D(INPUT1, b, f, w, z, y, x);
|
||||
#else
|
||||
#if INPUT1_DIMS == 4
|
||||
return INPUT1_GET_INDEX_SAFE(b, f, y, x);
|
||||
return INPUT1_GET_INDEX(b, f, y, x);
|
||||
#elif INPUT1_DIMS == 5
|
||||
return INPUT1_GET_INDEX_SAFE(b, f, z, y, x);
|
||||
return INPUT1_GET_INDEX(b, f, z, y, x);
|
||||
#elif INPUT1_DIMS == 6
|
||||
return INPUT1_GET_INDEX_SAFE(b, f, w, z, y, x);
|
||||
return INPUT1_GET_INDEX(b, f, w, z, y, x);
|
||||
#else
|
||||
# error sdpa_ref.cl : Unsupported input 1 format
|
||||
# error sdpa_opt.cl : Unsupported input 1 format
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
|
|
@ -77,13 +77,13 @@ inline uint FUNC(get_input2_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, ui
|
|||
return GET_DATA_INDEX_6D_SAFE(INPUT2, b, f, w, z, y, x);
|
||||
#else
|
||||
#if INPUT2_DIMS == 4
|
||||
return INPUT2_GET_INDEX_SAFE(b, f, y, x);
|
||||
return INPUT2_GET_INDEX(b, f, y, x);
|
||||
#elif INPUT2_DIMS == 5
|
||||
return INPUT2_GET_INDEX_SAFE(b, f, z, y, x);
|
||||
return INPUT2_GET_INDEX(b, f, z, y, x);
|
||||
#elif INPUT2_DIMS == 6
|
||||
return INPUT2_GET_INDEX_SAFE(b, f, w, z, y, x);
|
||||
return INPUT2_GET_INDEX(b, f, w, z, y, x);
|
||||
#else
|
||||
# error sdpa_ref.cl : Unsupported input 1 format
|
||||
# error sdpa_opt.cl : Unsupported input 1 format
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
|
|
@ -101,7 +101,7 @@ inline uint FUNC(get_bt_index_nt)(OPTIONAL_SHAPE_INFO_ARG uint b, uint f, uint w
|
|||
#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
|
||||
# error sdpa_opt.cl : Unsupported beam table format
|
||||
#endif
|
||||
}
|
||||
|
||||
|
|
@ -466,8 +466,8 @@ KERNEL(sdpa_opt)(
|
|||
|
||||
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);
|
||||
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)];
|
||||
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);
|
||||
|
|
@ -715,7 +715,7 @@ KERNEL(sdpa_opt)(
|
|||
#endif
|
||||
#if INPUT0_TYPE_SIZE == 2
|
||||
/* Adding this clamp improves performance for some reason */
|
||||
acc[i] = SOFTMAX_ACCUMULATOR_MIN_FUNC(SOFTMAX_ACCUMULATOR_MAX_FUNC(acc[i], INPUT0_VAL_MIN), INPUT0_VAL_MAX);
|
||||
acc[i] = INPUT0_MIN_FUNC(INPUT0_MAX_FUNC(acc[i], INPUT0_VAL_MIN), INPUT0_VAL_MAX);
|
||||
#endif
|
||||
if (seq_len + i >= partition_seq_len) {
|
||||
acc[i] = INPUT0_VAL_MIN;
|
||||
|
|
@ -749,7 +749,7 @@ KERNEL(sdpa_opt)(
|
|||
#define QK_ITERS_END QK_MAX_NUMS_PER_SG
|
||||
#endif
|
||||
|
||||
OUTPUT_TYPE qk_max[QK_MAX_NUMS_PER_SG];
|
||||
SOFTMAX_ACCUMULATOR_TYPE qk_max[QK_MAX_NUMS_PER_SG];
|
||||
for (uint i = 0; i < QK_MAX_NUMS_PER_SG; i++)
|
||||
qk_max[i] = SOFTMAX_ACCUMULATOR_VAL_MIN;
|
||||
|
||||
|
|
|
|||
|
|
@ -23,12 +23,15 @@ static size_t get_target_seq_len_block_size() {
|
|||
return block_size;
|
||||
}
|
||||
|
||||
|
||||
static size_t get_seq_len_partition_size() {
|
||||
const size_t seq_len = 256;
|
||||
return seq_len;
|
||||
}
|
||||
|
||||
static Datatype get_softmax_acc_type() {
|
||||
return Datatype::F32;
|
||||
}
|
||||
|
||||
static std::string GetKernelName(std::string base_name, KernelsTypes type, bool is_indirect) {
|
||||
auto kernel_name = base_name;
|
||||
if (is_indirect)
|
||||
|
|
@ -80,7 +83,7 @@ bool SDPAKernelOpt::Validate(const Params& p) const {
|
|||
JitConstants SDPAKernelOpt::GetJitConstants(const sdpa_params& params, size_t kernel_idx) const {
|
||||
auto jit = SDPAKernelBase::GetJitConstants(params);
|
||||
|
||||
const auto softmax_acc_dt = params.inputs[0].GetDType();
|
||||
const auto softmax_acc_dt = get_softmax_acc_type();
|
||||
jit.Merge(MakeTypeJitConstants(softmax_acc_dt, "SOFTMAX_ACCUMULATOR"));
|
||||
|
||||
const auto& config = params.conf;
|
||||
|
|
@ -232,7 +235,7 @@ void SDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) const {
|
|||
|
||||
auto num_of_partitions = CeilDiv(source_seq_len, get_seq_len_partition_size());
|
||||
|
||||
auto buf_dt_size = output.ElementSize();
|
||||
auto buf_dt_size = BytesPerElement(get_softmax_acc_type());
|
||||
auto buf_elements_count = (num_of_partitions == 1) ? 1 : output.LogicalSize() / head_size * num_of_partitions;
|
||||
auto buf_size = buf_elements_count * buf_dt_size;
|
||||
|
||||
|
|
@ -241,26 +244,26 @@ void SDPAKernelOpt::GetUpdateDispatchDataFunc(KernelData& kd) const {
|
|||
auto tmp_out_size = tmp_out_elements_count * tmp_out_dt_size;
|
||||
|
||||
auto dispatch_data1 = SetDefault(prim_params, 0);
|
||||
kernel_data.kernels[0].params.workGroups.global = dispatch_data1.gws;
|
||||
kernel_data.kernels[0].params.workGroups.local = dispatch_data1.lws;
|
||||
kernel_data.kernels[0].skip_execution = target_seq_len > 1;
|
||||
kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.global = dispatch_data1.gws;
|
||||
kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].params.workGroups.local = dispatch_data1.lws;
|
||||
kernel_data.kernels[KernelsTypes::SINGLE_TOKEN].skip_execution = target_seq_len > 1;
|
||||
|
||||
auto dispatch_data2 = SetDefault(prim_params, 1);
|
||||
kernel_data.kernels[1].params.workGroups.global = dispatch_data2.gws;
|
||||
kernel_data.kernels[1].params.workGroups.local = dispatch_data2.lws;
|
||||
kernel_data.kernels[1].skip_execution = target_seq_len == 1;
|
||||
kernel_data.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.global = dispatch_data2.gws;
|
||||
kernel_data.kernels[KernelsTypes::MULTI_TOKENS].params.workGroups.local = dispatch_data2.lws;
|
||||
kernel_data.kernels[KernelsTypes::MULTI_TOKENS].skip_execution = target_seq_len == 1;
|
||||
|
||||
ScalarDescriptor num_of_partitions_scalar;
|
||||
num_of_partitions_scalar.t = ScalarDescriptor::Types::UINT32;
|
||||
num_of_partitions_scalar.v.u32 = static_cast<uint32_t>(num_of_partitions);
|
||||
|
||||
auto dispatch_data3 = SetDefault(prim_params, 2);
|
||||
kernel_data.kernels[2].params.workGroups.global = dispatch_data3.gws;
|
||||
kernel_data.kernels[2].params.workGroups.local = dispatch_data3.lws;
|
||||
kernel_data.kernels[2].skip_execution = num_of_partitions == 1;
|
||||
kernel_data.kernels[KernelsTypes::FINALIZATION].params.workGroups.global = dispatch_data3.gws;
|
||||
kernel_data.kernels[KernelsTypes::FINALIZATION].params.workGroups.local = dispatch_data3.lws;
|
||||
kernel_data.kernels[KernelsTypes::FINALIZATION].skip_execution = num_of_partitions == 1;
|
||||
|
||||
kernel_data.kernels[2].params.scalars.clear();
|
||||
kernel_data.kernels[2].params.scalars.push_back(num_of_partitions_scalar);
|
||||
kernel_data.kernels[KernelsTypes::FINALIZATION].params.scalars.clear();
|
||||
kernel_data.kernels[KernelsTypes::FINALIZATION].params.scalars.push_back(num_of_partitions_scalar);
|
||||
|
||||
kernel_data.internalBufferSizes.clear();
|
||||
kernel_data.internalBufferSizes.push_back(buf_size);
|
||||
|
|
|
|||
|
|
@ -190,45 +190,27 @@ IndirectSDPAOpt::IndirectSDPAOpt() {
|
|||
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);
|
||||
OutputVector data_inputs;
|
||||
data_inputs.push_back(sdpa->get_input_node_shared_ptr(0)); // Q
|
||||
data_inputs.push_back(sdpa->get_input_node_shared_ptr(1)); // K
|
||||
data_inputs.push_back(sdpa->get_input_node_shared_ptr(2)); // V
|
||||
|
||||
if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
|
||||
data_inputs.push_back(sdpa->get_input_source_output(3));
|
||||
} 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);
|
||||
data_inputs.push_back(sdpa->get_input_source_output(3));
|
||||
data_inputs.push_back(sdpa->get_input_source_output(4));
|
||||
}
|
||||
|
||||
auto indirect_sdpa = std::make_shared<ov::intel_gpu::op::IndirectSDPA>(data_inputs,
|
||||
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());
|
||||
|
|
|
|||
|
|
@ -9,9 +9,7 @@ namespace ov {
|
|||
namespace intel_gpu {
|
||||
namespace op {
|
||||
|
||||
IndirectSDPA::IndirectSDPA(const ov::Output<Node>& Q,
|
||||
const ov::Output<Node>& K,
|
||||
const ov::Output<Node>& V,
|
||||
IndirectSDPA::IndirectSDPA(const OutputVector& data_inputs,
|
||||
const ov::Output<Node>& beam_table,
|
||||
const bool is_causal,
|
||||
const int64_t indirect_axis,
|
||||
|
|
@ -20,62 +18,28 @@ IndirectSDPA::IndirectSDPA(const ov::Output<Node>& Q,
|
|||
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)
|
||||
: ov::intel_gpu::op::SDPA(data_inputs, is_causal, order_q, order_k, order_v, order_out, 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);
|
||||
auto beam_table_idx = data_inputs.size();
|
||||
set_argument(beam_table_idx, 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);
|
||||
}
|
||||
// Exclude beam_table input
|
||||
OutputVector data_inputs(new_args.begin(), new_args.end() - 1);
|
||||
|
||||
return std::make_shared<IndirectSDPA>(data_inputs,
|
||||
new_args.back(),
|
||||
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() {
|
||||
|
|
|
|||
|
|
@ -14,65 +14,20 @@ namespace ov {
|
|||
namespace intel_gpu {
|
||||
namespace op {
|
||||
|
||||
SDPA::SDPA(const ov::Output<Node>& Q,
|
||||
const ov::Output<Node>& K,
|
||||
const ov::Output<Node>& V,
|
||||
SDPA::SDPA(const OutputVector& inputs,
|
||||
const bool is_causal,
|
||||
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 bool is_causal,
|
||||
const ov::element::Type output_type)
|
||||
: m_order_q(order_q)
|
||||
: m_is_causal(is_causal)
|
||||
, m_order_q(order_q)
|
||||
, m_order_k(order_k)
|
||||
, m_order_v(order_v)
|
||||
, m_order_out(order_out)
|
||||
, m_is_causal(is_causal)
|
||||
, m_output_type(output_type) {
|
||||
set_arguments({Q, K, V});
|
||||
set_causal(is_causal);
|
||||
validate_and_infer_types();
|
||||
}
|
||||
|
||||
SDPA::SDPA(const ov::Output<Node>& Q,
|
||||
const ov::Output<Node>& K,
|
||||
const ov::Output<Node>& V,
|
||||
const ov::Output<Node>& attn_mask,
|
||||
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 bool is_causal,
|
||||
const ov::element::Type output_type)
|
||||
: m_order_q(order_q)
|
||||
, m_order_k(order_k)
|
||||
, m_order_v(order_v)
|
||||
, m_order_out(order_out)
|
||||
, 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();
|
||||
}
|
||||
|
||||
SDPA::SDPA(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 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 bool is_causal,
|
||||
const ov::element::Type output_type)
|
||||
: m_order_q(order_q)
|
||||
, m_order_k(order_k)
|
||||
, m_order_v(order_v)
|
||||
, m_order_out(order_out)
|
||||
, m_is_causal(is_causal)
|
||||
, m_output_type(output_type) {
|
||||
set_arguments({Q, K, V, attn_mask, scale});
|
||||
set_arguments(inputs);
|
||||
set_causal(is_causal);
|
||||
validate_and_infer_types();
|
||||
}
|
||||
|
|
@ -80,16 +35,13 @@ SDPA::SDPA(const ov::Output<Node>& Q,
|
|||
std::shared_ptr<ov::Node> SDPA::clone_with_new_inputs(const ov::OutputVector& new_args) const {
|
||||
check_new_args_count(this, new_args);
|
||||
|
||||
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);
|
||||
}
|
||||
return std::make_shared<SDPA>(new_args,
|
||||
m_is_causal,
|
||||
m_order_q,
|
||||
m_order_k,
|
||||
m_order_v,
|
||||
m_order_out,
|
||||
m_output_type);
|
||||
}
|
||||
|
||||
void SDPA::validate_and_infer_types() {
|
||||
|
|
|
|||
|
|
@ -134,18 +134,20 @@ TransposeSDPAMatcher::TransposeSDPAMatcher() {
|
|||
auto input_k = ov::Output<Node>(pattern_map.at(input_k_m).get_node_shared_ptr(), input_k_output_idx);
|
||||
auto input_v = ov::Output<Node>(pattern_map.at(input_v_m).get_node_shared_ptr(), input_v_output_idx);
|
||||
|
||||
std::shared_ptr<op::SDPA> sdpa_new;
|
||||
if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) {
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_q, input_k, input_v, order_q, order_k, order_v, order_output, sdpa->get_causal());
|
||||
} else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
|
||||
auto attn_mask = sdpa->get_input_source_output(3);
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_q, input_k, input_v, attn_mask, order_q, order_k, order_v, order_output, sdpa->get_causal());
|
||||
OutputVector inputs;
|
||||
inputs.push_back(ov::Output<Node>(pattern_map.at(input_q_m).get_node_shared_ptr(), input_q_output_idx));
|
||||
inputs.push_back(ov::Output<Node>(pattern_map.at(input_k_m).get_node_shared_ptr(), input_k_output_idx));
|
||||
inputs.push_back(ov::Output<Node>(pattern_map.at(input_v_m).get_node_shared_ptr(), input_v_output_idx));
|
||||
|
||||
if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
|
||||
inputs.push_back(sdpa->get_input_source_output(3));
|
||||
} else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) {
|
||||
auto attn_mask = sdpa->get_input_source_output(3);
|
||||
auto scale = sdpa->get_input_source_output(4);
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_q, input_k, input_v, attn_mask, scale, order_q, order_k, order_v, order_output, sdpa->get_causal());
|
||||
inputs.push_back(sdpa->get_input_source_output(3));
|
||||
inputs.push_back(sdpa->get_input_source_output(4));
|
||||
}
|
||||
|
||||
auto sdpa_new = std::make_shared<op::SDPA>(inputs, sdpa->get_causal(), order_q, order_k, order_v, order_output);
|
||||
|
||||
sdpa_new->set_friendly_name(sdpa->get_friendly_name());
|
||||
ov::copy_runtime_info(m.get_matched_nodes(), sdpa_new);
|
||||
ov::replace_node(sdpa, sdpa_new);
|
||||
|
|
|
|||
|
|
@ -97,27 +97,25 @@ UnsqueezeBroadcastReshapeSDPAFusion::UnsqueezeBroadcastReshapeSDPAFusion() {
|
|||
return false;
|
||||
}
|
||||
|
||||
auto input_a = pattern_map.at(input_a_m).get_node_shared_ptr();
|
||||
auto input_b = pattern_map.at(input_b_m).get_node_shared_ptr();
|
||||
auto input_c = pattern_map.at(input_c_m).get_node_shared_ptr();
|
||||
OutputVector data_inputs;
|
||||
data_inputs.push_back(pattern_map.at(input_a_m).get_node_shared_ptr()); // Q
|
||||
data_inputs.push_back(pattern_map.at(input_b_m).get_node_shared_ptr()); // K
|
||||
data_inputs.push_back(pattern_map.at(input_c_m).get_node_shared_ptr()); // V
|
||||
|
||||
auto sdpa = std::dynamic_pointer_cast<op::SDPA>(m.get_match_root());
|
||||
if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
|
||||
data_inputs.push_back(sdpa->get_input_source_output(3)); // attn_mask
|
||||
} else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) {
|
||||
data_inputs.push_back(sdpa->get_input_source_output(3)); // attn_mask
|
||||
data_inputs.push_back(sdpa->get_input_source_output(4)); // scale
|
||||
}
|
||||
|
||||
auto order_a = sdpa->get_input0_transpose_order();
|
||||
auto order_b = sdpa->get_input1_transpose_order();
|
||||
auto order_c = sdpa->get_input2_transpose_order();
|
||||
auto order_d = sdpa->get_output_transpose_order();
|
||||
|
||||
std::shared_ptr<op::SDPA> sdpa_new;
|
||||
if (pattern_map.find(sdpa_without_attn_mask_m) != pattern_map.end()) {
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_a, input_b, input_c, order_a, order_b, order_c, order_d, sdpa->get_causal());
|
||||
} else if (pattern_map.find(sdpa_with_attn_mask_m) != pattern_map.end()) {
|
||||
auto attn_mask = sdpa->get_input_source_output(3);
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_a, input_b, input_c, attn_mask, order_a, order_b, order_c, order_d, sdpa->get_causal());
|
||||
} else if (pattern_map.find(sdpa_with_attn_mask_and_scale_m) != pattern_map.end()) {
|
||||
auto attn_mask = sdpa->get_input_source_output(3);
|
||||
auto scale = sdpa->get_input_source_output(4);
|
||||
sdpa_new = std::make_shared<op::SDPA>(input_a, input_b, input_c, attn_mask, scale, order_a, order_b, order_c, order_d, sdpa->get_causal());
|
||||
}
|
||||
auto sdpa_new = std::make_shared<op::SDPA>(data_inputs, sdpa->get_causal(), order_a, order_b, order_c, order_d);
|
||||
|
||||
sdpa_new->set_friendly_name(sdpa->get_friendly_name());
|
||||
ov::copy_runtime_info(m.get_matched_nodes(), sdpa_new);
|
||||
|
|
|
|||
|
|
@ -25,8 +25,7 @@ typedef std::tuple<ov::element::Type, // netPrecision
|
|||
std::vector<InputShape>, // shape
|
||||
bool, // is_causal
|
||||
bool, // has_attn
|
||||
bool, // has_scale
|
||||
std::string // targetDevice
|
||||
bool // has_scale
|
||||
> ScaledAttnGPUTestParams;
|
||||
|
||||
class ScaledAttnLayerGPUTest : public testing::WithParamInterface<ScaledAttnGPUTestParams>,
|
||||
|
|
@ -48,8 +47,7 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo
|
|||
bool is_causal;
|
||||
bool has_attn;
|
||||
bool has_scale;
|
||||
std::string targetDevice;
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = obj.param;
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = obj.param;
|
||||
|
||||
std::ostringstream result;
|
||||
result << "netPRC=" << inType << "_";
|
||||
|
|
@ -67,7 +65,6 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo
|
|||
result << "is_causal=" << is_causal << "_";
|
||||
result << "has_attn=" << has_attn << "_";
|
||||
result << "has_scale=" << has_scale << "_";
|
||||
result << "trgDev=" << targetDevice;
|
||||
|
||||
return result.str();
|
||||
}
|
||||
|
|
@ -75,7 +72,10 @@ std::string ScaledAttnLayerGPUTest::getTestCaseName(const testing::TestParamInfo
|
|||
void ScaledAttnLayerGPUTest::SetUp() {
|
||||
ov::element::Type inType;
|
||||
std::vector<InputShape> inputShapes;
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = this->GetParam();
|
||||
|
||||
targetDevice = ov::test::utils::DEVICE_GPU;
|
||||
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = this->GetParam();
|
||||
|
||||
init_input_shapes(inputShapes);
|
||||
ov::ParameterVector inputParams;
|
||||
|
|
@ -166,8 +166,7 @@ TEST_P(ScaledAttnLayerGPUTest, CompareWithRefs) {
|
|||
bool is_causal;
|
||||
bool has_attn;
|
||||
bool has_scale;
|
||||
std::string targetDevice;
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale, targetDevice) = this->GetParam();
|
||||
std::tie(inType, inputShapes, is_causal, has_attn, has_scale) = this->GetParam();
|
||||
run();
|
||||
}
|
||||
|
||||
|
|
@ -190,15 +189,15 @@ const std::vector<std::vector<InputShape>> shapes{
|
|||
{
|
||||
// q shape
|
||||
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128},
|
||||
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}}
|
||||
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 387, 128}}}
|
||||
},
|
||||
// kv shape
|
||||
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128},
|
||||
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}}
|
||||
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 387, 128}}}
|
||||
},
|
||||
// attn shape: [B, 1, -1, L0+L1]
|
||||
{ov::test::InputShape{ov::PartialShape{-1, 1, -1, -1},
|
||||
{ov::Shape{1, 1, 100, 100}, ov::Shape{1, 1, 1, 1}, ov::Shape{2, 1, 384, 384}}}
|
||||
{ov::Shape{1, 1, 100, 100}, ov::Shape{1, 1, 1, 1}, ov::Shape{2, 1, 387, 387}}}
|
||||
},
|
||||
},
|
||||
// heads number of kv is 1, attn mask: [B, H, L1, L0+L1]
|
||||
|
|
@ -222,8 +221,7 @@ const auto params = testing::Combine(testing::Values(ov::element::f16 /*, ov::el
|
|||
testing::ValuesIn(shapes),
|
||||
testing::Values(true, false),
|
||||
testing::Values(true, false),
|
||||
testing::Values(true, false),
|
||||
testing::Values(ov::test::utils::DEVICE_GPU));
|
||||
testing::Values(true, false));
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(smoke_ScaledAttn_GPU,
|
||||
ScaledAttnLayerGPUTest,
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@
|
|||
|
||||
#include "intel_gpu/op/indirect_gemm.hpp"
|
||||
#include "intel_gpu/op/gemm.hpp"
|
||||
#include "intel_gpu/op/indirect_sdpa.hpp"
|
||||
#include "intel_gpu/op/sdpa.hpp"
|
||||
#include "intel_gpu/op/read_value.hpp"
|
||||
#include "intel_gpu/op/kv_cache.hpp"
|
||||
|
||||
|
|
@ -132,3 +134,51 @@ TEST_F(TransformationTestsF, IndirectKVCache3) {
|
|||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TransformationTestsF, IndirectKVCache4) {
|
||||
std::vector<int64_t> in0_order = {0, 1, 2, 3};
|
||||
std::vector<int64_t> in1_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in2_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> out_order = {0, 1, 2, 3};
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v1"});
|
||||
auto parameter_key = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, -1, 32, 128});
|
||||
auto parameter_value = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, -1, 32, 128});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(value_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 1);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, parameter_key, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, parameter_value, value_variable, 2, ov::element::f32);
|
||||
auto sdpa_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, 32, -1, 128});
|
||||
auto inputs = ov::OutputVector{sdpa_q, key_cache, value_cache};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
auto result = std::make_shared<ov::op::v0::Result>(sdpa);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{parameter_key, parameter_value, beam_idx, sdpa_q});
|
||||
manager.register_pass<IndirectKVCache>();
|
||||
}
|
||||
{
|
||||
auto indirect_axis = 1;
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{1, -1, 32, 128}, ov::element::f32, "v1"});
|
||||
auto parameter_key = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, -1, 32, 128});
|
||||
auto parameter_value = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, -1, 32, 128});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(value_variable);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_past, parameter_key, beam_idx, key_variable, 2, indirect_axis, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_past, parameter_value, beam_idx, value_variable, 2, indirect_axis, ov::element::f32);
|
||||
auto sdpa_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1, 32, -1, 128});
|
||||
auto inputs = ov::OutputVector{sdpa_q, key_cache, value_cache};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::IndirectSDPA>(inputs, key_cache->output(1), is_causal, indirect_axis, in0_order, in1_order, in2_order, out_order);
|
||||
auto result = std::make_shared<ov::op::v0::Result>(sdpa);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::ResultVector{result}, ov::ParameterVector{parameter_key, parameter_value, beam_idx, sdpa_q});
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,11 +25,12 @@ namespace test {
|
|||
namespace intel_gpu {
|
||||
|
||||
TEST_F(TransformationTestsF, TranposeSDPAFusion1) {
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(input_a, input_b, input_c, true);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(input_a, input_b, input_c, is_causal);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
manager.register_pass<TransposeFusion>();
|
||||
|
|
@ -42,21 +43,22 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion1) {
|
|||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(input_a, input_b, input_c, order_a, order_b, order_c, order_output, true, ov::element::undefined );
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined );
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TransformationTestsF, TranposeSDPAFusion2) {
|
||||
TEST_F(TransformationTestsF, TransformationTestsF) {
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3});
|
||||
auto tranpose_a = std::make_shared<ov::op::v1::Transpose>(input_a, tranpose_a_const);
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, input_b, input_c, true);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, input_b, input_c, is_causal);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
manager.register_pass<TransposeFusion>();
|
||||
|
|
@ -69,7 +71,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion2) {
|
|||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(input_a, input_b, input_c, order_a, order_b, order_c, order_output, true, ov::element::undefined);
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
|
|
@ -77,6 +79,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion2) {
|
|||
}
|
||||
|
||||
TEST_F(TransformationTestsF, TranposeSDPAFusion3) {
|
||||
const bool is_causal = false;
|
||||
{
|
||||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3});
|
||||
|
|
@ -86,7 +89,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) {
|
|||
auto tranpose_b = std::make_shared<ov::op::v1::Transpose>(input_b, tranpose_b_const);
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, input_c, false);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, input_c, is_causal);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
manager.register_pass<TransposeFusion>();
|
||||
|
|
@ -99,7 +102,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) {
|
|||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(input_a, input_b, input_c, order_a, order_b, order_c, order_output, false, ov::element::undefined);
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
|
|
@ -107,6 +110,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion3) {
|
|||
}
|
||||
|
||||
TEST_F(TransformationTestsF, TranposeSDPAFusion4) {
|
||||
const bool is_causal = false;
|
||||
{
|
||||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3});
|
||||
|
|
@ -118,7 +122,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) {
|
|||
auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3});
|
||||
auto tranpose_c = std::make_shared<ov::op::v1::Transpose>(input_c, tranpose_c_const);
|
||||
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, false);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, is_causal);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
manager.register_pass<TransposeFusion>();
|
||||
|
|
@ -131,7 +135,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) {
|
|||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_b = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_c = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(input_a, input_b, input_c, order_a, order_b, order_c, order_output, false, ov::element::undefined);
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(ov::OutputVector{input_a, input_b, input_c}, is_causal, order_a, order_b, order_c, order_output, ov::element::undefined);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
|
|
@ -139,6 +143,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion4) {
|
|||
}
|
||||
|
||||
TEST_F(TransformationTestsF, TranposeSDPAFusion5) {
|
||||
const bool is_causal = false;
|
||||
{
|
||||
auto input_a = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto tranpose_a_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3});
|
||||
|
|
@ -150,7 +155,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion5) {
|
|||
auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {3, 2, 1, 0});
|
||||
auto tranpose_c = std::make_shared<ov::op::v1::Transpose>(input_c, tranpose_c_const);
|
||||
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, false);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, is_causal);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
manager.register_pass<TransposeFusion>();
|
||||
|
|
@ -166,7 +171,7 @@ TEST_F(TransformationTestsF, TranposeSDPAFusion5) {
|
|||
auto tranpose_c_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, {3, 2, 1, 0});
|
||||
auto tranpose_c = std::make_shared<ov::op::v1::Transpose>(input_c, tranpose_c_const);
|
||||
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, false);
|
||||
auto sdpa = std::make_shared<ov::op::v13::ScaledDotProductAttention>(tranpose_a, tranpose_b, tranpose_c, is_causal);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_a, input_b, input_c });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,233 @@
|
|||
// Copyright (C) 2024 Intel Corporation
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
//
|
||||
|
||||
#include "common_test_utils/ov_test_utils.hpp"
|
||||
|
||||
#include "openvino/core/model.hpp"
|
||||
#include "openvino/pass/manager.hpp"
|
||||
#include "openvino/op/abs.hpp"
|
||||
#include "openvino/op/broadcast.hpp"
|
||||
#include "openvino/op/constant.hpp"
|
||||
#include "openvino/op/gather.hpp"
|
||||
#include "openvino/op/reshape.hpp"
|
||||
#include "openvino/op/unsqueeze.hpp"
|
||||
|
||||
#include "intel_gpu/op/sdpa.hpp"
|
||||
#include "intel_gpu/op/read_value.hpp"
|
||||
#include "intel_gpu/op/kv_cache.hpp"
|
||||
|
||||
#include "plugin/transformations/unsqueeze_broadcast_reshape_sdpa_fusion.hpp"
|
||||
|
||||
#include <memory>
|
||||
|
||||
using namespace testing;
|
||||
using namespace ov::intel_gpu;
|
||||
|
||||
namespace ov {
|
||||
namespace test {
|
||||
namespace intel_gpu {
|
||||
|
||||
TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion1) {
|
||||
std::vector<int64_t> in0_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in1_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in2_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> out_order = {0, 1, 2, 3};
|
||||
std::vector<int64_t> axes_val = {-2};
|
||||
std::vector<int32_t> target_shape_kv = {1, 1, 1, 16, 1};
|
||||
std::vector<int64_t> pattern_shape = {0, 0, 32, 32};
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"});
|
||||
auto key_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto value_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, key_token_param, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, value_token_param, value_variable, 2, ov::element::f32);
|
||||
auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val);
|
||||
auto key_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(key_cache, axes);
|
||||
auto value_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(value_cache, axes);
|
||||
auto target_shape = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_kv);
|
||||
auto key_broadcast = std::make_shared<ov::op::v3::Broadcast>(key_unsqueeze, target_shape, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto value_broadcast = std::make_shared<ov::op::v3::Broadcast>(value_unsqueeze, target_shape, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape);
|
||||
auto key_reshape = std::make_shared<ov::op::v1::Reshape>(key_broadcast, pattern, true);
|
||||
auto value_reshape = std::make_shared<ov::op::v1::Reshape>(value_broadcast, pattern, true);
|
||||
auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx });
|
||||
manager.register_pass<UnsqueezeBroadcastReshapeSDPAFusion>();
|
||||
}
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"});
|
||||
auto key_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto value_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, key_token_param, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, value_token_param, value_variable, 2, ov::element::f32);
|
||||
auto inputs = ov::OutputVector{input_q, key_cache, value_cache};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion2) {
|
||||
std::vector<int64_t> in0_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in1_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in2_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> out_order = {0, 1, 2, 3};
|
||||
std::vector<int64_t> axes_val = {2};
|
||||
std::vector<int64_t> pattern_shape = {0, 32, -1, 32};
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v1"});
|
||||
auto key_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, 8, -1, 32});
|
||||
auto value_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, 8, -1, 32});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, key_token_param, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, value_token_param, value_variable, 2, ov::element::f32);
|
||||
auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val);
|
||||
auto key_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(key_cache, axes);
|
||||
auto value_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(value_cache, axes);
|
||||
auto abs_param = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::Shape{5});
|
||||
auto abs = std::make_shared<ov::op::v0::Abs>(abs_param);
|
||||
auto key_broadcast = std::make_shared<ov::op::v3::Broadcast>(key_unsqueeze, abs, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto value_broadcast = std::make_shared<ov::op::v3::Broadcast>(value_unsqueeze, abs, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape);
|
||||
auto key_reshape = std::make_shared<ov::op::v1::Reshape>(key_broadcast, pattern, true);
|
||||
auto value_reshape = std::make_shared<ov::op::v1::Reshape>(value_broadcast, pattern, true);
|
||||
auto inputs = ov::OutputVector{input_q, key_cache, value_cache};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx });
|
||||
manager.register_pass<UnsqueezeBroadcastReshapeSDPAFusion>();
|
||||
}
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, 8, -1, 32}, ov::element::f32, "v1"});
|
||||
auto key_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, 8, -1, 32});
|
||||
auto value_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, 8, -1, 32});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, key_token_param, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, value_token_param, value_variable, 2, ov::element::f32);
|
||||
auto inputs = ov::OutputVector{input_q, key_cache, value_cache};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model_ref = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx });
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion3) {
|
||||
std::vector<int64_t> in0_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in1_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in2_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> out_order = {0, 1, 2, 3};
|
||||
std::vector<int64_t> axes_val = {-2};
|
||||
std::vector<int32_t> target_shape_kv = {1, 1, 1, 16, 1};
|
||||
std::vector<int64_t> pattern_shape = {0, 0, -1, 32};
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto input_k = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, -1, 32});
|
||||
auto input_v = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, -1, 32});
|
||||
auto unsqueeze_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val);
|
||||
auto key_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(input_k, unsqueeze_const);
|
||||
auto value_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(input_v, unsqueeze_const);
|
||||
auto broadcast_const = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_kv);
|
||||
auto key_broadcast = std::make_shared<ov::op::v3::Broadcast>(key_unsqueeze, broadcast_const, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto value_broadcast = std::make_shared<ov::op::v3::Broadcast>(value_unsqueeze, broadcast_const, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto reshape_const = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape);
|
||||
auto key_reshape = std::make_shared<ov::op::v1::Reshape>(key_broadcast, reshape_const, true);
|
||||
auto value_reshape = std::make_shared<ov::op::v1::Reshape>(value_broadcast, reshape_const, true);
|
||||
auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, input_k, input_v });
|
||||
manager.register_pass<UnsqueezeBroadcastReshapeSDPAFusion>();
|
||||
}
|
||||
{
|
||||
model_ref = model->clone();
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(TransformationTestsF, UnsqueezeBroadReshapeSDPAFusion4) {
|
||||
std::vector<int64_t> in0_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in1_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> in2_order = {0, 2, 1, 3};
|
||||
std::vector<int64_t> out_order = {0, 1, 2, 3};
|
||||
std::vector<int64_t> axes_val = {-2};
|
||||
std::vector<int32_t> target_shape_k = {1, 1, 1, 16, 1};
|
||||
std::vector<int32_t> target_shape_v = {1, 1, 1, 32, 1};
|
||||
std::vector<int64_t> pattern_shape = {0, 0, 32, 32};
|
||||
const bool is_causal = true;
|
||||
{
|
||||
auto input_q = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape::dynamic(4));
|
||||
auto key_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v0"});
|
||||
auto value_variable = std::make_shared<ov::op::util::Variable>(ov::op::util::VariableInfo{{-1, -1, 2, 32}, ov::element::f32, "v1"});
|
||||
auto key_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto value_token_param = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{-1, -1, 2, 32});
|
||||
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});
|
||||
auto key_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto value_past = std::make_shared<ov::intel_gpu::op::ReadValue>(key_variable);
|
||||
auto axis = std::make_shared<ov::op::v0::Constant>(ov::element::i64, ov::Shape{}, 0);
|
||||
auto key_gather_past = std::make_shared<ov::op::v8::Gather>(key_past, beam_idx, axis);
|
||||
auto value_gather_past = std::make_shared<ov::op::v8::Gather>(value_past, beam_idx, axis);
|
||||
auto key_cache = std::make_shared<ov::intel_gpu::op::KVCache>(key_gather_past, key_token_param, key_variable, 2, ov::element::f32);
|
||||
auto value_cache = std::make_shared<ov::intel_gpu::op::KVCache>(value_gather_past, value_token_param, value_variable, 2, ov::element::f32);
|
||||
auto axes = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{1}, axes_val);
|
||||
auto key_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(key_cache, axes);
|
||||
auto value_unsqueeze = std::make_shared<ov::op::v0::Unsqueeze>(value_cache, axes);
|
||||
auto target_shape_key = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_k);
|
||||
auto target_shape_value = ov::op::v0::Constant::create(ov::element::i32, ov::Shape{5}, target_shape_v);
|
||||
auto key_broadcast = std::make_shared<ov::op::v3::Broadcast>(key_unsqueeze, target_shape_key, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto value_broadcast = std::make_shared<ov::op::v3::Broadcast>(value_unsqueeze, target_shape_value, ov::op::BroadcastType::BIDIRECTIONAL);
|
||||
auto pattern = ov::op::v0::Constant::create(ov::element::i64, ov::Shape{4}, pattern_shape);
|
||||
auto key_reshape = std::make_shared<ov::op::v1::Reshape>(key_broadcast, pattern, true);
|
||||
auto value_reshape = std::make_shared<ov::op::v1::Reshape>(value_broadcast, pattern, true);
|
||||
auto inputs = ov::OutputVector{input_q, key_reshape, value_reshape};
|
||||
auto sdpa = std::make_shared<ov::intel_gpu::op::SDPA>(inputs, is_causal, in0_order, in1_order, in2_order, out_order);
|
||||
|
||||
model = std::make_shared<ov::Model>(ov::NodeVector{ sdpa }, ov::ParameterVector{ input_q, key_token_param, value_token_param, beam_idx });
|
||||
manager.register_pass<UnsqueezeBroadcastReshapeSDPAFusion>();
|
||||
}
|
||||
{
|
||||
model_ref = model->clone();
|
||||
comparator.enable(FunctionsComparator::ATTRIBUTES);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace intel_gpu
|
||||
} // namespace test
|
||||
} // namespace ov
|
||||
Loading…
Reference in New Issue