[TRANSFORMATIONS] Align fused SDPA output by mryzhov · Pull Request #34097 · openvinotoolkit/openvino

mryzhov · 2026-02-12T15:54:29Z

Details:

This pull request enhances the SDPA Fusion transformation by improving shape alignment between the original and fused nodes, particularly for handling cases when the fused SDPA operation output changes the original output shape (2D->3D, 4D-3D)

SDPA Fusion Logic Improvements

Added a try_align_outputs helper function to insert a Squeeze node when the fused SDPA node's output shape does not match the original node's output, ensuring correct shape alignment after fusion.
Updated the SDPA fusion matcher to use try_align_outputs after creating the fused node, so that 2D inputs are properly handled and output shapes remain consistent.

Testing Enhancements

Extended the SDPA test utility class with squeeze and unsqueeze methods for more flexible test graph construction.
Added convenience methods to the SDPA class for unsqueezing Q, K, V inputs and squeezing SDPA outputs, supporting more detailed test scenarios.
Added new tests for SDPA fusion with 2D static and dynamic input shapes to ensure correct behavior and output shape alignment in these cases.

Tickets:

CVS-180217

Copilot

Pull request overview

Improve SDPA Fusion output shape alignment for 2D inputs by inserting a post-fusion squeeze when needed, and extend tests to cover 2D static/dynamic shapes.

Changes:

Added try_align_outputs helper to insert Squeeze when fused SDPA output shape differs from the original.
Updated SDPA fusion matcher to use try_align_outputs before replacing the matched subgraph.
Extended SDPA test utilities and added parametrized 2D-input SDPA fusion tests (static and dynamic).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp	Adds output-shape alignment logic (optional squeeze) and integrates it into the fusion replacement.
src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp	Adds squeeze/unsqueeze helpers and new parametrized tests for 2D input fusion scenarios.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp

…ion_test.cpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

rkazants · 2026-02-13T05:55:31Z

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp

        });
 };
+
+std::shared_ptr<ov::Node> try_align_outputs(const std::shared_ptr<ov::Node>& src_output,


can you please add a comment where we have a case of 2D tensors (Q, K, V)? Which model and why it happens?

The issue was identified in the GeekbenchAI model. It is not related to the SDPA inputs themselves; rather, it occurs during fusion of the SDPA subgraph. When the subgraph is replaced with the SDPA op, the resulting output shape can change in some cases, making it incompatible with the remainder of the model. The simplest way to reproduce the issue is with 2D inputs.
Another solution would be fixing the SDPA shape validation logic, @mitruska fyi

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 8 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-02-13T19:14:58Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

+}
+
+INSTANTIATE_TEST_SUITE_P(SDPAFusion,
+                         SDPAUnsquieezeOutput,


Typo in parameter reference: "Unsquieeze" should be "Unsqueeze".

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp

Copilot · 2026-02-13T19:14:59Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

    comparator.enable(FunctionsComparator::CmpValues::ATTRIBUTES);
 }

+class SDPASquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};


Typo in class name: "Squieeze" should be "Squeeze".

Copilot · 2026-02-13T19:14:59Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp


+class SDPASquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};
+
+TEST_P(SDPASquieezeOutput, SDPAFusionTest_SquieezeOutput) {


Typo in test name: "Squieeze" should be "Squeeze".

Suggested change

TEST_P(SDPASquieezeOutput, SDPAFusionTest_SquieezeOutput) {

TEST_P(SDPASquieezeOutput, SDPAFusionTest_SqueezeOutput) {

Copilot · 2026-02-13T19:14:59Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

+}
+
+INSTANTIATE_TEST_SUITE_P(SDPAFusion,
+                         SDPASquieezeOutput,


Typo in parameter reference: "Squieeze" should be "Squeeze".

Copilot · 2026-02-13T19:14:59Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

+class SDPAUnsquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};
+
+TEST_P(SDPAUnsquieezeOutput, SDPAFusionTest_UnsquieezeOutput) {


Typo in class name: "Unsquieeze" should be "Unsqueeze".

Suggested change

class SDPAUnsquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};

TEST_P(SDPAUnsquieezeOutput, SDPAFusionTest_UnsquieezeOutput) {

class SDPAUnsqueezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};

TEST_P(SDPAUnsqueezeOutput, SDPAFusionTest_UnsqueezeOutput) {

Copilot · 2026-02-13T19:15:00Z

src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp

+
+class SDPAUnsquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};
+
+TEST_P(SDPAUnsquieezeOutput, SDPAFusionTest_UnsquieezeOutput) {


Typo in test name: "Unsquieeze" should be "Unsqueeze".

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

align sdpa output

3ef857a

mryzhov requested a review from a team as a code owner February 12, 2026 15:54

github-actions bot added the category: transformations OpenVINO Runtime library - Transformations label Feb 12, 2026

mryzhov requested a review from mitruska February 12, 2026 15:58

extended tests

db339e4

mryzhov requested a review from Copilot February 12, 2026 16:12

Merge branch 'master' into sdpa_squeeze

e7c33de

Copilot AI reviewed Feb 12, 2026

View reviewed changes

Update src/common/transformations/tests/common_optimizations/sdpa_fus…

610c3af

…ion_test.cpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

rkazants reviewed Feb 13, 2026

View reviewed changes

mryzhov added 2 commits February 13, 2026 16:08

extend supported cases

1e015a3

Unsquieeze tests

19d967b

mryzhov requested a review from Copilot February 13, 2026 19:09

Copilot started reviewing on behalf of mryzhov February 13, 2026 19:10 View session

Copilot AI reviewed Feb 13, 2026

View reviewed changes

mryzhov and others added 2 commits February 13, 2026 20:15

Apply suggestion from @Copilot

5b8de14

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Merge branch 'master' into sdpa_squeeze

78751d5

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TRANSFORMATIONS] Align fused SDPA output#34097

[TRANSFORMATIONS] Align fused SDPA output#34097
mryzhov wants to merge 8 commits intoopenvinotoolkit:masterfrom
mryzhov:sdpa_squeeze

mryzhov commented Feb 12, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rkazants Feb 13, 2026

Uh oh!

mryzhov Feb 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Copilot AI Feb 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants


		class SDPASquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};

		TEST_P(SDPASquieezeOutput, SDPAFusionTest_SquieezeOutput) {

		class SDPAUnsquieezeOutput : public TransformationTestsF, public ::testing::WithParamInterface<PartialShape> {};

		TEST_P(SDPAUnsquieezeOutput, SDPAFusionTest_UnsquieezeOutput) {

Conversation

mryzhov commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Details:

SDPA Fusion Logic Improvements

Testing Enhancements

Tickets:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rkazants Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

mryzhov Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mryzhov commented Feb 12, 2026 •

edited

Loading