From 9f5117820fb0e25ca76a1c7d0ec5c4c6766c4c4e Mon Sep 17 00:00:00 2001 From: Rukhaiya2004 <162583766+Rukhaiya2004@users.noreply.github.com> Date: Tue, 23 Jun 2026 18:54:49 +0530 Subject: [PATCH] [HARDWARE][POWER] Enable fp16 support for PowerPC (#46135) Signed-off-by: Rukhaiya --- csrc/cpu/cpu_attn_impl.hpp | 2 - csrc/cpu/cpu_attn_vsx.hpp | 11 +- csrc/cpu/cpu_types_vsx.hpp | 222 ++++++++++++++++++++++++++++--------- csrc/cpu/mla_decode.cpp | 8 -- csrc/cpu/pos_encoding.cpp | 155 +++++++++++++++++++++++++- vllm/platforms/cpu.py | 2 +- 6 files changed, 332 insertions(+), 68 deletions(-) diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index d1b6c71c182..7b3757b313d 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -887,12 +887,10 @@ struct VecTypeTrait { using vec_t = vec_op::BF16Vec16; }; -#if !defined(__powerpc__) template <> struct VecTypeTrait { using vec_t = vec_op::FP16Vec16; }; -#endif template void print_logits(const char* name, T* ptr, int32_t row, int32_t col, diff --git a/csrc/cpu/cpu_attn_vsx.hpp b/csrc/cpu/cpu_attn_vsx.hpp index c7e1502bcb0..dd24b95ba02 100644 --- a/csrc/cpu/cpu_attn_vsx.hpp +++ b/csrc/cpu/cpu_attn_vsx.hpp @@ -50,7 +50,16 @@ FORCE_INLINE void load_row8_B_as_f32(const c10::BFloat16* p, b1 = (__vector float)vec_mergel(zeros, raw); } -// Note: c10::Half (FP16) is not supported on PowerPC architecture +// [3] Half (FP16) Specialization +template <> +FORCE_INLINE void load_row8_B_as_f32(const c10::Half* p, + __vector float& b0, + __vector float& b1) { + vec_op::FP16Vec8 fp16_vec(p); + vec_op::FP32Vec8 fp32_vec(fp16_vec); + b0 = fp32_vec.reg.val[0]; + b1 = fp32_vec.reg.val[1]; +} template FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4( diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index ba65e27a15e..2031e4c14a8 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -13,10 +13,10 @@ namespace vec_op { struct fp8_e4m3_tag {}; struct fp8_e5m2_tag {}; -// FIXME: FP16 is not fully supported in Torch-CPU -#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ - AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ - AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) @@ -34,6 +34,87 @@ struct fp8_e5m2_tag {}; #define FORCE_INLINE __attribute__((always_inline)) inline namespace { + +FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) { + const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000}; + const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00}; + const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF}; + const __vector unsigned int bias_adj = {112, 112, 112, 112}; + const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F}; + const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF}; + + __vector unsigned int s = (x & mask_sign) << 16; + __vector unsigned int e = (x & mask_exp) >> 10; + __vector unsigned int m = (x & mask_mant) << 13; + + __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16); + + __vector unsigned int e_normal = e + bias_adj; + e = vec_sel(e_normal, exp_max_fp32, is_nan_inf); + + return (__vector float)(s | (e << 23) | m); +} + +FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) { + __vector unsigned int in = (__vector unsigned int)f_in; + + const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000, + 0x80000000, 0x80000000}; + const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000, + 0x7F800000}; + const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF, + 0x007FFFFF, 0x007FFFFF}; + + const __vector signed int bias_adj = {112, 112, 112, 112}; + const __vector signed int zero = {0, 0, 0, 0}; + const __vector signed int max_exp = {31, 31, 31, 31}; + const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF}; + const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F}; + + __vector unsigned int s = (in & mask_sign_32) >> 16; + __vector unsigned int e_u = (in & mask_exp_32) >> 23; + + __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32); + + __vector signed int e_s = (__vector signed int)e_u; + e_s = vec_sub(e_s, bias_adj); + e_s = vec_max(e_s, zero); + e_s = vec_min(e_s, max_exp); + __vector unsigned int e_normal = (__vector unsigned int)e_s; + + __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf); + + const __vector unsigned int one_v = {1, 1, 1, 1}; + const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF}; + + __vector unsigned int round_bit = (in >> 12) & one_v; + __vector unsigned int sticky = in & mask_sticky; + __vector unsigned int m = (in & mask_mant_32) >> 13; + __vector unsigned int lsb = m & one_v; + + // Round up if: round_bit && (sticky || lsb) + __vector __bool int sticky_nonzero = + vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0}); + __vector __bool int lsb_set = vec_cmpeq(lsb, one_v); + __vector __bool int round_up = + vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set)); + + m = vec_sel(m, m + one_v, round_up); + + const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF}; + const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E}; + __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask); + __vector __bool int would_overflow_to_inf = + vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp)); + __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16); + e_final = vec_sel(e_final, e_inc, mant_overflows); + m = vec_and(m, mant_mask); + e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf); + m = vec_sel(m, mant_mask, would_overflow_to_inf); + + return s | (e_final << 10) | m; +} + template constexpr void unroll_loop_item(std::integer_sequence, F&& f) { (f(std::integral_constant{}), ...); @@ -89,6 +170,19 @@ struct BF16Vec8 : public Vec { } }; +struct FP16Vec8 : public Vec { + constexpr static int VEC_ELEM_NUM = 8; + + __vector signed short reg; + + explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {} + explicit FP16Vec8(const FP32Vec8&); + + void save(void* ptr) const { + *reinterpret_cast<__vector signed short*>(ptr) = reg; + } +}; + struct FP16Vec16 : public Vec { constexpr static int VEC_ELEM_NUM = 16; ss16x8x2_t reg; @@ -124,13 +218,11 @@ struct BF16Vec16 : public Vec { ss16x8x2_t reg; explicit BF16Vec16(const void* ptr) { - // Load 256 bits in two parts reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr); reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr); } explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {} - explicit BF16Vec16(const FP32Vec16&); void save(void* ptr) const { @@ -142,20 +234,16 @@ struct BF16Vec16 : public Vec { void save(void* ptr, const int elem_num) const { const int clamped_elem = std::max(0, std::min(elem_num, 16)); - // Calculate elements to store in each 128-bit part (8 elements each) const int elements_val0 = std::min(clamped_elem, 8); const int elements_val1 = std::max(clamped_elem - 8, 0); - // Convert elements to bytes (2 bytes per element) const size_t bytes_val0 = elements_val0 * sizeof(signed short); const size_t bytes_val1 = elements_val1 * sizeof(signed short); signed short* dest = static_cast(ptr); - // Store the first part using vec_xst_len if (bytes_val0 > 0) { vec_xst_len(reg.val[0], dest, bytes_val0); } - // Store the second part if needed if (bytes_val1 > 0) { vec_xst_len(reg.val[1], dest + elements_val0, bytes_val1); } @@ -238,6 +326,15 @@ struct FP32Vec8 : public Vec { reg.val[1] = (__vector float)vec_mergel(zero, v.reg); } + explicit FP32Vec8(const FP16Vec8& v) { + __vector unsigned short raw_u = (__vector unsigned short)v.reg; + __vector unsigned int raw_hi = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u); + __vector unsigned int raw_lo = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u); + reg.val[0] = fp16_to_fp32_bits(raw_hi); + reg.val[1] = fp16_to_fp32_bits(raw_lo); + } float reduce_sum() const { AliasReg ar; ar.reg = reg; @@ -410,8 +507,9 @@ struct FP32Vec16 : public Vec { reg.val[3] = vec_xl(48, ptr); } + explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {} + explicit FP32Vec16(const FP16Vec16&); explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {} - explicit FP32Vec16(f32x4x4_t data) : reg(data) {} explicit FP32Vec16(const FP32Vec16& data) { @@ -435,7 +533,6 @@ struct FP32Vec16 : public Vec { reg.val[3] = data.reg.val[1]; } - explicit FP32Vec16(const FP16Vec16& v); explicit FP32Vec16(const BF16Vec16& v) { reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]); reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]); @@ -502,28 +599,20 @@ struct FP32Vec16 : public Vec { FP32Vec16 max(const FP32Vec16& b, int elem_num) const { FP32Vec16 result; - // Create a vector of element indices for each chunk __vector unsigned int indices = {0, 1, 2, 3}; __vector unsigned int elem_num_vec = vec_splats(static_cast(elem_num)); - // Compute masks for each chunk - __vector unsigned int chunk_offset0 = {0, 0, 0, - 0}; // Chunk 0: Elements 0-3 - __vector unsigned int chunk_offset1 = {4, 4, 4, - 4}; // Chunk 1: Elements 4-7 - __vector unsigned int chunk_offset2 = {8, 8, 8, - 8}; // Chunk 2: Elements 8-11 - __vector unsigned int chunk_offset3 = {12, 12, 12, - 12}; // Chunk 3: Elements 12-15 + __vector unsigned int chunk_offset0 = {0, 0, 0, 0}; + __vector unsigned int chunk_offset1 = {4, 4, 4, 4}; + __vector unsigned int chunk_offset2 = {8, 8, 8, 8}; + __vector unsigned int chunk_offset3 = {12, 12, 12, 12}; - // Compute masks for each chunk __vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec); __vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec); __vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec); __vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec); - // Apply masks to compute the result for each chunk result.reg.val[0] = vec_sel(this->reg.val[0], vec_max(this->reg.val[0], b.reg.val[0]), mask0); result.reg.val[1] = vec_sel(this->reg.val[1], @@ -626,6 +715,16 @@ struct FP32Vec16 : public Vec { vec_xst(reg.val[3], 48, ptr); } + void save(c10::Half* ptr) const { + FP16Vec16 fp16_vec(*this); + fp16_vec.save(ptr); + } + + void save(c10::Half* ptr, const int elem_num) const { + FP16Vec16 fp16_vec(*this); + fp16_vec.save(ptr, elem_num); + } + void save(float* ptr, const int elem_num) const { const int elements_in_chunk1 = (elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0; @@ -659,7 +758,7 @@ struct FP32Vec16 : public Vec { }; struct INT8Vec16 : public Vec { - constexpr static int VEC_NUM_ELEM = 16; // 128 bits / 8 bits = 16 + constexpr static int VEC_NUM_ELEM = 16; union AliasReg { __vector signed char reg; @@ -707,6 +806,11 @@ struct VecType { using vec_type = BF16Vec8; }; +template <> +struct VecType { + using vec_type = FP16Vec8; +}; + template void storeFP32(float v, T* ptr) { *ptr = v; @@ -723,6 +827,15 @@ inline void storeFP32(float v, c10::BFloat16* ptr) { *ptr = *(v_ptr + 1); } +template <> +inline void storeFP32(float v, c10::Half* ptr) { + __vector float v_vec = {v, 0.0f, 0.0f, 0.0f}; + __vector unsigned int fp16_bits = fp32_to_fp16_bits(v_vec); + unsigned short result = + (unsigned short)((__vector unsigned short)fp16_bits)[0]; + *reinterpret_cast(ptr) = result; +} + #ifndef __VEC_CLASS_FP_NAN #define __VEC_CLASS_FP_NAN (1 << 6) #endif @@ -769,38 +882,39 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) { #endif } +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { + __vector unsigned int fp16_hi = fp32_to_fp16_bits(v.reg.val[0]); + __vector unsigned int fp16_lo = fp32_to_fp16_bits(v.reg.val[1]); + reg = (__vector signed short)vec_perm((__vector unsigned char)fp16_hi, + (__vector unsigned char)fp16_lo, omask); +} + inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { - alignas(16) float temp_fp32[16]; - alignas(16) c10::Half temp_fp16[16]; - - vec_xst(v.reg.val[0], 0, temp_fp32); - vec_xst(v.reg.val[1], 16, temp_fp32); - vec_xst(v.reg.val[2], 32, temp_fp32); - vec_xst(v.reg.val[3], 48, temp_fp32); - - for (int i = 0; i < 16; i++) { - temp_fp16[i] = c10::Half(temp_fp32[i]); - } - - reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16); - reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16); + __vector unsigned int fp16_0 = fp32_to_fp16_bits(v.reg.val[0]); + __vector unsigned int fp16_1 = fp32_to_fp16_bits(v.reg.val[1]); + __vector unsigned int fp16_2 = fp32_to_fp16_bits(v.reg.val[2]); + __vector unsigned int fp16_3 = fp32_to_fp16_bits(v.reg.val[3]); + reg.val[0] = (__vector signed short)vec_perm( + (__vector unsigned char)fp16_0, (__vector unsigned char)fp16_1, omask); + reg.val[1] = (__vector signed short)vec_perm( + (__vector unsigned char)fp16_2, (__vector unsigned char)fp16_3, omask); } inline FP32Vec16::FP32Vec16(const FP16Vec16& v) { - alignas(16) c10::Half temp_fp16[16]; - alignas(16) float temp_fp32[16]; - - vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16); - vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16); - - for (int i = 0; i < 16; i++) { - temp_fp32[i] = float(temp_fp16[i]); - } - - reg.val[0] = vec_xl(0, temp_fp32); - reg.val[1] = vec_xl(16, temp_fp32); - reg.val[2] = vec_xl(32, temp_fp32); - reg.val[3] = vec_xl(48, temp_fp32); + __vector unsigned short raw_u0 = (__vector unsigned short)v.reg.val[0]; + __vector unsigned short raw_u1 = (__vector unsigned short)v.reg.val[1]; + __vector unsigned int raw_hi0 = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u0); + __vector unsigned int raw_lo0 = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u0); + __vector unsigned int raw_hi1 = + (__vector unsigned int)vec_unpackh((__vector signed short)raw_u1); + __vector unsigned int raw_lo1 = + (__vector unsigned int)vec_unpackl((__vector signed short)raw_u1); + reg.val[0] = fp16_to_fp32_bits(raw_hi0); + reg.val[1] = fp16_to_fp32_bits(raw_lo0); + reg.val[2] = fp16_to_fp32_bits(raw_hi1); + reg.val[3] = fp16_to_fp32_bits(raw_lo1); } inline BF16Vec16::BF16Vec16(const FP32Vec16& v) { @@ -864,7 +978,6 @@ inline void prefetch(const void* addr) { struct INT8Vec64 { __vector signed char data[4]; - INT8Vec64() = default; explicit INT8Vec64(const int8_t* ptr) { @@ -900,5 +1013,4 @@ struct INT8Vec64 { void nt_save(int8_t* ptr) const { save(ptr); } }; } // namespace vec_op - #endif diff --git a/csrc/cpu/mla_decode.cpp b/csrc/cpu/mla_decode.cpp index 582c480c3be..3bd0d2e688f 100644 --- a/csrc/cpu/mla_decode.cpp +++ b/csrc/cpu/mla_decode.cpp @@ -18,17 +18,9 @@ struct KernelVecType { template <> struct KernelVecType { -#if defined(__powerpc64__) - // Power specific vector types - using qk_load_vec_type = vec_op::FP32Vec16; - using qk_vec_type = vec_op::FP32Vec16; - using v_load_vec_type = vec_op::FP32Vec16; -#else - // Fallback for other architectures, including x86 using qk_load_vec_type = vec_op::FP16Vec16; using qk_vec_type = vec_op::FP32Vec16; using v_load_vec_type = vec_op::FP16Vec16; -#endif }; #ifdef __AVX512BF16__ diff --git a/csrc/cpu/pos_encoding.cpp b/csrc/cpu/pos_encoding.cpp index 9f41e4e222b..b241918902e 100644 --- a/csrc/cpu/pos_encoding.cpp +++ b/csrc/cpu/pos_encoding.cpp @@ -1,4 +1,3 @@ - #include "cpu_types.hpp" namespace { @@ -97,6 +96,91 @@ void rotary_embedding_impl( } } +template <> +void rotary_embedding_impl( + const int64_t* __restrict__ positions, c10::Half* __restrict__ query, + c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache, + const int rot_dim, const int64_t query_stride, const int64_t key_stride, + const int num_heads, const int num_kv_heads, const int head_size, + const int num_tokens) { + using scalar_vec_t = vec_op::FP16Vec8; + constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); + + const int embed_dim = rot_dim / 2; + bool flag = (embed_dim % VEC_ELEM_NUM == 0); + const int loop_upper = flag ? embed_dim : embed_dim - VEC_ELEM_NUM; + + auto compute_loop = [&](const int64_t token_head, const c10::Half* cache_ptr, + c10::Half* qk) { + int j = 0; + for (; j < loop_upper; j += VEC_ELEM_NUM) { + const int rot_offset = j; + const int x_index = rot_offset; + const int y_index = embed_dim + rot_offset; + + const int64_t out_x = token_head + x_index; + const int64_t out_y = token_head + y_index; + + const vec_op::FP16Vec8 cos_fp16(cache_ptr + x_index); + const vec_op::FP16Vec8 sin_fp16(cache_ptr + y_index); + const vec_op::FP16Vec8 q_x_fp16(qk + out_x); + const vec_op::FP16Vec8 q_y_fp16(qk + out_y); + + const vec_op::FP32Vec8 fp32_cos(cos_fp16); + const vec_op::FP32Vec8 fp32_sin(sin_fp16); + const vec_op::FP32Vec8 fp32_q_x(q_x_fp16); + const vec_op::FP32Vec8 fp32_q_y(q_y_fp16); + + auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin; + auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin; + + vec_op::FP16Vec8(out1).save(qk + out_x); + vec_op::FP16Vec8(out2).save(qk + out_y); + } + if (!flag) { + for (; j < embed_dim; ++j) { + const int x_index = j; + const int y_index = embed_dim + j; + + const int64_t out_x = token_head + x_index; + const int64_t out_y = token_head + y_index; + + const float fp32_cos = static_cast(cache_ptr[x_index]); + const float fp32_sin = static_cast(cache_ptr[y_index]); + const float fp32_q_x = static_cast(qk[out_x]); + const float fp32_q_y = static_cast(qk[out_y]); + + qk[out_x] = + static_cast(fp32_q_x * fp32_cos - fp32_q_y * fp32_sin); + qk[out_y] = + static_cast(fp32_q_y * fp32_cos + fp32_q_x * fp32_sin); + } + } + }; + +#pragma omp parallel for + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + + for (int i = 0; i < num_heads; ++i) { + const int head_idx = i; + const int64_t token_head = + token_idx * query_stride + head_idx * head_size; + compute_loop(token_head, cache_ptr, query); + } + + if (key != nullptr) { + for (int i = 0; i < num_kv_heads; ++i) { + const int head_idx = i; + const int64_t token_head = + token_idx * key_stride + head_idx * head_size; + compute_loop(token_head, cache_ptr, key); + } + } + } +} + template void rotary_embedding_gptj_impl( const int64_t* __restrict__ positions, // [batch_size, seq_len] or @@ -174,6 +258,75 @@ void rotary_embedding_gptj_impl( } } } + +template <> +void rotary_embedding_gptj_impl( + const int64_t* __restrict__ positions, c10::Half* __restrict__ query, + c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache, + const int rot_dim, const int64_t query_stride, const int64_t key_stride, + const int num_heads, const int num_kv_heads, const int head_size, + const int num_tokens) { + const int embed_dim = rot_dim / 2; + +#pragma omp parallel for collapse(2) + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + for (int i = 0; i < num_heads; ++i) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + const c10::Half* cos_cache_ptr = cache_ptr; + const c10::Half* sin_cache_ptr = cache_ptr + embed_dim; + const int head_idx = i; + const int64_t token_head = + token_idx * query_stride + head_idx * head_size; + c10::Half* head_query = token_head + query; + for (int j = 0; j < embed_dim; j += 1) { + const int rot_offset = j; + const int x_index = 2 * rot_offset; + const int y_index = 2 * rot_offset + 1; + + const float cos = static_cast(cos_cache_ptr[rot_offset]); + const float sin = static_cast(sin_cache_ptr[rot_offset]); + + const float x = static_cast(head_query[x_index]); + const float y = static_cast(head_query[y_index]); + + head_query[x_index] = static_cast(x * cos - y * sin); + head_query[y_index] = static_cast(y * cos + x * sin); + } + } + } + + if (key == nullptr) { + return; + } + +#pragma omp parallel for collapse(2) + for (int token_idx = 0; token_idx < num_tokens; ++token_idx) { + for (int i = 0; i < num_kv_heads; ++i) { + int64_t pos = positions[token_idx]; + const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim; + const c10::Half* cos_cache_ptr = cache_ptr; + const c10::Half* sin_cache_ptr = cache_ptr + embed_dim; + const int head_idx = i; + const int64_t token_head = token_idx * key_stride + head_idx * head_size; + c10::Half* head_key = key + token_head; + for (int j = 0; j < embed_dim; j += 1) { + const int rot_offset = j; + const int x_index = 2 * rot_offset; + const int y_index = 2 * rot_offset + 1; + + const float cos = static_cast(cos_cache_ptr[rot_offset]); + const float sin = static_cast(sin_cache_ptr[rot_offset]); + + const float x = static_cast(head_key[x_index]); + const float y = static_cast(head_key[y_index]); + + head_key[x_index] = static_cast(x * cos - y * sin); + head_key[y_index] = static_cast(y * cos + x * sin); + } + } + } +} }; // namespace void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index b1414665869..c529af46df9 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -50,7 +50,7 @@ class CpuPlatform(Platform): @property def supported_dtypes(self) -> list[torch.dtype]: if self.get_cpu_architecture() == CpuArchEnum.POWERPC: - return [torch.bfloat16, torch.float32] + return [torch.bfloat16, torch.float32, torch.float16] elif self.get_cpu_architecture() == CpuArchEnum.ARM and sys.platform.startswith( "darwin" ):