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