|
| 1 | +// Copyright (C) 2018-2025 Intel Corporation |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | +// |
| 4 | + |
| 5 | +#include "transformations/op_conversions/group_query_attention_decomposition.hpp" |
| 6 | + |
| 7 | +#include <memory> |
| 8 | + |
| 9 | +#include "itt.hpp" |
| 10 | +#include "openvino/core/rt_info.hpp" |
| 11 | +#include "openvino/op/add.hpp" |
| 12 | +#include "openvino/op/concat.hpp" |
| 13 | +#include "openvino/op/constant.hpp" |
| 14 | +#include "openvino/op/convert.hpp" |
| 15 | +#include "openvino/op/gather.hpp" |
| 16 | +#include "openvino/op/greater.hpp" |
| 17 | +#include "openvino/op/multiply.hpp" |
| 18 | +#include "openvino/op/range.hpp" |
| 19 | +#include "openvino/op/reshape.hpp" |
| 20 | +#include "openvino/op/scaled_dot_product_attention.hpp" |
| 21 | +#include "openvino/op/select.hpp" |
| 22 | +#include "openvino/op/shape_of.hpp" |
| 23 | +#include "openvino/op/slice.hpp" |
| 24 | +#include "openvino/op/split.hpp" |
| 25 | +#include "openvino/op/subtract.hpp" |
| 26 | +#include "openvino/op/transpose.hpp" |
| 27 | +#include "openvino/op/unsqueeze.hpp" |
| 28 | +#include "openvino/pass/pattern/op/wrap_type.hpp" |
| 29 | + |
| 30 | +ov::pass::GroupQueryAttentionDecomposition::GroupQueryAttentionDecomposition() { |
| 31 | + MATCHER_SCOPE(GroupQeuryAttentionDecomposition); |
| 32 | + auto pattern_node = ov::pass::pattern::wrap_type<ov::op::internal::GroupQueryAttention>(); |
| 33 | + |
| 34 | + matcher_pass_callback callback = [OV_CAPTURE_CPY_AND_THIS](ov::pass::pattern::Matcher& m) { |
| 35 | + auto& pattern_to_output = m.get_pattern_value_map(); |
| 36 | + auto node = ov::as_type_ptr<ov::op::internal::GroupQueryAttention>( |
| 37 | + pattern_to_output.at(pattern_node).get_node_shared_ptr()); |
| 38 | + |
| 39 | + if (node == nullptr || transformation_callback(node)) { |
| 40 | + return false; |
| 41 | + } |
| 42 | + |
| 43 | + auto new_output_node = decompose(node); |
| 44 | + ov::replace_node(node, new_output_node); |
| 45 | + return true; |
| 46 | + }; |
| 47 | + |
| 48 | + auto m = std::make_shared<ov::pass::pattern::Matcher>(pattern_node, matcher_name); |
| 49 | + register_matcher(m, callback); |
| 50 | +} |
| 51 | + |
| 52 | +ov::OutputVector ov::pass::GroupQueryAttentionDecomposition::decompose( |
| 53 | + std::shared_ptr<ov::op::internal::GroupQueryAttention> node) { |
| 54 | + using namespace ov::op; |
| 55 | + |
| 56 | + const auto num_heads = node->get_num_heads(); |
| 57 | + const auto kv_num_heads = node->get_kv_num_heads(); |
| 58 | + const auto scale = node->get_scale(); |
| 59 | + const auto do_rotary = node->get_do_rotary(); |
| 60 | + const auto rotary_interleaved = node->get_rotary_interleaved(); |
| 61 | + // TODO: add softcap support |
| 62 | + |
| 63 | + auto Q = node->input_value(0); |
| 64 | + auto K = node->input_value(1); |
| 65 | + auto V = node->input_value(2); |
| 66 | + auto past_key = node->input_value(3); |
| 67 | + auto past_value = node->input_value(4); |
| 68 | + auto seqlens_k = node->input_value(5); |
| 69 | + auto cos_cache = node->input_value(6); |
| 70 | + auto sin_cache = node->input_value(7); |
| 71 | + |
| 72 | + // The length of all tokens (past + current) is `seqlens_k` + 1 |
| 73 | + // current = Q.shape[2], past = `seqlens_k` + 1 - current |
| 74 | + |
| 75 | + const auto T = Q.get_element_type(); |
| 76 | + const auto q_shape = register_new_node<v3::ShapeOf>(Q); |
| 77 | + const auto current_sequence_length = get_dimensions(q_shape, {2}); |
| 78 | + const auto head_size_node = get_dimensions(q_shape, {3}); |
| 79 | + |
| 80 | + auto zero = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{1}, {0})); |
| 81 | + auto one = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{1}, {1})); |
| 82 | + auto one_without_shape = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{}, {1})); |
| 83 | + auto two = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{1}, {2})); |
| 84 | + auto seqlens_elemi64 = register_new_node<v0::Convert>(seqlens_k, ov::element::i64); |
| 85 | + auto real_seqlens = register_new_node<v1::Add>(seqlens_elemi64, one); |
| 86 | + |
| 87 | + // Only consider batch is 1 |
| 88 | + auto seqlens_1d = register_new_node<v1::Reshape>(real_seqlens, one, false); |
| 89 | + auto past_sequence_length = register_new_node<v1::Subtract>(seqlens_1d, current_sequence_length); |
| 90 | + if (do_rotary) { |
| 91 | + Q = rotaryEmbedding(Q, |
| 92 | + past_sequence_length, |
| 93 | + seqlens_1d, |
| 94 | + cos_cache.get_node_shared_ptr(), |
| 95 | + sin_cache.get_node_shared_ptr(), |
| 96 | + head_size_node, |
| 97 | + rotary_interleaved); |
| 98 | + K = rotaryEmbedding(K, |
| 99 | + past_sequence_length, |
| 100 | + seqlens_1d, |
| 101 | + cos_cache.get_node_shared_ptr(), |
| 102 | + sin_cache.get_node_shared_ptr(), |
| 103 | + head_size_node, |
| 104 | + rotary_interleaved); |
| 105 | + } |
| 106 | + |
| 107 | + auto construct_kv_cache = [&](const ov::Output<ov::Node>& past, const ov::Output<ov::Node>& current) { |
| 108 | + auto past_datas = register_new_node<v8::Slice>(past, zero, past_sequence_length, one, two); |
| 109 | + auto curr_datas = register_new_node<v8::Slice>(current, zero, current_sequence_length, one, two); |
| 110 | + return register_new_node<v0::Concat>(ov::NodeVector{past_datas, curr_datas}, 2); |
| 111 | + }; |
| 112 | + K = construct_kv_cache(past_key, K); |
| 113 | + V = construct_kv_cache(past_value, V); |
| 114 | + auto present_k = K; |
| 115 | + auto present_v = V; |
| 116 | + |
| 117 | + const size_t kv_num_heads_factor = num_heads / kv_num_heads; |
| 118 | + if (kv_num_heads_factor > 1) { |
| 119 | + const auto kv_shape = register_new_node<v3::ShapeOf>(K); |
| 120 | + const auto kv_shape_prev_2 = get_dimensions(kv_shape, {0, 1}); |
| 121 | + const auto kv_shape_last_2 = get_dimensions(kv_shape, {2, 3}); |
| 122 | + auto new_kv_shape = register_new_node<v0::Concat>(ov::NodeVector{kv_shape_prev_2, one, kv_shape_last_2}, 0); |
| 123 | + K = register_new_node<v1::Reshape>(K, new_kv_shape, false); |
| 124 | + V = register_new_node<v1::Reshape>(V, new_kv_shape, false); |
| 125 | + K = register_new_node<v0::Concat>(ov::OutputVector(kv_num_heads_factor, K), 2); |
| 126 | + V = register_new_node<v0::Concat>(ov::OutputVector(kv_num_heads_factor, V), 2); |
| 127 | + const auto q_shape = register_new_node<v3::ShapeOf>(Q); |
| 128 | + const auto q_shape_prev_2 = get_dimensions(q_shape, {0, 1}); |
| 129 | + auto extended_kv_shape = register_new_node<v0::Concat>(ov::NodeVector{q_shape_prev_2, kv_shape_last_2}, 0); |
| 130 | + K = register_new_node<v1::Reshape>(K, extended_kv_shape, false); |
| 131 | + V = register_new_node<v1::Reshape>(V, extended_kv_shape, false); |
| 132 | + } |
| 133 | + |
| 134 | + // need to apply low-triangle mask to attention score. |
| 135 | + // two steps, construct the total_sequence x total_sequence triangle, then slice the current length |
| 136 | + auto seqlens_1d_scalar = register_new_node<v1::Reshape>(seqlens_1d, one_without_shape, false); |
| 137 | + std::shared_ptr<ov::Node> mask_per_line_node = |
| 138 | + register_new_node<v4::Range>(register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{}, {0})), |
| 139 | + seqlens_1d_scalar, |
| 140 | + one_without_shape, |
| 141 | + ov::element::i64); |
| 142 | + auto hori_range = register_new_node<v0::Unsqueeze>(mask_per_line_node, zero); |
| 143 | + auto vert_range = register_new_node<v0::Unsqueeze>(mask_per_line_node, one); |
| 144 | + auto triu = register_new_node<v1::Greater>(hori_range, vert_range); |
| 145 | + auto typed_zero = register_new_node(v0::Constant::create(T, ov::Shape{}, {0})); |
| 146 | + // cf. make_attention_mask@src\plugins\intel_gpu\tests\common\subgraphs_builders.hpp |
| 147 | + std::shared_ptr<ov::Node> minus_inf = nullptr; |
| 148 | + if (T == ov::element::f32) |
| 149 | + minus_inf = register_new_node(v0::Constant::create(T, ov::Shape{}, {-std::numeric_limits<float>::infinity()})); |
| 150 | + else if (T == ov::element::f16) |
| 151 | + minus_inf = |
| 152 | + register_new_node(v0::Constant::create(T, ov::Shape{}, {std::numeric_limits<ov::float16>::lowest()})); |
| 153 | + auto atten_mask = register_new_node<v1::Select>(triu, minus_inf, typed_zero); |
| 154 | + auto atten_mask_sliced = register_new_node<v8::Slice>(atten_mask, past_sequence_length, seqlens_1d, one, zero); |
| 155 | + |
| 156 | + std::shared_ptr<ov::Node> qga_output; |
| 157 | + if (scale != 0.0f) { |
| 158 | + auto scale_node = register_new_node(v0::Constant::create(T, Shape{}, {scale})); |
| 159 | + qga_output = register_new_node<v13::ScaledDotProductAttention>(Q, K, V, atten_mask_sliced, scale_node, false); |
| 160 | + } else { |
| 161 | + qga_output = register_new_node<v13::ScaledDotProductAttention>(Q, K, V, atten_mask_sliced, false); |
| 162 | + } |
| 163 | + |
| 164 | + // transpose the result from (batch_size, num_heads, sequence_length, head_size) |
| 165 | + // to (batch_size, sequence_length, num_heads * head_size) |
| 166 | + auto perm = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 2, 1, 3})); |
| 167 | + auto qga_output_transposed = register_new_node<v1::Transpose>(qga_output, perm); |
| 168 | + auto dim_merge_shape = register_new_node(v0::Constant::create(ov::element::i32, ov::Shape{3}, {0, 0, -1})); |
| 169 | + auto output = register_new_node<v1::Reshape>(qga_output_transposed, dim_merge_shape, true)->output(0); |
| 170 | + |
| 171 | + return {output, present_k, present_v}; |
| 172 | +} |
| 173 | + |
| 174 | +// make split functions is a copy-past from ONNX FE. TODO: move it to one place |
| 175 | +ov::OutputVector ov::pass::GroupQueryAttentionDecomposition::make_split(const ov::Output<ov::Node>& value, |
| 176 | + int64_t num_splits, |
| 177 | + int64_t axis) { |
| 178 | + using namespace ov::op; |
| 179 | + const auto axis_node = register_new_node(v0::Constant::create(ov::element::i64, ov::Shape{}, {axis})); |
| 180 | + const auto split = register_new_node<v1::Split>(value, axis_node, num_splits); |
| 181 | + |
| 182 | + return split->outputs(); |
| 183 | +} |
| 184 | + |
| 185 | +std::shared_ptr<ov::Node> ov::pass::GroupQueryAttentionDecomposition::get_dimensions( |
| 186 | + const std::shared_ptr<ov::op::v3::ShapeOf>& shape, |
| 187 | + const std::vector<int>& dims) { |
| 188 | + using namespace ov::op; |
| 189 | + const auto zero = v0::Constant::create(ov::element::i32, ov::Shape{}, {0}); |
| 190 | + const auto dims_const = v0::Constant::create(ov::element::i32, ov::Shape{dims.size()}, dims); |
| 191 | + return register_new_node<v8::Gather>(shape, dims_const, zero); |
| 192 | +} |
| 193 | + |
| 194 | +std::shared_ptr<ov::Node> ov::pass::GroupQueryAttentionDecomposition::get_dimensions( |
| 195 | + const std::shared_ptr<ov::Node>& node, |
| 196 | + const std::vector<int>& dims) { |
| 197 | + return get_dimensions(register_new_node<ov::op::v3::ShapeOf>(node), dims); |
| 198 | +} |
| 199 | + |
| 200 | +std::shared_ptr<ov::Node> ov::pass::GroupQueryAttentionDecomposition::rotaryEmbedding( |
| 201 | + ov::Output<ov::Node> input, |
| 202 | + ov::Output<ov::Node> past_seqlen, |
| 203 | + std::shared_ptr<ov::Node> seqlen_k, |
| 204 | + std::shared_ptr<ov::Node> cos_cache, |
| 205 | + std::shared_ptr<ov::Node> sin_cache, |
| 206 | + std::shared_ptr<ov::Node> dim_head_size, |
| 207 | + bool interleaved) { |
| 208 | + using namespace ov::op; |
| 209 | + auto zero = v0::Constant::create(ov::element::i64, ov::Shape{1}, {0}); |
| 210 | + auto one = v0::Constant::create(ov::element::i64, ov::Shape{1}, {1}); |
| 211 | + |
| 212 | + auto slice_cache_dim_shape = seqlen_k; |
| 213 | + |
| 214 | + auto cos = register_new_node<v8::Slice>(cos_cache, past_seqlen, slice_cache_dim_shape, one, zero); |
| 215 | + auto sin = register_new_node<v8::Slice>(sin_cache, past_seqlen, slice_cache_dim_shape, one, zero); |
| 216 | + |
| 217 | + if (interleaved) { |
| 218 | + auto two = v0::Constant::create(ov::element::i64, ov::Shape{1}, {2}); |
| 219 | + |
| 220 | + auto cache_shape = register_new_node<v3::ShapeOf>(cos_cache); |
| 221 | + auto cache_last_dim = get_dimensions(cos_cache, {-1}); |
| 222 | + |
| 223 | + auto input_shape = register_new_node<v3::ShapeOf>(input); |
| 224 | + |
| 225 | + auto dim_bns = get_dimensions(input_shape, {0, 1, 2}); |
| 226 | + std::shared_ptr<ov::Node> half_last_dim = cache_last_dim; |
| 227 | + |
| 228 | + auto negtive_one = v0::Constant::create(ov::element::i64, ov::Shape{1}, {-1}); |
| 229 | + auto split_input_shape = register_new_node<v0::Concat>(ov::NodeVector{dim_bns, half_last_dim, two}, 0); |
| 230 | + auto reshaped_input = register_new_node<v1::Reshape>(input, split_input_shape, false); |
| 231 | + |
| 232 | + auto in_split = make_split(reshaped_input, 2, -1); |
| 233 | + split_input_shape = register_new_node<v0::Concat>(ov::NodeVector{dim_bns, half_last_dim}, 0); |
| 234 | + auto in_split_0 = register_new_node<v1::Reshape>(in_split[0], split_input_shape, false); |
| 235 | + auto in_split_1 = register_new_node<v1::Reshape>(in_split[1], split_input_shape, false); |
| 236 | + |
| 237 | + auto res_0 = register_new_node<v1::Subtract>(register_new_node<v1::Multiply>(in_split_0, cos), |
| 238 | + register_new_node<v1::Multiply>(in_split_1, sin)); |
| 239 | + auto res_1 = register_new_node<v1::Add>(register_new_node<v1::Multiply>(in_split_0, sin), |
| 240 | + register_new_node<v1::Multiply>(in_split_1, cos)); |
| 241 | + |
| 242 | + split_input_shape = register_new_node<v0::Concat>(ov::NodeVector{dim_bns, half_last_dim, one}, 0); |
| 243 | + auto res_0_5d = register_new_node<v1::Reshape>(res_0, split_input_shape, false); |
| 244 | + auto res_1_5d = register_new_node<v1::Reshape>(res_1, split_input_shape, false); |
| 245 | + |
| 246 | + auto concat_ret = register_new_node<v0::Concat>(ov::NodeVector{res_0_5d, res_1_5d}, -1); |
| 247 | + return register_new_node<v1::Reshape>(concat_ret, input_shape, false); |
| 248 | + } else { |
| 249 | + auto in_split = make_split(input, 2, -1); |
| 250 | + auto res_0 = register_new_node<v1::Subtract>(register_new_node<v1::Multiply>(in_split[0], cos), |
| 251 | + register_new_node<v1::Multiply>(in_split[1], sin)); |
| 252 | + auto res_1 = register_new_node<v1::Add>(register_new_node<v1::Multiply>(in_split[0], sin), |
| 253 | + register_new_node<v1::Multiply>(in_split[1], cos)); |
| 254 | + |
| 255 | + return register_new_node<v0::Concat>(ov::NodeVector{res_0, res_1}, -1); |
| 256 | + } |
| 257 | +} |
0 commit comments