[CPU][Perf] refactor paged attention for Arm CPUs (#56045)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Fadi Arafeh committed
Sep 20, 2026 at 05:25 UTC
92b40f5a144763c8b2ef6ebe848793de634ada05
4 files changed
+401
-506
csrc/cpu/cpu_attn_impl.hpp
+36
-13
@@ -956,6 +956,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
956
std::printf("%s", ss.str().c_str());
957
}
958
959
+// This allows us to store probabilities in a packed format
960
+// for GEMM implementations that need packed A.
961
+template <typename attention_impl_t, typename = void>
962
+struct ProbabilityTokenStore {
963
+ using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
964
+ static constexpr int32_t TokenStride = 1;
965
+
966
+ FORCE_INLINE static void store_probabilities(
967
+ prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values,
968
+ const int32_t, const int64_t) {
969
+ using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
970
+ prob_buffer_vec_t output_vec(values);
971
+ output_vec.save(probability);
972
+ }
973
+};
974
+
975
+template <typename attention_impl_t>
976
+struct ProbabilityTokenStore<
977
+ attention_impl_t,
978
+ std::void_t<typename attention_impl_t::ProbabilityTokenStore>>
979
+ : attention_impl_t::ProbabilityTokenStore {};
980
+
981
template <typename attention_impl_t>
982
class AttentionMainLoop {
983
public:
@@ -966,6 +988,7 @@ class AttentionMainLoop {
988
using partial_output_buffer_t =
989
typename attention_impl_t::partial_output_buffer_t;
990
using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
991
+ using probability_token_store_t = ProbabilityTokenStore<attention_impl_t>;
992
993
static constexpr int64_t max_q_head_num_per_iter =
994
attention_impl_t::MaxQHeadNumPerIteration;
@@ -974,6 +997,8 @@ class AttentionMainLoop {
997
static constexpr int64_t headdim_alignment =
998
attention_impl_t::HeadDimAlignment;
999
static constexpr int64_t head_dim = attention_impl_t::HeadDim;
1000
+ static constexpr int32_t probability_token_stride =
1001
+ probability_token_store_t::TokenStride;
1002
static constexpr ISA ISAType = attention_impl_t::ISAType;
1003
static constexpr bool scale_on_logits =
1004
attention_impl_t::scale_on_logits; // apply scale on logits, otherwise
@@ -1152,8 +1177,6 @@ class AttentionMainLoop {
1177
? kv_tile_token_num
1178
: kv_tile_token_num *
1179
(sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
1155
- constexpr int64_t prob_buffer_elem_size =
1156
- prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);
1180
partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
1181
bool accum_c = !is_first_iter;
1182
for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
@@ -1187,9 +1210,7 @@ class AttentionMainLoop {
1210
remaining_group_num -= curr_group_num_in_block;
1211
curr_group_offset = 0;
1212
curr_group_num_in_block = token_group_num_per_block;
1190
- curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(
1191
- reinterpret_cast<uint8_t*>(curr_prob_buffer) +
1192
- curr_token_num * prob_buffer_elem_size);
1213
+ curr_prob_buffer += curr_token_num * probability_token_stride;
1214
curr_partial_q_buffer = partial_q_buffer;
1215
accum_c = true;
1216
}
@@ -1262,7 +1283,6 @@ class AttentionMainLoop {
1283
std::is_same_v<query_t, c10::Half>;
1284
#endif
1285
1265
- using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
1286
constexpr bool prequantize_probabilities = []() {
1287
if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
1288
return tile_gemm_t::prequantize_probabilities;
@@ -1331,12 +1351,6 @@ class AttentionMainLoop {
1351
vec = fast_exp(vec);
1352
}
1353
1334
- if constexpr (prequantize_probabilities) {
1335
- vec.save(curr_logits_buffer_iter);
1336
- } else {
1337
- prob_buffer_vec_t output_vec(vec);
1338
- output_vec.save(curr_prob_buffer_iter);
1339
- }
1354
#else
1355
vec.save(curr_logits_buffer_iter);
1356
for (int32_t k = 0; k < 16; ++k) {
@@ -1345,10 +1359,19 @@ class AttentionMainLoop {
1359
vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
1360
#endif
1361
1362
+ if constexpr (prequantize_probabilities) {
1363
+ vec.save(curr_logits_buffer_iter);
1364
+ } else {
1365
+ probability_token_store_t::store_probabilities(
1366
+ curr_prob_buffer_iter, vec, i,
1367
+ logits_buffer_stride * sizeof(logits_buffer_t) /
1368
+ sizeof(prob_buffer_t));
1369
+ }
1370
+
1371
sum_vec = sum_vec + vec;
1372
1373
curr_logits_buffer_iter += 16;
1351
- curr_prob_buffer_iter += 16;
1374
+ curr_prob_buffer_iter += 16 * probability_token_stride;
1375
}
1376
if constexpr (prequantize_probabilities) {
1377
tile_gemm_t::quantize_probability_row(
csrc/cpu/cpu_attn_neon.hpp
+2
-2
@@ -11,7 +11,7 @@ namespace cpu_attention {
11
12
namespace {
13
14
-#define BLOCK_SIZE_ALIGNMENT 32
14
+#define BLOCK_SIZE_ALIGNMENT 16
15
#define HEAD_SIZE_ALIGNMENT 32
16
#define MAX_Q_HEAD_NUM_PER_ITER 16
17
@@ -387,7 +387,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
387
};
388
389
#ifdef ARM_BF16_SUPPORT
390
-// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
390
+// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment.
391
template <int64_t head_dim>
392
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
393
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
csrc/cpu/cpu_attn_neon_bfmmla.hpp
+360
-487
@@ -5,432 +5,275 @@
5
6
#include "cpu_attn_impl.hpp"
7
8
+#include <arm_bf16.h>
9
#include <arm_neon.h>
10
+#include <c10/util/BFloat16.h>
11
12
+#include <algorithm>
13
#include <cstdint>
11
-#include <vector>
14
15
namespace cpu_attention {
16
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;
17
+class BfmmlaGemm {
18
+ public:
19
+ static constexpr int32_t KTile = 4;
20
+ static constexpr int32_t NTile = 8;
21
+ static constexpr int32_t MaxRows = 8;
22
+
23
+ FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a,
24
+ const c10::BFloat16* __restrict__ b,
25
+ float* __restrict__ c, const int32_t m,
26
+ const int32_t n, const int32_t k,
27
+ const int64_t a_pair_stride,
28
+ const int64_t b_n_group_stride,
29
+ const int64_t b_k_group_stride,
30
+ const int64_t ldc, const bool accumulate) {
31
+ const auto* a_ptr = reinterpret_cast<const bfloat16_t*>(a);
32
+ const auto* b_ptr = reinterpret_cast<const bfloat16_t*>(b);
33
+
34
+ for (int32_t n_idx = 0; n_idx < n; n_idx += 16) {
35
+ const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride;
36
+ float* c_panel = c + n_idx;
37
+
38
+ // Preserve this range so the inactive row pair is optimized away.
39
+ if (m <= 2) {
40
+ gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
41
+ b_n_group_stride, b_k_group_stride, ldc, accumulate);
42
+ } else if (m <= 4) {
43
+ gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
44
+ b_n_group_stride, b_k_group_stride, ldc, accumulate);
45
+ } else {
46
+ gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride,
47
+ ldc, accumulate);
48
+ gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k,
49
+ a_pair_stride, b_k_group_stride, ldc, accumulate);
50
+ }
51
}
52
}
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;
53
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;
54
+ private:
55
+ FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a,
56
+ const float32x4_t b) {
57
+ return vreinterpretq_f32_f64(
58
+ vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
59
}
60
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); \
61
+ FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a,
62
+ const float32x4_t b) {
63
+ return vreinterpretq_f32_f64(
64
+ vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
65
}
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);
66
+
67
+ FORCE_INLINE static void init_accumulators(
68
+ float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45,
69
+ float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc,
70
+ const int32_t rows, const bool accumulate) {
71
+ if (!accumulate || rows == 0) {
72
+ acc01 = vdupq_n_f32(0.0f);
73
+ acc23 = vdupq_n_f32(0.0f);
74
+ acc45 = vdupq_n_f32(0.0f);
75
+ acc67 = vdupq_n_f32(0.0f);
76
+ return;
77
}
78
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
79
+ const float32x4_t row0_0123 = vld1q_f32(c);
80
+ const float32x4_t row0_4567 = vld1q_f32(c + 4);
81
+ const float32x4_t row1_0123 =
82
+ (rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f);
83
+ const float32x4_t row1_4567 =
84
+ (rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f);
85
+
86
+ acc01 = zip_low_pairs(row0_0123, row1_0123);
87
+ acc23 = zip_high_pairs(row0_0123, row1_0123);
88
+ acc45 = zip_low_pairs(row0_4567, row1_4567);
89
+ acc67 = zip_high_pairs(row0_4567, row1_4567);
90
}
91
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
- }
92
+ FORCE_INLINE static void store_accumulators(
93
+ const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45,
94
+ const float32x4_t acc67, float* __restrict__ c, const int64_t ldc,
95
+ const int32_t rows) {
96
+ if (rows == 0) {
97
+ return;
98
+ }
99
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
- }
100
+ vst1q_f32(c, zip_low_pairs(acc01, acc23));
101
+ vst1q_f32(c + 4, zip_low_pairs(acc45, acc67));
102
+ if (rows == 2) {
103
+ vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23));
104
+ vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67));
105
}
106
}
107
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);
108
+ FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a,
109
+ const int32_t rows) {
110
+ if (rows == 0) {
111
+ return vdupq_n_bf16(bfloat16_t{});
112
+ }
113
+ // Packed A reserves both rows for an M tail.
114
+ return vld1q_bf16(a);
115
}
116
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
- }
117
+ FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a,
118
+ const bfloat16_t* __restrict__ b,
119
+ float* __restrict__ c, const int32_t m,
120
+ const int32_t k,
121
+ const int64_t a_pair_stride,
122
+ const int64_t b_n_group_stride,
123
+ const int64_t b_k_group_stride,
124
+ const int64_t ldc, const bool accumulate) {
125
+ const int32_t rows01 = std::min(2, std::max(0, m));
126
+ const int32_t rows23 = std::min(2, std::max(0, m - 2));
127
+ float32x4_t acc0101, acc0123, acc0145, acc0167;
128
+ float32x4_t acc2301, acc2323, acc2345, acc2367;
129
+ float32x4_t acc0189, acc011011, acc011213, acc011415;
130
+ float32x4_t acc2389, acc231011, acc231213, acc231415;
131
+ init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
132
+ accumulate);
133
+ init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
134
+ rows23, accumulate);
135
+ init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
136
+ rows01, accumulate);
137
+ init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8,
138
+ ldc, rows23, accumulate);
139
+
140
+ const bfloat16_t* a01 = a;
141
+ const bfloat16_t* a23 = a + a_pair_stride;
142
+ const bfloat16_t* b0 = b;
143
+ const bfloat16_t* b1 = b + b_n_group_stride;
144
+
145
+#pragma GCC unroll 4
146
+ for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
147
+ const bfloat16x8_t av01 = load_a_pair(a01, rows01);
148
+ const bfloat16x8_t av23 = load_a_pair(a23, rows23);
149
+ const bfloat16x8_t b01 = vld1q_bf16(b0);
150
+ const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile);
151
+ const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile);
152
+ const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile);
153
+ const bfloat16x8_t b89 = vld1q_bf16(b1);
154
+ const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile);
155
+ const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile);
156
+ const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile);
157
+
158
+ acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
159
+ acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
160
+ acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
161
+ acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
162
+ acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
163
+ acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
164
+ acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
165
+ acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
166
+ acc0189 = vbfmmlaq_f32(acc0189, av01, b89);
167
+ acc2389 = vbfmmlaq_f32(acc2389, av23, b89);
168
+ acc011011 = vbfmmlaq_f32(acc011011, av01, b1011);
169
+ acc231011 = vbfmmlaq_f32(acc231011, av23, b1011);
170
+ acc011213 = vbfmmlaq_f32(acc011213, av01, b1213);
171
+ acc231213 = vbfmmlaq_f32(acc231213, av23, b1213);
172
+ acc011415 = vbfmmlaq_f32(acc011415, av01, b1415);
173
+ acc231415 = vbfmmlaq_f32(acc231415, av23, b1415);
174
+
175
+ a01 += 2 * KTile;
176
+ a23 += 2 * KTile;
177
+ b0 += b_k_group_stride;
178
+ b1 += b_k_group_stride;
179
+ }
180
348
- PROCESS_RP_GROUP(4);
349
- PROCESS_RP_GROUP(2);
350
- PROCESS_RP_GROUP(1);
351
-#undef PROCESS_RP_GROUP
181
+ store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
182
+ store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
183
+ rows23);
184
+ store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
185
+ rows01);
186
+ store_accumulators(acc2389, acc231011, acc231213, acc231415,
187
+ c + 2 * ldc + 8, ldc, rows23);
188
}
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
189
+
190
+ FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a,
191
+ const bfloat16_t* __restrict__ b,
192
+ float* __restrict__ c, const int32_t m,
193
+ const int32_t k,
194
+ const int64_t a_pair_stride,
195
+ const int64_t b_k_group_stride,
196
+ const int64_t ldc, const bool accumulate) {
197
+ const int32_t rows01 = std::min(2, std::max(0, m));
198
+ const int32_t rows23 = std::min(2, std::max(0, m - 2));
199
+ const int32_t rows45 = std::min(2, std::max(0, m - 4));
200
+ const int32_t rows67 = std::min(2, std::max(0, m - 6));
201
+ float32x4_t acc0101, acc0123, acc0145, acc0167;
202
+ float32x4_t acc2301, acc2323, acc2345, acc2367;
203
+ float32x4_t acc4501, acc4523, acc4545, acc4567;
204
+ float32x4_t acc6701, acc6723, acc6745, acc6767;
205
+ init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
206
+ accumulate);
207
+ init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
208
+ rows23, accumulate);
209
+ init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
210
+ rows45, accumulate);
211
+ init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
212
+ rows67, accumulate);
213
+
214
+ const bfloat16_t* a01 = a;
215
+ const bfloat16_t* a23 = a + a_pair_stride;
216
+ const bfloat16_t* a45 = a + 2 * a_pair_stride;
217
+ const bfloat16_t* a67 = a + 3 * a_pair_stride;
218
+ const bfloat16_t* b_ptr = b;
219
+
220
+#pragma GCC unroll 4
221
+ for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
222
+ const bfloat16x8_t av01 = load_a_pair(a01, rows01);
223
+ const bfloat16x8_t av23 = load_a_pair(a23, rows23);
224
+ const bfloat16x8_t av45 = load_a_pair(a45, rows45);
225
+ const bfloat16x8_t av67 = load_a_pair(a67, rows67);
226
+ const bfloat16x8_t b01 = vld1q_bf16(b_ptr);
227
+ const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile);
228
+ const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile);
229
+ const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile);
230
+
231
+ acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
232
+ acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
233
+ acc4501 = vbfmmlaq_f32(acc4501, av45, b01);
234
+ acc6701 = vbfmmlaq_f32(acc6701, av67, b01);
235
+ acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
236
+ acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
237
+ acc4523 = vbfmmlaq_f32(acc4523, av45, b23);
238
+ acc6723 = vbfmmlaq_f32(acc6723, av67, b23);
239
+ acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
240
+ acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
241
+ acc4545 = vbfmmlaq_f32(acc4545, av45, b45);
242
+ acc6745 = vbfmmlaq_f32(acc6745, av67, b45);
243
+ acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
244
+ acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
245
+ acc4567 = vbfmmlaq_f32(acc4567, av45, b67);
246
+ acc6767 = vbfmmlaq_f32(acc6767, av67, b67);
247
+
248
+ a01 += 2 * KTile;
249
+ a23 += 2 * KTile;
250
+ a45 += 2 * KTile;
251
+ a67 += 2 * KTile;
252
+ b_ptr += b_k_group_stride;
253
}
254
+
255
+ store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
256
+ store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
257
+ rows23);
258
+ store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
259
+ rows45);
260
+ store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
261
+ rows67);
262
}
424
-}
263
+};
264
+
265
+namespace {
266
426
-#undef INIT_ACC_ROWPAIR_4
427
-#undef STORE_ACC_ROWPAIR_4
428
-#undef BFMMLA_COMPUTE_4
267
+constexpr int32_t TILE_K = BfmmlaGemm::KTile;
268
+constexpr int32_t TILE_COLS = 2;
269
+constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile;
270
+constexpr int32_t K_TOKENS_PER_GROUP = 8;
271
+constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;
272
+constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
273
+constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
274
275
} // namespace
276
432
-// TileGemm Adapter for Attention
433
-
277
template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
278
class TileGemmNEONBFMMLA {
279
public:
@@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA {
286
[[maybe_unused]] const int32_t block_size,
287
[[maybe_unused]] const int32_t dynamic_k_size,
288
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.
289
+ static_assert(BlockTokens % 16 == 0);
290
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.");
291
+ static_assert(head_dim_ct == HeadDim);
292
}
293
294
+ const auto* a = reinterpret_cast<const c10::BFloat16*>(a_tile);
295
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);
296
+ constexpr int64_t b_n_group_stride =
297
+ (HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE;
298
+
299
+ for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
300
+ const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
301
+ BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m,
302
+ BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride,
303
+ K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
304
+ }
305
} 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;
306
+ const int64_t b_n_group_stride =
307
+ (block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE;
308
+
309
+ for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
310
+ const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
311
+ BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m,
312
+ HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride,
313
+ K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
314
}
315
}
316
}
@@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA {
337
static constexpr int64_t HeadDim = head_dim;
338
static constexpr ISA ISAType = isa_type;
339
static constexpr bool scale_on_logits = false;
340
+ static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK;
341
+ static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K;
342
516
- static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
517
- static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
343
+ static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0);
344
+ static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
345
static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
346
347
public:
@@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA {
354
attention_iteration(CPU_ATTENTION_PARAMS);
355
}
356
357
+ struct ProbabilityTokenStore {
358
+ static constexpr int32_t TokenStride = 2;
359
+
360
+ FORCE_INLINE static void store_probabilities(
361
+ c10::BFloat16* __restrict__ probability,
362
+ const vec_op::FP32Vec16& values, const int32_t row,
363
+ const int64_t row_stride) {
364
+ const int32_t row_in_pair = row & 1;
365
+ auto* dst = reinterpret_cast<bfloat16_t*>(
366
+ probability + row_in_pair * (BfmmlaGemm::KTile - row_stride));
367
+ vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0]));
368
+ vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1]));
369
+ vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2]));
370
+ vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3]));
371
+ }
372
+ };
373
+
374
// Key cache stride per token group (TokenColumn layout; QK)
375
static constexpr int64_t k_cache_token_group_stride(
376
[[maybe_unused]] const int32_t block_size) {
377
static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
378
return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
535
- ((head_dim / TILE_K) * K_INNER_STRIDE);
379
+ ((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE);
380
}
381
382
// Value cache stride per token group (TokenRow layout; PV)
383
static constexpr int64_t v_cache_token_group_stride(
384
[[maybe_unused]] const int32_t block_size) {
385
static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
542
- return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
386
+ return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
387
}
388
389
// The stride to move to the "next" head_dim group
@@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA {
394
return head_dim * block_size;
395
}
396
553
- // Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
397
+ // Scale Q and write row pairs in BFMMLA reduction order.
398
static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
399
c10::BFloat16* __restrict__ q_buffer,
400
const int32_t q_num,
@@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA {
403
const int64_t q_head_stride, float scale) {
404
constexpr int32_t dim = static_cast<int32_t>(head_dim);
405
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);
406
+ const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{});
407
+ const int32_t row_num = q_num * q_heads_per_kv;
408
+
409
+ for (int32_t row = 0; row < row_num; row += 2) {
410
+ const int32_t q0 = row / q_heads_per_kv;
411
+ const int32_t h0 = row % q_heads_per_kv;
412
+ const auto* row0 = reinterpret_cast<const bfloat16_t*>(
413
+ src + q0 * q_num_stride + h0 * q_head_stride);
414
+ const bool has_row1 = row + 1 < row_num;
415
+ const int32_t q1 = (row + 1) / q_heads_per_kv;
416
+ const int32_t h1 = (row + 1) % q_heads_per_kv;
417
+ const auto* row1 = has_row1
418
+ ? reinterpret_cast<const bfloat16_t*>(
419
+ src + q1 * q_num_stride + h1 * q_head_stride)
420
+ : nullptr;
421
+ auto* dst = reinterpret_cast<bfloat16_t*>(q_buffer + row * head_dim);
422
+
423
+ for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) {
424
+ const bfloat16x8_t in0 = vld1q_bf16(row0 + k);
425
+ const bfloat16x4_t out0_lo =
426
+ vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec));
427
+ const bfloat16x4_t out0_hi =
428
+ vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec));
429
+
430
+ bfloat16x4_t out1_lo = zero;
431
+ bfloat16x4_t out1_hi = zero;
432
+ if (has_row1) {
433
+ const bfloat16x8_t in1 = vld1q_bf16(row1 + k);
434
+ out1_lo =
435
+ vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec));
436
+ out1_hi =
437
+ vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec));
438
}
439
+
440
+ vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo));
441
+ vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi));
442
}
443
}
444
}
@@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA {
446
public:
447
// Reshape and cache K/V into BFMMLA-optimized layouts
448
// K cache:
588
- // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
449
+ // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K]
450
+ // [K_CACHE_K_GROUP_STRIDE]
451
// - TokenColumn
452
// V cache:
591
- // [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
592
- // - TokenRows
453
+ // [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK]
454
+ // [VCacheKGroupStride]
455
static void reshape_and_cache(
456
const c10::BFloat16* __restrict__ key,
457
const c10::BFloat16* __restrict__ value,
@@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA {
466
const int64_t block_size,
467
[[maybe_unused]] const int64_t block_size_stride,
468
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
607
- const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
608
- const int64_t v_pair_stride =
609
- (block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
469
+ const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE;
470
+ const int64_t v_n_group_stride =
471
+ (block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
472
611
-#pragma omp parallel for
612
- for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
613
- for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
473
+#pragma omp parallel for collapse(2)
474
+ for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
475
+ for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
476
const int64_t pos = slot_mapping[token_idx];
477
if (pos < 0) continue;
478
@@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA {
498
key_base + block_in_block * k_block_stride;
499
500
for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
639
- uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
640
- block_base + hd4 * K_INNER_STRIDE +
641
- pair_in_block * V_INNER_STRIDE + lane_base);
642
- const uint16_t* src_u16 =
643
- reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
644
- vst1_u16(dst_u16, vld1_u16(src_u16));
501
+ bfloat16_t* dst = reinterpret_cast<bfloat16_t*>(
502
+ block_base + hd4 * K_CACHE_K_GROUP_STRIDE +
503
+ pair_in_block * B_COL_PAIR_STRIDE + lane_base);
504
+ const bfloat16_t* src =
505
+ reinterpret_cast<const bfloat16_t*>(key_src + hd4 * TILE_K);
506
+ vst1_bf16(dst, vld1_bf16(src));
507
}
508
}
509
@@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA {
517
value_cache + block_idx * num_blocks_stride +
518
head_idx * cache_head_num_stride;
519
658
- const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
520
+ const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK;
521
const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
522
661
- c10::BFloat16* __restrict row_block_base =
662
- value_base + row_block * V_INNER_STRIDE;
663
-
664
- for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
665
- c10::BFloat16* __restrict dst_val =
666
- row_block_base + hd2 * v_pair_stride;
667
-
668
- const uint16_t* src_u16 =
669
- reinterpret_cast<const uint16_t*>(value_src);
670
- uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
671
- dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
672
- dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
673
- src_u16[hd2 * TILE_COLS + 1];
523
+ const auto* src = reinterpret_cast<const bfloat16_t*>(value_src);
524
+ auto* dst = reinterpret_cast<bfloat16_t*>(value_base);
525
+ for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) {
526
+ const bfloat16x8_t values =
527
+ vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK);
528
+ const bfloat16x4_t low = vget_low_bf16(values);
529
+ const bfloat16x4_t high = vget_high_bf16(values);
530
+ bfloat16_t* group =
531
+ dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride;
532
+
533
+ vst1_lane_bf16(group + lane, low, 0);
534
+ vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1);
535
+ vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2);
536
+ vst1_lane_bf16(
537
+ group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low,
538
+ 3);
539
+ vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0);
540
+ vst1_lane_bf16(
541
+ group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
542
+ high, 1);
543
+ vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2);
544
+ vst1_lane_bf16(
545
+ group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
546
+ high, 3);
547
}
548
}
549
}
@@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA {
553
554
} // namespace cpu_attention
555
683
-#endif // CPU_ATTN_ASIMD_BFMMLA_HPP
556
+#endif // CPU_ATTN_NEON_BFMMLA_HPP
vllm/v1/attention/backends/cpu_attn.py
+3
-4
@@ -512,11 +512,10 @@ def _get_attn_isa(
512
)
513
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
514
return "amx"
515
+ elif supports_arm:
516
+ return "neon"
517
elif block_size % 32 == 0:
516
- if supports_arm:
517
- # support ARM NEON FMLA and BFMMLA (bf16) for block size 32
518
- return "neon"
519
- elif supports_riscv and _riscv_supports_rvv():
518
+ if supports_riscv and _riscv_supports_rvv():
519
return "rvv"
520
elif supports_vxe:
521
return "vxe"