| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | #ifndef CPU_ATTN_NEON_BFMMLA_HPP |
| 4 | #define CPU_ATTN_NEON_BFMMLA_HPP |
| 5 | |
| 6 | #include "cpu_attn_impl.hpp" |
| 7 | |
| 8 | #include <arm_bf16.h> |
| 9 | #include <arm_neon.h> |
| 10 | #include <c10/util/BFloat16.h> |
| 11 | |
| 12 | #include <algorithm> |
| 13 | #include <cstdint> |
| 14 | |
| 15 | namespace cpu_attention { |
| 16 | |
| 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 | } |
| 53 | |
| 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 | |
| 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 | } |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | } |
| 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 | } |
| 263 | }; |
| 264 | |
| 265 | namespace { |
| 266 | |
| 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 | |
| 277 | template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim> |
| 278 | class TileGemmNEONBFMMLA { |
| 279 | public: |
| 280 | template <AttentionGemmPhase phase, int32_t head_dim_ct> |
| 281 | FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile, |
| 282 | kv_cache_t* __restrict__ b_tile, |
| 283 | float* __restrict__ c_tile, const int64_t lda, |
| 284 | [[maybe_unused]] const int64_t ldb, |
| 285 | const int64_t ldc, |
| 286 | [[maybe_unused]] const int32_t block_size, |
| 287 | [[maybe_unused]] const int32_t dynamic_k_size, |
| 288 | const bool accum_c) { |
| 289 | static_assert(BlockTokens % 16 == 0); |
| 290 | if constexpr (head_dim_ct >= 0) { |
| 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) { |
| 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 { |
| 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 | } |
| 317 | }; |
| 318 | |
| 319 | // Shared ASIMD BFMMLA implementation (BF16 only). The block size alignment and |
| 320 | // ISA tag are template parameters so we can reuse the same kernels for |
| 321 | // different NEON configurations. |
| 322 | template <int64_t block_size_alignment, ISA isa_type, int64_t head_dim> |
| 323 | class AttentionImplNEONBFMMLA { |
| 324 | public: |
| 325 | using query_t = c10::BFloat16; |
| 326 | using q_buffer_t = c10::BFloat16; |
| 327 | using kv_cache_t = c10::BFloat16; |
| 328 | using logits_buffer_t = float; |
| 329 | using partial_output_buffer_t = float; |
| 330 | using prob_buffer_t = c10::BFloat16; |
| 331 | |
| 332 | static constexpr int64_t BlockSizeAlignment = block_size_alignment; |
| 333 | // HeadDimAlignment equals head_dim so that the PV phase processes |
| 334 | // the full head dimension in a single gemm call. |
| 335 | static constexpr int64_t HeadDimAlignment = head_dim; |
| 336 | static constexpr int64_t MaxQHeadNumPerIteration = 16; |
| 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 | |
| 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: |
| 348 | template <template <typename tile_gemm_t> typename attention> |
| 349 | FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) { |
| 350 | attention< |
| 351 | TileGemmNEONBFMMLA<kv_cache_t, static_cast<int32_t>(BlockSizeAlignment), |
| 352 | static_cast<int32_t>(HeadDim)>> |
| 353 | attention_iteration; |
| 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) * |
| 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); |
| 386 | return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * VCacheKGroupStride; |
| 387 | } |
| 388 | |
| 389 | // The stride to move to the "next" head_dim group |
| 390 | // is the full V cache size per head, since HeadDimAlignment == head_dim. |
| 391 | // Hence, the stride is not used in this case |
| 392 | static constexpr int64_t v_cache_head_group_stride( |
| 393 | [[maybe_unused]] const int32_t block_size) { |
| 394 | return head_dim * block_size; |
| 395 | } |
| 396 | |
| 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, |
| 401 | const int32_t q_heads_per_kv, |
| 402 | const int64_t q_num_stride, |
| 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); |
| 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 | } |
| 445 | |
| 446 | public: |
| 447 | // Reshape and cache K/V into BFMMLA-optimized layouts |
| 448 | // K cache: |
| 449 | // [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K] |
| 450 | // [K_CACHE_K_GROUP_STRIDE] |
| 451 | // - TokenColumn |
| 452 | // V cache: |
| 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, |
| 458 | c10::BFloat16* __restrict__ key_cache, |
| 459 | c10::BFloat16* __restrict__ value_cache, |
| 460 | const int64_t* __restrict__ slot_mapping, const int64_t token_num, |
| 461 | const int64_t key_token_num_stride, const int64_t value_token_num_stride, |
| 462 | const int64_t head_num, const int64_t key_head_num_stride, |
| 463 | const int64_t value_head_num_stride, |
| 464 | [[maybe_unused]] const int64_t num_blocks, |
| 465 | const int64_t num_blocks_stride, const int64_t cache_head_num_stride, |
| 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) { |
| 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 | |
| 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 | |
| 479 | const int64_t block_idx = pos / block_size; |
| 480 | const int64_t block_offset = pos % block_size; |
| 481 | |
| 482 | // Key cache: TokenColumn QK |
| 483 | { |
| 484 | const c10::BFloat16* __restrict key_src = |
| 485 | key + token_idx * key_token_num_stride + |
| 486 | head_idx * key_head_num_stride; |
| 487 | |
| 488 | c10::BFloat16* __restrict key_base = key_cache + |
| 489 | block_idx * num_blocks_stride + |
| 490 | head_idx * cache_head_num_stride; |
| 491 | |
| 492 | const int64_t block_in_block = block_offset / K_TOKENS_PER_GROUP; |
| 493 | const int64_t pair_in_block = |
| 494 | (block_offset % K_TOKENS_PER_GROUP) / TILE_COLS; |
| 495 | const int64_t lane_base = (block_offset & 1) ? TILE_K : 0; |
| 496 | |
| 497 | c10::BFloat16* __restrict block_base = |
| 498 | key_base + block_in_block * k_block_stride; |
| 499 | |
| 500 | for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) { |
| 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 | |
| 510 | // Value cache: TokenRow PV |
| 511 | { |
| 512 | const c10::BFloat16* __restrict value_src = |
| 513 | value + token_idx * value_token_num_stride + |
| 514 | head_idx * value_head_num_stride; |
| 515 | |
| 516 | c10::BFloat16* __restrict value_base = |
| 517 | value_cache + block_idx * num_blocks_stride + |
| 518 | head_idx * cache_head_num_stride; |
| 519 | |
| 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 | |
| 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 | } |
| 550 | } |
| 551 | } |
| 552 | }; |
| 553 | |
| 554 | } // namespace cpu_attention |
| 555 | |
| 556 | #endif // CPU_ATTN_NEON_BFMMLA_HPP |