[Transformations] Fuse split-attention sub-graph into v13::ScaledDotProductAttention

zlma7001 · zlma7001 · commit af1fab475efb · 2026-06-11T16:16:57.000+08:00
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.
diff --git a/src/common/transformations/include/transformations/common_optimizations/sdpa_fusion.hpp b/src/common/transformations/include/transformations/common_optimizations/sdpa_fusion.hpp
@@ -121,6 +121,83 @@ class TRANSFORMATIONS_API SDPAFusionMatcherSinks : public ov::pass::MatcherPass
     SDPAFusionMatcherSinks();
 };
 
+/// This pass fuses a "split-attention" sub-graph into a single ScaledDotProductAttention.
+/// The pattern is produced by some Gemma models which compute attention against past KV
+/// and current KV separately and merge via Concat/VariadicSplit.
+///
+/// Before:
+///       ┌─────────┐     ┌─────┐     ┌───────┐
+///       │ K_cache │     │  Q  │     │ K_new │
+///       └────┬────┘     └──┬──┘     └───┬───┘
+///            │          ┌──┴──┐         │
+///            │       ┌──┘     └──┐      │
+///            │       │           │      │
+///       ┌────┴───────┴─┐    ┌────┴──────┴┐
+///       │    MatMul    │    │   MatMul   │
+///       └──────┬───────┘    └──────┬─────┘
+///              │ qk_cache          │ qk_new
+///              └──────────┐   ┌────┘
+///                         │   │
+///                     ┌───┴───┴───┐
+///                     │  Concat   │ (axis=-1)
+///                     └─────┬─────┘
+///                           │            ┌─────────────┐
+///                        ┌──┴──┐         │AttentionMask│
+///                        │ Add │<────────┤             │
+///                        └──┬──┘         └─────────────┘
+///                           │
+///                      ┌────┴────┐
+///                      │ Softmax │ (axis=-1)
+///                      └────┬────┘
+///                           │
+///                  ┌────────┴────────┐
+///                  │  VariadicSplit  │ (axis=-1)
+///                  └───┬─────────┬───┘
+///                   pc │         │ pn
+///       ┌─────────┐    │         │    ┌───────┐
+///       │ V_cache │    │         │    │ V_new │
+///       └────┬────┘    │         │    └───┬───┘
+///            │      ┌──┘         └──┐     │
+///            │      │               │     │
+///       ┌────┴──────┴┐          ┌───┴─────┴──┐
+///       │   MatMul   │          │   MatMul   │
+///       └──────┬─────┘          └──────┬─────┘
+///              │ attn_c                │ attn_n
+///              └──────────┐   ┌────────┘
+///                         │   │
+///                       ┌─┴───┴─┐
+///                       │  Add  │
+///                       └───┬───┘
+///                           │
+///                       ┌───┴───┐
+///                       │ Output│
+///                       └───────┘
+///
+/// After:
+///     ┌───────┐  ┌──────────────────────┐  ┌──────────────────────┐  ┌─────────────┐
+///     │   Q   │  │Concat(K_cache,K_new) │  │Concat(V_cache,V_new) │  │AttentionMask│
+///     └───┬───┘  └──────────┬───────────┘  └──────────┬───────────┘  └──────┬──────┘
+///         │                 │                         │                     │
+///     ┌───┴─────────────────┴─────────────────────────┴─────────────────────┴─┐
+///     │                     ScaledDotProductAttention                         │
+///     └─────────────────────────────────┬─────────────────────────────────────┘
+///                                       │
+///                                  ┌────┴────┐
+///                                  │  Output │
+///                                  └─────────┘
+///
+/// K_cache and K_new are concatenated along the sequence axis (and similarly for V).
+/// When the matmul transpose flags indicate that K or V are stored with seq as the
+/// innermost dimension (i.e. [B,H,D,S]), explicit Transpose nodes are inserted to bring
+/// them to the canonical [B,H,S,D] layout required by v13::ScaledDotProductAttention.
+/// Scale is set to 1.0 because the matched pattern does not contain a separate scale node;
+/// any required scaling is assumed to be pre-applied to Q by the model.
+class TRANSFORMATIONS_API SDPASplitAttentionFusionMatcher : public ov::pass::MatcherPass {
+public:
+    OPENVINO_MATCHER_PASS_RTTI("SDPASplitAttentionFusionMatcher", "0");
+    SDPASplitAttentionFusionMatcher();
+};
+
 // Temporary wrapper to enable Symbolic infrastructure inside.
 class TRANSFORMATIONS_API SDPAFusion : public ov::pass::ModelPass {
 public:
diff --git a/src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp b/src/common/transformations/src/transformations/common_optimizations/sdpa_fusion.cpp
@@ -26,6 +26,7 @@
 #include "openvino/op/subtract.hpp"
 #include "openvino/op/transpose.hpp"
 #include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/variadic_split.hpp"
 #include "openvino/pass/pattern/op/optional.hpp"
 #include "openvino/pass/pattern/op/pattern.hpp"
 #include "openvino/pass/pattern/op/wrap_type.hpp"
@@ -152,6 +153,7 @@ bool SDPAFusion::run_on_model(const std::shared_ptr<ov::Model>& model) {
     ov::pass::SymbolicOptimizations symbolic_optimizations(false, get_pass_config());
     auto symbolic_ctx_manager = symbolic_optimizations.get_manager();
     symbolic_ctx_manager->register_pass<ov::pass::SDPAFusionMatcher>();
+    symbolic_ctx_manager->register_pass<ov::pass::SDPASplitAttentionFusionMatcher>();
     symbolic_ctx_manager->register_pass<ov::pass::SDPAFusionMatcherSinks>();
     symbolic_ctx_manager->register_pass<ov::pass::SDPAReshapeFusion>();
     return symbolic_optimizations.run_on_model(model);
@@ -632,4 +634,196 @@ SDPAFusionMatcherSinks::SDPAFusionMatcherSinks() {
     this->register_matcher(m, callback);
 }
 
+SDPASplitAttentionFusionMatcher::SDPASplitAttentionFusionMatcher() {
+    MATCHER_SCOPE(SDPASplitAttentionFusionMatcher);
+
+    // Match the "split-attention" pattern; see header documentation.
+    auto q_input = any_input();
+    auto k_cache_input = any_input();
+    auto k_new_input = any_input();
+    auto v_cache_input = any_input();
+    auto v_new_input = any_input();
+    auto mask_input = any_input();
+
+    auto qk_cache = wrap_type<v0::MatMul>({q_input, k_cache_input}, consumers_count(1));
+    auto qk_new = wrap_type<v0::MatMul>({q_input, k_new_input}, consumers_count(1));
+
+    auto scores_concat = wrap_type<v0::Concat>({qk_cache, qk_new}, consumers_count(1));
+    auto masked_scores = wrap_type<v1::Add>({scores_concat, mask_input}, consumers_count(1));
+    auto softmax = wrap_type<v8::Softmax>({masked_scores}, consumers_count(1));
+
+    auto split_axis = any_input();
+    auto split_sizes = any_input();
+    auto vsplit = wrap_type<v1::VariadicSplit>({softmax, split_axis, split_sizes});
+
+    // VariadicSplit has multiple outputs; the v0::MatMul wrap cannot easily express
+    // "use output 0/1 of vsplit", so the first MatMul input is matched as any_input and
+    // verified explicitly in the callback below.
+    auto attn_cache = wrap_type<v0::MatMul>({any_input(), v_cache_input});
+    auto attn_new = wrap_type<v0::MatMul>({any_input(), v_new_input});
+    auto output_add = wrap_type<v1::Add>({attn_cache, attn_new});
+
+    ov::matcher_pass_callback callback = [=](Matcher& m) {
+        if (transformation_callback(m.get_match_root())) {
+            return false;
+        }
+
+        auto add_node = ov::as_type_ptr<v1::Add>(m.get_match_root());
+        if (!add_node) {
+            return false;
+        }
+
+        auto matmul_cache = ov::as_type_ptr<v0::MatMul>(add_node->input_value(0).get_node_shared_ptr());
+        auto matmul_new = ov::as_type_ptr<v0::MatMul>(add_node->input_value(1).get_node_shared_ptr());
+        if (!matmul_cache || !matmul_new) {
+            return false;
+        }
+
+        // Both V matmuls' first input must come from the same VariadicSplit (output 0 and 1).
+        auto cache_src = matmul_cache->input_value(0);
+        auto new_src = matmul_new->input_value(0);
+        auto vsplit_node = ov::as_type_ptr<v1::VariadicSplit>(cache_src.get_node_shared_ptr());
+        if (!vsplit_node ||
+            vsplit_node != ov::as_type_ptr<v1::VariadicSplit>(new_src.get_node_shared_ptr())) {
+            return false;
+        }
+        if (cache_src.get_index() != 0 || new_src.get_index() != 1) {
+            return false;
+        }
+
+        auto softmax_node = ov::as_type_ptr<v8::Softmax>(vsplit_node->input_value(0).get_node_shared_ptr());
+        if (!softmax_node) {
+            return false;
+        }
+
+        // Softmax must be on the last axis.
+        auto sm_rank = softmax_node->get_output_partial_shape(0).rank();
+        if (sm_rank.is_dynamic()) {
+            return false;
+        }
+        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;
+        }
+
+        auto mask_add_node = ov::as_type_ptr<v1::Add>(softmax_node->input_value(0).get_node_shared_ptr());
+        if (!mask_add_node) {
+            return false;
+        }
+
+        auto concat_node = ov::as_type_ptr<v0::Concat>(mask_add_node->input_value(0).get_node_shared_ptr());
+        if (!concat_node || concat_node->get_input_size() != 2) {
+            return false;
+        }
+        auto concat_rank = concat_node->get_output_partial_shape(0).rank();
+        if (concat_rank.is_dynamic()) {
+            return false;
+        }
+        auto concat_axis = concat_node->get_axis();
+        if (concat_axis < 0) {
+            concat_axis += concat_rank.get_length();
+       }
+        if (concat_axis != concat_rank.get_length() - 1) {
+            return false;
+        }
+
+        auto qk_cache_node = ov::as_type_ptr<v0::MatMul>(concat_node->input_value(0).get_node_shared_ptr());
+        auto qk_new_node = ov::as_type_ptr<v0::MatMul>(concat_node->input_value(1).get_node_shared_ptr());
+        if (!qk_cache_node || !qk_new_node) {
+            return false;
+        }
+
+        // Both QK matmuls must share the same Q and agree on transpose flags.
+        if (qk_cache_node->input_value(0) != qk_new_node->input_value(0)) {
+            return false;
+        }
+        if (qk_cache_node->get_transpose_a() || qk_new_node->get_transpose_a()) {
+            return false;
+        }
+        if (qk_cache_node->get_transpose_b() != qk_new_node->get_transpose_b()) {
+            return false;
+        }
+        if (matmul_cache->get_transpose_a() || matmul_new->get_transpose_a()) {
+            return false;
+        }
+        if (matmul_cache->get_transpose_b() != matmul_new->get_transpose_b()) {
+            return false;
+        }
+
+        auto Q = qk_cache_node->input_value(0);
+        auto K_cache = qk_cache_node->input_value(1);
+        auto K_new = qk_new_node->input_value(1);
+        auto V_cache = matmul_cache->input_value(1);
+        auto V_new = matmul_new->input_value(1);
+        auto mask_value = mask_add_node->input_value(1);
+
+        // Require static 4D ranks for predictable layout reasoning.
+        for (const auto& in : {Q, K_cache, K_new, V_cache, V_new}) {
+            const auto& ps = in.get_partial_shape();
+            if (ps.rank().is_dynamic() || ps.rank().get_length() != 4) {
+                return false;
+            }
+        }
+
+        // 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.
+        //   attn transpose_b == false; V stored as [B,H,S,D] (standard).
+        const bool k_is_transposed = !qk_cache_node->get_transpose_b();
+        const bool v_is_transposed = matmul_cache->get_transpose_b();
+
+        // Concat along the physical sequence axis of each input, then add an explicit
+        // Transpose to bring it to the [B,H,S,D] layout expected by v13::SDPA.
+        const int64_t k_concat_axis = k_is_transposed ? 3 : 2;
+        const int64_t v_concat_axis = v_is_transposed ? 3 : 2;
+
+        auto k_concat = std::make_shared<v0::Concat>(ov::OutputVector{K_cache, K_new}, k_concat_axis);
+        auto v_concat = std::make_shared<v0::Concat>(ov::OutputVector{V_cache, V_new}, v_concat_axis);
+
+        std::shared_ptr<ov::Node> k_input = k_concat;
+        std::shared_ptr<ov::Node> v_input = v_concat;
+
+        ov::NodeVector new_nodes{k_concat, v_concat};
+        auto add_transpose = [&new_nodes](std::shared_ptr<ov::Node>& target) {
+            auto perm = v0::Constant::create(ov::element::i64, ov::Shape{4}, {0, 1, 3, 2});
+            auto transpose = std::make_shared<v1::Transpose>(target, perm);
+            new_nodes.push_back(perm);
+            new_nodes.push_back(transpose);
+            target = transpose;
+        };
+        if (k_is_transposed) {
+            add_transpose(k_input);
+        }
+        if (v_is_transposed) {
+            add_transpose(v_input);
+        }
+
+        // Scale = 1.0: the matched pattern has no explicit scale node; any scaling is
+        // assumed to be baked into Q by the model (e.g. Gemma's query_pre_attn_scalar).
+        auto scale_const = v0::Constant::create(Q.get_element_type(), ov::Shape{}, {1.0f});
+        new_nodes.push_back(scale_const);
+
+        auto sdpa = std::make_shared<v13::ScaledDotProductAttention>(Q,
+                                                                     k_input,
+                                                                     v_input,
+                                                                     mask_value,
+                                                                     scale_const,
+                                                                     /*is_causal=*/false);
+        new_nodes.push_back(sdpa);
+
+        sdpa->set_friendly_name(add_node->get_friendly_name());
+        ov::copy_runtime_info(m.get_matched_nodes(), new_nodes);
+        ov::replace_node(add_node, sdpa);
+
+        return true;
+    };
+
+    auto m = std::make_shared<Matcher>(output_add, matcher_name);
+    this->register_matcher(m, callback);
+}
+
 }  // namespace ov::pass
diff --git a/src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp b/src/common/transformations/tests/common_optimizations/sdpa_fusion_test.cpp
