[GPU][CM] optimize CM VLSDPA and CM PA performance.

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

@@ -58,24 +58,16 @@ std::shared_ptr<ov::Model> build_target_model(const string& mask_name) {
     k->set_friendly_name("k");
     v->set_friendly_name("v");
 
-    auto transpose_q = std::make_shared<Transpose>(q, Constant::create(element::i64, Shape{3}, {1, 0, 2}));
-    auto transpose_k = std::make_shared<Transpose>(k, Constant::create(element::i64, Shape{3}, {1, 0, 2}));
-    auto transpose_v = std::make_shared<Transpose>(v, Constant::create(element::i64, Shape{3}, {1, 0, 2}));
-    transpose_q->set_friendly_name("transpose_q");
-    transpose_k->set_friendly_name("transpose_k");
-    transpose_v->set_friendly_name("transpose_v");
-
     auto cuseq_mask = std::make_shared<Parameter>(element::i32, PartialShape{-1});
     cuseq_mask->set_friendly_name(mask_name);
     cuseq_mask->get_output_tensor(0).set_names({mask_name});
 
+    const std::vector<int64_t> order{1, 0, 2};
     auto vlsdpa =
-        std::make_shared<ov::op::internal::VLSDPA>(OutputVector{transpose_q, transpose_k, transpose_v, cuseq_mask});
-
-    auto transpose_o = std::make_shared<Transpose>(vlsdpa, Constant::create(element::i64, Shape{3}, {1, 0, 2}));
-    transpose_o->set_friendly_name("transpose_o");
+        std::make_shared<ov::op::internal::VLSDPA>(OutputVector{q, k, v, cuseq_mask}, order, order, order, order);
+    vlsdpa->set_friendly_name("transpose_o");
 
-    return std::make_shared<ov::Model>(OutputVector{transpose_o}, ParameterVector{q, k, v, cuseq_mask});
+    return std::make_shared<ov::Model>(OutputVector{vlsdpa}, ParameterVector{q, k, v, cuseq_mask});
 }
 };  // namespace
 

src/core/src/pass/sdpa_to_vlsdpa.cpp

@@ -9,7 +9,11 @@
 #include <utility>
 
 #include "openvino/cc/pass/itt.hpp"
+#include "openvino/core/graph_util.hpp"
+#include "openvino/core/rt_info.hpp"
+#include "openvino/op/constant.hpp"
 #include "openvino/op/scaled_dot_product_attention.hpp"
+#include "openvino/op/transpose.hpp"
 #include "openvino/op/util/node_util.hpp"
 #include "openvino/pass/manager.hpp"
 #include "ov_ops/vl_sdpa.hpp"
@@ -85,14 +89,74 @@ bool SDPAToVLSDPA::run_on_model(const std::shared_ptr<ov::Model>& model) {
 
                 const auto sdpa_consumers = sdpa->get_output_target_inputs(0);
                 const auto new_args = sdpa->input_values();
-                OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param};
+
+                // Try to fuse the {1,0,2} Transposes surrounding SDPA into the VLSDPA op
+                // at op-creation time. Doing it here (rather than as a later plugin-side
+                // matcher on `Transpose->VLSDPA->Transpose`) is race-free w.r.t. common
+                // optimizations that may convert Transpose->Reshape when one of the
+                // permuted dims degenerates to size 1 (e.g. num_head==1). Constraints
+                // mirror `VLSDPA::validate_and_infer_types`: all four orders must be
+                // equal and, when non-empty, equal to {1,0,2}.
+                const std::vector<int64_t> expected_order{1, 0, 2};
+                auto try_peel_input_transpose =
+                    [&expected_order](ov::Output<ov::Node>& value) -> std::shared_ptr<v1::Transpose> {
+                    auto transpose = ov::as_type_ptr<v1::Transpose>(value.get_node_shared_ptr());
+                    if (!transpose)
+                        return nullptr;
+                    if (transpose->get_output_target_inputs(0).size() != 1)
+                        return nullptr;
+                    auto order_const = ov::as_type_ptr<v0::Constant>(transpose->input_value(1).get_node_shared_ptr());
+                    if (!order_const || ov::shape_size(order_const->get_shape()) != expected_order.size())
+                        return nullptr;
+                    if (order_const->cast_vector<int64_t>() != expected_order)
+                        return nullptr;
+                    value = transpose->input_value(0);
+                    return transpose;
+                };
+
+                ov::Output<ov::Node> q_in = new_args.at(0);
+                ov::Output<ov::Node> k_in = new_args.at(1);
+                ov::Output<ov::Node> v_in = new_args.at(2);
+                auto tq = try_peel_input_transpose(q_in);
+                auto tk = try_peel_input_transpose(k_in);
+                auto tv = try_peel_input_transpose(v_in);
+
+                std::shared_ptr<v1::Transpose> to;
+                if (sdpa_consumers.size() == 1) {
+                    auto candidate_out =
+                        ov::as_type_ptr<v1::Transpose>(sdpa_consumers.begin()->get_node()->shared_from_this());
+                    if (candidate_out) {
+                        auto out_order_const =
+                            ov::as_type_ptr<v0::Constant>(candidate_out->input_value(1).get_node_shared_ptr());
+                        if (out_order_const && ov::shape_size(out_order_const->get_shape()) == expected_order.size() &&
+                            out_order_const->cast_vector<int64_t>() == expected_order) {
+                            to = candidate_out;
+                        }
+                    }
+                }
+
+                const bool fuse_transposes = tq && tk && tv && to;
 
                 std::shared_ptr<op::internal::VLSDPA> vl_sdpa;
