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

Merge 8c867fba67 into 43de73cbd6

This commit is contained in:
Luo Cheng 2024-06-12 05:36:57 +00:00 committed by GitHub
parent 43de73cbd6 8c867fba67
commit ecb954224c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 120 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp
View File

@ -53,7 +53,10 @@ StatefulSDPAFusion::StatefulSDPAFusion() {
auto concat_k = makePattern<opset1::Concat>({gather_input_k, cur_k}, {{"axis", axis_seq_len}});
auto concat_v = makePattern<opset1::Concat>({gather_input_v, cur_v}, {{"axis", axis_seq_len}});
auto multi_query_bcst = [](std::shared_ptr<Node> kv) {
std::shared_ptr<Node> reshape_k, reshape_v, unsqueeze_k, unsqueeze_v;
std::shared_ptr<Node> computed_bcst_k, computed_bcst_v, multiply_k, multiply_v;
std::shared_ptr<Node> mq_reshape_k, mq_reshape_v;
auto multi_query_bcst = [](const std::shared_ptr<Node>& kv) {
auto reshape_kv = wrap_type<opset6::Reshape>({kv, any_input()});
auto unsqueeze_kv = makePattern<opset1::Unsqueeze>({kv, any_input()});
@ -70,11 +73,14 @@ StatefulSDPAFusion::StatefulSDPAFusion() {
any_input(), any_input()}, {{"mode", "numpy"}});
auto multiply_kv = wrap_type<opset6::Multiply>({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst});
return wrap_type<opset6::Reshape>({multiply_kv, any_input()});
auto result = wrap_type<opset6::Reshape>({multiply_kv, any_input()});
return std::make_tuple(result, reshape_kv, unsqueeze_kv, computed_bcst, multiply_kv);
};
auto present_k = concat_k | multi_query_bcst(concat_k);
auto present_v = concat_v | multi_query_bcst(concat_v);
std::tie(mq_reshape_k, reshape_k, unsqueeze_k, computed_bcst_k, multiply_k) = multi_query_bcst(concat_k);
std::tie(mq_reshape_v, reshape_v, unsqueeze_v, computed_bcst_v, multiply_v) = multi_query_bcst(concat_v);
auto present_k = concat_k | mq_reshape_k;
auto present_v = concat_v | mq_reshape_v;
// canonical q/k/v shape definition: [B,H,...L,S]
auto sdp0 = makePattern<opset13::ScaledDotProductAttention>({cur_q, present_k, present_v});
@ -181,12 +187,31 @@ StatefulSDPAFusion::StatefulSDPAFusion() {
if (past_v_node->get_variable_id() != assign_v_node->get_variable_id())
return false;
auto is_optional_one_child = [&pattern_map] (const std::vector<std::shared_ptr<Node>>& nodes) {
for (auto&& node : nodes) {
if (pattern_map.count(node)) {
auto p = pattern_map.at(node).get_node_shared_ptr();
if (p->get_output_target_inputs(0).size() != 1)
return false;
}
}
return true;
};
if (!is_optional_one_child({convert_past_k, convert_past_v, transpose_q, transpose_k, transpose_v,
reshape_k, unsqueeze_k, computed_bcst_k, multiply_k,
reshape_v, unsqueeze_v, computed_bcst_v, multiply_v,
mq_reshape_k, mq_reshape_v})) {
return false;
}
// past_k & past_v must be reordered by same beam_idx
const auto gather_k_node =
ov::as_type_ptr<opset8::Gather>(pattern_map.at(gather_input_k).get_node_shared_ptr());
const auto gather_v_node =
ov::as_type_ptr<opset8::Gather>(pattern_map.at(gather_input_v).get_node_shared_ptr());
if (gather_k_node->input_value(1) != gather_v_node->input_value(1)) {
if (gather_k_node->input_value(1) != gather_v_node->input_value(1) ||
gather_k_node->get_output_target_inputs(0).size() != 1 ||
gather_v_node->get_output_target_inputs(0).size() != 1) {
return false;
}

44
src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_sdp.cpp
View File

@ -66,7 +66,6 @@ public:
void SetUp() override {
ElementType inType;
std::vector<InputShape> inputShapes;
bool hasShapeOf;
std::tie(inType, inputShapes, hasShapeOf) = this->GetParam();
targetDevice = ov::test::utils::DEVICE_CPU;
rel_threshold = 1e-2f;
@ -93,16 +92,21 @@ public:
ov::op::util::VariableInfo{inputDynamicShapes[1], inType, "pastv"});
auto pastv = std::make_shared<ov::op::v6::ReadValue>(inputParams[3], var_v);
pastv->set_friendly_name("pastv_r");
std::shared_ptr<Node> pastk_shapeof, pastv_shapeof;
if (hasShapeOf) {
pastk_shapeof = std::make_shared<ov::op::v0::ShapeOf>(pastk);
pastv_shapeof = std::make_shared<ov::op::v0::ShapeOf>(pastv);
}
auto beam_idx = std::make_shared<ov::op::v0::Parameter>(ElementType::i32, ov::PartialShape{-1});
beam_idx->set_friendly_name("beam_idx");
inputParams.push_back(beam_idx);
auto gatherK = std::make_shared<ov::op::v8::Gather>(pastk, beam_idx, op::v0::Constant::create(ElementType::i32, {1}, {0}));
auto gatherV = std::make_shared<ov::op::v8::Gather>(pastv, beam_idx, op::v0::Constant::create(ElementType::i32, {1}, {0}));
auto gatherK = std::make_shared<ov::op::v8::Gather>(pastk, beam_idx, op::v0::Constant::create(ElementType::i32, {}, {0}));
auto gatherV = std::make_shared<ov::op::v8::Gather>(pastv, beam_idx, op::v0::Constant::create(ElementType::i32, {}, {0}));
std::shared_ptr<Node> shapeof_k, shapeof_v;
// test special case:
// ReadValue->Gather->Concat->SDPA...
// |
// ShapeOf...
// The transformation 'SimplifyGatherShapeOf' will move ShapeOf to be the child of ReadValue
if (hasShapeOf) {
shapeof_k = std::make_shared<ov::op::v0::ShapeOf>(gatherK);
shapeof_v = std::make_shared<ov::op::v0::ShapeOf>(gatherV);
}
auto concatK = std::make_shared<ov::op::v0::Concat>(OutputVector{gatherK, inputParams[1]}, 2);
auto concatV = std::make_shared<ov::op::v0::Concat>(OutputVector{gatherV, inputParams[2]}, 2);
auto sdp = std::make_shared<ov::opset13::ScaledDotProductAttention>(inputParams[0], concatK, concatV, false);
@ -115,8 +119,8 @@ public:
ResultVector results{std::make_shared<ov::op::v0::Result>(add)};
if (hasShapeOf) {
results.push_back(std::make_shared<ov::op::v0::Result>(pastk_shapeof));
results.push_back(std::make_shared<ov::op::v0::Result>(pastv_shapeof));
results.push_back(std::make_shared<ov::op::v0::Result>(shapeof_k));
results.push_back(std::make_shared<ov::op::v0::Result>(shapeof_v));
}
SinkVector sinks{pastk_assign, pastv_assign};
function = std::make_shared<ov::Model>(results, sinks, inputParams, "ConcatSDP");
@ -202,14 +206,26 @@ public:
return outputs;
}
bool hasShapeOf;
};
TEST_P(ConcatSDPTest, CompareWithRefs) {
auto actualOutputs = run_test(function);
CheckNumberOfNodesWithType(compiledModel, "ScaledDotProductAttention", 1);
CheckNumberOfNodesWithType(compiledModel, "Concatenation", 0);
CheckNumberOfNodesWithType(compiledModel, "Reorder", 0);
CheckNumberOfNodesWithType(compiledModel, "Gather", 0);
if (!hasShapeOf) {
CheckNumberOfNodesWithType(compiledModel, "ScaledDotProductAttention", 1);
CheckNumberOfNodesWithType(compiledModel, "Concatenation", 0);
CheckNumberOfNodesWithType(compiledModel, "Reorder", 0);
CheckNumberOfNodesWithType(compiledModel, "Gather", 0);
} else {
// SimplifyGatherShapeOf will generate subgraph which contains Gather/Concat/Reorder, so could not check the number to confirm it's expected.
// W/ or w/o fusion the SDPA name is the same, but if fused the output number should be 3(data+pastk+pastv)
for (const auto& node : compiledModel.get_runtime_model()->get_ops()) {
if (node->get_rt_info().find(ov::exec_model_info::LAYER_TYPE)->second.as<std::string>() == "ScaledDotProductAttention") {
ASSERT_EQ(3, node->get_output_size()) << "ScaledDotProductAttention should be fused";
}
}
}
auto expectedOutputs = run_test(functionRefs);
CheckNumberOfNodesWithType(compiledModel, "ScaledDotProductAttention", 0);
for (size_t i = 0; i < actualOutputs.size(); i++) {

61
src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp
View File

@ -24,7 +24,22 @@ using namespace ov;
using namespace ov::intel_cpu;
using namespace ov::gen_pattern;
static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, bool isRef = false, bool hasConvert = false, bool hasMultiquery = false) {
namespace {
enum InsertPoint {
At_None,
At_Convert,
At_Gather,
At_MQ_Unsqueeze,
At_MQ_Broadcast,
At_MQ_Multiply,
At_MQ_Reshape,
At_End
};
} // namespace
static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, bool isRef = false, bool hasConvert = false, bool hasMultiquery = false,
InsertPoint at = InsertPoint::At_None) {
std::shared_ptr<ov::op::v0::Result> insert_result;
auto q = std::make_shared<ov::op::v0::Parameter>(element::f32, inputShape);
auto kvInputShape = inputShape;
if (hasMultiquery) {
@ -44,6 +59,10 @@ static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, b
if (hasConvert) {
pastk = std::make_shared<ov::op::v0::Convert>(pastk, element::f32);
pastv = std::make_shared<ov::op::v0::Convert>(pastv, element::f32);
if (at == InsertPoint::At_Convert) {
// insert one should be enough to cover the 'convert' case
insert_result = std::make_shared<ov::op::v0::Result>(pastk);
}
}
if (isRef) {
ov::intel_cpu::ScaledDotProductAttentionWithKVCache::Config config;
@ -54,6 +73,9 @@ static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, b
concatV = new_node->output(2);
} else {
pastk = std::make_shared<ov::op::v8::Gather>(pastk, beam_idx, op::v0::Constant::create(element::i32, {1}, {0}));
if (at == InsertPoint::At_Gather) {
insert_result = std::make_shared<ov::op::v0::Result>(pastk);
}
pastv = std::make_shared<ov::op::v8::Gather>(pastv, beam_idx, op::v0::Constant::create(element::i32, {1}, {0}));
concatK = std::make_shared<ov::op::v0::Concat>(OutputVector{pastk, k}, 2);
concatV = std::make_shared<ov::op::v0::Concat>(OutputVector{pastv, v}, 2);
@ -62,6 +84,9 @@ static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, b
auto beam_idx_shape = makeOP<ov::op::TypeRelaxed<opset1::ShapeOf>>({beam_idx},
{{"type_relax", true}, {"input_data_types", {}}, {"output_data_types", {element::i32}}});
auto unsqueeze_concat = makeOP<opset1::Unsqueeze>({conat, 2});
if (at == InsertPoint::At_MQ_Unsqueeze && !insert_result) {
insert_result = std::make_shared<ov::op::v0::Result>(unsqueeze_concat);
}
auto concat_shape = makeOP<ov::op::TypeRelaxed<opset1::ShapeOf>>(
{conat},
{{"type_relax", true}, {"input_data_types", {}}, {"output_data_types", {element::i32}}});
@ -72,10 +97,19 @@ static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, b
auto expand_ones = makeOP<opset1::Broadcast>({{1.0f},
expand_Abs,
axis_mapping}, {{"mode", "numpy"}});
if (at == InsertPoint::At_MQ_Broadcast && !insert_result) {
insert_result = std::make_shared<ov::op::v0::Result>(expand_ones);
}
auto expand_Broadcast = makeOP<opset1::Multiply>({unsqueeze_concat,
expand_ones}, {{"auto_broadcast", "numpy"}});
if (at == InsertPoint::At_MQ_Multiply && !insert_result) {
insert_result = std::make_shared<ov::op::v0::Result>(expand_Broadcast);
}
auto expected_shape = makeOP<opset1::Concat>({beam_idx_shape, {inputShape[1]}, gather_ls}, {{"axis", 0}});
auto reshape_Reshape = makeOP<opset1::Reshape>({expand_Broadcast, expected_shape}, {{"special_zero", false}});
if (at == InsertPoint::At_MQ_Reshape && !insert_result) {
insert_result = std::make_shared<ov::op::v0::Result>(reshape_Reshape);
}
return reshape_Reshape;
};
sdp = std::make_shared<ov::opset13::ScaledDotProductAttention>(q, make_multi_query(concatK), make_multi_query(concatV), false);
@ -92,6 +126,8 @@ static std::shared_ptr<ov::Model> makeSDPA(const ov::PartialShape& inputShape, b
auto add = std::make_shared<op::v1::Add>(sdp, op::v0::Constant::create(element::f32, {1}, {1.0f}));
ResultVector results{std::make_shared<ov::op::v0::Result>(add)};
if (insert_result)
results.push_back(insert_result);
SinkVector sinks{pastk_assign, pastv_assign};
return std::make_shared<Model>(results, sinks, ParameterVector{q, k, v, init, beam_idx}, "ConcatSDP");
}
@ -158,3 +194,26 @@ TEST(TransformationTests, StateConcatSDPAMixtral) {
ASSERT_TRUE(res.first) << res.second;
}
}
TEST(TransformationTests, StateConcatSDPAWithExtraNode) {
// when some unexpected extra nodes exist in SDPA, the fusion should fail
std::shared_ptr<ov::Model> f(nullptr), f_ref(nullptr);
{
using namespace ov;
auto inputShape = ov::PartialShape{-1, 32, -1, 64};
size_t i = static_cast<size_t>(InsertPoint::At_None) + 1;
size_t end = static_cast<size_t>(InsertPoint::At_End);
// check each position
for (; i < end; i++) {
InsertPoint at = static_cast<InsertPoint>(i);
f = makeSDPA(inputShape, false, true, true, at);
f_ref = makeSDPA(inputShape, false, true, true, at);
pass::Manager m;
m.register_pass<ov::pass::InitNodeInfo>();
m.register_pass<StatefulSDPAFusion>();
m.run_passes(f);
auto res = compare_functions(f, f_ref);
ASSERT_TRUE(res.first) << res.second;
}
}
}