| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | #ifndef CPU_ATTN_VSX_HPP |
| 4 | #define CPU_ATTN_VSX_HPP |
| 5 | |
| 6 | #include "cpu_attn_impl.hpp" |
| 7 | #include <altivec.h> |
| 8 | #include <type_traits> |
| 9 | |
| 10 | namespace cpu_attention { |
| 11 | |
| 12 | namespace { |
| 13 | |
| 14 | // ppc64le Vector = 16 bytes (128 bits) |
| 15 | #define BLOCK_SIZE_ALIGNMENT 32 |
| 16 | #define HEAD_SIZE_ALIGNMENT 32 |
| 17 | #define MAX_Q_HEAD_NUM_PER_ITER 16 |
| 18 | |
| 19 | template <typename kv_cache_t> |
| 20 | FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0, |
| 21 | __vector float& b1); |
| 22 | |
| 23 | // [1] Float Specialization |
| 24 | template <> |
| 25 | FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0, |
| 26 | __vector float& b1) { |
| 27 | b0 = vec_xl(0, const_cast<float*>(p)); |
| 28 | b1 = vec_xl(0, const_cast<float*>(p + 4)); |
| 29 | } |
| 30 | |
| 31 | // [2] BFloat16 Specialization (Little Endian ppc64le) |
| 32 | // On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32. |
| 33 | // Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)] |
| 34 | // We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed. |
| 35 | // vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]} |
| 36 | // So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair: |
| 37 | // uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs |
| 38 | // uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits |
| 39 | // Cast to float32 gives exactly (bf16_bits << 16) per element. |
| 40 | template <> |
| 41 | FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p, |
| 42 | __vector float& b0, |
| 43 | __vector float& b1) { |
| 44 | __vector unsigned short raw = vec_xl( |
| 45 | 0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p))); |
| 46 | __vector unsigned short zeros = vec_splat_u16(0); |
| 47 | |
| 48 | // LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32 |
| 49 | b0 = (__vector float)vec_mergeh(zeros, raw); |
| 50 | b1 = (__vector float)vec_mergel(zeros, raw); |
| 51 | } |
| 52 | |
| 53 | // [3] Half (FP16) Specialization |
| 54 | template <> |
| 55 | FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p, |
| 56 | __vector float& b0, |
| 57 | __vector float& b1) { |
| 58 | vec_op::FP16Vec8 fp16_vec(p); |
| 59 | vec_op::FP32Vec8 fp32_vec(fp16_vec); |
| 60 | b0 = fp32_vec.reg.val[0]; |
| 61 | b1 = fp32_vec.reg.val[1]; |
| 62 | } |
| 63 | |
| 64 | template <int32_t M, typename kv_cache_t> |
| 65 | FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4( |
| 66 | const float* __restrict A, // [M x K] |
| 67 | const kv_cache_t* __restrict B, // [K x 8] |
| 68 | float* __restrict C, // [M x 8] |
| 69 | int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) { |
| 70 | static_assert(1 <= M && M <= 8, "M must be in [1,8]"); |
| 71 | |
| 72 | #define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7) |
| 73 | #define IF_M(i) if constexpr (M > (i)) |
| 74 | |
| 75 | // 1. Define A pointers |
| 76 | #define DECL_A(i) const float* a##i = A + (i) * lda; |
| 77 | ROWS_APPLY(DECL_A) |
| 78 | #undef DECL_A |
| 79 | |
| 80 | // 2. Define Accumulators (2 vectors covers 8 columns) |
| 81 | #define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1; |
| 82 | ROWS_APPLY(DECL_ACC) |
| 83 | #undef DECL_ACC |
| 84 | |
| 85 | // 3. Initialize Accumulators (Load C or Zero) |
| 86 | #define INIT_ACC(i) \ |
| 87 | IF_M(i) { \ |
| 88 | if (accumulate) { \ |
| 89 | acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \ |
| 90 | acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \ |
| 91 | } else { \ |
| 92 | acc##i##_0 = vec_splats(0.0f); \ |
| 93 | acc##i##_1 = vec_splats(0.0f); \ |
| 94 | } \ |
| 95 | } |
| 96 | ROWS_APPLY(INIT_ACC) |
| 97 | #undef INIT_ACC |
| 98 | |
| 99 | int32_t k = 0; |
| 100 | |
| 101 | for (; k + 3 < K; k += 4) { |
| 102 | // Load 4 values of A for each Row M: A[k...k+3] |
| 103 | #define LOAD_A4(i) \ |
| 104 | __vector float a##i##v; \ |
| 105 | IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k)); |
| 106 | ROWS_APPLY(LOAD_A4) |
| 107 | #undef LOAD_A4 |
| 108 | |
| 109 | // FMA for specific lane L of A |
| 110 | // ppc64le: vec_madd(b, vec_splat(a, lane), acc) |
| 111 | #define FMAS_LANE(i, aiv, L) \ |
| 112 | IF_M(i) { \ |
| 113 | __vector float a_broad = vec_splat(aiv, L); \ |
| 114 | acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \ |
| 115 | acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \ |
| 116 | } |
| 117 | |
| 118 | // Unroll K=0..3 |
| 119 | { |
| 120 | __vector float b0, b1; |
| 121 | load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1); |
| 122 | #define STEP_K0(i) FMAS_LANE(i, a##i##v, 0) |
| 123 | ROWS_APPLY(STEP_K0) |
| 124 | #undef STEP_K0 |
| 125 | } |
| 126 | { |
| 127 | __vector float b0, b1; |
| 128 | load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1); |
| 129 | #define STEP_K1(i) FMAS_LANE(i, a##i##v, 1) |
| 130 | ROWS_APPLY(STEP_K1) |
| 131 | #undef STEP_K1 |
| 132 | } |
| 133 | { |
| 134 | __vector float b0, b1; |
| 135 | load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1); |
| 136 | #define STEP_K2(i) FMAS_LANE(i, a##i##v, 2) |
| 137 | ROWS_APPLY(STEP_K2) |
| 138 | #undef STEP_K2 |
| 139 | } |
| 140 | { |
| 141 | __vector float b0, b1; |
| 142 | load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1); |
| 143 | #define STEP_K3(i) FMAS_LANE(i, a##i##v, 3) |
| 144 | ROWS_APPLY(STEP_K3) |
| 145 | #undef STEP_K3 |
| 146 | } |
| 147 | #undef FMAS_LANE |
| 148 | } |
| 149 | |
| 150 | for (; k < K; ++k) { |
| 151 | __vector float b0, b1; |
| 152 | load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1); |
| 153 | #define TAIL_ROW(i) \ |
| 154 | IF_M(i) { \ |
| 155 | __vector float ai = vec_splats(*(a##i + k)); \ |
| 156 | acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \ |
| 157 | acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \ |
| 158 | } |
| 159 | ROWS_APPLY(TAIL_ROW) |
| 160 | #undef TAIL_ROW |
| 161 | } |
| 162 | |
| 163 | #define STORE_ROW(i) \ |
| 164 | IF_M(i) { \ |
| 165 | vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \ |
| 166 | vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \ |
| 167 | } |
| 168 | ROWS_APPLY(STORE_ROW) |
| 169 | #undef STORE_ROW |
| 170 | |
| 171 | #undef ROWS_APPLY |
| 172 | #undef IF_M |
| 173 | } |
| 174 | |
| 175 | template <int32_t N, typename kv_cache_t> |
| 176 | FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A, |
| 177 | const kv_cache_t* __restrict B, |
| 178 | float* __restrict C, int32_t M, |
| 179 | int32_t K, int64_t lda, |
| 180 | int64_t ldb, int64_t ldc, |
| 181 | bool accumulate) { |
| 182 | static_assert(N % 8 == 0, "N must be a multiple of 8"); |
| 183 | for (int32_t m = 0; m < M;) { |
| 184 | int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1; |
| 185 | const float* Ab = A + m * lda; |
| 186 | float* Cb = C + m * ldc; |
| 187 | |
| 188 | for (int32_t n = 0; n < N; n += 8) { |
| 189 | const kv_cache_t* Bn = B + n; |
| 190 | float* Cn = Cb + n; |
| 191 | switch (mb) { |
| 192 | case 8: |
| 193 | gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, |
| 194 | K, accumulate); |
| 195 | break; |
| 196 | case 4: |
| 197 | gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, |
| 198 | K, accumulate); |
| 199 | break; |
| 200 | case 2: |
| 201 | gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, |
| 202 | K, accumulate); |
| 203 | break; |
| 204 | default: |
| 205 | gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, |
| 206 | K, accumulate); |
| 207 | break; |
| 208 | } |
| 209 | } |
| 210 | m += mb; |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | template <typename kv_cache_t> |
| 215 | class TileGemmPPC64 { |
| 216 | public: |
| 217 | template <AttentionGemmPhase phase, int32_t k_size> |
| 218 | FORCE_INLINE static void gemm(const int32_t m_size, |
| 219 | float* __restrict__ a_tile, |
| 220 | kv_cache_t* __restrict__ b_tile, |
| 221 | float* __restrict__ c_tile, const int64_t lda, |
| 222 | const int64_t ldb, const int64_t ldc, |
| 223 | const int32_t block_size, |
| 224 | const int32_t dynamic_k_size, |
| 225 | const bool accum_c) { |
| 226 | if constexpr (phase == AttentionGemmPhase::QK) { |
| 227 | gemm_macro_ppc64le_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>( |
| 228 | a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c); |
| 229 | } else { |
| 230 | gemm_macro_ppc64le_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>( |
| 231 | a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc, |
| 232 | accum_c); |
| 233 | } |
| 234 | } |
| 235 | }; |
| 236 | |
| 237 | } // namespace |
| 238 | |
| 239 | template <typename scalar_t, int64_t head_dim> |
| 240 | class AttentionImpl<ISA::VSX, scalar_t, head_dim> { |
| 241 | public: |
| 242 | using query_t = scalar_t; |
| 243 | using q_buffer_t = float; |
| 244 | using kv_cache_t = scalar_t; |
| 245 | using logits_buffer_t = float; |
| 246 | using partial_output_buffer_t = float; |
| 247 | using prob_buffer_t = float; |
| 248 | |
| 249 | constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT; |
| 250 | constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT; |
| 251 | constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER; |
| 252 | constexpr static int64_t HeadDim = head_dim; |
| 253 | constexpr static ISA ISAType = ISA::VSX; |
| 254 | constexpr static bool scale_on_logits = |
| 255 | false; // Scale is applied to Q during copy |
| 256 | |
| 257 | public: |
| 258 | AttentionImpl() {} |
| 259 | |
| 260 | template <template <typename tile_gemm_t> typename attention> |
| 261 | FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) { |
| 262 | attention<TileGemmPPC64<kv_cache_t>> attention_iteration; |
| 263 | attention_iteration(CPU_ATTENTION_PARAMS); |
| 264 | } |
| 265 | |
| 266 | // Strides for Memory Layout |
| 267 | constexpr static int64_t k_cache_token_group_stride( |
| 268 | const int32_t block_size) { |
| 269 | return BlockSizeAlignment; // [head_dim, block_size] layout |
| 270 | } |
| 271 | |
| 272 | constexpr static int64_t v_cache_token_group_stride( |
| 273 | const int32_t block_size) { |
| 274 | return head_dim * BlockSizeAlignment; |
| 275 | } |
| 276 | |
| 277 | constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) { |
| 278 | return HeadDimAlignment; |
| 279 | } |
| 280 | |
| 281 | static void copy_q_heads_tile(scalar_t* __restrict__ src, |
| 282 | float* __restrict__ q_buffer, |
| 283 | const int32_t q_num, |
| 284 | const int32_t q_heads_per_kv, |
| 285 | const int64_t q_num_stride, |
| 286 | const int64_t q_head_stride, float scale) { |
| 287 | __vector float scale_vec = vec_splats(scale); |
| 288 | constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value; |
| 289 | |
| 290 | for (int32_t i = 0; i < q_num; ++i) { |
| 291 | for (int32_t h = 0; h < q_heads_per_kv; ++h) { |
| 292 | scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride; |
| 293 | float* curr_dst = |
| 294 | q_buffer + i * q_heads_per_kv * head_dim + h * head_dim; |
| 295 | |
| 296 | int32_t d = 0; |
| 297 | for (; d <= head_dim - 8; d += 8) { |
| 298 | __vector float v0, v1; |
| 299 | load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1); |
| 300 | |
| 301 | v0 = vec_mul(v0, scale_vec); |
| 302 | v1 = vec_mul(v1, scale_vec); |
| 303 | |
| 304 | vec_xst(v0, 0, curr_dst + d); |
| 305 | vec_xst(v1, 0, curr_dst + d + 4); |
| 306 | } |
| 307 | |
| 308 | for (; d < head_dim; ++d) { |
| 309 | float val = static_cast<float>(curr_src[d]); |
| 310 | curr_dst[d] = val * scale; |
| 311 | } |
| 312 | } |
| 313 | } |
| 314 | } |
| 315 | |
| 316 | static void reshape_and_cache( |
| 317 | const scalar_t* __restrict__ key, const scalar_t* __restrict__ value, |
| 318 | scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache, |
| 319 | const int64_t* __restrict__ slot_mapping, const int64_t token_num, |
| 320 | const int64_t key_token_num_stride, const int64_t value_token_num_stride, |
| 321 | const int64_t head_num, const int64_t key_head_num_stride, |
| 322 | const int64_t value_head_num_stride, const int64_t num_blocks, |
| 323 | const int64_t num_blocks_stride, const int64_t cache_head_num_stride, |
| 324 | const int64_t block_size, const int64_t block_size_stride, |
| 325 | const float k_inv = 0.0f, const float v_inv = 0.0f) { |
| 326 | #pragma omp parallel for collapse(2) |
| 327 | for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) { |
| 328 | for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) { |
| 329 | const int64_t pos = slot_mapping[token_idx]; |
| 330 | if (pos < 0) continue; |
| 331 | |
| 332 | const int64_t block_idx = pos / block_size; |
| 333 | const int64_t block_offset = pos % block_size; |
| 334 | |
| 335 | { |
| 336 | const scalar_t* key_src = key + token_idx * key_token_num_stride + |
| 337 | head_idx * key_head_num_stride; |
| 338 | scalar_t* key_dst = key_cache + block_idx * num_blocks_stride + |
| 339 | head_idx * cache_head_num_stride + block_offset; |
| 340 | |
| 341 | for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) { |
| 342 | key_dst[j] = key_src[i]; |
| 343 | } |
| 344 | } |
| 345 | |
| 346 | { |
| 347 | const scalar_t* val_src = value + token_idx * value_token_num_stride + |
| 348 | head_idx * value_head_num_stride; |
| 349 | scalar_t* val_dst = value_cache + block_idx * num_blocks_stride + |
| 350 | head_idx * cache_head_num_stride + |
| 351 | block_offset * head_dim; |
| 352 | |
| 353 | std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim); |
| 354 | } |
| 355 | } |
| 356 | } |
| 357 | } |
| 358 | }; |
| 359 | |
| 360 | } // namespace cpu_attention |
| 361 | |
| 362 | #undef BLOCK_SIZE_ALIGNMENT |
| 363 | #undef HEAD_SIZE_ALIGNMENT |
| 364 | #undef MAX_Q_HEAD_NUM_PER_ITER |
| 365 | |
| 366 | #endif // CPU_ATTN_VSX_HPP |