[Transformations] Fuse split-attention sub-graph into v13::ScaledDotProductAttention#36153
[Transformations] Fuse split-attention sub-graph into v13::ScaledDotProductAttention#36153zlma7001 wants to merge 1 commit into
Conversation
github-actions
Bot
added
the
category: transformations
|
build_jenkins |
|
build_jenkins |
|
Patch updated for code style requirement, PTAL |
|
Just thinking out loud, wouldn't it be simpler to transform the graph in such a way that it's going to be matched by SDPAFusion? Such that in your transformation you won't fuse SDPA, but rather canonicalize the graph around it. I believe, that might make the code much simpler. What do you think? |
|
build_jenkins |
There was a problem hiding this comment.
Pull request overview
This PR extends the common SDPA fusion transformations to recognize and fuse a “split-attention” attention pattern (separate cache/new KV attention merged via Concat → Softmax → VariadicSplit → 2×MatMul → Add) into a single ov::op::v13::ScaledDotProductAttention, enabling optimized plugin SDPA kernels for affected Gemma-style graphs.
Changes:
- Register a new
SDPASplitAttentionFusionMatcherfromSDPAFusion::run_on_model. - Implement the matcher to concatenate
K_cache/K_newandV_cache/V_newalong the implied sequence axis and insert(0,1,3,2)transposes when inputs are stored as[B,H,D,S]. - Add unit tests covering standard layout, pre-transposed K, pre-transposed V, and a negative case (Softmax axis not trailing).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp | Adds and registers SDPASplitAttentionFusionMatcher that fuses the split-attention subgraph into v13::ScaledDotProductAttention. |
| src/common/transformations/include/transformations/common_optimizations/sdpa_fusion.hpp | Declares the new matcher pass and documents the fused pattern. |
| src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp | Adds regression tests validating fusion behavior for split-attention patterns and a no-fusion negative case. |
I considered that option. Conceptually it is attractive, but for this split-attention pattern I do not think it would make the implementation simpler. The graph is not just a syntactic variant of the existing SDPAFusion pattern: it computes QK and PV separately for cache/new branches, then uses Concat/Softmax/VariadicSplit/Add to preserve the cache/new partitioning. A canonicalization pass would still need to prove the same constraints currently checked by this matcher: matching concat/split axes and order, compatible split lengths, equivalent mask placement, and equivalent scale handling. It would also need to materialize a temporary canonical MatMul/Softmax/MatMul form and rely on a later pass to fuse it, which makes pass ordering and fallback behavior more delicate. |
|
build_jenkins |
|
Rebased to re-run CI, seems the previous failures were unrelated to this PR. |
|
build_jenkins |
1 similar comment
|
build_jenkins |
|
Rebased to retrigger ci, thanks! |
|
build_jenkins |
|
build_jenkins |
|
Thanks @CuriousPanCake for the review, PTAL |
|
build_jenkins |
…roductAttention
Some Gemma models compute attention against the persistent KV cache and
the current step's KV separately, merging the results via Concat /
VariadicSplit:
qk_cache = MatMul(Q, K_cache)
qk_new = MatMul(Q, K_new)
scores = Concat(qk_cache, qk_new, axis=-1)
masked = Add(scores, mask)
probs = Softmax(masked, axis=-1)
[pc, pn] = VariadicSplit(probs, axis=-1, sizes=[S_cache, S_new])
attn_c = MatMul(pc, V_cache)
attn_n = MatMul(pn, V_new)
output = Add(attn_c, attn_n)
The existing SDPAFusionMatcher does not handle this shape, so the layer
runs as ~12 unfused primitives and never reaches the plugins' optimized
SDPA kernel.
Add SDPASplitAttentionFusionMatcher (registered from SDPAFusion::run_on_model)
which concatenates K_cache+K_new and V_cache+V_new along the sequence
axis implied by each MatMul's transpose_b, inserts a Transpose(0,1,3,2)
when K or V is stored as [B,H,D,S] to restore the [B,H,S,D] layout
required by v13::SDPA, and uses scale=1.0 (any scaling is assumed
pre-baked into Q, e.g. Gemma's query_pre_attn_scalar). Restricted to
static 4-D ranks; other shapes fall through unchanged.
Tests cover the standard [B,H,S,D] layout, the pre-transposed K and V
variants, and a negative case where Softmax is on a non-trailing axis.
|
Thanks Andrii! PTAL |
| SDPASplitAttentionFusionMatcher::SDPASplitAttentionFusionMatcher() { | ||
| MATCHER_SCOPE(SDPASplitAttentionFusionMatcher); | ||
|
|
||
| // Match the "split-attention" pattern; see header documentation. |
There was a problem hiding this comment.
| // Match the "split-attention" pattern; see header documentation. |
| auto add_node = pattern_map.at(output_add).get_node_shared_ptr(); | ||
| auto matmul_cache = ov::as_type_ptr<v0::MatMul>(pattern_map.at(attn_cache).get_node_shared_ptr()); | ||
| auto matmul_new = ov::as_type_ptr<v0::MatMul>(pattern_map.at(attn_new).get_node_shared_ptr()); | ||
| auto vsplit_node = ov::as_type_ptr<v1::VariadicSplit>(pattern_map.at(vsplit_out0).get_node_shared_ptr()); |
There was a problem hiding this comment.
Since you're casting using ov::as_type_ptr(), check they're all not nullptr. Use one if statement for this.
| auto split_axis_value = axis_vec[0]; | ||
| if (split_axis_value < 0) { | ||
| split_axis_value += split_rank.get_length(); | ||
| } |
There was a problem hiding this comment.
Check for some normalization utils in OpenVINO. They should cover this functionality
| return false; | ||
| } | ||
|
|
||
| auto softmax_node = ov::as_type_ptr<v8::Softmax>(pattern_map.at(softmax).get_node_shared_ptr()); |
There was a problem hiding this comment.
Check for nullptr.
| auto sm_axis = softmax_node->get_axis(); | ||
| if (sm_axis < 0) { | ||
| sm_axis += sm_rank.get_length(); | ||
| } | ||
| if (sm_axis != sm_rank.get_length() - 1) { | ||
| return false; | ||
| } |
There was a problem hiding this comment.
Same comment as above regarding normalization
| if (concat_rank.is_dynamic()) { | ||
| return false; | ||
| } | ||
| const auto concat_axis = ov::util::try_normalize_axis(concat_node->get_axis(), concat_rank, *concat_node); |
There was a problem hiding this comment.
Why are you using try_normalize_axis instead of just normalize_axis? Not a request to change it, just asking to understand.
| std::shared_ptr<ov::Node> k_input = k_concat; | ||
| std::shared_ptr<ov::Node> v_input = v_concat; |
There was a problem hiding this comment.
Why do we need these variables?
| // Interpret MatMul transpose flags as a physical-layout hint: | ||
| // qk transpose_b == true; K stored as [B,H,S,D] (standard); MatMul performs Q ┬╖ K^T. | ||
| // qk transpose_b == false; K stored as [B,H,D,S] (already K^T in memory). | ||
| // attn transpose_b == true; V stored as [B,H,D,S]; MatMul performs probs ┬╖ V^T. |
There was a problem hiding this comment.
Looks like this comment hasn't been rendered correctly at the end of the line
Some Gemma models compute attention against the persistent KV cache and the current step's KV separately, merging the results via Concat / VariadicSplit:
The existing SDPAFusionMatcher does not handle this shape, so the layer runs as ~12 unfused primitives and never reaches the plugins' optimized SDPA kernel.
Add SDPASplitAttentionFusionMatcher (registered from SDPAFusion::run_on_model) which concatenates K_cache+K_new and V_cache+V_new along the sequence axis implied by each MatMul's transpose_b, inserts a Transpose(0,1,3,2) when K or V is stored as [B,H,D,S] to restore the [B,H,S,D] layout required by v13::SDPA, and uses scale=1.0 (any scaling is assumed pre-baked into Q, e.g. Gemma's query_pre_attn_scalar). Restricted to static 4-D ranks; other shapes fall through unchanged.
Tests cover the standard [B,H,S,D] layout, the pre-transposed K and V variants, and a negative case where Softmax is on a non-trailing axis.
Tickets: