| 1 | #include "cpu_types.hpp" |
| 2 | #include <float.h> |
| 3 | |
| 4 | namespace { |
| 5 | template <typename scalar_t> |
| 6 | struct KernelVecType { |
| 7 | using qk_load_vec_type = void; |
| 8 | using qk_vec_type = void; |
| 9 | using v_load_vec_type = void; |
| 10 | }; |
| 11 | |
| 12 | template <> |
| 13 | struct KernelVecType<float> { |
| 14 | using qk_load_vec_type = vec_op::FP32Vec16; |
| 15 | using qk_vec_type = vec_op::FP32Vec16; |
| 16 | using v_load_vec_type = vec_op::FP32Vec16; |
| 17 | }; |
| 18 | |
| 19 | template <> |
| 20 | struct KernelVecType<c10::Half> { |
| 21 | using qk_load_vec_type = vec_op::FP16Vec16; |
| 22 | using qk_vec_type = vec_op::FP32Vec16; |
| 23 | using v_load_vec_type = vec_op::FP16Vec16; |
| 24 | }; |
| 25 | |
| 26 | #ifdef __AVX512BF16__ |
| 27 | template <> |
| 28 | struct KernelVecType<c10::BFloat16> { |
| 29 | using qk_load_vec_type = vec_op::BF16Vec32; |
| 30 | using qk_vec_type = vec_op::BF16Vec32; |
| 31 | using v_load_vec_type = vec_op::BF16Vec16; |
| 32 | }; |
| 33 | #else |
| 34 | template <> |
| 35 | struct KernelVecType<c10::BFloat16> { |
| 36 | using qk_load_vec_type = vec_op::BF16Vec16; |
| 37 | using qk_vec_type = vec_op::FP32Vec16; |
| 38 | using v_load_vec_type = vec_op::BF16Vec16; |
| 39 | }; |
| 40 | #endif |
| 41 | |
| 42 | template <int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE, int HEAD_UNROLL, |
| 43 | typename qk_vec_type> |
| 44 | void mla_decode_block_head( |
| 45 | const qk_vec_type* __restrict__ q_vecs, // [HEAD_UNROLL, head_dim] |
| 46 | const qk_vec_type* __restrict__ k_vecs, // [block_size, head_dim] |
| 47 | const vec_op::FP32Vec16* __restrict v_vecs_f32, // [block_size, v_head_dim] |
| 48 | float* __restrict__ acc_out, // [HEAD_UNROLL, v_head_dim] |
| 49 | float* __restrict__ acc_lse, // [HEAD_UNROLL] |
| 50 | const float scale, const int num_tokens) { |
| 51 | using f32_vec_type = vec_op::FP32Vec16; |
| 52 | constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM; |
| 53 | constexpr int V_NUM_ELEM = f32_vec_type::VEC_ELEM_NUM; |
| 54 | |
| 55 | float logits[BLOCK_SIZE][HEAD_UNROLL] = {}; // initialize to zeros |
| 56 | float max_val[HEAD_UNROLL]; |
| 57 | std::fill(max_val, max_val + HEAD_UNROLL, -FLT_MAX); |
| 58 | |
| 59 | f32_vec_type acc_vec[BLOCK_SIZE][HEAD_UNROLL]; |
| 60 | // Explicitly zero-initialize the accumulator. On x86 the default |
| 61 | // ctor of FP32Vec16 already sets zero, but on ARM (and any backend |
| 62 | // whose Vectorized<float> default ctor is trivial) the array holds |
| 63 | // undefined values and the fma below accumulates garbage. |
| 64 | for (int b = 0; b < BLOCK_SIZE; ++b) |
| 65 | for (int u = 0; u < HEAD_UNROLL; ++u) acc_vec[b][u] = f32_vec_type{0.0f}; |
| 66 | for (int i = 0; i < HEAD_DIM; i += QK_NUM_ELEM) { |
| 67 | // load to registers |
| 68 | qk_vec_type q_vec[HEAD_UNROLL]; |
| 69 | |
| 70 | #pragma unroll |
| 71 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) |
| 72 | q_vec[unroll] = |
| 73 | qk_vec_type{q_vecs[(i + unroll * HEAD_DIM) / QK_NUM_ELEM]}; |
| 74 | |
| 75 | for (int block_offset = 0; block_offset < num_tokens; ++block_offset) { |
| 76 | qk_vec_type k_vec(k_vecs[(block_offset * HEAD_DIM + i) / QK_NUM_ELEM]); |
| 77 | |
| 78 | #pragma unroll |
| 79 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) |
| 80 | vec_op::fma(acc_vec[block_offset][unroll], q_vec[unroll], k_vec); |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | for (int block_offset = 0; block_offset < num_tokens; ++block_offset) { |
| 85 | #pragma unroll |
| 86 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) { |
| 87 | const float acc = acc_vec[block_offset][unroll].reduce_sum() * scale; |
| 88 | logits[block_offset][unroll] = acc; |
| 89 | max_val[unroll] = std::max(max_val[unroll], acc); |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | float sum_exp[HEAD_UNROLL] = {}; |
| 94 | for (int block_offset = 0; block_offset < num_tokens; ++block_offset) { |
| 95 | #pragma unroll |
| 96 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) { |
| 97 | const float val = |
| 98 | std::exp(logits[block_offset][unroll] - max_val[unroll]); |
| 99 | logits[block_offset][unroll] = val; |
| 100 | sum_exp[unroll] += val; |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | f32_vec_type this_out[V_HEAD_DIM / V_NUM_ELEM][HEAD_UNROLL]; |
| 105 | // Same zero-initialization concern as acc_vec above. |
| 106 | for (int i = 0; i < V_HEAD_DIM / V_NUM_ELEM; ++i) |
| 107 | for (int u = 0; u < HEAD_UNROLL; ++u) this_out[i][u] = f32_vec_type{0.0f}; |
| 108 | |
| 109 | for (int block_offset = 0; block_offset < num_tokens; ++block_offset) { |
| 110 | // load to registers |
| 111 | f32_vec_type scale_[HEAD_UNROLL]; |
| 112 | |
| 113 | #pragma unroll |
| 114 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) |
| 115 | scale_[unroll] = |
| 116 | f32_vec_type{logits[block_offset][unroll] / sum_exp[unroll]}; |
| 117 | |
| 118 | for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) { |
| 119 | f32_vec_type v_vec( |
| 120 | v_vecs_f32[(block_offset * HEAD_DIM + i) / V_NUM_ELEM]); |
| 121 | |
| 122 | #pragma unroll |
| 123 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) |
| 124 | vec_op::fma(this_out[i / V_NUM_ELEM][unroll], v_vec, scale_[unroll]); |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | // merge attention state |
| 129 | // section 2.2 in https://arxiv.org/pdf/2501.01005 |
| 130 | f32_vec_type prev_scale[HEAD_UNROLL]; |
| 131 | f32_vec_type curr_scale[HEAD_UNROLL]; |
| 132 | |
| 133 | #pragma unroll |
| 134 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) { |
| 135 | const float prev_lse = acc_lse[unroll]; |
| 136 | const float curr_lse = std::log(sum_exp[unroll]) + |
| 137 | max_val[unroll]; // add back max_val to get true lse |
| 138 | // softmax trick |
| 139 | const float max_lse = std::max(prev_lse, curr_lse); |
| 140 | const float prev_sum_exp = std::exp(prev_lse - max_lse); |
| 141 | const float curr_sum_exp = std::exp(curr_lse - max_lse); |
| 142 | |
| 143 | const float new_sum_exp = prev_sum_exp + curr_sum_exp; |
| 144 | acc_lse[unroll] = std::log(new_sum_exp) + max_lse; |
| 145 | |
| 146 | prev_scale[unroll] = f32_vec_type{prev_sum_exp / new_sum_exp}; |
| 147 | curr_scale[unroll] = f32_vec_type{curr_sum_exp / new_sum_exp}; |
| 148 | } |
| 149 | |
| 150 | for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) { |
| 151 | #pragma unroll |
| 152 | for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) { |
| 153 | f32_vec_type o_vec(acc_out + i + V_HEAD_DIM * unroll); |
| 154 | o_vec = o_vec * prev_scale[unroll] + |
| 155 | this_out[i / V_NUM_ELEM][unroll] * curr_scale[unroll]; |
| 156 | o_vec.save(acc_out + i + V_HEAD_DIM * unroll); |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | q_vecs += HEAD_DIM / QK_NUM_ELEM * HEAD_UNROLL; |
| 161 | acc_out += V_HEAD_DIM * HEAD_UNROLL; |
| 162 | } |
| 163 | |
| 164 | template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE, |
| 165 | typename qk_vec_type> |
| 166 | void mla_decode_block( |
| 167 | const qk_vec_type* __restrict__ q_vecs, // [num_heads, head_dim] |
| 168 | const scalar_t* __restrict__ kv_cache, // [block_size, head_dim] |
| 169 | float* __restrict__ acc_out, // [num_heads, v_head_dim] |
| 170 | float* __restrict__ acc_lse, // [num_heads] |
| 171 | const int num_heads, const float scale, const int num_tokens) { |
| 172 | using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type; |
| 173 | static_assert( |
| 174 | std::is_same<qk_vec_type, |
| 175 | typename KernelVecType<scalar_t>::qk_vec_type>::value); |
| 176 | using v_load_vec_type = typename KernelVecType<scalar_t>::v_load_vec_type; |
| 177 | using f32_vec_type = vec_op::FP32Vec16; |
| 178 | static_assert(qk_load_vec_type::VEC_ELEM_NUM == qk_vec_type::VEC_ELEM_NUM); |
| 179 | static_assert(v_load_vec_type::VEC_ELEM_NUM == f32_vec_type::VEC_ELEM_NUM); |
| 180 | constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM; |
| 181 | constexpr int V_NUM_ELEM = v_load_vec_type::VEC_ELEM_NUM; |
| 182 | |
| 183 | const qk_vec_type* k_vecs; |
| 184 | const f32_vec_type* v_vecs_f32; |
| 185 | float* kv_cache_f32 = nullptr; |
| 186 | |
| 187 | if constexpr (!std::is_same<scalar_t, float>::value) { |
| 188 | // convert KV cache block to FP32 to reuse it across query heads and |
| 189 | // attn @ V computation, since FP16/BF16->FP32 is expensive. |
| 190 | // TODO: move malloc outside of this fn to reuse across iterations. |
| 191 | const int nbytes = BLOCK_SIZE * HEAD_DIM * sizeof(float); |
| 192 | kv_cache_f32 = static_cast<float*>(std::aligned_alloc(64, nbytes)); |
| 193 | |
| 194 | for (int block_offset = 0; block_offset < num_tokens; ++block_offset) |
| 195 | for (int i = 0; i < HEAD_DIM; i += V_NUM_ELEM) { |
| 196 | v_load_vec_type kv_load_vec(kv_cache + block_offset * HEAD_DIM + i); |
| 197 | f32_vec_type kv_vec_f32(kv_load_vec); |
| 198 | kv_vec_f32.save(kv_cache_f32 + block_offset * HEAD_DIM + i); |
| 199 | } |
| 200 | |
| 201 | if constexpr (std::is_same<qk_load_vec_type, qk_vec_type>::value) { |
| 202 | // for AVX512_BF16, Q @ K.T uses BF16 for K (no conversion) |
| 203 | // NOTE: in this case, we only need to convert the V section to FP32. |
| 204 | // But for simplicity, we will convert the whole KV block to FP32. |
| 205 | k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache); |
| 206 | } else { |
| 207 | k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache_f32); |
| 208 | } |
| 209 | |
| 210 | // attn @ V always use FP32 for V, since attn is FP32. |
| 211 | v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache_f32); |
| 212 | |
| 213 | } else { |
| 214 | // KV cache is FP32. don't need to do anything. |
| 215 | k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache); |
| 216 | v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache); |
| 217 | } |
| 218 | |
| 219 | // compute 2 heads at the same time to improve ILP and |
| 220 | // take advantage of register cache for K and V. |
| 221 | constexpr int HEAD_UNROLL = 2; |
| 222 | for (int iter = 0; iter < num_heads / HEAD_UNROLL; ++iter) { |
| 223 | mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, HEAD_UNROLL>( |
| 224 | q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens); |
| 225 | |
| 226 | q_vecs += HEAD_UNROLL * HEAD_DIM / QK_NUM_ELEM; |
| 227 | acc_out += HEAD_UNROLL * V_HEAD_DIM; |
| 228 | acc_lse += HEAD_UNROLL; |
| 229 | } |
| 230 | |
| 231 | // take care of the remaining heads |
| 232 | for (int iter = 0; iter < num_heads % HEAD_UNROLL; ++iter) { |
| 233 | mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, 1>( |
| 234 | q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens); |
| 235 | |
| 236 | q_vecs += HEAD_DIM / QK_NUM_ELEM; |
| 237 | acc_out += V_HEAD_DIM; |
| 238 | acc_lse += 1; |
| 239 | } |
| 240 | |
| 241 | if (kv_cache_f32 != nullptr) { |
| 242 | std::free(kv_cache_f32); |
| 243 | } |
| 244 | } |
| 245 | } // namespace |
| 246 | |
| 247 | template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE> |
| 248 | void mla_decode_kvcache_cpu_impl( |
| 249 | scalar_t* __restrict__ out, // [num_seqs, num_heads, v_head_dim] |
| 250 | const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_dim] |
| 251 | const scalar_t* __restrict__ kv_cache, // [num_blocks, block_size, |
| 252 | // head_dim] |
| 253 | const int num_heads, const float scale, |
| 254 | const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] |
| 255 | const int* __restrict__ seq_lens, // [num_seqs] |
| 256 | const int max_num_blocks_per_seq, const int o_stride, const int q_stride, |
| 257 | const int kv_stride, const int num_seqs) { |
| 258 | using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type; |
| 259 | using qk_vec_type = typename KernelVecType<scalar_t>::qk_vec_type; |
| 260 | constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM; |
| 261 | |
| 262 | // shared across threads |
| 263 | const int max_threads = cpu_utils::get_max_threads(); |
| 264 | const int acc_out_nbytes = |
| 265 | max_threads * num_heads * V_HEAD_DIM * sizeof(float); |
| 266 | float* acc_out = static_cast<float*>(std::aligned_alloc(64, acc_out_nbytes)); |
| 267 | std::vector<float> acc_lse(max_threads * num_heads); |
| 268 | |
| 269 | // allocate memory to pre-convert query to FP32 later |
| 270 | float* q_f32; |
| 271 | constexpr bool PRE_CONVERT_QUERY = |
| 272 | !std::is_same<scalar_t, float>::value && |
| 273 | std::is_same<qk_vec_type, vec_op::FP32Vec16>::value; |
| 274 | if constexpr (PRE_CONVERT_QUERY) { |
| 275 | const int q_f32_nbytes = num_heads * HEAD_DIM * sizeof(float); |
| 276 | q_f32 = static_cast<float*>(std::aligned_alloc(64, q_f32_nbytes)); |
| 277 | } |
| 278 | |
| 279 | #pragma omp parallel |
| 280 | { |
| 281 | const int num_threads = omp_get_num_threads(); |
| 282 | const int thread_id = omp_get_thread_num(); |
| 283 | float* __restrict__ acc_out_thread = |
| 284 | acc_out + thread_id * num_heads * V_HEAD_DIM; |
| 285 | float* __restrict__ acc_lse_thread = acc_lse.data() + thread_id * num_heads; |
| 286 | |
| 287 | for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) { |
| 288 | // reset accumulator |
| 289 | std::fill(acc_out_thread, acc_out_thread + num_heads * V_HEAD_DIM, 0.0f); |
| 290 | std::fill(acc_lse_thread, acc_lse_thread + num_heads, -FLT_MAX); |
| 291 | |
| 292 | const int seq_len = seq_lens[seq_idx]; |
| 293 | const int block_num = (seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE; |
| 294 | const int last_block_size = seq_len - (block_num - 1) * BLOCK_SIZE; |
| 295 | |
| 296 | const qk_vec_type* q_vecs; |
| 297 | if constexpr (PRE_CONVERT_QUERY) { |
| 298 | // pre-convert query to FP32 since FP16/BF16->FP32 is slow. |
| 299 | #pragma omp for |
| 300 | for (int i = 0; i < num_heads * HEAD_DIM; i += QK_NUM_ELEM) { |
| 301 | qk_load_vec_type q_load_vec(q + seq_idx * q_stride + i); |
| 302 | qk_vec_type q_vec(q_load_vec); |
| 303 | q_vec.save(q_f32 + i); |
| 304 | } |
| 305 | q_vecs = reinterpret_cast<const qk_vec_type*>(q_f32); |
| 306 | } else { |
| 307 | q_vecs = reinterpret_cast<const qk_vec_type*>(q + seq_idx * q_stride); |
| 308 | } |
| 309 | |
| 310 | #pragma omp for |
| 311 | for (int block_idx = 0; block_idx < block_num; ++block_idx) { |
| 312 | const int physical_block_idx = |
| 313 | block_tables[seq_idx * max_num_blocks_per_seq + block_idx]; |
| 314 | const int num_tokens = |
| 315 | block_idx < block_num - 1 ? BLOCK_SIZE : last_block_size; |
| 316 | |
| 317 | mla_decode_block<scalar_t, HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE>( |
| 318 | q_vecs, kv_cache + physical_block_idx * kv_stride, acc_out_thread, |
| 319 | acc_lse_thread, num_heads, scale, num_tokens); |
| 320 | } |
| 321 | |
| 322 | // merge attention states across threads |
| 323 | // section 2.2 in https://arxiv.org/pdf/2501.01005 |
| 324 | // each thread is responsible for 1 head |
| 325 | #pragma omp for |
| 326 | for (int head_idx = 0; head_idx < num_heads; ++head_idx) { |
| 327 | float* acc_lse_head = acc_lse.data() + head_idx; |
| 328 | float* acc_out_head = acc_out + head_idx * V_HEAD_DIM; |
| 329 | |
| 330 | float max_val = -FLT_MAX; |
| 331 | for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) { |
| 332 | max_val = std::max(max_val, acc_lse_head[thread_id_ * num_heads]); |
| 333 | } |
| 334 | |
| 335 | float sum_exp = 0.0f; |
| 336 | for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) { |
| 337 | float val = std::exp(acc_lse_head[thread_id_ * num_heads] - max_val); |
| 338 | acc_lse_head[thread_id_ * num_heads] = val; |
| 339 | sum_exp += val; |
| 340 | } |
| 341 | |
| 342 | float inv_sum = 1.0f / sum_exp; |
| 343 | float out_head[V_HEAD_DIM] = {}; |
| 344 | for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) { |
| 345 | float scale_ = acc_lse_head[thread_id_ * num_heads] * inv_sum; |
| 346 | for (int i = 0; i < V_HEAD_DIM; ++i) { |
| 347 | out_head[i] += |
| 348 | acc_out_head[thread_id_ * num_heads * V_HEAD_DIM + i] * scale_; |
| 349 | } |
| 350 | } |
| 351 | |
| 352 | for (int i = 0; i < V_HEAD_DIM; ++i) { |
| 353 | vec_op::storeFP32(out_head[i], out + seq_idx * o_stride + |
| 354 | head_idx * V_HEAD_DIM + i); |
| 355 | } |
| 356 | } |
| 357 | } |
| 358 | } |
| 359 | if (PRE_CONVERT_QUERY) { |
| 360 | std::free(q_f32); |
| 361 | } |
| 362 | std::free(acc_out); |
| 363 | } |
| 364 | |
| 365 | void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query, |
| 366 | torch::Tensor& kv_cache, double scale, |
| 367 | torch::Tensor& block_tables, torch::Tensor& seq_lens) { |
| 368 | const int num_seqs = query.size(0); |
| 369 | const int num_heads = query.size(1); |
| 370 | const int head_dim = query.size(2); |
| 371 | const int block_size = kv_cache.size(1); |
| 372 | const int v_head_dim = out.size(2); |
| 373 | |
| 374 | const int max_num_blocks_per_seq = block_tables.size(1); |
| 375 | const int o_stride = out.stride(0); |
| 376 | const int q_stride = query.stride(0); |
| 377 | const int kv_stride = kv_cache.stride(0); |
| 378 | |
| 379 | VLLM_DISPATCH_FLOATING_TYPES( |
| 380 | query.scalar_type(), "mla_decode_kvcache_cpu_impl", [&] { |
| 381 | CPU_KERNEL_GUARD_IN(mla_decode_kvcache_cpu_impl) |
| 382 | if (head_dim == 576 && v_head_dim == 512 && block_size == 16) |
| 383 | mla_decode_kvcache_cpu_impl<scalar_t, 576, 512, 16>( |
| 384 | out.data_ptr<scalar_t>(), query.data_ptr<scalar_t>(), |
| 385 | kv_cache.data_ptr<scalar_t>(), num_heads, scale, |
| 386 | block_tables.data_ptr<int>(), seq_lens.data_ptr<int>(), |
| 387 | max_num_blocks_per_seq, o_stride, q_stride, kv_stride, num_seqs); |
| 388 | else |
| 389 | TORCH_CHECK(false, "Unsupported block size: ", block_size); |
| 390 | CPU_KERNEL_GUARD_OUT(mla_decode_kvcache_cpu_impl) |
| 391 | }); |
| 392 | } |