Add cache rotation inputs and CPU kernel implementation for cache rotation #27088

vshampor · 2024-10-16T13:04:23Z

Tickets:
153783

slyalin · 2024-10-28T07:13:59Z

src/core/src/op/paged_attention.cpp

+                          get_input_size() == 15,
+                          "PagedAttensionExtension expects 15 inputs, but it has ",


Doesn't look as optional inputs. According to the spec they could be omitted. If you replace by get_input_size() == 13 || get_input_size() == 15 it wouldn't be a big code modification but unlock a bit of flexibility in the transition period where various mixes of main ov and genai may happen. As we keep PA op internal and not very particular on op version numbering, then a bit of backward compatibility care would be nice.

Done, but the alibi parameter doesn't seem to follow that approach

slyalin · 2024-10-28T07:17:35Z

...mon/transformations/src/transformations/sdpa_to_paged_attention/state_management_pattern.cpp

+            pa_arguments.insert(pa_arguments.begin() + 13, v0::Constant::create(element::f32, Shape{0}, {}));
+            pa_arguments.insert(pa_arguments.begin() + 14, v0::Constant::create(element::i32, Shape{0}, {}));


If you make these inputs really optional, these two lines are not required.

slyalin · 2024-10-28T07:19:57Z

src/core/src/op/paged_attention.cpp

+            get_input_partial_shape(13).rank().is_dynamic() ||
+            get_input_partial_shape(13).rank().get_length() == 0 ||
+            get_input_partial_shape(13).rank().get_length() == 1,
+            "Input `rotation_coefficients` should either have an empty shape or rank 1, but it has rank ",


Suggested change

"Input `rotation_coefficients` should either have an empty shape or rank 1, but it has rank ",

"Input `rotation_coefficients` should either have rank 1 or omitted, but it has rank ",

"Empty" shape means [0] here, which have rank 1.

slyalin · 2024-10-28T07:20:39Z

src/core/src/op/paged_attention.cpp

