[perf][cpu] Accelerate paged attention GEMMs (QK, PV) on Arm CPUs with NEON

[fadara01]

[perf][cpu] Accelerate paged attention GEMMs (QK, PV) on Arm CPUs with NEON

NEON is an Arm SIMD instruction set extension, compulsory since armv8-a

Purpose

Fixes #28981 for Arm CPUs
Related to #23934 (since it enables attention sinks for Arm)

PR #27954 added cpu_attention_with_kv_cache which supports chucked prefill, prefix caching, SWA, alibi, softcap and sinks.

However, it's currently disabled for the prefill phase on Arm CPUs because it's slower than torch.sdpa for relatively long prefills. Hence chunked prefill, prefix caching, sinks, etc remained disabled on Arm.

This PR accelerates cpu_attention_with_kv_cache on Arm CPUs by introducing NEON accelerated GEMMs (enabled with ISA::NEON) for QK and PV. With the new GEMMs, performance of cpu_attention_with_kv_cache is similar to torch.sdpa for long prefills, which allows us to enable cpu_attention_with_kv_cache for prefill path on Arm and therefore enable chunked prefill, prefix caching, sinks, alibi, softcap, etc.

Performance

Uplift with ISA::NEON vs ISA::VEC:
For batch size = 64, query tokens = kv tokens = 512, q heads = 32, kv heads = 8, head size = 128, block size = 128: using ISA::NEON for cpu_attention_with_kv_cache accelerates prefill attention by 2x compared to the current state with ISA::VEC

For the throughput benchmark below on Arm Neoverse-V2, using cpu_attention_with_kv_cache for prefills and decodes: ISA::NEON yields ~ %13 higher throughput than ISA::VEC and similar throughput to using torch.sdpa for prefill.

export VLLM_CPU_OMP_THREADS_BIND=0-63
export LD_PRELOAD="/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4:/usr/lib/aarch64-linux-gnu/libgomp.so.1"
export VLLM_TARGET_DEVICE=cpu
export VLLM_CPU_KVCACHE_SPACE=64
vllm bench throughput \
  --num-prompts 128 \
  --seed 0 \
  --dataset-name sharegpt \
  --input-len 1024 \
  --output-len 128 \
  --max-model-len 2048 \
  --max-num-batched-tokens 8192 \
  --model  meta-llama/Llama-3.1-8B-Instruct \
  --load-format dummy

Test Plan

./run-cpu-test-arm.sh
which includes tests/kernels/attention/test_cpu_attn.py

Test Result

All tests pass

Future Work

This PR enables a solid reference path for attention GEMMs on Arm CPUs.
Future PRs will accelerate attention further by introducing faster/vectorized exp implementations and leveraging bfmmla/bfdot for QK, PV on Arm CPUs with bf16 support, which should yield another 2x speedup for attention.

---

Essential Elements of an Effective PR Description Checklist

• [Y] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• [Y] The test plan, such as providing test command.
• [Y] The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces NEON acceleration for Arm CPU Attention GEMMs, which is a significant performance improvement. The changes are well-structured, with a new cpu_attn_neon.hpp file containing the optimized kernels and modifications in other files to integrate the new ISA path. The NEON implementation itself is solid, using intrinsics and unrolling to achieve better performance. I've found an important issue regarding naming clarity in the new NEON implementation that should be addressed to improve maintainability.

[fadara01]

@mgoin @bigPYJ1151

Can you guys have a look? this deprecates torch.sdpa prefill path for Arm CPUs and enables chunked prefill, prefix caching, sinks, etc

[fadara01]

Hmmm CI used to pass, current failures are unrelated.

[mgoin]

This works great on M1 Mac, validated with a Qwen3-0.6B eval on GSM8k. LGTM!