[CPU] Fuse SDPA and Concat as early as possible

[luo-cheng2021]

Details:

• Move StatefulSDPAFusion before CommonOptimizations
• ...

Tickets:

• 158738

[usstq]

LGTM!

[dmitry-gorokhov]

Any details on that?
I recall that pass was added for Whisper model. So how do we guarantee we don't bring regression to Whisper?

[luo-cheng2021]

Currently no LLM models will hit SDPAFuseTransposeReshape including Whisper so there is no regression, checked on WW52 models. The pattern would work for customized model which should not upstream yet. From @maxnick's comment, there is a plan to use snippets to cover the pattern. So, should we leave the case to snippets or use this transformation to cover?

[dmitry-gorokhov]

This customized pattern is expected to become default: huggingface/optimum-intel#1078. Export part has been waiting runtime optimizations actually and SDPAFuseTransposeReshape is part of these optimizations https://jira.devtools.intel.com/browse/CVS-153616.
Snippets are responsible for SDPA patterns w/o states (like regular transformers or diffusers). SDPA with states (LLMs. Whisper) should be processed via custom SDPA node as of now.

[luo-cheng2021]

SDPAFuseTransposeReshape only works with SDPA when fusing to SDPA+concat failed(if success the node type will be ScaledDotProductAttentionWithKVCache):

openvino/src/plugins/intel_cpu/src/transformations/cpu_opset/x64/pass/sdpa_fuse_transpose_reshape.cpp

Lines 54 to 55 in 548786a

	auto sdpa_node =
	wrap_type<op::v13::ScaledDotProductAttention>({q_transpose_node, k_transpose_node, v_transpose_node});

I suppose huggingface/optimum-intel#1078 will change current Whisper model into stateful model, so, we may need to change SDPAFuseTransposeReshape to support ScaledDotProductAttentionWithKVCache or add the function into StatefulSDPAFusion. If my understanding is correct, I will create a ticket to track this.

[dmitry-gorokhov]

@xipingyan Could you please clarify how SDPAFuseTransposeReshape worked for Whisper model? Have StatefulSDPAFusion falied for this model before?

[xipingyan]

hi @dmitry-gorokhov @luo-cheng2021 ,
Actually, Whisper model has 2 kinds of SDPA OPS. First one is mapped to ScaledDotProductAttentionWithKVCache, second one is mapped to ScaledDotProductAttention. In my case, SDPAFuseTransposeReshape only works for second one. So @luo-cheng2021 's suggestion "change SDPAFuseTransposeReshape to support ScaledDotProductAttentionWithKVCache" will not work.

@dmitry-gorokhov , I just aligned with @luo-cheng2021 , we need to double confirm whether current master has supported dynamic shape's snippet,

If yes, this pattern SDPAFuseTransposeReshape can be removed.
If no, maybe we can try to merge Reshape+Transpose before SDPA into init graph of readvalue, take it as temp solution, the snippet with dynamic shape still is our final target.

[dmitry-gorokhov]

Disscussed offline.
Agreed to updated conditions for SDPAFuseTransposeReshape applicability: it should check that Reshape op goes after ReadValue. That will help to limit the pass implication on Whisper model only and avoif negative perf impact on SD and others.
Once Snippets will support such patterns with dynamic shapes, SDPAFuseTransposeReshape might be fully removed.

[luo-cheng2021]

Done, no models in ww52 will hit SDPAFuseTransposeReshape, only customized Whisper will hit it.