diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index e9ddaacea5..4ade47eff3 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -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 +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::vec_t; + prob_buffer_vec_t output_vec(values); + output_vec.save(probability); + } +}; + +template +struct ProbabilityTokenStore< + attention_impl_t, + std::void_t> + : attention_impl_t::ProbabilityTokenStore {}; + template 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; 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( - reinterpret_cast(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; #endif - using prob_buffer_vec_t = typename VecTypeTrait::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( diff --git a/csrc/cpu/cpu_attn_neon.hpp b/csrc/cpu/cpu_attn_neon.hpp index db4c5df2e8..eed7499444 100644 --- a/csrc/cpu/cpu_attn_neon.hpp +++ b/csrc/cpu/cpu_attn_neon.hpp @@ -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 { }; #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 class AttentionImpl : public AttentionImplNEONBFMMLA #include +#include +#include #include -#include 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(a); + const auto* b_ptr = reinterpret_cast(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(r0); - const uint16_t* s1 = reinterpret_cast(r1); - uint16_t* d = reinterpret_cast(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 -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 -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 -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 -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 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(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( \ - 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( \ - Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \ - } else { \ - gemm_rowpairs_x8_bfmmla_neon( \ - 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 -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(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(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 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(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( - reinterpret_cast(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( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 32, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 128: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 128, lda, ldc, b_pair_stride, 0, accum_c); - break; - case 256: - gemm_macro_neon_bfmmla( - reinterpret_cast(a_tile), b_tile, c_tile, - m_size, 256, lda, ldc, b_pair_stride, 0, accum_c); - break; - default: - gemm_macro_neon_bfmmla( - reinterpret_cast(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( + 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(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( + 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( + src + q1 * q_num_stride + h1 * q_head_stride) + : nullptr; + auto* dst = reinterpret_cast(q_buffer + row * head_dim); - for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) { - bfloat16x8_t in8 = - vld1q_bf16(reinterpret_cast(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(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( - block_base + hd4 * K_INNER_STRIDE + - pair_in_block * V_INNER_STRIDE + lane_base); - const uint16_t* src_u16 = - reinterpret_cast(key_src + hd4 * TILE_K); - vst1_u16(dst_u16, vld1_u16(src_u16)); + bfloat16_t* dst = reinterpret_cast( + block_base + hd4 * K_CACHE_K_GROUP_STRIDE + + pair_in_block * B_COL_PAIR_STRIDE + lane_base); + const bfloat16_t* src = + reinterpret_cast(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(value_src); + auto* dst = reinterpret_cast(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(value_src); - uint16_t* dst_u16 = reinterpret_cast(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 diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b6864897bf..24be5b8ec9 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -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"