iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CPU][Perf] refactor paged attention for Arm CPUs (#56045)

Browse Source 
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>
This commit is contained in:
Fadi Arafeh
2026-09-20 12:25:13 +08:00
committed by GitHub
co-authored by Codex Li, Jiang <jiang1.li@intel.com>
parent 70fc7bd695
commit 92b40f5a14
4 changed files with 414 additions and 519 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/cpu_attn_impl.hpp
+36 -13
View File
@@ -956,6 +956,28 @@ void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
std::printf("%s", ss.str().c_str());
}
// This allows us to store probabilities in a packed format
// for GEMM implementations that need packed A.
template <typename attention_impl_t, typename = void>
struct ProbabilityTokenStore {
using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
static constexpr int32_t TokenStride = 1;
FORCE_INLINE static void store_probabilities(
prob_buffer_t* __restrict__ probability, const vec_op::FP32Vec16& values,
const int32_t, const int64_t) {
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
prob_buffer_vec_t output_vec(values);
output_vec.save(probability);
}
};
template <typename attention_impl_t>
struct ProbabilityTokenStore<
attention_impl_t,
std::void_t<typename attention_impl_t::ProbabilityTokenStore>>
: attention_impl_t::ProbabilityTokenStore {};
template <typename attention_impl_t>
class AttentionMainLoop {
public:
@@ -966,6 +988,7 @@ class AttentionMainLoop {
using partial_output_buffer_t =
typename attention_impl_t::partial_output_buffer_t;
using prob_buffer_t = typename attention_impl_t::prob_buffer_t;
using probability_token_store_t = ProbabilityTokenStore<attention_impl_t>;
static constexpr int64_t max_q_head_num_per_iter =
attention_impl_t::MaxQHeadNumPerIteration;
@@ -974,6 +997,8 @@ class AttentionMainLoop {
static constexpr int64_t headdim_alignment =
attention_impl_t::HeadDimAlignment;
static constexpr int64_t head_dim = attention_impl_t::HeadDim;
static constexpr int32_t probability_token_stride =
probability_token_store_t::TokenStride;
static constexpr ISA ISAType = attention_impl_t::ISAType;
static constexpr bool scale_on_logits =
attention_impl_t::scale_on_logits; // apply scale on logits, otherwise
@@ -1152,8 +1177,6 @@ class AttentionMainLoop {
? kv_tile_token_num
: kv_tile_token_num *
(sizeof(logits_buffer_t) / sizeof(prob_buffer_t));
constexpr int64_t prob_buffer_elem_size =
prequantize_probabilities ? sizeof(uint8_t) : sizeof(prob_buffer_t);
partial_output_buffer_t* curr_partial_q_buffer = partial_q_buffer;
bool accum_c = !is_first_iter;
for (int32_t block_idx = start_block_idx; block_idx < end_block_idx;
@@ -1187,9 +1210,7 @@ class AttentionMainLoop {
remaining_group_num -= curr_group_num_in_block;
curr_group_offset = 0;
curr_group_num_in_block = token_group_num_per_block;
curr_prob_buffer = reinterpret_cast<pv_prob_buffer_t*>(
reinterpret_cast<uint8_t*>(curr_prob_buffer) +
curr_token_num * prob_buffer_elem_size);
curr_prob_buffer += curr_token_num * probability_token_stride;
curr_partial_q_buffer = partial_q_buffer;
accum_c = true;
}
@@ -1262,7 +1283,6 @@ class AttentionMainLoop {
std::is_same_v<query_t, c10::Half>;
#endif
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
constexpr bool prequantize_probabilities = []() {
if constexpr (requires { tile_gemm_t::prequantize_probabilities; }) {
return tile_gemm_t::prequantize_probabilities;
@@ -1331,12 +1351,6 @@ class AttentionMainLoop {
vec = fast_exp(vec);
}
if constexpr (prequantize_probabilities) {
vec.save(curr_logits_buffer_iter);
} else {
prob_buffer_vec_t output_vec(vec);
output_vec.save(curr_prob_buffer_iter);
}
#else
vec.save(curr_logits_buffer_iter);
for (int32_t k = 0; k < 16; ++k) {
@@ -1345,10 +1359,19 @@ class AttentionMainLoop {
vec = vec_op::FP32Vec16(curr_logits_buffer_iter);
#endif
if constexpr (prequantize_probabilities) {
vec.save(curr_logits_buffer_iter);
} else {
probability_token_store_t::store_probabilities(
curr_prob_buffer_iter, vec, i,
logits_buffer_stride * sizeof(logits_buffer_t) /
sizeof(prob_buffer_t));
}
sum_vec = sum_vec + vec;
curr_logits_buffer_iter += 16;
curr_prob_buffer_iter += 16;
curr_prob_buffer_iter += 16 * probability_token_stride;
}
if constexpr (prequantize_probabilities) {
tile_gemm_t::quantize_probability_row(
csrc/cpu/cpu_attn_neon.hpp
+2 -2
View File
@@ -11,7 +11,7 @@ namespace cpu_attention {
namespace {
#define BLOCK_SIZE_ALIGNMENT 32
#define BLOCK_SIZE_ALIGNMENT 16
#define HEAD_SIZE_ALIGNMENT 32
#define MAX_Q_HEAD_NUM_PER_ITER 16
@@ -387,7 +387,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
};
#ifdef ARM_BF16_SUPPORT
// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
// For BF16 on Arm, reuse the BFMMLA kernels with the same token alignment.
template <int64_t head_dim>
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
csrc/cpu/cpu_attn_neon_bfmmla.hpp
+373 -500
View File
@@ -5,432 +5,275 @@
#include "cpu_attn_impl.hpp"
#include <arm_bf16.h>
#include <arm_neon.h>
#include <c10/util/BFloat16.h>
#include <algorithm>
#include <cstdint>
#include <vector>
namespace cpu_attention {
class BfmmlaGemm {
public:
static constexpr int32_t KTile = 4;
static constexpr int32_t NTile = 8;
static constexpr int32_t MaxRows = 8;
FORCE_INLINE static void gemm(const c10::BFloat16* __restrict__ a,
const c10::BFloat16* __restrict__ b,
float* __restrict__ c, const int32_t m,
const int32_t n, const int32_t k,
const int64_t a_pair_stride,
const int64_t b_n_group_stride,
const int64_t b_k_group_stride,
const int64_t ldc, const bool accumulate) {
const auto* a_ptr = reinterpret_cast<const bfloat16_t*>(a);
const auto* b_ptr = reinterpret_cast<const bfloat16_t*>(b);
for (int32_t n_idx = 0; n_idx < n; n_idx += 16) {
const auto* b_panel = b_ptr + (n_idx / NTile) * b_n_group_stride;
float* c_panel = c + n_idx;
// Preserve this range so the inactive row pair is optimized away.
if (m <= 2) {
gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
b_n_group_stride, b_k_group_stride, ldc, accumulate);
} else if (m <= 4) {
gemm_4x16(a_ptr, b_panel, c_panel, m, k, a_pair_stride,
b_n_group_stride, b_k_group_stride, ldc, accumulate);
} else {
gemm_8x8(a_ptr, b_panel, c_panel, m, k, a_pair_stride, b_k_group_stride,
ldc, accumulate);
gemm_8x8(a_ptr, b_panel + b_n_group_stride, c_panel + NTile, m, k,
a_pair_stride, b_k_group_stride, ldc, accumulate);
}
}
}
private:
FORCE_INLINE static float32x4_t zip_low_pairs(const float32x4_t a,
const float32x4_t b) {
return vreinterpretq_f32_f64(
vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
}
FORCE_INLINE static float32x4_t zip_high_pairs(const float32x4_t a,
const float32x4_t b) {
return vreinterpretq_f32_f64(
vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
}
FORCE_INLINE static void init_accumulators(
float32x4_t& acc01, float32x4_t& acc23, float32x4_t& acc45,
float32x4_t& acc67, const float* __restrict__ c, const int64_t ldc,
const int32_t rows, const bool accumulate) {
if (!accumulate || rows == 0) {
acc01 = vdupq_n_f32(0.0f);
acc23 = vdupq_n_f32(0.0f);
acc45 = vdupq_n_f32(0.0f);
acc67 = vdupq_n_f32(0.0f);
return;
}
const float32x4_t row0_0123 = vld1q_f32(c);
const float32x4_t row0_4567 = vld1q_f32(c + 4);
const float32x4_t row1_0123 =
(rows == 2) ? vld1q_f32(c + ldc) : vdupq_n_f32(0.0f);
const float32x4_t row1_4567 =
(rows == 2) ? vld1q_f32(c + ldc + 4) : vdupq_n_f32(0.0f);
acc01 = zip_low_pairs(row0_0123, row1_0123);
acc23 = zip_high_pairs(row0_0123, row1_0123);
acc45 = zip_low_pairs(row0_4567, row1_4567);
acc67 = zip_high_pairs(row0_4567, row1_4567);
}
FORCE_INLINE static void store_accumulators(
const float32x4_t acc01, const float32x4_t acc23, const float32x4_t acc45,
const float32x4_t acc67, float* __restrict__ c, const int64_t ldc,
const int32_t rows) {
if (rows == 0) {
return;
}
vst1q_f32(c, zip_low_pairs(acc01, acc23));
vst1q_f32(c + 4, zip_low_pairs(acc45, acc67));
if (rows == 2) {
vst1q_f32(c + ldc, zip_high_pairs(acc01, acc23));
vst1q_f32(c + ldc + 4, zip_high_pairs(acc45, acc67));
}
}
FORCE_INLINE static bfloat16x8_t load_a_pair(const bfloat16_t* __restrict__ a,
const int32_t rows) {
if (rows == 0) {
return vdupq_n_bf16(bfloat16_t{});
}
// Packed A reserves both rows for an M tail.
return vld1q_bf16(a);
}
FORCE_INLINE static void gemm_4x16(const bfloat16_t* __restrict__ a,
const bfloat16_t* __restrict__ b,
float* __restrict__ c, const int32_t m,
const int32_t k,
const int64_t a_pair_stride,
const int64_t b_n_group_stride,
const int64_t b_k_group_stride,
const int64_t ldc, const bool accumulate) {
const int32_t rows01 = std::min(2, std::max(0, m));
const int32_t rows23 = std::min(2, std::max(0, m - 2));
float32x4_t acc0101, acc0123, acc0145, acc0167;
float32x4_t acc2301, acc2323, acc2345, acc2367;
float32x4_t acc0189, acc011011, acc011213, acc011415;
float32x4_t acc2389, acc231011, acc231213, acc231415;
init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
accumulate);
init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
rows23, accumulate);
init_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
rows01, accumulate);
init_accumulators(acc2389, acc231011, acc231213, acc231415, c + 2 * ldc + 8,
ldc, rows23, accumulate);
const bfloat16_t* a01 = a;
const bfloat16_t* a23 = a + a_pair_stride;
const bfloat16_t* b0 = b;
const bfloat16_t* b1 = b + b_n_group_stride;
#pragma GCC unroll 4
for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
const bfloat16x8_t av01 = load_a_pair(a01, rows01);
const bfloat16x8_t av23 = load_a_pair(a23, rows23);
const bfloat16x8_t b01 = vld1q_bf16(b0);
const bfloat16x8_t b23 = vld1q_bf16(b0 + NTile);
const bfloat16x8_t b45 = vld1q_bf16(b0 + 2 * NTile);
const bfloat16x8_t b67 = vld1q_bf16(b0 + 3 * NTile);
const bfloat16x8_t b89 = vld1q_bf16(b1);
const bfloat16x8_t b1011 = vld1q_bf16(b1 + NTile);
const bfloat16x8_t b1213 = vld1q_bf16(b1 + 2 * NTile);
const bfloat16x8_t b1415 = vld1q_bf16(b1 + 3 * NTile);
acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
acc0189 = vbfmmlaq_f32(acc0189, av01, b89);
acc2389 = vbfmmlaq_f32(acc2389, av23, b89);
acc011011 = vbfmmlaq_f32(acc011011, av01, b1011);
acc231011 = vbfmmlaq_f32(acc231011, av23, b1011);
acc011213 = vbfmmlaq_f32(acc011213, av01, b1213);
acc231213 = vbfmmlaq_f32(acc231213, av23, b1213);
acc011415 = vbfmmlaq_f32(acc011415, av01, b1415);
acc231415 = vbfmmlaq_f32(acc231415, av23, b1415);
a01 += 2 * KTile;
a23 += 2 * KTile;
b0 += b_k_group_stride;
b1 += b_k_group_stride;
}
store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
rows23);
store_accumulators(acc0189, acc011011, acc011213, acc011415, c + 8, ldc,
rows01);
store_accumulators(acc2389, acc231011, acc231213, acc231415,
c + 2 * ldc + 8, ldc, rows23);
}
FORCE_INLINE static void gemm_8x8(const bfloat16_t* __restrict__ a,
const bfloat16_t* __restrict__ b,
float* __restrict__ c, const int32_t m,
const int32_t k,
const int64_t a_pair_stride,
const int64_t b_k_group_stride,
const int64_t ldc, const bool accumulate) {
const int32_t rows01 = std::min(2, std::max(0, m));
const int32_t rows23 = std::min(2, std::max(0, m - 2));
const int32_t rows45 = std::min(2, std::max(0, m - 4));
const int32_t rows67 = std::min(2, std::max(0, m - 6));
float32x4_t acc0101, acc0123, acc0145, acc0167;
float32x4_t acc2301, acc2323, acc2345, acc2367;
float32x4_t acc4501, acc4523, acc4545, acc4567;
float32x4_t acc6701, acc6723, acc6745, acc6767;
init_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01,
accumulate);
init_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
rows23, accumulate);
init_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
rows45, accumulate);
init_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
rows67, accumulate);
const bfloat16_t* a01 = a;
const bfloat16_t* a23 = a + a_pair_stride;
const bfloat16_t* a45 = a + 2 * a_pair_stride;
const bfloat16_t* a67 = a + 3 * a_pair_stride;
const bfloat16_t* b_ptr = b;
#pragma GCC unroll 4
for (int32_t k_idx = 0; k_idx < k; k_idx += KTile) {
const bfloat16x8_t av01 = load_a_pair(a01, rows01);
const bfloat16x8_t av23 = load_a_pair(a23, rows23);
const bfloat16x8_t av45 = load_a_pair(a45, rows45);
const bfloat16x8_t av67 = load_a_pair(a67, rows67);
const bfloat16x8_t b01 = vld1q_bf16(b_ptr);
const bfloat16x8_t b23 = vld1q_bf16(b_ptr + NTile);
const bfloat16x8_t b45 = vld1q_bf16(b_ptr + 2 * NTile);
const bfloat16x8_t b67 = vld1q_bf16(b_ptr + 3 * NTile);
acc0101 = vbfmmlaq_f32(acc0101, av01, b01);
acc2301 = vbfmmlaq_f32(acc2301, av23, b01);
acc4501 = vbfmmlaq_f32(acc4501, av45, b01);
acc6701 = vbfmmlaq_f32(acc6701, av67, b01);
acc0123 = vbfmmlaq_f32(acc0123, av01, b23);
acc2323 = vbfmmlaq_f32(acc2323, av23, b23);
acc4523 = vbfmmlaq_f32(acc4523, av45, b23);
acc6723 = vbfmmlaq_f32(acc6723, av67, b23);
acc0145 = vbfmmlaq_f32(acc0145, av01, b45);
acc2345 = vbfmmlaq_f32(acc2345, av23, b45);
acc4545 = vbfmmlaq_f32(acc4545, av45, b45);
acc6745 = vbfmmlaq_f32(acc6745, av67, b45);
acc0167 = vbfmmlaq_f32(acc0167, av01, b67);
acc2367 = vbfmmlaq_f32(acc2367, av23, b67);
acc4567 = vbfmmlaq_f32(acc4567, av45, b67);
acc6767 = vbfmmlaq_f32(acc6767, av67, b67);
a01 += 2 * KTile;
a23 += 2 * KTile;
a45 += 2 * KTile;
a67 += 2 * KTile;
b_ptr += b_k_group_stride;
}
store_accumulators(acc0101, acc0123, acc0145, acc0167, c, ldc, rows01);
store_accumulators(acc2301, acc2323, acc2345, acc2367, c + 2 * ldc, ldc,
rows23);
store_accumulators(acc4501, acc4523, acc4545, acc4567, c + 4 * ldc, ldc,
rows45);
store_accumulators(acc6701, acc6723, acc6745, acc6767, c + 6 * ldc, ldc,
rows67);
}
};
namespace {
// BFMMLA tile dimensions
constexpr int32_t TILE_ROWS = 2; // M dimension
constexpr int32_t TILE_K = 4; // K reduction
constexpr int32_t TILE_COLS = 2; // N dimension (column-pair)
// Derived constants
constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs
constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache
constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block
constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing
// Matrix Packing and Accumulator
// Reshape two rows of Q into BFMMLA-friendly interleaved
// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
const c10::BFloat16* __restrict r1,
c10::BFloat16* __restrict dst,
int32_t K) {
const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
uint16_t* d = reinterpret_cast<uint16_t*>(dst);
// Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
int32_t k = 0;
for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
}
// Handle K tail: pack remaining elements with zero-padding
const int32_t tail = K - k;
if (tail > 0) {
// Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
// pad]
for (int32_t t = 0; t < tail; ++t) {
d[t] = s0[k + t];
d[t + TILE_K] = s1[k + t];
}
// Zero-pad the rest
for (int32_t t = tail; t < TILE_K; ++t) {
d[t] = 0;
d[t + TILE_K] = 0;
}
}
}
// 2x2 accumulator load/store with compile-time row count
template <int32_t m_rows>
FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
static_assert(m_rows == 1 || m_rows == 2);
float32x2_t row0 = vld1_f32(base + col_off);
float32x2_t row1 =
(m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
return vcombine_f32(row0, row1);
}
template <int32_t m_rows>
FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
int col_off) {
static_assert(m_rows == 1 || m_rows == 2);
vst1_f32(base + col_off, vget_low_f32(acc));
if constexpr (m_rows == 2) {
vst1_f32(base + ldc + col_off, vget_high_f32(acc));
}
}
// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
do { \
if (accum) { \
if (m_rows == 2) { \
a0 = load_acc_2x2<2>(Crow, ldc, 0); \
a1 = load_acc_2x2<2>(Crow, ldc, 2); \
a2 = load_acc_2x2<2>(Crow, ldc, 4); \
a3 = load_acc_2x2<2>(Crow, ldc, 6); \
} else { \
a0 = load_acc_2x2<1>(Crow, ldc, 0); \
a1 = load_acc_2x2<1>(Crow, ldc, 2); \
a2 = load_acc_2x2<1>(Crow, ldc, 4); \
a3 = load_acc_2x2<1>(Crow, ldc, 6); \
} \
} else { \
a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \
} \
} while (0)
// Store 4 column-pair accumulators back to C matrix
#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
do { \
if (m_rows == 2) { \
store_acc_2x2<2>(a0, Crow, ldc, 0); \
store_acc_2x2<2>(a1, Crow, ldc, 2); \
store_acc_2x2<2>(a2, Crow, ldc, 4); \
store_acc_2x2<2>(a3, Crow, ldc, 6); \
} else { \
store_acc_2x2<1>(a0, Crow, ldc, 0); \
store_acc_2x2<1>(a1, Crow, ldc, 2); \
store_acc_2x2<1>(a2, Crow, ldc, 4); \
store_acc_2x2<1>(a3, Crow, ldc, 6); \
} \
} while (0)
// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
do { \
r0 = vbfmmlaq_f32(r0, a, b0); \
r1 = vbfmmlaq_f32(r1, a, b1); \
r2 = vbfmmlaq_f32(r2, a, b2); \
r3 = vbfmmlaq_f32(r3, a, b3); \
} while (0)
// Micro-kernel: updates a small fixed tile using BFMMLA.
// RP = number of row-pairs (1,2,4)
// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
// - AttentionGemmPhase::PV -> token-row layout (P @ V)
// K_static < 0 enables runtime K (PV only)
template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
const bfloat16_t* const* __restrict A_packed_rp,
const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
int32_t K_runtime = 0) {
static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
static_assert(K_static < 0 || K_static % TILE_K == 0,
"K must be divisible by TILE_K");
static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
"Runtime K only supported for PV");
constexpr bool runtime_k = (K_static < 0);
const int32_t K_iters =
runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
if (!runtime_k) {
// Help the compiler fold away unused K_runtime when K is compile-time
(void)K_runtime;
}
auto* C_al = C;
const auto* B_al = B_blk;
// Setup A pointers
const bfloat16_t* a_ptr[4] = {
A_packed_rp[0],
(RP >= 2) ? A_packed_rp[1] : nullptr,
(RP >= 4) ? A_packed_rp[2] : nullptr,
(RP >= 4) ? A_packed_rp[3] : nullptr,
};
// Setup B pointers based on layout
const bfloat16_t* b_ptr[4];
if constexpr (phase == AttentionGemmPhase::PV) {
b_ptr[0] = B_blk + 0 * b_stride;
b_ptr[1] = B_blk + 1 * b_stride;
b_ptr[2] = B_blk + 2 * b_stride;
b_ptr[3] = B_blk + 3 * b_stride;
}
float32x4_t acc[4][4];
// Initialize accumulators
#define INIT_RP(rp) \
if constexpr (RP > rp) { \
INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
}
INIT_RP(0);
INIT_RP(1);
INIT_RP(2);
INIT_RP(3);
#undef INIT_RP
// Main compute loop
for (int32_t ki = 0; ki < K_iters; ++ki) {
bfloat16x8_t b0, b1, b2, b3;
if constexpr (phase == AttentionGemmPhase::PV) {
b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
} else {
const bfloat16_t* b_base = B_al + ki * b_stride;
b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
}
#define COMPUTE_RP(rp) \
if constexpr (RP > rp) { \
bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \
b1, b2, b3); \
}
COMPUTE_RP(0);
COMPUTE_RP(1);
COMPUTE_RP(2);
COMPUTE_RP(3);
#undef COMPUTE_RP
}
// K tail for runtime PV: fallback path
if constexpr (runtime_k) {
if (K_tail > 0) {
const int32_t tail_offset = K_iters * V_INNER_STRIDE;
const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
for (int32_t kt = 0; kt < K_tail; ++kt) {
float32x4_t b_vecs[4];
for (int32_t p = 0; p < 4; ++p) {
const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
const float b0 = vcvtah_f32_bf16(bp[0]);
const float b1 = vcvtah_f32_bf16(bp[1]);
const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
b_vecs[p] = vcombine_f32(b_pair, b_pair);
}
#define TAIL_RP(rp) \
if constexpr (RP > rp) { \
const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \
float a_row0 = vcvtah_f32_bf16(ap[kt]); \
float a_row1 = \
(m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
const float32x4_t a_vec = \
vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \
for (int32_t p = 0; p < 4; ++p) { \
acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \
} \
}
TAIL_RP(0);
TAIL_RP(1);
TAIL_RP(2);
TAIL_RP(3);
#undef TAIL_RP
}
}
}
// Store results
#define STORE_RP(rp) \
if constexpr (RP > rp) { \
STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \
}
STORE_RP(0);
STORE_RP(1);
STORE_RP(2);
STORE_RP(3);
#undef STORE_RP
}
// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
// runtime K (PV only).
template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
FORCE_INLINE void gemm_packA_compute_MB_xN(
const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
"N must be a multiple of OUTPUT_COLS_PER_BLOCK");
static_assert(K_static < 0 || K_static % TILE_K == 0,
"K must be divisible by TILE_K");
static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
"Runtime K only supported for PV");
constexpr bool runtime_k = (K_static < 0);
const int32_t K_val = runtime_k ? K_runtime : K_static;
// Keep small packs on-stack to avoid heap churn
constexpr int32_t STACK_PACK_STRIDE =
(1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
const int32_t pack_stride =
runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
: (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
std::vector<c10::BFloat16> A_packed_heap;
c10::BFloat16* A_packed =
(pack_stride <= STACK_PACK_STRIDE)
? A_packed_stack
: (A_packed_heap.resize(ROW_PAIRS * pack_stride),
A_packed_heap.data());
for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
const int32_t m = rp * TILE_ROWS;
const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
const c10::BFloat16* A0 = A + m * lda;
const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
}
for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
const c10::BFloat16* B_blk_c10 =
(phase == AttentionGemmPhase::PV)
? (B + (n / TILE_COLS) * b_layout_stride)
: (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
// Process row-pairs in groups of 4, 2, then 1
int32_t row_pair_idx = 0;
#define PROCESS_RP_GROUP(group_size) \
for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \
row_pair_idx += group_size) { \
const bfloat16_t* Ap[group_size]; \
int32_t mr[group_size]; \
for (int32_t i = 0; i < group_size; ++i) { \
Ap[i] = reinterpret_cast<const bfloat16_t*>( \
A_packed + (row_pair_idx + i) * pack_stride); \
mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
} \
float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \
if constexpr (runtime_k) { \
gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>( \
Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \
} else { \
gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>( \
Ap, mr, B_blk, C_blk, ldc, accumulate, \
(phase == AttentionGemmPhase::PV) ? b_layout_stride \
: b_reduction_stride); \
} \
}
PROCESS_RP_GROUP(4);
PROCESS_RP_GROUP(2);
PROCESS_RP_GROUP(1);
#undef PROCESS_RP_GROUP
}
}
// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
// (runtime K path is only supported for PV).
template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
FORCE_INLINE void gemm_macro_neon_bfmmla(
const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
"N must be a multiple of OUTPUT_COLS_PER_BLOCK");
if constexpr (K >= 0) {
static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
for (int32_t m = 0; m < M;) {
const int32_t rem = M - m;
const c10::BFloat16* A_blk = A + m * lda;
float* C_blk = C + m * ldc;
#define DISPATCH_MB(mb) \
gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
b_layout_stride, \
b_reduction_stride, accumulate)
if (rem >= 8) {
DISPATCH_MB(8);
m += 8;
} else if (rem >= 4) {
DISPATCH_MB(4);
m += 4;
} else if (rem >= 2) {
DISPATCH_MB(2);
m += 2;
} else {
DISPATCH_MB(1);
m += 1;
}
#undef DISPATCH_MB
}
} else {
static_assert(phase == AttentionGemmPhase::PV,
"Runtime K specialization only supported for PV.");
const int32_t K_val = K_runtime;
for (int32_t m = 0; m < M;) {
const int32_t rem = M - m;
const c10::BFloat16* A_blk = A + m * lda;
float* C_blk = C + m * ldc;
#define DISPATCH_MB_RUNTIME(mb) \
gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
b_layout_stride, \
b_reduction_stride, accumulate)
if (rem >= 8) {
DISPATCH_MB_RUNTIME(8);
m += 8;
} else if (rem >= 4) {
DISPATCH_MB_RUNTIME(4);
m += 4;
} else if (rem >= 2) {
DISPATCH_MB_RUNTIME(2);
m += 2;
} else {
DISPATCH_MB_RUNTIME(1);
m += 1;
}
#undef DISPATCH_MB_RUNTIME
}
}
}
#undef INIT_ACC_ROWPAIR_4
#undef STORE_ACC_ROWPAIR_4
#undef BFMMLA_COMPUTE_4
constexpr int32_t TILE_K = BfmmlaGemm::KTile;
constexpr int32_t TILE_COLS = 2;
constexpr int32_t OUTPUT_COLS_PER_BLOCK = BfmmlaGemm::NTile;
constexpr int32_t K_TOKENS_PER_GROUP = 8;
constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4;
constexpr int32_t K_CACHE_K_GROUP_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
constexpr int32_t B_COL_PAIR_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
} // namespace
// TileGemm Adapter for Attention
template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
class TileGemmNEONBFMMLA {
public:
@@ -443,49 +286,31 @@ class TileGemmNEONBFMMLA {
[[maybe_unused]] const int32_t block_size,
[[maybe_unused]] const int32_t dynamic_k_size,
const bool accum_c) {
static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
// BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
// API parity with other tile_gemm implementations.
static_assert(BlockTokens % 16 == 0);
if constexpr (head_dim_ct >= 0) {
static_assert(head_dim_ct == HeadDim,
"BFMMLA expects head_dim_ct to match HeadDim; PV passes "
"-1 for API parity.");
static_assert(head_dim_ct == HeadDim);
}
const auto* a = reinterpret_cast<const c10::BFloat16*>(a_tile);
if constexpr (phase == AttentionGemmPhase::QK) {
const int64_t b_reduction_stride = K_INNER_STRIDE;
const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
constexpr int64_t b_n_group_stride =
(HeadDim / BfmmlaGemm::KTile) * K_CACHE_K_GROUP_STRIDE;
gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
accum_c);
for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
BfmmlaGemm::gemm(a + row * HeadDim, b_tile, c_tile + row * ldc, panel_m,
BlockTokens, HeadDim, 2 * HeadDim, b_n_group_stride,
K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
}
} else {
const int64_t b_pair_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
const int64_t b_n_group_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * K_CACHE_K_GROUP_STRIDE;
// PV gemm with runtime K specialization
switch (dynamic_k_size) {
case 32:
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
break;
case 128:
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
break;
case 256:
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
break;
default:
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
break;
for (int32_t row = 0; row < m_size; row += BfmmlaGemm::MaxRows) {
const int32_t panel_m = std::min(BfmmlaGemm::MaxRows, m_size - row);
BfmmlaGemm::gemm(a + row * lda, b_tile, c_tile + row * ldc, panel_m,
HeadDim, dynamic_k_size, 2 * lda, b_n_group_stride,
K_CACHE_K_GROUP_STRIDE, ldc, accum_c);
}
}
}
@@ -512,9 +337,11 @@ class AttentionImplNEONBFMMLA {
static constexpr int64_t HeadDim = head_dim;
static constexpr ISA ISAType = isa_type;
static constexpr bool scale_on_logits = false;
static constexpr int64_t VCacheNGroup = OUTPUT_COLS_PER_BLOCK;
static constexpr int64_t VCacheKGroupStride = VCacheNGroup * TILE_K;
static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
static_assert(HeadDim % (2 * OUTPUT_COLS_PER_BLOCK) == 0);
static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
public:
@@ -527,19 +354,36 @@ class AttentionImplNEONBFMMLA {
attention_iteration(CPU_ATTENTION_PARAMS);
}
struct ProbabilityTokenStore {
static constexpr int32_t TokenStride = 2;
FORCE_INLINE static void store_probabilities(
c10::BFloat16* __restrict__ probability,
const vec_op::FP32Vec16& values, const int32_t row,
const int64_t row_stride) {
const int32_t row_in_pair = row & 1;
auto* dst = reinterpret_cast<bfloat16_t*>(
probability + row_in_pair * (BfmmlaGemm::KTile - row_stride));
vst1_bf16(dst, vcvt_bf16_f32(values.reg.val[0]));
vst1_bf16(dst + 2 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[1]));
vst1_bf16(dst + 4 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[2]));
vst1_bf16(dst + 6 * BfmmlaGemm::KTile, vcvt_bf16_f32(values.reg.val[3]));
}
};
// Key cache stride per token group (TokenColumn layout; QK)
static constexpr int64_t k_cache_token_group_stride(
[[maybe_unused]] const int32_t block_size) {
static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
((head_dim / TILE_K) * K_INNER_STRIDE);
((head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE);
}
// Value cache stride per token group (TokenRow layout; PV)
static constexpr int64_t v_cache_token_group_stride(
[[maybe_unused]] const int32_t block_size) {
static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
}
// The stride to move to the "next" head_dim group
@@ -550,7 +394,7 @@ class AttentionImplNEONBFMMLA {
return head_dim * block_size;
}
// Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
// Scale Q and write row pairs in BFMMLA reduction order.
static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
c10::BFloat16* __restrict__ q_buffer,
const int32_t q_num,
@@ -559,25 +403,42 @@ class AttentionImplNEONBFMMLA {
const int64_t q_head_stride, float scale) {
constexpr int32_t dim = static_cast<int32_t>(head_dim);
const float32x4_t scale_vec = vdupq_n_f32(scale);
const bfloat16x4_t zero = vdup_n_bf16(bfloat16_t{});
const int32_t row_num = q_num * q_heads_per_kv;
for (int32_t qi = 0; qi < q_num; ++qi) {
for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
c10::BFloat16* __restrict__ curr_q =
src + qi * q_num_stride + hi * q_head_stride;
c10::BFloat16* __restrict__ dst =
q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
for (int32_t row = 0; row < row_num; row += 2) {
const int32_t q0 = row / q_heads_per_kv;
const int32_t h0 = row % q_heads_per_kv;
const auto* row0 = reinterpret_cast<const bfloat16_t*>(
src + q0 * q_num_stride + h0 * q_head_stride);
const bool has_row1 = row + 1 < row_num;
const int32_t q1 = (row + 1) / q_heads_per_kv;
const int32_t h1 = (row + 1) % q_heads_per_kv;
const auto* row1 = has_row1
? reinterpret_cast<const bfloat16_t*>(
src + q1 * q_num_stride + h1 * q_head_stride)
: nullptr;
auto* dst = reinterpret_cast<bfloat16_t*>(q_buffer + row * head_dim);
for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
bfloat16x8_t in8 =
vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
for (int32_t k = 0; k < dim; k += OUTPUT_COLS_PER_BLOCK) {
const bfloat16x8_t in0 = vld1q_bf16(row0 + k);
const bfloat16x4_t out0_lo =
vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in0), scale_vec));
const bfloat16x4_t out0_hi =
vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in0), scale_vec));
bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
bfloat16x4_t out1_lo = zero;
bfloat16x4_t out1_hi = zero;
if (has_row1) {
const bfloat16x8_t in1 = vld1q_bf16(row1 + k);
out1_lo =
vcvt_bf16_f32(vmulq_f32(vcvtq_low_f32_bf16(in1), scale_vec));
out1_hi =
vcvt_bf16_f32(vmulq_f32(vcvtq_high_f32_bf16(in1), scale_vec));
}
vst1q_bf16(dst + 2 * k, vcombine_bf16(out0_lo, out1_lo));
vst1q_bf16(dst + 2 * k + 8, vcombine_bf16(out0_hi, out1_hi));
}
}
}
@@ -585,11 +446,12 @@ class AttentionImplNEONBFMMLA {
public:
// Reshape and cache K/V into BFMMLA-optimized layouts
// K cache:
// [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
// [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K]
// [K_CACHE_K_GROUP_STRIDE]
// - TokenColumn
// V cache:
// [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
// - TokenRows
// [head_dim/VCacheNGroup][block_size/V_TOKENS_PER_ROW_BLOCK]
// [VCacheKGroupStride]
static void reshape_and_cache(
const c10::BFloat16* __restrict__ key,
const c10::BFloat16* __restrict__ value,
@@ -604,13 +466,13 @@ class AttentionImplNEONBFMMLA {
const int64_t block_size,
[[maybe_unused]] const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
const int64_t v_pair_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
const int64_t k_block_stride = (head_dim / TILE_K) * K_CACHE_K_GROUP_STRIDE;
const int64_t v_n_group_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride;
#pragma omp parallel for
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
const int64_t pos = slot_mapping[token_idx];
if (pos < 0) continue;
@@ -636,12 +498,12 @@ class AttentionImplNEONBFMMLA {
key_base + block_in_block * k_block_stride;
for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
block_base + hd4 * K_INNER_STRIDE +
pair_in_block * V_INNER_STRIDE + lane_base);
const uint16_t* src_u16 =
reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
vst1_u16(dst_u16, vld1_u16(src_u16));
bfloat16_t* dst = reinterpret_cast<bfloat16_t*>(
block_base + hd4 * K_CACHE_K_GROUP_STRIDE +
pair_in_block * B_COL_PAIR_STRIDE + lane_base);
const bfloat16_t* src =
reinterpret_cast<const bfloat16_t*>(key_src + hd4 * TILE_K);
vst1_bf16(dst, vld1_bf16(src));
}
}
@@ -655,22 +517,33 @@ class AttentionImplNEONBFMMLA {
value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride;
const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
const int64_t token_group = block_offset / V_TOKENS_PER_ROW_BLOCK;
const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
c10::BFloat16* __restrict row_block_base =
value_base + row_block * V_INNER_STRIDE;
const auto* src = reinterpret_cast<const bfloat16_t*>(value_src);
auto* dst = reinterpret_cast<bfloat16_t*>(value_base);
for (int64_t hd8 = 0; hd8 < head_dim / OUTPUT_COLS_PER_BLOCK; ++hd8) {
const bfloat16x8_t values =
vld1q_bf16(src + hd8 * OUTPUT_COLS_PER_BLOCK);
const bfloat16x4_t low = vget_low_bf16(values);
const bfloat16x4_t high = vget_high_bf16(values);
bfloat16_t* group =
dst + hd8 * v_n_group_stride + token_group * VCacheKGroupStride;
for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
c10::BFloat16* __restrict dst_val =
row_block_base + hd2 * v_pair_stride;
const uint16_t* src_u16 =
reinterpret_cast<const uint16_t*>(value_src);
uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
src_u16[hd2 * TILE_COLS + 1];
vst1_lane_bf16(group + lane, low, 0);
vst1_lane_bf16(group + V_TOKENS_PER_ROW_BLOCK + lane, low, 1);
vst1_lane_bf16(group + B_COL_PAIR_STRIDE + lane, low, 2);
vst1_lane_bf16(
group + B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane, low,
3);
vst1_lane_bf16(group + 2 * B_COL_PAIR_STRIDE + lane, high, 0);
vst1_lane_bf16(
group + 2 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
high, 1);
vst1_lane_bf16(group + 3 * B_COL_PAIR_STRIDE + lane, high, 2);
vst1_lane_bf16(
group + 3 * B_COL_PAIR_STRIDE + V_TOKENS_PER_ROW_BLOCK + lane,
high, 3);
}
}
}
@@ -680,4 +553,4 @@ class AttentionImplNEONBFMMLA {
} // namespace cpu_attention
#endif // CPU_ATTN_ASIMD_BFMMLA_HPP
#endif // CPU_ATTN_NEON_BFMMLA_HPP
vllm/v1/attention/backends/cpu_attn.py
+3 -4
View File
@@ -512,11 +512,10 @@ def _get_attn_isa(
)
if supports_amx and dtype in (torch.bfloat16,) and block_size % 32 == 0:
return "amx"
elif supports_arm:
return "neon"
elif block_size % 32 == 0:
if supports_arm:
# support ARM NEON FMLA and BFMMLA (bf16) for block size 32
return "neon"
elif supports_riscv and _riscv_supports_rvv():
if supports_riscv and _riscv_supports_rvv():
return "rvv"
elif supports_vxe:
return "vxe"