[Feature Request]: Optimize PagedAttention operation on aarch64 HW

[dmitry-gorokhov]

Request Description

PagedAttention operation is already implemented in bounds of CPU plugin using C++ and optimized for x64 using avx2/avx512 instrinsics.
The request is to optimize PA operation for aarch64 using NEON/SVE extensions.

Please refer to SDPA optimization using NEON for reference.
How to build OV on ARM: https://github.com/openvinotoolkit/openvino/blob/master/docs/dev/build.md

Feature Use Case

PagedAttention operation implements attention algo required for workloads like continuous batching or speculative decoding. PagedAttention is used as basic attention block in VLLM OpenVINO backend and under OpenVINO GenAI API (for some use-cases). PA operation might take significant resources for execution (especially for long contexts), so its optimization is crucial for overall LLM based workloads.

Issue submission checklist

• The feature request or improvement must be related to OpenVINO

[rkazants]

@dmitry-gorokhov, can we add good-first-issue label here?

[samkitshah1262]

Hi @dmitry-gorokhov @rkazants @wenjiew , is anyone working on this issue ? If not can I take it up ?

[mlukasze]

it's your now, have fun :)

[samkitshah1262]

Hi , currently the paged attention executor logic depends on brgemm (only supports x86-64) , so to extend the support for ARM-64 (aarach-64) , which library should we use. following are the prospects - ARM compute library(ACL) , OpenBLAS, Eigen, TFlite.

[dmitry-gorokhov]

Hi @samkitshah1262. OneDNN has Brgemm block implemented for aarch64. The problem is it supports only SVE512/SVE256.
So in order to cover Neon powered devices (like Apple silicon) we need to use another backend. I would propose to take a look on ACL NEGEMM kernel. It was already successfully applied for the same purpose inside SDPA op: #25183

[samkitshah1262]

@dmitry-gorokhov Implementation for PA on ARM using NEON and ACL is done , except the MHAHelper class as the ACL Gemm kernal does not support stride (yet). Could you please review the current changes and provide direction. Thanks!

[ashwins990]

Hi @dmitry-gorokhov,
I am working towards enabling f16 precision for pagedAttention for ARM.
To enable f16 pagedattention node, I made the below changes:

• setting hints::inference_precsion to f16
• forcing f16 infer_precision in file intel_cpu/src/graph.cpp
• made changes to transformation_pipeline.cpp

I found that pagedAttention is still falling back to f32 precision. Could you please refer me to some documentation regarding the same. Thank you.

[dmitry-gorokhov]

@ashwins990 Have you modified this condition

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

Lines 221 to 223 in bd04c5a

	} else if (rtPrecision == ov::element::f16 && ov::with_cpu_x86_avx512_core_fp16()) {
	rtPrecision = ov::element::f16;
	} else {

? It explicitly check for avx512 isa.
You need to do smt like that https://github.com/openvinotoolkit/openvino/blob/master/src/plugins/intel_cpu/src/nodes/scaled_attn.cpp#L1855-L1856

[ashwins990]

Hi @dmitry-gorokhov

Is there a way to perform GEMM f16 for ARM for small matrices as NGEMM ACL have very significant overhead for small matrices.?

Here are the performance analysis comparision.
Model :: LLAMA2-7B on Graviton 3E

Experiment 1 : keeping prompt tokens less than 32, and generation tokens 256:

We get 70% improvement over f32 implementation in total throughput

Experiment 2 : server benchmark on ShareGPT dataset

{prompt and generation tokens are 1:1 ratio approx }

Variant	Total Throughput	throughput %	TTFT	TPOT
f32 with BRGEMM	327	reference	reference	reference
f16 with ACL	232	-29%	-650%	-354%
f16 with custom gemm	318	-3%	-8%	-2%

Analysis

The two GEMM operation performed are q[32 * 128] X k[ 128 X 32 ] and w[32 * 32] X v[ 32 X 128 ].
ACL NGEMM used to carry these multiplication carries the penalty of initialization [esp. for small martrix ] when done multiple times.
Therefore a very basic matrix multiplication is able to produce better results than NGEMM in this case.

[dmitry-gorokhov]

@ashwins990 I would recommend to create PR with current ACL NEGEMM based impl. so we could help with the review and investigation of the possible solutions.
AFAIK, NEGEMM is recommended ACL primitive for matrix multiplication, so the key is to avoid initialization overhead on each iteration (if possible).
Alternative solution might be using KleidiAI microkernels library, which integration we are finalizing now.