[TRANSFORMATIONS][GPU] Adjust SDPA Fusion pass #29575

merezman · 2025-03-19T12:12:50Z

Details:

Adjusted SDPA fusion pass to be able to handle patterns observed in many customer models, that requires to include reshapes to the pattern that changes 4D tensors to 3D or 2D. It handles scales that can be attached in different places, and inner reshapes ie. reshape of softmax output.

vladimir-paramuzov · 2025-03-20T05:39:27Z

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

+    auto q_base = makePattern(ov::Rank(4));
+    auto q_shape = ov::pass::pattern::any_input();
+    auto q_reshaped = ov::pass::pattern::wrap_type<ov::op::v1::Reshape>({q_base, q_shape});
+    auto q = q_reshaped | q_base;


I guess pattern::optional can be used here to simplify that part:

auto q_base = makePattern(ov::Rank(4)); auto q = optional<ov::op::v1::Reshape>({q_base, any_input()});

I used this way for later checks

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

Lines 146 to 150 in d32a498

// make sure that if inputs are reshaped the output is reshaped back

bool inputs_reshaped = pattern_map.count(q_reshaped) > 0 && pattern_map.count(k_reshaped) > 0 && pattern_map.count(v_reshaped) > 0;

bool output_reshaped = pattern_map.count(qkv_reshaped) > 0;

if (inputs_reshaped && !output_reshaped || !inputs_reshaped && output_reshaped)

return false;

vladimir-paramuzov · 2025-03-20T05:39:42Z

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

+
+    // Optional k scale
+    auto attn_scale = ov::pass::pattern::any_input();
+    // auto k_opt_scaled = optional<ov::op::v1::Multiply>({k, attn_scale});


remove unused code

vladimir-paramuzov · 2025-03-20T05:40:48Z

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

-    auto softmax = makePattern<ov::op::v8::Softmax>({optional_add_mask}, {{"axis", "-1"}});
-    auto qkv = makePattern<ov::op::v0::MatMul>({softmax, v}, {{"transpose_a", false}, {"transpose_b", false}});
+    // Optional reshape befor adding mask
+    auto qk_opt_scaled_pre_mask_shape = ov::pass::pattern::any_input();


We have using ov::pass::pattern in the beginning of the function, so you can omit explicit namespace specification

vladimir-paramuzov · 2025-03-20T05:45:04Z

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

+    auto qk_opt_scaled_pre_mask_reshaped = ov::pass::pattern::wrap_type<ov::op::v1::Reshape>({qk_opt_scaled, qk_opt_scaled_pre_mask_shape});
+    auto qk_opt_scaled_pre_mask_opt_reshaped = qk_opt_scaled_pre_mask_reshaped | qk_opt_scaled;
+    // Optional mask add
+    auto qk_opt_scaled_mask_added = makePattern<ov::op::v1::Add>({qk_opt_scaled_pre_mask_opt_reshaped, mask});


I think makePattern shall be replaced here with wrap_type

vladimir-paramuzov · 2025-03-20T05:47:37Z

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

+    auto k_opt_transposed_scaled = ov::pass::pattern::wrap_type<ov::op::v1::Multiply>({k_opt_transposed, attn_scale});
+    auto k_opt_transposed_opt_scaled = k_opt_transposed_scaled | k_opt_transposed;


I think optional can be used here too.

It is also used for checks

vladimir-paramuzov · 2025-03-20T05:54:01Z

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

+            ov::Output<ov::Node> qk_out;
+            if (pattern_map.count(qk_opt_scaled_pre_mask_reshaped) > 0)
+                qk_out = pattern_map.at(qk_opt_scaled_pre_mask_reshaped);
+            else if (pattern_map.count(qk_unsqueeze) > 0)
+                qk_out = pattern_map.at(qk_unsqueeze);
+            else if (pattern_map.count(qk) > 0)
+                qk_out = pattern_map.at(qk);
+            else
+                return false;


Probably can be simplified by taking first input of qk_opt_scaled_mask_added node.

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

src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/sdpa_kernel_ref.cpp

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

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

sshlyapn

In general looks good to me

itikhono · 2025-03-27T12:43:52Z

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

-    auto v = makePattern(ov::Rank(4));
+    auto q_base = makePattern(ov::Rank(4));
+    auto q_shape = any_input();
+    auto q_reshaped = wrap_type<ov::op::v1::Reshape>({q_base, q_shape});


minor: can we use
pattern::optional for v1::Reshape?

openvino/src/core/include/openvino/pass/pattern/op/optional.hpp

Line 47 in 94bc742

