keep PA sink logic focus on constant with shape {1, H, 1, 1} to align with transofrm process

Author: liubo-intel

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa.cpp

@@ -926,7 +926,7 @@ struct MHAHelper {
                 // sink processing is independent of sliding_window size
                 float* sink = nullptr;
                 if (sinks) {
-                    sink = &sinks.at<float>({batch_in_seq, h, m, 0}, true);
+                    sink = &sinks.at<float>({0, h, 0, 0}, true);
                 }
                 if (_sliding_window) {
                     size_t start_idx = 0;
@@ -1227,7 +1227,6 @@ struct MHAHelper {
                             size_t cur_kv_len,
                             const PlainTensor& alibi_slopes,
                             float* score_output,
-                            size_t batch_in_seq,
                             const PlainTensor& sinks) {
 #    if defined(OPENVINO_ARCH_X86_64)
         if (any_of(_fastpath_valid_prec, ov::element::bf16, ov::element::f16)) {
@@ -1288,7 +1287,7 @@ struct MHAHelper {
                 }
                 float* sink = nullptr;
                 if (sinks) {
-                    sink = &sinks.at<float>({batch_in_seq, h, pq, 0}, true);
+                    sink = &sinks.at<float>({0, h, 0, 0}, true);
                 }
                 if (_sliding_window) {
                     size_t start_idx = 0;
@@ -1498,7 +1497,7 @@ struct MHAHelper {
             }
             float* sink = nullptr;
             if (sinks) {
-                sink = &sinks.at<float>({b, h, pq, 0}, true);
+                sink = &sinks.at<float>({0, h, 0, 0}, true);
             }
             if (_sliding_window) {
                 size_t start_idx = 0;
@@ -1808,7 +1807,6 @@ struct MHA {
                     cur_kv_len,
                     alibi_slopes,
                     score_output,
-                    batch_in_seq,
                     sinks);
             } else {
                 const auto batch_in_reorder = item.batch_in_reorder;

src/plugins/intel_cpu/src/nodes/paged_attn.cpp

@@ -27,6 +27,7 @@
 #include "openvino/core/type/element_type.hpp"
 #include "openvino/runtime/system_conf.hpp"
 #include "shape_inference/shape_inference_internal_dyn.hpp"
+#include "transformations/utils/utils.hpp"
 #include "utils/general_utils.h"
 
 #if defined(OPENVINO_ARCH_X86) || defined(OPENVINO_ARCH_X86_64) || defined(OPENVINO_ARCH_ARM64)
@@ -302,7 +303,20 @@ bool PagedAttention::isSupportedOperation(const std::shared_ptr<const ov::Node>&
         }
         auto orgInput = static_cast<int>(op->get_input_size());
         if (op->get_type_name() == std::string("PagedAttentionExtension") &&
-            orgInput == PagedAttentionExecutor::ID_SLIDING_WINDOW + 1) {
+            orgInput == PagedAttentionExecutor::ID_SINKS + 1) {
+            if (!ov::op::util::is_on_path<ov::op::v0::Constant>(op->input_value(PagedAttentionExecutor::ID_SINKS))) {
+                errorMessage = "Only Constant operation on sink input is supported";
+                return false;
+            }
+#if defined(OPENVINO_ARCH_ARM64)
+            // ARM platform doesn't support non-empty sink input yet
+            // Check if sink input is non-empty (shape size > 0)
+            const auto& sink_shape = op->get_input_partial_shape(PagedAttentionExecutor::ID_SINKS);
+            if (sink_shape.is_static() && ov::shape_size(sink_shape.to_shape()) > 0) {
+                errorMessage = "PagedAttentionExtension with non-empty sink input is not supported on ARM platform";
+                return false;
+            }
+#endif
             return true;
         }
     } catch (...) {

src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/x64/paged_attn.cpp

@@ -138,20 +138,23 @@ class PagedAttnTestBase : public testing::WithParamInterface<PagedAttnTestParams
             std::make_shared<ov::op::v0::Constant>(ov::element::i32, Shape{}, std::vector<int32_t>{64});
         auto xattention_stride =
             std::make_shared<ov::op::v0::Constant>(ov::element::i32, Shape{}, std::vector<int32_t>{8});
-        // Create sink_input parameter for testing - shape [1, num_heads, 1, 1] as per PagedAttentionExecutor::ID_SINKS
-        // PagedAttentionExtension always expects 21 inputs, so we must always include sinks parameter
-        auto sinks = make_param(PartialShape{1, head_num, 1, 1}, data_type, "sinks");
-
-        ParameterVector params = {q,
-                                  k,
-                                  v,
-                                  key_cache,
-                                  value_cache,
-                                  past_lens,
-                                  subsequence_begins,
-                                  block_indices,
-                                  block_indices_begins,
-                                  sinks};
+        // Create sink input as Constant (not Parameter) for testing
+        // PagedAttentionExtension requires sink input to be Constant
+        // Use shape [1, head_num, 1, 1] when use_sink_input=true, or empty shape [0] when false
+        std::shared_ptr<ov::op::v0::Constant> sinks;
+        if (use_sink_input) {
+            // Create real sink tokens for testing sink functionality
+            std::vector<float> sink_data(static_cast<size_t>(head_num), 0.1f);
+            sinks = std::make_shared<ov::op::v0::Constant>(data_type,
+                                                           Shape{1, static_cast<size_t>(head_num), 1, 1},
+                                                           sink_data);
+        } else {
+            // Create empty sink (matching SDPA->PA transformation behavior when no sink)
+            sinks = std::make_shared<ov::op::v0::Constant>(data_type, Shape{0}, std::vector<float>{});
+        }
+
+        ParameterVector params =
+            {q, k, v, key_cache, value_cache, past_lens, subsequence_begins, block_indices, block_indices_begins};
         OutputVector paged_attn_inputs = {q,
                                           k,
                                           v,
@@ -461,10 +464,7 @@ class PagedAttnTestBase : public testing::WithParamInterface<PagedAttnTestParams
             inputs.insert({function->get_parameters()[7], block_indices});
             inputs.insert({function->get_parameters()[8], block_indices_begins});
 
-            // Create sink_input data - shape [1, num_heads, 1, 1] as per PagedAttentionExtensor specification
-            // Always create the sink input data since PagedAttentionExtension expects 21 inputs
-            // The value will be ignored when use_sink_input=false
-            create_input(function->get_parameters()[9], {1, qkv_shape[2], 1, 1}, use_sink_input ? 0.1f : 0.0f);
+            // Note: sink is a Constant in the model, not a Parameter, so no need to provide input for it
 
             past_len_count += static_cast<int32_t>(qkv_shape[0]);
 
@@ -655,21 +655,22 @@ INSTANTIATE_TEST_SUITE_P(smoke_PagedAttnVSSDPATest,
                          ::testing::Combine(::testing::Values(ElementType::f32, ElementType::bf16),
                                             ::testing::ValuesIn(inputShapeAndReorders),
                                             ::testing::Values(true, false),
+                                            // TODO: Xattn should not direcctly compare with SDPA/decomposed Matmul
+                                            // which not contain sparse logics
                                             ::testing::Values(true, false),
-                                            ::testing::Values(true, false),
-                                            ::testing::Values(0),  // sliding_window = 0
+                                            ::testing::Values(false),
+                                            ::testing::Values(0),
                                             ::testing::ValuesIn(additional_configs)),
                          PagedAttnTestBase::getTestCaseName);
 
-// Sliding window test with same shapes as normal test
-INSTANTIATE_TEST_SUITE_P(smoke_PagedAttnVSSDPATest_WithSlidingWindow,
+INSTANTIATE_TEST_SUITE_P(smoke_PagedAttnVSSDPATest_WithSlidingWindowAndSinks,
                          PagedAttnVSSDPATest,
                          ::testing::Combine(::testing::Values(ElementType::f32),
                                             ::testing::ValuesIn(inputShapeAndReorders),
-                                            ::testing::Values(false),  // extendBlockIndices
-                                            ::testing::Values(false),  // enableXattn
-                                            ::testing::Values(true),   // sinkInput
-                                            ::testing::Values(8),      // sliding_window = 8
+                                            ::testing::Values(false),        // extendBlockIndices
+                                            ::testing::Values(false),        // enableXattn
+                                            ::testing::Values(true, false),  // sinkInput
+                                            ::testing::Values(0, 8),         // sliding_window = 8
                                             ::testing::Values(ov::AnyMap{
                                                 {ov::intel_cpu::enable_sage_attn.name(), false}})),
                          PagedAttnTestBase::getTestCaseName);