-                vl_sdpa = std::make_shared<op::internal::VLSDPA>(inputs);
-                vl_sdpa->set_friendly_name(sdpa->get_friendly_name());
+                if (fuse_transposes) {
+                    OutputVector inputs{q_in, k_in, v_in, cu_seqlens_param};
+                    vl_sdpa = std::make_shared<op::internal::VLSDPA>(inputs,
+                                                                     expected_order,
+                                                                     expected_order,
+                                                                     expected_order,
+                                                                     expected_order);
+                    vl_sdpa->set_friendly_name(to->get_friendly_name());
+                    ov::copy_runtime_info({tq, tk, tv, sdpa->shared_from_this(), to}, vl_sdpa);
+                    ov::replace_node(to, vl_sdpa);
+                } else {
+                    OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param};
+                    vl_sdpa = std::make_shared<op::internal::VLSDPA>(inputs);
+                    vl_sdpa->set_friendly_name(sdpa->get_friendly_name());
+                    ov::copy_runtime_info(sdpa->shared_from_this(), vl_sdpa);
 
-                for (auto& consumer : sdpa_consumers)
-                    consumer.replace_source_output(vl_sdpa);
+                    for (auto& consumer : sdpa_consumers)
+                        consumer.replace_source_output(vl_sdpa);
+                }
             }
         }
     }

src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm

