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 8, 19:00 UTC
[cpu][performance] CPU Paged Attention NEON BFMMLA BF16 Implementation (#32263)
Signed-off-by: Gassan <gassan.salama@arm.com>
Gassan Salama committed
Feb 6, 2026 at 07:01 UTC
1363e3d6d5659b58376fa5284afc2c8be548cc9d
4 files changed
+704
-4
csrc/cpu/cpu_attn_impl.hpp
+2
-1
@@ -1107,7 +1107,8 @@ class AttentionMainLoop {
1107
if (sliding_window_left != -1) {
1108
pos = std::max(pos, curr_token_pos - sliding_window_left);
1109
}
1110
- return pos;
1110
+ // Clamp to tile end to avoid OOB when window starts past the tile
1111
+ return std::min(pos, kv_tile_end_pos);
1112
}();
1113
1114
int32_t right_kv_pos = [&]() {
csrc/cpu/cpu_attn_neon.hpp
+17
-2
@@ -4,6 +4,9 @@
4
#include "cpu_attn_impl.hpp"
5
#include <arm_neon.h>
6
#include <type_traits>
7
+#ifdef ARM_BF16_SUPPORT
8
+ #include "cpu_attn_neon_bfmmla.hpp"
9
+#endif
10
namespace cpu_attention {
11
12
namespace {
@@ -57,7 +60,7 @@ FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
60
#endif
61
}
62
60
-// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with NEON FMLAs
63
+// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with ASIMD FMLAs
64
// #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
65
// #FMLAs = (K // 4) * (4 * 2 * M)
66
// We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
@@ -381,6 +384,18 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
384
}
385
}
386
};
387
+
388
+#ifdef ARM_BF16_SUPPORT
389
+// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
390
+template <int64_t head_dim>
391
+class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
392
+ : public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
393
+ head_dim> {};
394
+#endif
395
} // namespace cpu_attention
396
386
-#endif // #ifndef CPU_ATTN_NEON_HPP
397
+#undef BLOCK_SIZE_ALIGNMENT
398
+#undef HEAD_SIZE_ALIGNMENT
399
+#undef MAX_Q_HEAD_NUM_PER_ITER
400
+
401
+#endif // #ifndef CPU_ATTN_ASIMD_HPP
csrc/cpu/cpu_attn_neon_bfmmla.hpp
new
+682
@@ -0,0 +1,682 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#ifndef CPU_ATTN_NEON_BFMMLA_HPP
4
+#define CPU_ATTN_NEON_BFMMLA_HPP
5
+
6
+#include "cpu_attn_impl.hpp"
7
+
8
+#include <arm_neon.h>
9
+
10
+#include <cstdint>
11
+#include <vector>
12
+
13
+namespace cpu_attention {
14
+
15
+namespace {
16
+
17
+// BFMMLA tile dimensions
18
+constexpr int32_t TILE_ROWS = 2; // M dimension
19
+constexpr int32_t TILE_K = 4; // K reduction
20
+constexpr int32_t TILE_COLS = 2; // N dimension (column-pair)
21
+
22
+// Derived constants
23
+constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs
24
+constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache
25
+constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block
26
+constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
27
+constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
28
+constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing
29
+
30
+// Matrix Packing and Accumulator
31
+// Reshape two rows of Q into BFMMLA-friendly interleaved
32
+// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
33
+// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
34
+// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
35
+FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
36
+ const c10::BFloat16* __restrict r1,
37
+ c10::BFloat16* __restrict dst,
38
+ int32_t K) {
39
+ const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
40
+ const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
41
+ uint16_t* d = reinterpret_cast<uint16_t*>(dst);
42
+
43
+ // Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
44
+ int32_t k = 0;
45
+ for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
46
+ vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
47
+ }
48
+
49
+ // Handle K tail: pack remaining elements with zero-padding
50
+ const int32_t tail = K - k;
51
+ if (tail > 0) {
52
+ // Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
53
+ // pad]
54
+ for (int32_t t = 0; t < tail; ++t) {
55
+ d[t] = s0[k + t];
56
+ d[t + TILE_K] = s1[k + t];
57
+ }
58
+ // Zero-pad the rest
59
+ for (int32_t t = tail; t < TILE_K; ++t) {
60
+ d[t] = 0;
61
+ d[t + TILE_K] = 0;
62
+ }
63
+ }
64
+}
65
+
66
+// 2x2 accumulator load/store with compile-time row count
67
+template <int32_t m_rows>
68
+FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
69
+ static_assert(m_rows == 1 || m_rows == 2);
70
+ float32x2_t row0 = vld1_f32(base + col_off);
71
+ float32x2_t row1 =
72
+ (m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
73
+ return vcombine_f32(row0, row1);
74
+}
75
+
76
+template <int32_t m_rows>
77
+FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
78
+ int col_off) {
79
+ static_assert(m_rows == 1 || m_rows == 2);
80
+ vst1_f32(base + col_off, vget_low_f32(acc));
81
+ if constexpr (m_rows == 2) {
82
+ vst1_f32(base + ldc + col_off, vget_high_f32(acc));
83
+ }
84
+}
85
+
86
+// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
87
+#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
88
+ do { \
89
+ if (accum) { \
90
+ if (m_rows == 2) { \
91
+ a0 = load_acc_2x2<2>(Crow, ldc, 0); \
92
+ a1 = load_acc_2x2<2>(Crow, ldc, 2); \
93
+ a2 = load_acc_2x2<2>(Crow, ldc, 4); \
94
+ a3 = load_acc_2x2<2>(Crow, ldc, 6); \
95
+ } else { \
96
+ a0 = load_acc_2x2<1>(Crow, ldc, 0); \
97
+ a1 = load_acc_2x2<1>(Crow, ldc, 2); \
98
+ a2 = load_acc_2x2<1>(Crow, ldc, 4); \
99
+ a3 = load_acc_2x2<1>(Crow, ldc, 6); \
100
+ } \
101
+ } else { \
102
+ a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \
103
+ } \
104
+ } while (0)
105
+
106
+// Store 4 column-pair accumulators back to C matrix
107
+#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
108
+ do { \
109
+ if (m_rows == 2) { \
110
+ store_acc_2x2<2>(a0, Crow, ldc, 0); \
111
+ store_acc_2x2<2>(a1, Crow, ldc, 2); \
112
+ store_acc_2x2<2>(a2, Crow, ldc, 4); \
113
+ store_acc_2x2<2>(a3, Crow, ldc, 6); \
114
+ } else { \
115
+ store_acc_2x2<1>(a0, Crow, ldc, 0); \
116
+ store_acc_2x2<1>(a1, Crow, ldc, 2); \
117
+ store_acc_2x2<1>(a2, Crow, ldc, 4); \
118
+ store_acc_2x2<1>(a3, Crow, ldc, 6); \
119
+ } \
120
+ } while (0)
121
+
122
+// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
123
+#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
124
+ do { \
125
+ r0 = vbfmmlaq_f32(r0, a, b0); \
126
+ r1 = vbfmmlaq_f32(r1, a, b1); \
127
+ r2 = vbfmmlaq_f32(r2, a, b2); \
128
+ r3 = vbfmmlaq_f32(r3, a, b3); \
129
+ } while (0)
130
+
131
+// Micro-kernel: updates a small fixed tile using BFMMLA.
132
+// RP = number of row-pairs (1,2,4)
133
+// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
134
+// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
135
+// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
136
+// - AttentionGemmPhase::PV -> token-row layout (P @ V)
137
+// K_static < 0 enables runtime K (PV only)
138
+template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
139
+FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
140
+ const bfloat16_t* const* __restrict A_packed_rp,
141
+ const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
142
+ float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
143
+ int32_t K_runtime = 0) {
144
+ static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
145
+ static_assert(K_static < 0 || K_static % TILE_K == 0,
146
+ "K must be divisible by TILE_K");
147
+ static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
148
+ "Runtime K only supported for PV");
149
+
150
+ constexpr bool runtime_k = (K_static < 0);
151
+ const int32_t K_iters =
152
+ runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
153
+ const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
154
+
155
+ if (!runtime_k) {
156
+ // Help the compiler fold away unused K_runtime when K is compile-time
157
+ (void)K_runtime;
158
+ }
159
+
160
+ auto* C_al = C;
161
+ const auto* B_al = B_blk;
162
+
163
+ // Setup A pointers
164
+ const bfloat16_t* a_ptr[4] = {
165
+ A_packed_rp[0],
166
+ (RP >= 2) ? A_packed_rp[1] : nullptr,
167
+ (RP >= 4) ? A_packed_rp[2] : nullptr,
168
+ (RP >= 4) ? A_packed_rp[3] : nullptr,
169
+ };
170
+
171
+ // Setup B pointers based on layout
172
+ const bfloat16_t* b_ptr[4];
173
+ if constexpr (phase == AttentionGemmPhase::PV) {
174
+ b_ptr[0] = B_blk + 0 * b_stride;
175
+ b_ptr[1] = B_blk + 1 * b_stride;
176
+ b_ptr[2] = B_blk + 2 * b_stride;
177
+ b_ptr[3] = B_blk + 3 * b_stride;
178
+ }
179
+
180
+ float32x4_t acc[4][4];
181
+
182
+// Initialize accumulators
183
+#define INIT_RP(rp) \
184
+ if constexpr (RP > rp) { \
185
+ INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
186
+ C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
187
+ }
188
+ INIT_RP(0);
189
+ INIT_RP(1);
190
+ INIT_RP(2);
191
+ INIT_RP(3);
192
+#undef INIT_RP
193
+
194
+ // Main compute loop
195
+ for (int32_t ki = 0; ki < K_iters; ++ki) {
196
+ bfloat16x8_t b0, b1, b2, b3;
197
+ if constexpr (phase == AttentionGemmPhase::PV) {
198
+ b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
199
+ b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
200
+ b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
201
+ b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
202
+ } else {
203
+ const bfloat16_t* b_base = B_al + ki * b_stride;
204
+ b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
205
+ b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
206
+ b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
207
+ b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
208
+ }
209
+
210
+#define COMPUTE_RP(rp) \
211
+ if constexpr (RP > rp) { \
212
+ bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
213
+ BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \
214
+ b1, b2, b3); \
215
+ }
216
+ COMPUTE_RP(0);
217
+ COMPUTE_RP(1);
218
+ COMPUTE_RP(2);
219
+ COMPUTE_RP(3);
220
+#undef COMPUTE_RP
221
+ }
222
+
223
+ // K tail for runtime PV: fallback path
224
+ if constexpr (runtime_k) {
225
+ if (K_tail > 0) {
226
+ const int32_t tail_offset = K_iters * V_INNER_STRIDE;
227
+ const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
228
+ for (int32_t kt = 0; kt < K_tail; ++kt) {
229
+ float32x4_t b_vecs[4];
230
+ for (int32_t p = 0; p < 4; ++p) {
231
+ const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
232
+ const float b0 = vcvtah_f32_bf16(bp[0]);
233
+ const float b1 = vcvtah_f32_bf16(bp[1]);
234
+ const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
235
+ b_vecs[p] = vcombine_f32(b_pair, b_pair);
236
+ }
237
+
238
+#define TAIL_RP(rp) \
239
+ if constexpr (RP > rp) { \
240
+ const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \
241
+ float a_row0 = vcvtah_f32_bf16(ap[kt]); \
242
+ float a_row1 = \
243
+ (m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
244
+ const float32x4_t a_vec = \
245
+ vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \
246
+ for (int32_t p = 0; p < 4; ++p) { \
247
+ acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \
248
+ } \
249
+ }
250
+ TAIL_RP(0);
251
+ TAIL_RP(1);
252
+ TAIL_RP(2);
253
+ TAIL_RP(3);
254
+#undef TAIL_RP
255
+ }
256
+ }
257
+ }
258
+
259
+ // Store results
260
+#define STORE_RP(rp) \
261
+ if constexpr (RP > rp) { \
262
+ STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
263
+ C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \
264
+ }
265
+ STORE_RP(0);
266
+ STORE_RP(1);
267
+ STORE_RP(2);
268
+ STORE_RP(3);
269
+#undef STORE_RP
270
+}
271
+
272
+// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
273
+// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
274
+// runtime K (PV only).
275
+template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
276
+FORCE_INLINE void gemm_packA_compute_MB_xN(
277
+ const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
278
+ float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
279
+ int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
280
+ static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
281
+ static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
282
+ "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
283
+ static_assert(K_static < 0 || K_static % TILE_K == 0,
284
+ "K must be divisible by TILE_K");
285
+ static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
286
+ "Runtime K only supported for PV");
287
+
288
+ constexpr bool runtime_k = (K_static < 0);
289
+ const int32_t K_val = runtime_k ? K_runtime : K_static;
290
+
291
+ // Keep small packs on-stack to avoid heap churn
292
+ constexpr int32_t STACK_PACK_STRIDE =
293
+ (1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
294
+
295
+ constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
296
+ const int32_t pack_stride =
297
+ runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
298
+ : (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
299
+
300
+ alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
301
+ std::vector<c10::BFloat16> A_packed_heap;
302
+ c10::BFloat16* A_packed =
303
+ (pack_stride <= STACK_PACK_STRIDE)
304
+ ? A_packed_stack
305
+ : (A_packed_heap.resize(ROW_PAIRS * pack_stride),
306
+ A_packed_heap.data());
307
+
308
+ for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
309
+ const int32_t m = rp * TILE_ROWS;
310
+ const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
311
+ const c10::BFloat16* A0 = A + m * lda;
312
+ const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
313
+ reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
314
+ }
315
+
316
+ for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
317
+ const c10::BFloat16* B_blk_c10 =
318
+ (phase == AttentionGemmPhase::PV)
319
+ ? (B + (n / TILE_COLS) * b_layout_stride)
320
+ : (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
321
+ const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
322
+
323
+ // Process row-pairs in groups of 4, 2, then 1
324
+ int32_t row_pair_idx = 0;
325
+
326
+#define PROCESS_RP_GROUP(group_size) \
327
+ for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \
328
+ row_pair_idx += group_size) { \
329
+ const bfloat16_t* Ap[group_size]; \
330
+ int32_t mr[group_size]; \
331
+ for (int32_t i = 0; i < group_size; ++i) { \
332
+ Ap[i] = reinterpret_cast<const bfloat16_t*>( \
333
+ A_packed + (row_pair_idx + i) * pack_stride); \
334
+ mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
335
+ } \
336
+ float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \
337
+ if constexpr (runtime_k) { \
338
+ gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>( \
339
+ Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \
340
+ } else { \
341
+ gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>( \
342
+ Ap, mr, B_blk, C_blk, ldc, accumulate, \
343
+ (phase == AttentionGemmPhase::PV) ? b_layout_stride \
344
+ : b_reduction_stride); \
345
+ } \
346
+ }
347
+
348
+ PROCESS_RP_GROUP(4);
349
+ PROCESS_RP_GROUP(2);
350
+ PROCESS_RP_GROUP(1);
351
+#undef PROCESS_RP_GROUP
352
+ }
353
+}
354
+
355
+// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
356
+// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
357
+// (runtime K path is only supported for PV).
358
+template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
359
+FORCE_INLINE void gemm_macro_neon_bfmmla(
360
+ const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
361
+ float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
362
+ int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
363
+ static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
364
+ "N must be a multiple of OUTPUT_COLS_PER_BLOCK");
365
+
366
+ if constexpr (K >= 0) {
367
+ static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
368
+ for (int32_t m = 0; m < M;) {
369
+ const int32_t rem = M - m;
370
+ const c10::BFloat16* A_blk = A + m * lda;
371
+ float* C_blk = C + m * ldc;
372
+
373
+#define DISPATCH_MB(mb) \
374
+ gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
375
+ b_layout_stride, \
376
+ b_reduction_stride, accumulate)
377
+
378
+ if (rem >= 8) {
379
+ DISPATCH_MB(8);
380
+ m += 8;
381
+ } else if (rem >= 4) {
382
+ DISPATCH_MB(4);
383
+ m += 4;
384
+ } else if (rem >= 2) {
385
+ DISPATCH_MB(2);
386
+ m += 2;
387
+ } else {
388
+ DISPATCH_MB(1);
389
+ m += 1;
390
+ }
391
+#undef DISPATCH_MB
392
+ }
393
+ } else {
394
+ static_assert(phase == AttentionGemmPhase::PV,
395
+ "Runtime K specialization only supported for PV.");
396
+ const int32_t K_val = K_runtime;
397
+
398
+ for (int32_t m = 0; m < M;) {
399
+ const int32_t rem = M - m;
400
+ const c10::BFloat16* A_blk = A + m * lda;
401
+ float* C_blk = C + m * ldc;
402
+
403
+#define DISPATCH_MB_RUNTIME(mb) \
404
+ gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
405
+ b_layout_stride, \
406
+ b_reduction_stride, accumulate)
407
+
408
+ if (rem >= 8) {
409
+ DISPATCH_MB_RUNTIME(8);
410
+ m += 8;
411
+ } else if (rem >= 4) {
412
+ DISPATCH_MB_RUNTIME(4);
413
+ m += 4;
414
+ } else if (rem >= 2) {
415
+ DISPATCH_MB_RUNTIME(2);
416
+ m += 2;
417
+ } else {
418
+ DISPATCH_MB_RUNTIME(1);
419
+ m += 1;
420
+ }
421
+#undef DISPATCH_MB_RUNTIME
422
+ }
423
+ }
424
+}
425
+
426
+#undef INIT_ACC_ROWPAIR_4
427
+#undef STORE_ACC_ROWPAIR_4
428
+#undef BFMMLA_COMPUTE_4
429
+
430
+} // namespace
431
+
432
+// TileGemm Adapter for Attention
433
+
434
+template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
435
+class TileGemmNEONBFMMLA {
436
+ public:
437
+ template <AttentionGemmPhase phase, int32_t head_dim_ct>
438
+ FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
439
+ kv_cache_t* __restrict__ b_tile,
440
+ float* __restrict__ c_tile, const int64_t lda,
441
+ [[maybe_unused]] const int64_t ldb,
442
+ const int64_t ldc,
443
+ [[maybe_unused]] const int32_t block_size,
444
+ [[maybe_unused]] const int32_t dynamic_k_size,
445
+ const bool accum_c) {
446
+ static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
447
+ // BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
448
+ // API parity with other tile_gemm implementations.
449
+ if constexpr (head_dim_ct >= 0) {
450
+ static_assert(head_dim_ct == HeadDim,
451
+ "BFMMLA expects head_dim_ct to match HeadDim; PV passes "
452
+ "-1 for API parity.");
453
+ }
454
+
455
+ if constexpr (phase == AttentionGemmPhase::QK) {
456
+ const int64_t b_reduction_stride = K_INNER_STRIDE;
457
+ const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
458
+
459
+ gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
460
+ reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
461
+ m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
462
+ accum_c);
463
+ } else {
464
+ const int64_t b_pair_stride =
465
+ (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
466
+
467
+ // PV gemm with runtime K specialization
468
+ switch (dynamic_k_size) {
469
+ case 32:
470
+ gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
471
+ reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
472
+ m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
473
+ break;
474
+ case 128:
475
+ gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
476
+ reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
477
+ m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
478
+ break;
479
+ case 256:
480
+ gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
481
+ reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
482
+ m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
483
+ break;
484
+ default:
485
+ gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
486
+ reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
487
+ m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
488
+ break;
489
+ }
490
+ }
491
+ }
492
+};
493
+
494
+// Shared ASIMD BFMMLA implementation (BF16 only). The block size alignment and
495
+// ISA tag are template parameters so we can reuse the same kernels for
496
+// different NEON configurations.
497
+template <int64_t block_size_alignment, ISA isa_type, int64_t head_dim>
498
+class AttentionImplNEONBFMMLA {
499
+ public:
500
+ using query_t = c10::BFloat16;
501
+ using q_buffer_t = c10::BFloat16;
502
+ using kv_cache_t = c10::BFloat16;
503
+ using logits_buffer_t = float;
504
+ using partial_output_buffer_t = float;
505
+ using prob_buffer_t = c10::BFloat16;
506
+
507
+ static constexpr int64_t BlockSizeAlignment = block_size_alignment;
508
+ // HeadDimAlignment equals head_dim so that the PV phase processes
509
+ // the full head dimension in a single gemm call.
510
+ static constexpr int64_t HeadDimAlignment = head_dim;
511
+ static constexpr int64_t MaxQHeadNumPerIteration = 16;
512
+ static constexpr int64_t HeadDim = head_dim;
513
+ static constexpr ISA ISAType = isa_type;
514
+ static constexpr bool scale_on_logits = false;
515
+
516
+ static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
517
+ static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
518
+ static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
519
+
520
+ public:
521
+ template <template <typename tile_gemm_t> typename attention>
522
+ FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
523
+ attention<
524
+ TileGemmNEONBFMMLA<kv_cache_t, static_cast<int32_t>(BlockSizeAlignment),
525
+ static_cast<int32_t>(HeadDim)>>
526
+ attention_iteration;
527
+ attention_iteration(CPU_ATTENTION_PARAMS);
528
+ }
529
+
530
+ // Key cache stride per token group (TokenColumn layout; QK)
531
+ static constexpr int64_t k_cache_token_group_stride(
532
+ [[maybe_unused]] const int32_t block_size) {
533
+ static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
534
+ return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
535
+ ((head_dim / TILE_K) * K_INNER_STRIDE);
536
+ }
537
+
538
+ // Value cache stride per token group (TokenRow layout; PV)
539
+ static constexpr int64_t v_cache_token_group_stride(
540
+ [[maybe_unused]] const int32_t block_size) {
541
+ static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
542
+ return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
543
+ }
544
+
545
+ // The stride to move to the "next" head_dim group
546
+ // is the full V cache size per head, since HeadDimAlignment == head_dim.
547
+ // Hence, the stride is not used in this case
548
+ static constexpr int64_t v_cache_head_group_stride(
549
+ [[maybe_unused]] const int32_t block_size) {
550
+ return head_dim * block_size;
551
+ }
552
+
553
+ // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
554
+ static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
555
+ c10::BFloat16* __restrict__ q_buffer,
556
+ const int32_t q_num,
557
+ const int32_t q_heads_per_kv,
558
+ const int64_t q_num_stride,
559
+ const int64_t q_head_stride, float scale) {
560
+ constexpr int32_t dim = static_cast<int32_t>(head_dim);
561
+ const float32x4_t scale_vec = vdupq_n_f32(scale);
562
+
563
+ for (int32_t qi = 0; qi < q_num; ++qi) {
564
+ for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
565
+ c10::BFloat16* __restrict__ curr_q =
566
+ src + qi * q_num_stride + hi * q_head_stride;
567
+ c10::BFloat16* __restrict__ dst =
568
+ q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
569
+
570
+ for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
571
+ bfloat16x8_t in8 =
572
+ vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
573
+ float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
574
+ float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
575
+
576
+ bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
577
+ bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
578
+ bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
579
+ vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
580
+ }
581
+ }
582
+ }
583
+ }
584
+
585
+ public:
586
+ // Reshape and cache K/V into BFMMLA-optimized layouts
587
+ // K cache:
588
+ // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
589
+ // - TokenColumn
590
+ // V cache:
591
+ // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
592
+ // - TokenRows
593
+ static void reshape_and_cache(
594
+ const c10::BFloat16* __restrict__ key,
595
+ const c10::BFloat16* __restrict__ value,
596
+ c10::BFloat16* __restrict__ key_cache,
597
+ c10::BFloat16* __restrict__ value_cache,
598
+ const int64_t* __restrict__ slot_mapping, const int64_t token_num,
599
+ const int64_t key_token_num_stride, const int64_t value_token_num_stride,
600
+ const int64_t head_num, const int64_t key_head_num_stride,
601
+ const int64_t value_head_num_stride,
602
+ [[maybe_unused]] const int64_t num_blocks,
603
+ const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
604
+ const int64_t block_size,
605
+ [[maybe_unused]] const int64_t block_size_stride) {
606
+ const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
607
+ const int64_t v_pair_stride =
608
+ (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
609
+
610
+#pragma omp parallel for
611
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
612
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
613
+ const int64_t pos = slot_mapping[token_idx];
614
+ if (pos < 0) continue;
615
+
616
+ const int64_t block_idx = pos / block_size;
617
+ const int64_t block_offset = pos % block_size;
618
+
619
+ // Key cache: TokenColumn QK
620
+ {
621
+ const c10::BFloat16* __restrict key_src =
622
+ key + token_idx * key_token_num_stride +
623
+ head_idx * key_head_num_stride;
624
+
625
+ c10::BFloat16* __restrict key_base = key_cache +
626
+ block_idx * num_blocks_stride +
627
+ head_idx * cache_head_num_stride;
628
+
629
+ const int64_t block_in_block = block_offset / K_TOKENS_PER_GROUP;
630
+ const int64_t pair_in_block =
631
+ (block_offset % K_TOKENS_PER_GROUP) / TILE_COLS;
632
+ const int64_t lane_base = (block_offset & 1) ? TILE_K : 0;
633
+
634
+ c10::BFloat16* __restrict block_base =
635
+ key_base + block_in_block * k_block_stride;
636
+
637
+ for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
638
+ uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
639
+ block_base + hd4 * K_INNER_STRIDE +
640
+ pair_in_block * V_INNER_STRIDE + lane_base);
641
+ const uint16_t* src_u16 =
642
+ reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
643
+ vst1_u16(dst_u16, vld1_u16(src_u16));
644
+ }
645
+ }
646
+
647
+ // Value cache: TokenRow PV
648
+ {
649
+ const c10::BFloat16* __restrict value_src =
650
+ value + token_idx * value_token_num_stride +
651
+ head_idx * value_head_num_stride;
652
+
653
+ c10::BFloat16* __restrict value_base =
654
+ value_cache + block_idx * num_blocks_stride +
655
+ head_idx * cache_head_num_stride;
656
+
657
+ const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
658
+ const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
659
+
660
+ c10::BFloat16* __restrict row_block_base =
661
+ value_base + row_block * V_INNER_STRIDE;
662
+
663
+ for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
664
+ c10::BFloat16* __restrict dst_val =
665
+ row_block_base + hd2 * v_pair_stride;
666
+
667
+ const uint16_t* src_u16 =
668
+ reinterpret_cast<const uint16_t*>(value_src);
669
+ uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
670
+ dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
671
+ dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
672
+ src_u16[hd2 * TILE_COLS + 1];
673
+ }
674
+ }
675
+ }
676
+ }
677
+ }
678
+};
679
+
680
+} // namespace cpu_attention
681
+
682
+#endif // CPU_ATTN_ASIMD_BFMMLA_HPP
vllm/v1/attention/backends/cpu_attn.py
+3
-1
@@ -487,10 +487,12 @@ def _get_attn_isa(
487
if head_size is not None and head_size % 32 != 0 and head_size % 16 == 0:
488
return "vec16"
489
supports_amx = torch._C._cpu._is_amx_tile_supported()
490
+ supports_arm = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
491
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
492
return "amx"
493
elif block_size % 32 == 0:
493
- if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
494
+ if supports_arm:
495
+ # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
496
return "neon"
497
else:
498
return "vec"