From 13c40a49baf28c8ed284dad6e04870ecbafa2fda Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 5 Jun 2024 05:30:31 +0200 Subject: [PATCH 1/4] add more check for sdpa pattern match --- .../common/pass/stateful_sdpa_fusion.cpp | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp index cb3fc689923..8f6b3b65772 100644 --- a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp @@ -181,12 +181,26 @@ 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::shared_ptr& node) { + 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) || !is_optional_one_child(convert_past_v) || + !is_optional_one_child(transpose_q) || !is_optional_one_child(transpose_k) || !is_optional_one_child(transpose_v)) + return false; + // past_k & past_v must be reordered by same beam_idx const auto gather_k_node = ov::as_type_ptr(pattern_map.at(gather_input_k).get_node_shared_ptr()); const auto gather_v_node = ov::as_type_ptr(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; } From d035478dd8d2662fc58bd08fc3f9f2aeb652bfb3 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Fri, 7 Jun 2024 03:10:06 +0200 Subject: [PATCH 2/4] more check for multi query --- .../common/pass/stateful_sdpa_fusion.cpp | 37 ++++++++++++------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp index 8f6b3b65772..c153df2f20d 100644 --- a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp @@ -53,9 +53,12 @@ StatefulSDPAFusion::StatefulSDPAFusion() { auto concat_k = makePattern({gather_input_k, cur_k}, {{"axis", axis_seq_len}}); auto concat_v = makePattern({gather_input_v, cur_v}, {{"axis", axis_seq_len}}); - auto multi_query_bcst = [](std::shared_ptr kv) { - auto reshape_kv = wrap_type({kv, any_input()}); - auto unsqueeze_kv = makePattern({kv, any_input()}); + std::shared_ptr reshape_k, reshape_v, unsqueeze_k, unsqueeze_v; + std::shared_ptr computed_bcst_k, computed_bcst_v, multiply_k, multiply_v; + auto multi_query_bcst = [](const std::shared_ptr& kv, std::shared_ptr& reshape_kv, + std::shared_ptr& unsqueeze_kv, std::shared_ptr& computed_bcst, std::shared_ptr& multiply_kv) { + reshape_kv = wrap_type({kv, any_input()}); + unsqueeze_kv = makePattern({kv, any_input()}); auto check_one = [] (Output output) -> bool { auto node = std::dynamic_pointer_cast(output.get_node_shared_ptr()); @@ -66,15 +69,17 @@ StatefulSDPAFusion::StatefulSDPAFusion() { }; auto constant_bcst = wrap_type(check_one); - auto computed_bcst = makePattern({wrap_type(check_one), + computed_bcst = makePattern({wrap_type(check_one), any_input(), any_input()}, {{"mode", "numpy"}}); - auto multiply_kv = wrap_type({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst}); + multiply_kv = wrap_type({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst}); return wrap_type({multiply_kv, any_input()}); }; - auto present_k = concat_k | multi_query_bcst(concat_k); - auto present_v = concat_v | multi_query_bcst(concat_v); + auto mq_reshape_k = multi_query_bcst(concat_k, reshape_k, unsqueeze_k, computed_bcst_k, multiply_k); + auto mq_reshape_v = multi_query_bcst(concat_v, reshape_v, unsqueeze_v, computed_bcst_v, multiply_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({cur_q, present_k, present_v}); @@ -181,16 +186,20 @@ 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::shared_ptr& node) { - 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; + auto is_optional_one_child = [&pattern_map] (const std::vector>& 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) || !is_optional_one_child(convert_past_v) || - !is_optional_one_child(transpose_q) || !is_optional_one_child(transpose_k) || !is_optional_one_child(transpose_v)) + 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 From a099ba34c163b5446994f3cccbdba18a48e0f199 Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Tue, 11 Jun 2024 09:07:47 +0200 Subject: [PATCH 3/4] apply review comments --- .../common/pass/stateful_sdpa_fusion.cpp | 22 ++++--- .../transformations/state_concat_sdpa.cpp | 61 ++++++++++++++++++- 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp index c153df2f20d..c2765b889c6 100644 --- a/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp +++ b/src/plugins/intel_cpu/src/transformations/cpu_opset/common/pass/stateful_sdpa_fusion.cpp @@ -55,10 +55,10 @@ StatefulSDPAFusion::StatefulSDPAFusion() { std::shared_ptr reshape_k, reshape_v, unsqueeze_k, unsqueeze_v; std::shared_ptr computed_bcst_k, computed_bcst_v, multiply_k, multiply_v; - auto multi_query_bcst = [](const std::shared_ptr& kv, std::shared_ptr& reshape_kv, - std::shared_ptr& unsqueeze_kv, std::shared_ptr& computed_bcst, std::shared_ptr& multiply_kv) { - reshape_kv = wrap_type({kv, any_input()}); - unsqueeze_kv = makePattern({kv, any_input()}); + std::shared_ptr mq_reshape_k, mq_reshape_v; + auto multi_query_bcst = [](const std::shared_ptr& kv) { + auto reshape_kv = wrap_type({kv, any_input()}); + auto unsqueeze_kv = makePattern({kv, any_input()}); auto check_one = [] (Output output) -> bool { auto node = std::dynamic_pointer_cast(output.get_node_shared_ptr()); @@ -69,15 +69,16 @@ StatefulSDPAFusion::StatefulSDPAFusion() { }; auto constant_bcst = wrap_type(check_one); - computed_bcst = makePattern({wrap_type(check_one), + auto computed_bcst = makePattern({wrap_type(check_one), any_input(), any_input()}, {{"mode", "numpy"}}); - multiply_kv = wrap_type({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst}); - return wrap_type({multiply_kv, any_input()}); + auto multiply_kv = wrap_type({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst}); + auto result = wrap_type({multiply_kv, any_input()}); + return std::make_tuple(result, reshape_kv, unsqueeze_kv, computed_bcst, multiply_kv); }; - auto mq_reshape_k = multi_query_bcst(concat_k, reshape_k, unsqueeze_k, computed_bcst_k, multiply_k); - auto mq_reshape_v = multi_query_bcst(concat_v, reshape_v, unsqueeze_v, computed_bcst_v, multiply_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; @@ -199,8 +200,9 @@ StatefulSDPAFusion::StatefulSDPAFusion() { 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})) + mq_reshape_k, mq_reshape_v})) { return false; + } // past_k & past_v must be reordered by same beam_idx const auto gather_k_node = diff --git a/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp b/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp index 36fc7d50a86..a91de574807 100644 --- a/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp +++ b/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp @@ -24,7 +24,22 @@ using namespace ov; using namespace ov::intel_cpu; using namespace ov::gen_pattern; -static std::shared_ptr 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 makeSDPA(const ov::PartialShape& inputShape, bool isRef = false, bool hasConvert = false, bool hasMultiquery = false, + InsertPoint at = InsertPoint::At_None) { + std::shared_ptr insert_result; auto q = std::make_shared(element::f32, inputShape); auto kvInputShape = inputShape; if (hasMultiquery) { @@ -44,6 +59,10 @@ static std::shared_ptr makeSDPA(const ov::PartialShape& inputShape, b if (hasConvert) { pastk = std::make_shared(pastk, element::f32); pastv = std::make_shared(pastv, element::f32); + if (at == InsertPoint::At_Convert) { + // insert one should be enough to cover the 'convert' case + insert_result = std::make_shared(pastk); + } } if (isRef) { ov::intel_cpu::ScaledDotProductAttentionWithKVCache::Config config; @@ -54,6 +73,9 @@ static std::shared_ptr makeSDPA(const ov::PartialShape& inputShape, b concatV = new_node->output(2); } else { pastk = std::make_shared(pastk, beam_idx, op::v0::Constant::create(element::i32, {1}, {0})); + if (at == InsertPoint::At_Gather) { + insert_result = std::make_shared(pastk); + } pastv = std::make_shared(pastv, beam_idx, op::v0::Constant::create(element::i32, {1}, {0})); concatK = std::make_shared(OutputVector{pastk, k}, 2); concatV = std::make_shared(OutputVector{pastv, v}, 2); @@ -62,6 +84,9 @@ static std::shared_ptr makeSDPA(const ov::PartialShape& inputShape, b auto beam_idx_shape = makeOP>({beam_idx}, {{"type_relax", true}, {"input_data_types", {}}, {"output_data_types", {element::i32}}}); auto unsqueeze_concat = makeOP({conat, 2}); + if (at == InsertPoint::At_MQ_Unsqueeze && !insert_result) { + insert_result = std::make_shared(unsqueeze_concat); + } auto concat_shape = makeOP>( {conat}, {{"type_relax", true}, {"input_data_types", {}}, {"output_data_types", {element::i32}}}); @@ -72,10 +97,19 @@ static std::shared_ptr makeSDPA(const ov::PartialShape& inputShape, b auto expand_ones = makeOP({{1.0f}, expand_Abs, axis_mapping}, {{"mode", "numpy"}}); + if (at == InsertPoint::At_MQ_Broadcast && !insert_result) { + insert_result = std::make_shared(expand_ones); + } auto expand_Broadcast = makeOP({unsqueeze_concat, expand_ones}, {{"auto_broadcast", "numpy"}}); + if (at == InsertPoint::At_MQ_Multiply && !insert_result) { + insert_result = std::make_shared(expand_Broadcast); + } auto expected_shape = makeOP({beam_idx_shape, {inputShape[1]}, gather_ls}, {{"axis", 0}}); auto reshape_Reshape = makeOP({expand_Broadcast, expected_shape}, {{"special_zero", false}}); + if (at == InsertPoint::At_MQ_Reshape && !insert_result) { + insert_result = std::make_shared(reshape_Reshape); + } return reshape_Reshape; }; sdp = std::make_shared(q, make_multi_query(concatK), make_multi_query(concatV), false); @@ -92,6 +126,8 @@ static std::shared_ptr makeSDPA(const ov::PartialShape& inputShape, b auto add = std::make_shared(sdp, op::v0::Constant::create(element::f32, {1}, {1.0f})); ResultVector results{std::make_shared(add)}; + if (insert_result) + results.push_back(insert_result); SinkVector sinks{pastk_assign, pastv_assign}; return std::make_shared(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 be fail + std::shared_ptr f(nullptr), f_ref(nullptr); + { + using namespace ov; + auto inputShape = ov::PartialShape{-1, 32, -1, 64}; + size_t i = static_cast(InsertPoint::At_None) + 1; + size_t end = static_cast(InsertPoint::At_End); + // check each position + for (; i < end; i++) { + InsertPoint at = static_cast(i); + f = makeSDPA(inputShape, false, true, true, at); + f_ref = makeSDPA(inputShape, false, true, true, at); + pass::Manager m; + m.register_pass(); + m.register_pass(); + m.run_passes(f); + auto res = compare_functions(f, f_ref); + ASSERT_TRUE(res.first) << res.second; + } + } +} \ No newline at end of file From abfafe0369f7ba35cbe6de946b0d344870f0a91a Mon Sep 17 00:00:00 2001 From: Luo Cheng Date: Wed, 12 Jun 2024 05:11:17 +0200 Subject: [PATCH 4/4] 1, apply review comment 2, cover test: read->gather->shapeof+concat --- .../custom/subgraph_tests/src/concat_sdp.cpp | 44 +++++++++++++------ .../transformations/state_concat_sdpa.cpp | 2 +- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_sdp.cpp b/src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_sdp.cpp index 9d8182ef320..bec3a2bbd2c 100644 --- a/src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_sdp.cpp +++ b/src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/concat_sdp.cpp @@ -66,7 +66,6 @@ public: void SetUp() override { ElementType inType; std::vector 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(inputParams[3], var_v); pastv->set_friendly_name("pastv_r"); - std::shared_ptr pastk_shapeof, pastv_shapeof; - if (hasShapeOf) { - pastk_shapeof = std::make_shared(pastk); - pastv_shapeof = std::make_shared(pastv); - } auto beam_idx = std::make_shared(ElementType::i32, ov::PartialShape{-1}); beam_idx->set_friendly_name("beam_idx"); inputParams.push_back(beam_idx); - auto gatherK = std::make_shared(pastk, beam_idx, op::v0::Constant::create(ElementType::i32, {1}, {0})); - auto gatherV = std::make_shared(pastv, beam_idx, op::v0::Constant::create(ElementType::i32, {1}, {0})); + auto gatherK = std::make_shared(pastk, beam_idx, op::v0::Constant::create(ElementType::i32, {}, {0})); + auto gatherV = std::make_shared(pastv, beam_idx, op::v0::Constant::create(ElementType::i32, {}, {0})); + std::shared_ptr 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(gatherK); + shapeof_v = std::make_shared(gatherV); + } auto concatK = std::make_shared(OutputVector{gatherK, inputParams[1]}, 2); auto concatV = std::make_shared(OutputVector{gatherV, inputParams[2]}, 2); auto sdp = std::make_shared(inputParams[0], concatK, concatV, false); @@ -115,8 +119,8 @@ public: ResultVector results{std::make_shared(add)}; if (hasShapeOf) { - results.push_back(std::make_shared(pastk_shapeof)); - results.push_back(std::make_shared(pastv_shapeof)); + results.push_back(std::make_shared(shapeof_k)); + results.push_back(std::make_shared(shapeof_v)); } SinkVector sinks{pastk_assign, pastv_assign}; function = std::make_shared(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() == "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++) { diff --git a/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp b/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp index a91de574807..ca55337fc3d 100644 --- a/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp +++ b/src/plugins/intel_cpu/tests/unit/transformations/state_concat_sdpa.cpp @@ -196,7 +196,7 @@ TEST(TransformationTests, StateConcatSDPAMixtral) { } TEST(TransformationTests, StateConcatSDPAWithExtraNode) { - // when some unexpected extra nodes exist in SDPA, the fusion should be fail + // when some unexpected extra nodes exist in SDPA, the fusion should fail std::shared_ptr f(nullptr), f_ref(nullptr); { using namespace ov;