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..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 @@ -53,7 +53,10 @@ 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) { + std::shared_ptr reshape_k, reshape_v, unsqueeze_k, unsqueeze_v; + std::shared_ptr computed_bcst_k, computed_bcst_v, multiply_k, multiply_v; + 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()}); @@ -70,11 +73,14 @@ StatefulSDPAFusion::StatefulSDPAFusion() { any_input(), any_input()}, {{"mode", "numpy"}}); auto multiply_kv = wrap_type({reshape_kv | unsqueeze_kv, constant_bcst | computed_bcst}); - return wrap_type({multiply_kv, any_input()}); + auto result = wrap_type({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({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>& 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(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; } 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 36fc7d50a86..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 @@ -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 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