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

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

Tickets:
153783

src/core/src/op/paged_attention.cpp

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.

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

@luo-cheng2021 Please review CPU PA changes.

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/common.hpp

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

luo-cheng2021 · 2024-11-18T06:54:51Z

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/cache_rotation.hpp

+                CT cache_value_1 = *cache_value_1_ptr;
+
+                *cache_value_0_ptr = cache_value_0 * rotation_value_cos - cache_value_1 * rotation_value_sin;
+                *cache_value_1_ptr = cache_value_0 * rotation_value_sin + cache_value_1 * rotation_value_cos;


Is the algorithm same with the following code?

openvino/src/plugins/intel_cpu/src/nodes/rope.cpp

Lines 158 to 161 in c4d6d2b

auto src0 = src[i];

auto src1 = src[i + half_rotary_dims];

dst[i] = cos[i] * src0 - sin[i] * src1;

dst[i + half_rotary_dims] = cos[i + half_rotary_dims] * src1 + sin[i + half_rotary_dims] * src0;

If so, the following code can be used as reference:

openvino/src/plugins/intel_cpu/src/nodes/rope.cpp

Lines 35 to 102 in c4d6d2b

static std::shared_ptr<kernel::JitKernelBase> createJitKernel(const jit_rotary_compile_params& param, bool check_vec_size2 = false) {

std::shared_ptr<kernel::JitKernelBase> res;

MAYBE_UNUSED(param);

MAYBE_UNUSED(check_vec_size2);

#if defined(OPENVINO_ARCH_X86_64)

if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core)) {

bool flag = true;

if (check_vec_size2) {

auto vec_size = jit_rotary_kernel<dnnl::impl::cpu::x64::avx512_core>::vec_size;

if (param.rotary_ndims % (vec_size * 2) != 0)

flag = false;

}

if (flag)

res = std::make_shared<jit_rotary_kernel<dnnl::impl::cpu::x64::avx512_core>>(param);

} else if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx2)) {

bool flag = true;

if (check_vec_size2) {

auto vec_size = jit_rotary_kernel<dnnl::impl::cpu::x64::avx2>::vec_size;

if (param.rotary_ndims % (vec_size * 2) != 0)

flag = false;

}

if (flag)

res = std::make_shared<jit_rotary_kernel<dnnl::impl::cpu::x64::avx2>>(param);

}

if (res)

res->create_kernel();

#endif // OPENVINO_ARCH_X86_64

return res;

}

static void execJitKernel(const std::shared_ptr<kernel::JitKernelBase>& ker, const void* src, void* dst, const float* cos, const float* sin) {

MAYBE_UNUSED(ker);

MAYBE_UNUSED(src);

MAYBE_UNUSED(dst);

MAYBE_UNUSED(cos);

MAYBE_UNUSED(sin);

#if defined(OPENVINO_ARCH_X86_64)

jit_rotary_call_args call_args;

call_args.src = src;

call_args.cos = cos;

call_args.sin = sin;

call_args.dst = dst;

(*ker)(&call_args);

#endif // OPENVINO_ARCH_X86_64

}

template <typename T>

struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {

const op::internal::RoPE::Config& m_config;

std::shared_ptr<kernel::JitKernelBase> m_rotaryKernel;

RoPEExecutorRotateHalf(const op::internal::RoPE::Config& config) : m_config(config) {

jit_rotary_compile_params jcp;

jcp.src_prc = precision_of<T>::value;

jcp.dst_prc = precision_of<T>::value;

jcp.rotary_ndims = config.rotary_ndims;

jcp.interleave = false;

m_rotaryKernel = createJitKernel(jcp);

}

I have already written and tested my implementation, besides, the code you've sent me probably cannot be reused without modifications or bulky instantiations.

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/cache_rotation.hpp

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

p-durandin

OK from GPU point of view

### Details: - This PR adds cache rotation support for PagedAttention and related tests - Should be merged after #27088 ### Tickets: - *ticket-id* --------- Co-authored-by: Vasily Shamporov <[email protected]> Co-authored-by: Pavel Durandin <[email protected]> Co-authored-by: cecilia peng <[email protected]>

…ation (openvinotoolkit#27088) Tickets: 153783

### Details: - This PR adds cache rotation support for PagedAttention and related tests - Should be merged after openvinotoolkit#27088 ### Tickets: - *ticket-id* --------- Co-authored-by: Vasily Shamporov <[email protected]> Co-authored-by: Pavel Durandin <[email protected]> Co-authored-by: cecilia peng <[email protected]>

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 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

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

dmitry-gorokhov assigned luo-cheng2021 Nov 15, 2024

luo-cheng2021 requested changes Nov 18, 2024

View reviewed changes

luo-cheng2021 reviewed Nov 18, 2024

View reviewed changes

src/plugins/intel_cpu/src/nodes/kernels/scaled_attn/executor_pa.cpp Outdated Show resolved Hide resolved

vshampor requested a review from slyalin November 18, 2024 10:11

vshampor force-pushed the token_rotation branch from afa851c to 8a355f3 Compare November 18, 2024 18:10

github-actions bot added category: Python API OpenVINO Python bindings category: TF FE OpenVINO TensorFlow FrontEnd category: PyTorch FE OpenVINO PyTorch Frontend category: JAX FE OpenVINO JAX FrontEnd labels Nov 18, 2024

vshampor force-pushed the token_rotation branch from 8a355f3 to a33f255 Compare November 19, 2024 10:13

github-actions bot added category: CI OpenVINO public CI github_actions Pull requests that update GitHub Actions code category: NPU OpenVINO NPU plugin labels Nov 19, 2024

vshampor force-pushed the token_rotation branch from 138de47 to a0818c2 Compare November 19, 2024 12:12

vshampor requested a review from luo-cheng2021 November 19, 2024 12:20

vshampor marked this pull request as ready for review November 19, 2024 12:22

vshampor requested review from a team as code owners November 19, 2024 12:22

vshampor added 18 commits January 9, 2025 11:19

Apply comments

439c2dc

Fix build on Win

88aeb6b

Fix PA build

2efde73

Add NeoX?

97506d1

Refactor to delta + LUT

bc0ea4c

Format

51e3638

Debug prints

75aa521

Add shape to rotation deltas

8e33270

Add test

7328f7c

Remove debug prints

b875546

Format

bea5d62

Fix PA tests

3f948be

Fix typo

ade9e62

Skip AVX2 vectorized test if no AVX2 on host

c36d183

Remove extra ifdef

94643bb

Fix GPU names

77fcc65

Fix rebase

7760293

Fix GPU

80449bc

vshampor force-pushed the token_rotation branch from 9306c4e to 80449bc Compare January 9, 2025 10:19

p-durandin approved these changes Jan 9, 2025

View reviewed changes

vshampor enabled auto-merge January 9, 2025 11:53

vshampor added this pull request to the merge queue Jan 9, 2025

Merged via the queue into openvinotoolkit:master with commit 98192e9 Jan 9, 2025
187 checks passed

vshampor deleted the token_rotation branch January 9, 2025 14:08

MirceaDan99 pushed a commit to MirceaDan99/openvino that referenced this pull request Jan 22, 2025

Add cache rotation inputs and CPU kernel implementation for cache rot…

5994651

…ation (openvinotoolkit#27088) Tickets: 153783

ilya-lavrenov mentioned this pull request Jan 23, 2025

Fixed macOS x86_64 vectorized CPU tests #28644

Draft

t-jankowski mentioned this pull request Jan 24, 2025

Fix macOS x86_64 vectorized CPU tests #28656

Draft

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 •

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

luo-cheng2021 Nov 18, 2024

vshampor Nov 19, 2024

p-durandin left a comment

		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 ",

	auto src0 = src[i];
	auto src1 = src[i + half_rotary_dims];
	dst[i] = cos[i] * src0 - sin[i] * src1;
	dst[i + half_rotary_dims] = cos[i + half_rotary_dims] * src1 + sin[i + half_rotary_dims] * src0;

	static std::shared_ptr<kernel::JitKernelBase> createJitKernel(const jit_rotary_compile_params& param, bool check_vec_size2 = false) {
	std::shared_ptr<kernel::JitKernelBase> res;

	MAYBE_UNUSED(param);
	MAYBE_UNUSED(check_vec_size2);

	#if defined(OPENVINO_ARCH_X86_64)

	if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx512_core)) {
	bool flag = true;
	if (check_vec_size2) {
	auto vec_size = jit_rotary_kernel<dnnl::impl::cpu::x64::avx512_core>::vec_size;
	if (param.rotary_ndims % (vec_size * 2) != 0)
	flag = false;
	}
	if (flag)
	res = std::make_shared<jit_rotary_kernel<dnnl::impl::cpu::x64::avx512_core>>(param);
	} else if (dnnl::impl::cpu::x64::mayiuse(dnnl::impl::cpu::x64::avx2)) {
	bool flag = true;
	if (check_vec_size2) {
	auto vec_size = jit_rotary_kernel<dnnl::impl::cpu::x64::avx2>::vec_size;
	if (param.rotary_ndims % (vec_size * 2) != 0)
	flag = false;
	}
	if (flag)
	res = std::make_shared<jit_rotary_kernel<dnnl::impl::cpu::x64::avx2>>(param);
	}

	if (res)
	res->create_kernel();

	#endif // OPENVINO_ARCH_X86_64

	return res;
	}

	static void execJitKernel(const std::shared_ptr<kernel::JitKernelBase>& ker, const void* src, void* dst, const float* cos, const float* sin) {
	MAYBE_UNUSED(ker);
	MAYBE_UNUSED(src);
	MAYBE_UNUSED(dst);
	MAYBE_UNUSED(cos);
	MAYBE_UNUSED(sin);

	#if defined(OPENVINO_ARCH_X86_64)

	jit_rotary_call_args call_args;
	call_args.src = src;
	call_args.cos = cos;
	call_args.sin = sin;
	call_args.dst = dst;
	(*ker)(&call_args);

	#endif // OPENVINO_ARCH_X86_64
	}

	template <typename T>
	struct RoPE::RoPEExecutorRotateHalf : public RoPE::Executor {
	const op::internal::RoPE::Config& m_config;
	std::shared_ptr<kernel::JitKernelBase> m_rotaryKernel;

	RoPEExecutorRotateHalf(const op::internal::RoPE::Config& config) : m_config(config) {
	jit_rotary_compile_params jcp;
	jcp.src_prc = precision_of<T>::value;
	jcp.dst_prc = precision_of<T>::value;
	jcp.rotary_ndims = config.rotary_ndims;
	jcp.interleave = false;
	m_rotaryKernel = createJitKernel(jcp);
	}

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

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

Conversation

vshampor commented Oct 16, 2024

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

p-durandin left a comment

Choose a reason for hiding this comment

slyalin Oct 28, 2024 •

edited

Loading

slyalin Oct 28, 2024 •

edited

Loading