@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 730bd35378
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 16:00 UTC

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

Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>

Fadi Arafeh committed Nov 22, 2025 at 17:04 UTC 730bd35378bf2a5b56b6d3a45be28b3092d26519
5 files changed +416 -5
csrc/cpu/cpu_attn.cpp
+17
@@ -13,6 +13,18 @@
13 #define AMX_DISPATCH(...) case cpu_attention::ISA::AMX:
14 #endif
15
16 +#ifdef __aarch64__
17 + #include "cpu_attn_neon.hpp"
18 + #define NEON_DISPATCH(...) \
19 + case cpu_attention::ISA::NEON: { \
20 + using attn_impl = cpu_attention::AttentionImpl<cpu_attention::ISA::NEON, \
21 + scalar_t, head_dim>; \
22 + return __VA_ARGS__(); \
23 + }
24 +#else
25 + #define NEON_DISPATCH(...) case cpu_attention::ISA::NEON:
26 +#endif // #ifdef __aarch64__
27 +
28 #define CPU_ATTN_DISPATCH_CASE(HEAD_DIM, ...) \
29 case HEAD_DIM: { \
30 constexpr size_t head_dim = HEAD_DIM; \
@@ -41,6 +53,7 @@
53 [&] { \
54 switch (ISA_TYPE) { \
55 AMX_DISPATCH(__VA_ARGS__) \
56 + NEON_DISPATCH(__VA_ARGS__) \
57 case cpu_attention::ISA::VEC: { \
58 using attn_impl = \
59 cpu_attention::AttentionImpl<cpu_attention::ISA::VEC, scalar_t, \
@@ -73,6 +86,8 @@ torch::Tensor get_scheduler_metadata(
86 isa = cpu_attention::ISA::VEC;
87 } else if (isa_hint == "vec16") {
88 isa = cpu_attention::ISA::VEC16;
89 + } else if (isa_hint == "neon") {
90 + isa = cpu_attention::ISA::NEON;
91 } else {
92 TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
93 }
@@ -158,6 +173,8 @@ void cpu_attn_reshape_and_cache(
173 return cpu_attention::ISA::VEC;
174 } else if (isa == "vec16") {
175 return cpu_attention::ISA::VEC16;
176 + } else if (isa == "neon") {
177 + return cpu_attention::ISA::NEON;
178 } else {
179 TORCH_CHECK(false, "Invalid ISA type: " + isa);
180 }
csrc/cpu/cpu_attn_impl.hpp
+7 -1
@@ -14,7 +14,7 @@
14 #include "utils.hpp"
15
16 namespace cpu_attention {
17 -enum class ISA { AMX, VEC, VEC16 };
17 +enum class ISA { AMX, VEC, VEC16, NEON };
18
19 template <ISA isa, typename scalar_t, int64_t head_dim>
20 class AttentionImpl {};
@@ -143,6 +143,12 @@ struct AttentionMetadata {
143 case ISA::VEC:
144 ss << "VEC, ";
145 break;
146 + case ISA::VEC16:
147 + ss << "VEC16, ";
148 + break;
149 + case ISA::NEON:
150 + ss << "NEON, ";
151 + break;
152 }
153 ss << "workitem_group_num: " << workitem_group_num
154 << ", reduction_item_num: " << reduction_item_num
csrc/cpu/cpu_attn_neon.hpp new
+386
@@ -0,0 +1,386 @@
1 +#ifndef CPU_ATTN_NEON_HPP
2 +#define CPU_ATTN_NEON_HPP
3 +
4 +#include "cpu_attn_impl.hpp"
5 +#include <arm_neon.h>
6 +#include <type_traits>
7 +namespace cpu_attention {
8 +
9 +namespace {
10 +
11 +#define BLOCK_SIZE_ALIGNMENT 32
12 +#define HEAD_SIZE_ALIGNMENT 32
13 +#define MAX_Q_HEAD_NUM_PER_ITER 16
14 +
15 +// These do not use vectorized class for loading / converting
16 +// because csrc/cpu/cpu_types_arm.hpp does not have fallback options
17 +// for vec_op::BF16Vec* / vec_op::BF16Vec* on Arm HW that
18 +// doesn't support BF16.
19 +// We don't use vec_op::FP32Vec* or vec_op::FP16Vec* for consistency.
20 +template <typename kv_cache_t>
21 +FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, float32x4_t& b0,
22 + float32x4_t& b1);
23 +
24 +template <>
25 +FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, float32x4_t& b0,
26 + float32x4_t& b1) {
27 + b0 = vld1q_f32(p + 0);
28 + b1 = vld1q_f32(p + 4);
29 +}
30 +
31 +template <>
32 +FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
33 + float32x4_t& b0,
34 + float32x4_t& b1) {
35 + const float16_t* h = reinterpret_cast<const float16_t*>(p);
36 + float16x8_t v = vld1q_f16(h);
37 + b0 = vcvt_f32_f16(vget_low_f16(v));
38 + b1 = vcvt_f32_f16(vget_high_f16(v));
39 +}
40 +
41 +template <>
42 +FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
43 + float32x4_t& b0,
44 + float32x4_t& b1) {
45 + const uint16_t* u = reinterpret_cast<const uint16_t*>(p);
46 +#ifdef ARM_BF16_SUPPORT
47 + uint16x8_t u0 = vld1q_u16(u);
48 + bfloat16x8_t bf0 = vreinterpretq_bf16_u16(u0);
49 + b0 = vcvtq_low_f32_bf16(bf0);
50 + b1 = vcvtq_high_f32_bf16(bf0);
51 +#else
52 + uint16x8_t x0 = vld1q_u16(u);
53 + uint32x4_t lo = vshlq_n_u32(vmovl_u16(vget_low_u16(x0)), 16);
54 + uint32x4_t hi = vshlq_n_u32(vmovl_u16(vget_high_u16(x0)), 16);
55 + b0 = vreinterpretq_f32_u32(lo);
56 + b1 = vreinterpretq_f32_u32(hi);
57 +#endif
58 +}
59 +
60 +// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs
61 +// #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
62 +// #FMLAs = (K // 4) * (4 * 2 * M)
63 +// We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
64 +template <int32_t M, typename kv_cache_t>
65 +FORCE_INLINE void gemm_micro_neon_fmla_Mx8_Ku4(
66 + const float* __restrict A, // [M x K],
67 + const kv_cache_t* __restrict B, // [K x 8],
68 + float* __restrict C, // [M x 8],
69 + int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
70 + // kernel supports max M of 8, as it'd spill for larger M
71 + static_assert(1 <= M && M <= 8, "M must be in [1,8]");
72 +
73 +// helpers for per-M codegen
74 +#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
75 +#define IF_M(i) if constexpr (M > (i))
76 +
77 + // A row base pointers
78 +#define DECL_A(i) const float* a##i = A + (i) * lda;
79 + ROWS_APPLY(DECL_A)
80 +#undef DECL_A
81 +
82 + // declare 2 accumulators per row of M
83 +#define DECL_ACC(i) float32x4_t acc##i##_0, acc##i##_1;
84 + ROWS_APPLY(DECL_ACC)
85 +#undef DECL_ACC
86 +
87 + // initialize accumulators
88 +#define INIT_ACC(i) \
89 + IF_M(i) { \
90 + if (accumulate) { \
91 + acc##i##_0 = vld1q_f32(C + (i) * ldc + 0); \
92 + acc##i##_1 = vld1q_f32(C + (i) * ldc + 4); \
93 + } else { \
94 + acc##i##_0 = vdupq_n_f32(0.f); \
95 + acc##i##_1 = vdupq_n_f32(0.f); \
96 + } \
97 + }
98 + ROWS_APPLY(INIT_ACC)
99 +#undef INIT_ACC
100 +
101 + int32_t k = 0;
102 +
103 + // K unrolled by 4
104 + for (; k + 3 < K; k += 4) {
105 + // load A[k..k+3] for each active row (M)
106 +#define LOAD_A4(i) \
107 + float32x4_t a##i##v; \
108 + IF_M(i) a##i##v = vld1q_f32(a##i + k);
109 + ROWS_APPLY(LOAD_A4)
110 +#undef LOAD_A4
111 +
112 + // helper: FMA lane L from aiv
113 +#define FMAS_LANE(i, aiv, L) \
114 + IF_M(i) { \
115 + acc##i##_0 = vfmaq_laneq_f32(acc##i##_0, b0, aiv, L); \
116 + acc##i##_1 = vfmaq_laneq_f32(acc##i##_1, b1, aiv, L); \
117 + }
118 +
119 + // k + 0
120 + {
121 + float32x4_t b0, b1;
122 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
123 +#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
124 + ROWS_APPLY(STEP_K0)
125 +#undef STEP_K0
126 + }
127 + // k + 1
128 + {
129 + float32x4_t b0, b1;
130 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
131 +#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
132 + ROWS_APPLY(STEP_K1)
133 +#undef STEP_K1
134 + }
135 + // k + 2
136 + {
137 + float32x4_t b0, b1;
138 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
139 +#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
140 + ROWS_APPLY(STEP_K2)
141 +#undef STEP_K2
142 + }
143 + // k + 3
144 + {
145 + float32x4_t b0, b1;
146 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
147 +#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
148 + ROWS_APPLY(STEP_K3)
149 +#undef STEP_K3
150 + }
151 +#undef FMAS_LANE
152 + }
153 +
154 + // K tail
155 + for (; k < K; ++k) {
156 + float32x4_t b0, b1;
157 + load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
158 +#define TAIL_ROW(i) \
159 + IF_M(i) { \
160 + float32x4_t ai = vdupq_n_f32(*(a##i + k)); \
161 + acc##i##_0 = vfmaq_f32(acc##i##_0, b0, ai); \
162 + acc##i##_1 = vfmaq_f32(acc##i##_1, b1, ai); \
163 + }
164 + ROWS_APPLY(TAIL_ROW)
165 +#undef TAIL_ROW
166 + }
167 +
168 + // store accumulators to C
169 +#define STORE_ROW(i) \
170 + IF_M(i) { \
171 + vst1q_f32(C + (i) * ldc + 0, acc##i##_0); \
172 + vst1q_f32(C + (i) * ldc + 4, acc##i##_1); \
173 + }
174 + ROWS_APPLY(STORE_ROW)
175 +#undef STORE_ROW
176 +
177 +#undef ROWS_APPLY
178 +#undef IF_M
179 +}
180 +
181 +template <int32_t N, typename kv_cache_t>
182 +FORCE_INLINE void gemm_macro_neon_fmla_Mx8_Ku4(const float* __restrict A,
183 + const kv_cache_t* __restrict B,
184 + float* __restrict C, int32_t M,
185 + int32_t K, int64_t lda,
186 + int64_t ldb, int64_t ldc,
187 + bool accumulate) {
188 + // micro kernel is Mx8
189 + static_assert(N % 8 == 0, "N must be a multiple of 8");
190 + for (int32_t m = 0; m < M;) {
191 + int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
192 + const float* Ab = A + m * lda;
193 + float* Cb = C + m * ldc;
194 +
195 + for (int32_t n = 0; n < N; n += 8) {
196 + const kv_cache_t* Bn = B + n;
197 + float* Cn = Cb + n;
198 + switch (mb) {
199 + case 8:
200 + gemm_micro_neon_fmla_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
201 + K, accumulate);
202 + break;
203 + case 4:
204 + gemm_micro_neon_fmla_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
205 + K, accumulate);
206 + break;
207 + case 2:
208 + gemm_micro_neon_fmla_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
209 + K, accumulate);
210 + break;
211 + default:
212 + gemm_micro_neon_fmla_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
213 + K, accumulate);
214 + break;
215 + }
216 + }
217 + // no tail loop for N as it's guaranteed to be a multiple of 8
218 + m += mb;
219 + }
220 +}
221 +
222 +template <typename kv_cache_t>
223 +class TileGemmNeonFMLA {
224 + public:
225 + template <AttentionGemmPhase phase, int32_t k_size>
226 + FORCE_INLINE static void gemm(const int32_t m_size,
227 + float* __restrict__ a_tile,
228 + kv_cache_t* __restrict__ b_tile,
229 + float* __restrict__ c_tile, const int64_t lda,
230 + const int64_t ldb, const int64_t ldc,
231 + const int32_t block_size,
232 + const int32_t dynamic_k_size,
233 + const bool accum_c) {
234 + if constexpr (phase == AttentionGemmPhase::QK) {
235 + gemm_macro_neon_fmla_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
236 + a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
237 + } else {
238 + gemm_macro_neon_fmla_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
239 + a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
240 + accum_c);
241 + }
242 + }
243 +};
244 +
245 +} // namespace
246 +
247 +// this is similar to "ISA::VEC" at the moment
248 +template <typename scalar_t, int64_t head_dim>
249 +class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
250 + public:
251 + using query_t = scalar_t;
252 + using q_buffer_t = float;
253 + using kv_cache_t = scalar_t;
254 + using logits_buffer_t = float;
255 + using partial_output_buffer_t = float;
256 + using prob_buffer_t = float;
257 +
258 + constexpr static int64_t BlockSizeAlignment =
259 + BLOCK_SIZE_ALIGNMENT; // KV token num unit of QK and PV phases
260 + constexpr static int64_t HeadDimAlignment =
261 + HEAD_SIZE_ALIGNMENT; // headdim num unit of PV phase
262 + constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
263 + constexpr static int64_t HeadDim = head_dim;
264 + constexpr static ISA ISAType = ISA::NEON;
265 + constexpr static bool scale_on_logits = false; // apply scale on q_buffer
266 +
267 + static_assert(HeadDim % HeadDimAlignment == 0);
268 + // the gemm micro kernel is Mx8
269 + static_assert(HeadDimAlignment % 8 == 0);
270 + static_assert(BlockSizeAlignment % 8 == 0);
271 +
272 + public:
273 + template <template <typename tile_gemm_t> typename attention>
274 + FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
275 + attention<TileGemmNeonFMLA<kv_cache_t>> attention_iteration;
276 + attention_iteration(CPU_ATTENTION_PARAMS);
277 + }
278 +
279 + // k_cache_token_group_stride: stride of K cache when move to next
280 + // BlockSizeAlignment tokens in a block
281 + constexpr static int64_t k_cache_token_group_stride(
282 + const int32_t block_size) {
283 + return BlockSizeAlignment; // layout of k_cache block is [head_dim,
284 + // block_size], row-major
285 + }
286 +
287 + // v_cache_token_group_stride: stride of V cache when move to next
288 + // BlockSizeAlignment tokens in a block
289 + constexpr static int64_t v_cache_token_group_stride(
290 + const int32_t block_size) {
291 + return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
292 + // head_dim], row-major
293 + }
294 +
295 + // v_cache_head_group_stride: stride of V cache when move to next
296 + // HeadDimAlignment head dims in a block
297 + constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
298 + return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
299 + // row-major
300 + }
301 +
302 + // Copy q to q_buffer and cast it to fp32
303 + static void copy_q_heads_tile(
304 + scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
305 + float* __restrict__ q_buffer, const int32_t q_num,
306 + const int32_t q_heads_per_kv, const int64_t q_num_stride,
307 + const int64_t q_head_stride, float scale) {
308 + static_assert(head_dim % 16 == 0);
309 + constexpr int32_t unroll_size = head_dim / 16;
310 + using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
311 +
312 + vec_op::FP32Vec16 scale_vec(scale);
313 + for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
314 + for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
315 + scalar_t* __restrict__ curr_q =
316 + src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
317 + float* __restrict__ curr_q_buffer =
318 + q_buffer + q_num_idx * q_heads_per_kv * head_dim +
319 + q_head_idx * head_dim;
320 +
321 + vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
322 + load_vec_t vec(curr_q);
323 + vec_op::FP32Vec16 fp32_vec(vec);
324 + fp32_vec = fp32_vec * scale_vec;
325 + fp32_vec.save(curr_q_buffer);
326 +
327 + curr_q += 16;
328 + curr_q_buffer += 16;
329 + });
330 + }
331 + }
332 + }
333 +
334 + // reshape K as column-major and V as row-major
335 + static void reshape_and_cache(
336 + const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
337 + scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
338 + const int64_t* __restrict__ slot_mapping, const int64_t token_num,
339 + const int64_t key_token_num_stride, const int64_t value_token_num_stride,
340 + const int64_t head_num, const int64_t key_head_num_stride,
341 + const int64_t value_head_num_stride, const int64_t num_blocks,
342 + const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
343 + const int64_t block_size, const int64_t block_size_stride) {
344 +#pragma omp parallel for collapse(2)
345 + for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
346 + for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
347 + const int64_t pos = slot_mapping[token_idx];
348 + if (pos < 0) {
349 + // skip
350 + continue;
351 + }
352 +
353 + const int64_t block_idx = pos / block_size;
354 + const int64_t block_offset = pos % block_size;
355 + {
356 + // Write Key
357 + const scalar_t* key_start_ptr = key +
358 + token_idx * key_token_num_stride +
359 + head_idx * key_head_num_stride;
360 + scalar_t* key_cache_start_ptr =
361 + key_cache + block_idx * num_blocks_stride +
362 + head_idx * cache_head_num_stride + block_offset;
363 +
364 +#pragma GCC unroll 8
365 + for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
366 + key_cache_start_ptr[j] = key_start_ptr[i];
367 + }
368 + }
369 + {
370 + // Write Value
371 + const scalar_t* value_start_ptr = value +
372 + token_idx * value_token_num_stride +
373 + head_idx * value_head_num_stride;
374 + scalar_t* value_cache_start_ptr =
375 + value_cache + block_idx * num_blocks_stride +
376 + head_idx * cache_head_num_stride + block_offset * head_dim;
377 + std::memcpy(value_cache_start_ptr, value_start_ptr,
378 + sizeof(scalar_t) * head_dim);
379 + }
380 + }
381 + }
382 + }
383 +};
384 +} // namespace cpu_attention
385 +
386 +#endif // #ifndef CPU_ATTN_NEON_HPP
vllm/engine/arg_utils.py
+1 -2
@@ -1392,11 +1392,10 @@ class EngineArgs:
1392 # Set default arguments for V1 Engine.
1393 self._set_default_args(usage_context, model_config)
1394 # Disable chunked prefill and prefix caching for:
1395 - # POWER (ppc64le)/ARM/s390x/RISCV CPUs in V1
1395 + # POWER (ppc64le)/s390x/RISCV CPUs in V1
1396 if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
1397 CpuArchEnum.POWERPC,
1398 CpuArchEnum.S390X,
1399 - CpuArchEnum.ARM,
1399 CpuArchEnum.RISCV,
1400 ):
1401 logger.info(
vllm/v1/attention/backends/cpu_attn.py
+5 -2
@@ -25,7 +25,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec
25
26 logger = init_logger(__name__)
27
28 -_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86,)
28 +_CPU_ARCH_PREFER_MIXED_BATCH = (CpuArchEnum.X86, CpuArchEnum.ARM)
29
30
31 class CPUAttentionBackend(AttentionBackend):
@@ -491,6 +491,9 @@ def _get_attn_isa(dtype: torch.dtype, block_size: int) -> str:
491 if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
492 return "amx"
493 elif block_size % 32 == 0:
494 - return "vec"
494 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
495 + return "neon"
496 + else:
497 + return "vec"
498 else:
499 return "vec16"