class OPENVINO_API Optional : public Pattern {

it should work as pattern::Or inside.

the same for K, V
and probably for QK

I used this way to be able to use those checks

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

Lines 146 to 150 in d32a498

// make sure that if inputs are reshaped the output is reshaped back

bool inputs_reshaped = pattern_map.count(q_reshaped) > 0 && pattern_map.count(k_reshaped) > 0 && pattern_map.count(v_reshaped) > 0;

bool output_reshaped = pattern_map.count(qkv_reshaped) > 0;

if (inputs_reshaped && !output_reshaped || !inputs_reshaped && output_reshaped)

return false;

I believe it should work the same way with auto q_reshaped = optional<Reshape>(in1, in2)> etc could you double check?
As I can see, we have a test for this

openvino/src/core/tests/pattern.cpp

Line 579 in 94bc742

EXPECT_EQ(pattern_val_mp.count(pattern_convert), 0);

Right, seems that it works

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

merezman added 18 commits January 16, 2025 08:20

Adjust SDPA to machine translation patterns

c4dedb8

Change Slice to StridedSlice

2574662

move SDPAFusion matcher

f037d77

Simplify SDPAFusion

e656f23

Undo optionals

b4a2962

Merge branch 'master' into merezman/sdpa_fusion_adjustment

79f061a

Add scale and mask checks

bef454f

Fix output reshape detection

4c26485

Add transformations tests

e5d7eca

Merge branch 'master' into merezman/sdpa_fusion

dca5cae

Remove concat

4f710b9

Add functional test

955bc00

Add reshaped inputs

3860b7e

Transformation tests update

abc7c92

Adjust test input generator

0e106eb

Merge branch 'master' into merezman/sdpa_fusion

47fe489

Correct naming

2e24efc

Remove debug prints

d32a498

github-actions bot added category: GPU OpenVINO GPU plugin category: transformations OpenVINO Runtime library - Transformations labels Mar 19, 2025

sys-openvino-ci added the ExternalIntelPR External contributor from Intel label Mar 19, 2025

vladimir-paramuzov reviewed Mar 20, 2025

View reviewed changes

merezman added 3 commits March 20, 2025 07:04

Format

ef46558

Simplify getting tensor shape

fc5a541

Format

e9843af

merezman requested a review from vladimir-paramuzov March 20, 2025 12:25

sshlyapn reviewed Mar 21, 2025

View reviewed changes

Fix SDPAKernelRef

df3daf1

merezman requested a review from sshlyapn March 21, 2025 10:51

Format

cf3ea82

merezman marked this pull request as ready for review March 24, 2025 12:06

merezman requested review from a team as code owners March 24, 2025 12:06

merezman requested review from itikhono and removed request for a team March 24, 2025 12:06

sshlyapn reviewed Mar 27, 2025

View reviewed changes

sshlyapn approved these changes Mar 27, 2025

View reviewed changes

Update sdpa_kernel_ref check

43776f7

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Outdated Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Outdated Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Outdated Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Outdated Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Outdated Show resolved Hide resolved

itikhono reviewed Mar 27, 2025

View reviewed changes

src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp Show resolved Hide resolved

itikhono requested review from CuriousPanCake and removed request for vladimir-paramuzov March 27, 2025 13:21

merezman added 2 commits March 27, 2025 14:38

Update scale and mask checks

c6ee970

Add optional nodes

890bb42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TRANSFORMATIONS][GPU] Adjust SDPA Fusion pass #29575

[TRANSFORMATIONS][GPU] Adjust SDPA Fusion pass #29575

merezman commented Mar 19, 2025 •

edited

Loading

vladimir-paramuzov Mar 20, 2025

merezman Mar 20, 2025 •

edited

Loading

vladimir-paramuzov Mar 20, 2025

vladimir-paramuzov Mar 20, 2025

vladimir-paramuzov Mar 20, 2025

vladimir-paramuzov Mar 20, 2025

merezman Mar 20, 2025

vladimir-paramuzov Mar 20, 2025

sshlyapn left a comment

itikhono Mar 27, 2025 •

edited

Loading

merezman Mar 27, 2025

itikhono Mar 27, 2025 •

edited

Loading

merezman Mar 28, 2025

	// make sure that if inputs are reshaped the output is reshaped back
	bool inputs_reshaped = pattern_map.count(q_reshaped) > 0 && pattern_map.count(k_reshaped) > 0 && pattern_map.count(v_reshaped) > 0;
	bool output_reshaped = pattern_map.count(qkv_reshaped) > 0;
	if (inputs_reshaped && !output_reshaped \|\| !inputs_reshaped && output_reshaped)
	return false;

		auto k_opt_transposed_scaled = ov::pass::pattern::wrap_type<ov::op::v1::Multiply>({k_opt_transposed, attn_scale});
		auto k_opt_transposed_opt_scaled = k_opt_transposed_scaled \| k_opt_transposed;

[TRANSFORMATIONS][GPU] Adjust SDPA Fusion pass #29575

Are you sure you want to change the base?

[TRANSFORMATIONS][GPU] Adjust SDPA Fusion pass #29575

Conversation

merezman commented Mar 19, 2025 • edited Loading

Details:

Choose a reason for hiding this comment

merezman Mar 20, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sshlyapn left a comment

Choose a reason for hiding this comment

itikhono Mar 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

itikhono Mar 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merezman commented Mar 19, 2025 •

edited

Loading

merezman Mar 20, 2025 •

edited

Loading

itikhono Mar 27, 2025 •

edited

Loading

itikhono Mar 27, 2025 •

edited

Loading