@@ -15,7 +15,7 @@
  * limitations under the License.
  *******************************************************************************/
 namespace KERNEL_NAME {
-    
+
 #include "cm_sdpa_common.hpp"
 
 #ifdef CM_HAS_LSC_UNTYPED_2D
@@ -36,10 +36,32 @@ _GENX_MAIN_ void KERNEL_NAME(
     SurfaceIndex query [[type("buffer_t")]],
     SurfaceIndex key [[type("buffer_t")]],
     SurfaceIndex value [[type("buffer_t")]],
-    SurfaceIndex output [[type("buffer_t")]],    
+    SurfaceIndex output [[type("buffer_t")]],
 #endif
     int* cu_seqlens [[type("svmptr_t")]],
-    int  need_wg_mapping
+    int  need_wg_mapping,
+    // Memory-layout contract for Q/K/V inputs (must be preserved by upstream passes):
+    //   Only the token (outer) axis is allowed to be discontinuous. The kernel treats
+    //   each input as a 3-D tensor [tokens, heads, head_size] and hardcodes the two
+    //   inner strides:
+    //     - per-head stride  == head_size (heads contiguous within a token)
+    //     - per-element stride == 1        (head_size dim is contiguous)
+    //   Only the outer per-token stride is layout-driven via {q,k,v}_token_pitch below.
+    //   If prepare_buffer_fusing ever propagates padding onto the head or head_size axis
+    //   of a vl_sdpa input layout, this kernel will silently address the wrong memory.
+    //
+    // Element-count offsets from the SVM base to the slice start of each input, taken
+    // from input_layout.get_linear_offset(). Non-zero when an in-place crop (e.g. the
+    // TransposeSplitMatcher axis=1 pattern) leaves the SVM base pointing at the packed
+    // parent buffer and shifts the logical view via layout padding on the token axis.
+    int  token_offset_q,
+    int  token_offset_k,
+    int  token_offset_v,
+    // Per-input token strides in elements, from input_layout.get_pitches()[0]. Allows
+    // mixed layouts (e.g. Q/K contiguous post-RoPE, V still viewing the padded parent).
+    int  q_token_pitch,
+    int  k_token_pitch,
+    int  v_token_pitch
     ) {
     constexpr int num_heads = CMFLA_NUM_HEADS;
     constexpr int head_size = CMFLA_HEAD_SIZE;
@@ -107,8 +129,12 @@ _GENX_MAIN_ void KERNEL_NAME(
         kv_start = cu_seqlens[wg_id];
         kv_seq_len = cu_seqlens[wg_id + 1] - kv_start;
     }
-    uint qo_offset = (q_start*num_heads + h)*head_size;
-    uint kv_offset = (kv_start*num_kv_heads + hkv)*head_size;
+    // Input offsets are computed with per-input runtime token pitches.
+    // Output is always contiguous [q_len, num_heads, head_size].
+    uint q_in_offset = q_start * q_token_pitch + h * head_size;
+    uint k_in_offset = kv_start * k_token_pitch + hkv * head_size;
+    uint v_in_offset = kv_start * v_token_pitch + hkv * head_size;
+    uint o_offset = (q_start * num_heads + h) * head_size;
 
 #if USE_LSC == 1
 #if USE_LSC_PREFETCH
@@ -118,10 +144,13 @@ _GENX_MAIN_ void KERNEL_NAME(
                                 kv_seq_len,
                                 q_len,
                                 kv_seq_len, //kv_len,
-                                reinterpret_cast<svmptr_t>(query + qo_offset),
-                                reinterpret_cast<svmptr_t>(key + kv_offset),
-                                reinterpret_cast<svmptr_t>(value + kv_offset),
-                                reinterpret_cast<svmptr_t>(output + qo_offset));
+                                q_token_pitch * sizeof(half),
+                                k_token_pitch * sizeof(half),
+                                v_token_pitch * sizeof(half),
+                                reinterpret_cast<svmptr_t>((query + token_offset_q) + q_in_offset),
+                                reinterpret_cast<svmptr_t>((key + token_offset_k) + k_in_offset),
+                                reinterpret_cast<svmptr_t>((value + token_offset_v) + v_in_offset),
+                                reinterpret_cast<svmptr_t>(output + o_offset));
 #else
     sdpa_kernel_lsc<false, num_heads, num_kv_heads, head_size>(
                                 slm_K,
@@ -132,10 +161,13 @@ _GENX_MAIN_ void KERNEL_NAME(
                                 kv_seq_len, //kv_stop,
                                 q_len, //q_len,
                                 kv_seq_len, //kv_len,
-                                reinterpret_cast<svmptr_t>(query + qo_offset),
-                                reinterpret_cast<svmptr_t>(key + kv_offset),
-                                reinterpret_cast<svmptr_t>(value + kv_offset),
-                                reinterpret_cast<svmptr_t>(output + qo_offset));
+                                q_token_pitch * sizeof(half),
+                                k_token_pitch * sizeof(half),
+                                v_token_pitch * sizeof(half),
+                                reinterpret_cast<svmptr_t>((query + token_offset_q) + q_in_offset),
+                                reinterpret_cast<svmptr_t>((key + token_offset_k) + k_in_offset),
+                                reinterpret_cast<svmptr_t>((value + token_offset_v) + v_in_offset),
+                                reinterpret_cast<svmptr_t>(output + o_offset));
 #endif
 #else
     sdpa_kernel<false, num_heads, num_kv_heads, head_size, 0>(
@@ -151,10 +183,13 @@ _GENX_MAIN_ void KERNEL_NAME(
                                 key,
                                 value,
                                 output,
-                                qo_offset * sizeof(half),
-                                kv_offset * sizeof(half),
-                                kv_offset * sizeof(half),
-                                qo_offset * sizeof(half)
+                                q_token_pitch * sizeof(half),
+                                k_token_pitch * sizeof(half),
+                                v_token_pitch * sizeof(half),
+                                (token_offset_q + q_in_offset) * sizeof(half),
+                                (token_offset_k + k_in_offset) * sizeof(half),
+                                (token_offset_v + v_in_offset) * sizeof(half),
+                                o_offset * sizeof(half)
                                 );
 #endif
 }