+    NODE_VALIDATION_CHECK(
+            this,
+            get_input_partial_shape(13).rank().is_dynamic() ||
+            get_input_partial_shape(13).rank().get_length() == 0 ||


Suggested change

get_input_partial_shape(13).rank().get_length() == 0 ||

slyalin · 2024-10-28T07:23:08Z

src/core/src/op/paged_attention.cpp

+            get_input_partial_shape(14).rank().get_length() == 0 ||
+            get_input_partial_shape(14).rank().get_length() == 1,
+            "Input `rotated_block_indices` should either have an empty shape or rank 1 but it has rank ",


The same comment are applicable here as for input 13 above.

slyalin · 2024-10-28T07:27:03Z

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

@@ -1576,6 +1591,11 @@ struct AttentionExecutor : public PagedAttentionExecutor {
        if (alibi_slopes) {
            alibi_slopes.assert_dims({H});
        }
+
+        if (rotated_block_indices) {
+            // Rotation, and cache eviction, is limited to cases when Q, K and V embedding sizes are equal, e.g. S == Sv


We already have cases where they are not: minicpm-3

Removed - realized that we don't need that limitation for cache rotation since we only rotate the K values

slyalin · 2024-10-28T07:30:07Z

src/plugins/intel_gpu/src/plugin/ops/paged_attention.cpp

@@ -58,6 +59,10 @@ static void CreatePagedAttentionExtensionOp(ProgramBuilder& p, const std::shared
    OPENVINO_ASSERT(alibi_const != nullptr);
    prim.has_alibi = ov::shape_size(alibi_const->get_output_shape(0)) > 0;

+    std::shared_ptr<ov::op::v0::Constant> rotation_coefficients_const = std::dynamic_pointer_cast<ov::op::v0::Constant>(op->get_input_node_shared_ptr(rotation_coefficients_idx));
+    OPENVINO_ASSERT(rotation_coefficients_const != nullptr);
+    prim.has_rotation_coefficients = ov::shape_size(alibi_const->get_output_shape(0)) > 0;


alibi_const shouldn't be used here -- bad copy&paste?

Fixed, thanks.

# Conflicts: # src/core/src/op/paged_attention.cpp

dmitry-gorokhov · 2024-11-15T07:52:52Z

@luo-cheng2021 Please review CPU PA changes.

github-actions bot added category: Core OpenVINO Core (aka ngraph) category: GPU OpenVINO GPU plugin category: CPU OpenVINO CPU plugin category: transformations OpenVINO Runtime library - Transformations category: CPP API OpenVINO CPP API bindings labels Oct 16, 2024

vshampor force-pushed the token_rotation branch from 2a172b2 to c071571 Compare October 26, 2024 00:52

slyalin requested changes Oct 28, 2024

View reviewed changes

github-actions bot added the category: build OpenVINO cmake script / infra label Oct 30, 2024

vshampor added 12 commits November 11, 2024 20:28

Add cache rotation inputs in transformations, CPU and GPU plugins

eb19902

# Conflicts: # src/core/src/op/paged_attention.cpp

Fix missing internal dimension

9302ada

Fixes

65c26e9

Apply comments

8de94e4

Use llama-style rotation

f6e04a8

Implement AVX512 version w/ intrinsics

870fb6d

Refactor the kernel

a2c2d0a

Add AVX2 impl

072c13b

Add SW impl reference unit test

452736d

Add HW impl block tests

929e5ba

Add chunk tests

c97cd62

Fix tests

7722675

vshampor changed the title ~~Add cache rotation inputs~~ Add cache rotation inputs and CPU kernel implementation for cache rotation Nov 12, 2024

vshampor mentioned this pull request Nov 12, 2024

Token rotation openvinotoolkit/openvino.genai#987

Open

Add testing for chunk rotation in FP16 and BF16

ed46cfe

vshampor force-pushed the token_rotation branch from d90e212 to ed46cfe Compare November 12, 2024 20:30

Fix bugs

8af57a6

dmitry-gorokhov assigned luo-cheng2021 Nov 15, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add cache rotation inputs and CPU kernel implementation for cache rotation #27088

Add cache rotation inputs and CPU kernel implementation for cache rotation #27088

vshampor commented Oct 16, 2024

slyalin Oct 28, 2024

vshampor Oct 30, 2024

slyalin Oct 28, 2024 •

edited

Loading

vshampor Oct 30, 2024

slyalin Oct 28, 2024

vshampor Oct 30, 2024

slyalin Oct 28, 2024

vshampor Oct 30, 2024

slyalin Oct 28, 2024

vshampor Oct 30, 2024

slyalin Oct 28, 2024 •

edited

Loading

vshampor Oct 30, 2024

slyalin Oct 28, 2024

vshampor Oct 30, 2024

dmitry-gorokhov commented Nov 15, 2024

		get_input_size() == 15,
		"PagedAttensionExtension expects 15 inputs, but it has ",

		pa_arguments.insert(pa_arguments.begin() + 13, v0::Constant::create(element::f32, Shape{0}, {}));
		pa_arguments.insert(pa_arguments.begin() + 14, v0::Constant::create(element::i32, Shape{0}, {}));

	"Input `rotation_coefficients` should either have an empty shape or rank 1, but it has rank ",
	"Input `rotation_coefficients` should either have rank 1 or omitted, but it has rank ",

Add cache rotation inputs and CPU kernel implementation for cache rotation #27088

Are you sure you want to change the base?

Add cache rotation inputs and CPU kernel implementation for cache rotation #27088

Conversation

vshampor commented Oct 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slyalin Oct 28, 2024 • 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

Choose a reason for hiding this comment

slyalin Oct 28, 2024 • 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

dmitry-gorokhov commented Nov 15, 2024

slyalin Oct 28, 2024 •

edited

Loading

slyalin Oct 28, 2024 •

edited

Loading