[HARDWARE][POWER] Enable fp16 support for PowerPC (#46135)
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com>
Rukhaiya2004 committed
Jun 23, 2026 at 18:54 UTC
9f5117820fb0e25ca76a1c7d0ec5c4c6766c4c4e
6 files changed
+333
-69
csrc/cpu/cpu_attn_impl.hpp
-2
@@ -887,12 +887,10 @@ struct VecTypeTrait<c10::BFloat16> {
887
using vec_t = vec_op::BF16Vec16;
888
};
889
890
-#if !defined(__powerpc__)
890
template <>
891
struct VecTypeTrait<c10::Half> {
892
using vec_t = vec_op::FP16Vec16;
893
};
895
-#endif
894
895
template <typename T>
896
void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
csrc/cpu/cpu_attn_vsx.hpp
+10
-1
@@ -50,7 +50,16 @@ FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
50
b1 = (__vector float)vec_mergel(zeros, raw);
51
}
52
53
-// Note: c10::Half (FP16) is not supported on PowerPC architecture
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(
csrc/cpu/cpu_types_vsx.hpp
+168
-56
@@ -13,10 +13,10 @@ namespace vec_op {
13
struct fp8_e4m3_tag {};
14
struct fp8_e5m2_tag {};
15
16
-// FIXME: FP16 is not fully supported in Torch-CPU
17
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
18
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
19
- AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
16
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
17
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
18
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
19
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
20
21
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
22
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -34,6 +34,87 @@ struct fp8_e5m2_tag {};
34
#define FORCE_INLINE __attribute__((always_inline)) inline
35
36
namespace {
37
+
38
+FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
39
+ const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
40
+ const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
41
+ const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
42
+ const __vector unsigned int bias_adj = {112, 112, 112, 112};
43
+ const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
44
+ const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
45
+
46
+ __vector unsigned int s = (x & mask_sign) << 16;
47
+ __vector unsigned int e = (x & mask_exp) >> 10;
48
+ __vector unsigned int m = (x & mask_mant) << 13;
49
+
50
+ __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
51
+
52
+ __vector unsigned int e_normal = e + bias_adj;
53
+ e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
54
+
55
+ return (__vector float)(s | (e << 23) | m);
56
+}
57
+
58
+FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
59
+ __vector unsigned int in = (__vector unsigned int)f_in;
60
+
61
+ const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
62
+ 0x80000000, 0x80000000};
63
+ const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
64
+ 0x7F800000};
65
+ const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
66
+ 0x007FFFFF, 0x007FFFFF};
67
+
68
+ const __vector signed int bias_adj = {112, 112, 112, 112};
69
+ const __vector signed int zero = {0, 0, 0, 0};
70
+ const __vector signed int max_exp = {31, 31, 31, 31};
71
+ const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
72
+ const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
73
+
74
+ __vector unsigned int s = (in & mask_sign_32) >> 16;
75
+ __vector unsigned int e_u = (in & mask_exp_32) >> 23;
76
+
77
+ __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
78
+
79
+ __vector signed int e_s = (__vector signed int)e_u;
80
+ e_s = vec_sub(e_s, bias_adj);
81
+ e_s = vec_max(e_s, zero);
82
+ e_s = vec_min(e_s, max_exp);
83
+ __vector unsigned int e_normal = (__vector unsigned int)e_s;
84
+
85
+ __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
86
+
87
+ const __vector unsigned int one_v = {1, 1, 1, 1};
88
+ const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
89
+
90
+ __vector unsigned int round_bit = (in >> 12) & one_v;
91
+ __vector unsigned int sticky = in & mask_sticky;
92
+ __vector unsigned int m = (in & mask_mant_32) >> 13;
93
+ __vector unsigned int lsb = m & one_v;
94
+
95
+ // Round up if: round_bit && (sticky || lsb)
96
+ __vector __bool int sticky_nonzero =
97
+ vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
98
+ __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
99
+ __vector __bool int round_up =
100
+ vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
101
+
102
+ m = vec_sel(m, m + one_v, round_up);
103
+
104
+ const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
105
+ const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
106
+ __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
107
+ __vector __bool int would_overflow_to_inf =
108
+ vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
109
+ __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
110
+ e_final = vec_sel(e_final, e_inc, mant_overflows);
111
+ m = vec_and(m, mant_mask);
112
+ e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
113
+ m = vec_sel(m, mant_mask, would_overflow_to_inf);
114
+
115
+ return s | (e_final << 10) | m;
116
+}
117
+
118
template <typename T, T... indexes, typename F>
119
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
120
(f(std::integral_constant<T, indexes>{}), ...);
@@ -89,6 +170,19 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
170
}
171
};
172
173
+struct FP16Vec8 : public Vec<FP16Vec8> {
174
+ constexpr static int VEC_ELEM_NUM = 8;
175
+
176
+ __vector signed short reg;
177
+
178
+ explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
179
+ explicit FP16Vec8(const FP32Vec8&);
180
+
181
+ void save(void* ptr) const {
182
+ *reinterpret_cast<__vector signed short*>(ptr) = reg;
183
+ }
184
+};
185
+
186
struct FP16Vec16 : public Vec<FP16Vec16> {
187
constexpr static int VEC_ELEM_NUM = 16;
188
ss16x8x2_t reg;
@@ -124,13 +218,11 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
218
ss16x8x2_t reg;
219
220
explicit BF16Vec16(const void* ptr) {
127
- // Load 256 bits in two parts
221
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
222
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
223
}
224
225
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
133
-
226
explicit BF16Vec16(const FP32Vec16&);
227
228
void save(void* ptr) const {
@@ -142,20 +234,16 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
234
void save(void* ptr, const int elem_num) const {
235
const int clamped_elem = std::max(0, std::min(elem_num, 16));
236
145
- // Calculate elements to store in each 128-bit part (8 elements each)
237
const int elements_val0 = std::min(clamped_elem, 8);
238
const int elements_val1 = std::max(clamped_elem - 8, 0);
239
149
- // Convert elements to bytes (2 bytes per element)
240
const size_t bytes_val0 = elements_val0 * sizeof(signed short);
241
const size_t bytes_val1 = elements_val1 * sizeof(signed short);
242
243
signed short* dest = static_cast<signed short*>(ptr);
154
- // Store the first part using vec_xst_len
244
if (bytes_val0 > 0) {
245
vec_xst_len(reg.val[0], dest, bytes_val0);
246
}
158
- // Store the second part if needed
247
if (bytes_val1 > 0) {
248
vec_xst_len(reg.val[1], dest + elements_val0, bytes_val1);
249
}
@@ -238,6 +326,15 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
326
reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
327
}
328
329
+ explicit FP32Vec8(const FP16Vec8& v) {
330
+ __vector unsigned short raw_u = (__vector unsigned short)v.reg;
331
+ __vector unsigned int raw_hi =
332
+ (__vector unsigned int)vec_unpackh((__vector signed short)raw_u);
333
+ __vector unsigned int raw_lo =
334
+ (__vector unsigned int)vec_unpackl((__vector signed short)raw_u);
335
+ reg.val[0] = fp16_to_fp32_bits(raw_hi);
336
+ reg.val[1] = fp16_to_fp32_bits(raw_lo);
337
+ }
338
float reduce_sum() const {
339
AliasReg ar;
340
ar.reg = reg;
@@ -410,8 +507,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
507
reg.val[3] = vec_xl(48, ptr);
508
}
509
510
+ explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {}
511
+ explicit FP32Vec16(const FP16Vec16&);
512
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
414
-
513
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
514
515
explicit FP32Vec16(const FP32Vec16& data) {
@@ -435,7 +533,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
533
reg.val[3] = data.reg.val[1];
534
}
535
438
- explicit FP32Vec16(const FP16Vec16& v);
536
explicit FP32Vec16(const BF16Vec16& v) {
537
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
538
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -502,28 +599,20 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
599
FP32Vec16 max(const FP32Vec16& b, int elem_num) const {
600
FP32Vec16 result;
601
505
- // Create a vector of element indices for each chunk
602
__vector unsigned int indices = {0, 1, 2, 3};
603
__vector unsigned int elem_num_vec =
604
vec_splats(static_cast<unsigned int>(elem_num));
605
510
- // Compute masks for each chunk
511
- __vector unsigned int chunk_offset0 = {0, 0, 0,
512
- 0}; // Chunk 0: Elements 0-3
513
- __vector unsigned int chunk_offset1 = {4, 4, 4,
514
- 4}; // Chunk 1: Elements 4-7
515
- __vector unsigned int chunk_offset2 = {8, 8, 8,
516
- 8}; // Chunk 2: Elements 8-11
517
- __vector unsigned int chunk_offset3 = {12, 12, 12,
518
- 12}; // Chunk 3: Elements 12-15
519
-
520
- // Compute masks for each chunk
606
+ __vector unsigned int chunk_offset0 = {0, 0, 0, 0};
607
+ __vector unsigned int chunk_offset1 = {4, 4, 4, 4};
608
+ __vector unsigned int chunk_offset2 = {8, 8, 8, 8};
609
+ __vector unsigned int chunk_offset3 = {12, 12, 12, 12};
610
+
611
__vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec);
612
__vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec);
613
__vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec);
614
__vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec);
615
526
- // Apply masks to compute the result for each chunk
616
result.reg.val[0] = vec_sel(this->reg.val[0],
617
vec_max(this->reg.val[0], b.reg.val[0]), mask0);
618
result.reg.val[1] = vec_sel(this->reg.val[1],
@@ -626,6 +715,16 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
715
vec_xst(reg.val[3], 48, ptr);
716
}
717
718
+ void save(c10::Half* ptr) const {
719
+ FP16Vec16 fp16_vec(*this);
720
+ fp16_vec.save(ptr);
721
+ }
722
+
723
+ void save(c10::Half* ptr, const int elem_num) const {
724
+ FP16Vec16 fp16_vec(*this);
725
+ fp16_vec.save(ptr, elem_num);
726
+ }
727
+
728
void save(float* ptr, const int elem_num) const {
729
const int elements_in_chunk1 =
730
(elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0;
@@ -659,7 +758,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
758
};
759
760
struct INT8Vec16 : public Vec<INT8Vec16> {
662
- constexpr static int VEC_NUM_ELEM = 16; // 128 bits / 8 bits = 16
761
+ constexpr static int VEC_NUM_ELEM = 16;
762
763
union AliasReg {
764
__vector signed char reg;
@@ -707,6 +806,11 @@ struct VecType<c10::BFloat16> {
806
using vec_type = BF16Vec8;
807
};
808
809
+template <>
810
+struct VecType<c10::Half> {
811
+ using vec_type = FP16Vec8;
812
+};
813
+
814
template <typename T>
815
void storeFP32(float v, T* ptr) {
816
*ptr = v;
@@ -723,6 +827,15 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
827
*ptr = *(v_ptr + 1);
828
}
829
830
+template <>
831
+inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
832
+ __vector float v_vec = {v, 0.0f, 0.0f, 0.0f};
833
+ __vector unsigned int fp16_bits = fp32_to_fp16_bits(v_vec);
834
+ unsigned short result =
835
+ (unsigned short)((__vector unsigned short)fp16_bits)[0];
836
+ *reinterpret_cast<unsigned short*>(ptr) = result;
837
+}
838
+
839
#ifndef __VEC_CLASS_FP_NAN
840
#define __VEC_CLASS_FP_NAN (1 << 6)
841
#endif
@@ -769,38 +882,39 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
882
#endif
883
}
884
772
-inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
773
- alignas(16) float temp_fp32[16];
774
- alignas(16) c10::Half temp_fp16[16];
775
-
776
- vec_xst(v.reg.val[0], 0, temp_fp32);
777
- vec_xst(v.reg.val[1], 16, temp_fp32);
778
- vec_xst(v.reg.val[2], 32, temp_fp32);
779
- vec_xst(v.reg.val[3], 48, temp_fp32);
780
-
781
- for (int i = 0; i < 16; i++) {
782
- temp_fp16[i] = c10::Half(temp_fp32[i]);
783
- }
885
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
886
+ __vector unsigned int fp16_hi = fp32_to_fp16_bits(v.reg.val[0]);
887
+ __vector unsigned int fp16_lo = fp32_to_fp16_bits(v.reg.val[1]);
888
+ reg = (__vector signed short)vec_perm((__vector unsigned char)fp16_hi,
889
+ (__vector unsigned char)fp16_lo, omask);
890
+}
891
785
- reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
786
- reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
892
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
893
+ __vector unsigned int fp16_0 = fp32_to_fp16_bits(v.reg.val[0]);
894
+ __vector unsigned int fp16_1 = fp32_to_fp16_bits(v.reg.val[1]);
895
+ __vector unsigned int fp16_2 = fp32_to_fp16_bits(v.reg.val[2]);
896
+ __vector unsigned int fp16_3 = fp32_to_fp16_bits(v.reg.val[3]);
897
+ reg.val[0] = (__vector signed short)vec_perm(
898
+ (__vector unsigned char)fp16_0, (__vector unsigned char)fp16_1, omask);
899
+ reg.val[1] = (__vector signed short)vec_perm(
900
+ (__vector unsigned char)fp16_2, (__vector unsigned char)fp16_3, omask);
901
}
902
903
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
790
- alignas(16) c10::Half temp_fp16[16];
791
- alignas(16) float temp_fp32[16];
792
-
793
- vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
794
- vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
795
-
796
- for (int i = 0; i < 16; i++) {
797
- temp_fp32[i] = float(temp_fp16[i]);
798
- }
799
-
800
- reg.val[0] = vec_xl(0, temp_fp32);
801
- reg.val[1] = vec_xl(16, temp_fp32);
802
- reg.val[2] = vec_xl(32, temp_fp32);
803
- reg.val[3] = vec_xl(48, temp_fp32);
904
+ __vector unsigned short raw_u0 = (__vector unsigned short)v.reg.val[0];
905
+ __vector unsigned short raw_u1 = (__vector unsigned short)v.reg.val[1];
906
+ __vector unsigned int raw_hi0 =
907
+ (__vector unsigned int)vec_unpackh((__vector signed short)raw_u0);
908
+ __vector unsigned int raw_lo0 =
909
+ (__vector unsigned int)vec_unpackl((__vector signed short)raw_u0);
910
+ __vector unsigned int raw_hi1 =
911
+ (__vector unsigned int)vec_unpackh((__vector signed short)raw_u1);
912
+ __vector unsigned int raw_lo1 =
913
+ (__vector unsigned int)vec_unpackl((__vector signed short)raw_u1);
914
+ reg.val[0] = fp16_to_fp32_bits(raw_hi0);
915
+ reg.val[1] = fp16_to_fp32_bits(raw_lo0);
916
+ reg.val[2] = fp16_to_fp32_bits(raw_hi1);
917
+ reg.val[3] = fp16_to_fp32_bits(raw_lo1);
918
}
919
920
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
@@ -864,7 +978,6 @@ inline void prefetch(const void* addr) {
978
979
struct INT8Vec64 {
980
__vector signed char data[4];
867
-
981
INT8Vec64() = default;
982
983
explicit INT8Vec64(const int8_t* ptr) {
@@ -900,5 +1013,4 @@ struct INT8Vec64 {
1013
void nt_save(int8_t* ptr) const { save(ptr); }
1014
};
1015
} // namespace vec_op
903
-
1016
#endif
csrc/cpu/mla_decode.cpp
-8
@@ -18,17 +18,9 @@ struct KernelVecType<float> {
18
19
template <>
20
struct KernelVecType<c10::Half> {
21
-#if defined(__powerpc64__)
22
- // Power specific vector types
23
- using qk_load_vec_type = vec_op::FP32Vec16;
24
- using qk_vec_type = vec_op::FP32Vec16;
25
- using v_load_vec_type = vec_op::FP32Vec16;
26
-#else
27
- // Fallback for other architectures, including x86
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;
31
-#endif
24
};
25
26
#ifdef __AVX512BF16__
csrc/cpu/pos_encoding.cpp
+154
-1
@@ -1,4 +1,3 @@
1
-
1
#include "cpu_types.hpp"
2
3
namespace {
@@ -97,6 +96,91 @@ void rotary_embedding_impl(
96
}
97
}
98
99
+template <>
100
+void rotary_embedding_impl<c10::Half>(
101
+ const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
102
+ c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
103
+ const int rot_dim, const int64_t query_stride, const int64_t key_stride,
104
+ const int num_heads, const int num_kv_heads, const int head_size,
105
+ const int num_tokens) {
106
+ using scalar_vec_t = vec_op::FP16Vec8;
107
+ constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
108
+
109
+ const int embed_dim = rot_dim / 2;
110
+ bool flag = (embed_dim % VEC_ELEM_NUM == 0);
111
+ const int loop_upper = flag ? embed_dim : embed_dim - VEC_ELEM_NUM;
112
+
113
+ auto compute_loop = [&](const int64_t token_head, const c10::Half* cache_ptr,
114
+ c10::Half* qk) {
115
+ int j = 0;
116
+ for (; j < loop_upper; j += VEC_ELEM_NUM) {
117
+ const int rot_offset = j;
118
+ const int x_index = rot_offset;
119
+ const int y_index = embed_dim + rot_offset;
120
+
121
+ const int64_t out_x = token_head + x_index;
122
+ const int64_t out_y = token_head + y_index;
123
+
124
+ const vec_op::FP16Vec8 cos_fp16(cache_ptr + x_index);
125
+ const vec_op::FP16Vec8 sin_fp16(cache_ptr + y_index);
126
+ const vec_op::FP16Vec8 q_x_fp16(qk + out_x);
127
+ const vec_op::FP16Vec8 q_y_fp16(qk + out_y);
128
+
129
+ const vec_op::FP32Vec8 fp32_cos(cos_fp16);
130
+ const vec_op::FP32Vec8 fp32_sin(sin_fp16);
131
+ const vec_op::FP32Vec8 fp32_q_x(q_x_fp16);
132
+ const vec_op::FP32Vec8 fp32_q_y(q_y_fp16);
133
+
134
+ auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin;
135
+ auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin;
136
+
137
+ vec_op::FP16Vec8(out1).save(qk + out_x);
138
+ vec_op::FP16Vec8(out2).save(qk + out_y);
139
+ }
140
+ if (!flag) {
141
+ for (; j < embed_dim; ++j) {
142
+ const int x_index = j;
143
+ const int y_index = embed_dim + j;
144
+
145
+ const int64_t out_x = token_head + x_index;
146
+ const int64_t out_y = token_head + y_index;
147
+
148
+ const float fp32_cos = static_cast<float>(cache_ptr[x_index]);
149
+ const float fp32_sin = static_cast<float>(cache_ptr[y_index]);
150
+ const float fp32_q_x = static_cast<float>(qk[out_x]);
151
+ const float fp32_q_y = static_cast<float>(qk[out_y]);
152
+
153
+ qk[out_x] =
154
+ static_cast<c10::Half>(fp32_q_x * fp32_cos - fp32_q_y * fp32_sin);
155
+ qk[out_y] =
156
+ static_cast<c10::Half>(fp32_q_y * fp32_cos + fp32_q_x * fp32_sin);
157
+ }
158
+ }
159
+ };
160
+
161
+#pragma omp parallel for
162
+ for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
163
+ int64_t pos = positions[token_idx];
164
+ const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
165
+
166
+ for (int i = 0; i < num_heads; ++i) {
167
+ const int head_idx = i;
168
+ const int64_t token_head =
169
+ token_idx * query_stride + head_idx * head_size;
170
+ compute_loop(token_head, cache_ptr, query);
171
+ }
172
+
173
+ if (key != nullptr) {
174
+ for (int i = 0; i < num_kv_heads; ++i) {
175
+ const int head_idx = i;
176
+ const int64_t token_head =
177
+ token_idx * key_stride + head_idx * head_size;
178
+ compute_loop(token_head, cache_ptr, key);
179
+ }
180
+ }
181
+ }
182
+}
183
+
184
template <typename scalar_t>
185
void rotary_embedding_gptj_impl(
186
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
@@ -174,6 +258,75 @@ void rotary_embedding_gptj_impl(
258
}
259
}
260
}
261
+
262
+template <>
263
+void rotary_embedding_gptj_impl<c10::Half>(
264
+ const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
265
+ c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
266
+ const int rot_dim, const int64_t query_stride, const int64_t key_stride,
267
+ const int num_heads, const int num_kv_heads, const int head_size,
268
+ const int num_tokens) {
269
+ const int embed_dim = rot_dim / 2;
270
+
271
+#pragma omp parallel for collapse(2)
272
+ for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
273
+ for (int i = 0; i < num_heads; ++i) {
274
+ int64_t pos = positions[token_idx];
275
+ const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
276
+ const c10::Half* cos_cache_ptr = cache_ptr;
277
+ const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
278
+ const int head_idx = i;
279
+ const int64_t token_head =
280
+ token_idx * query_stride + head_idx * head_size;
281
+ c10::Half* head_query = token_head + query;
282
+ for (int j = 0; j < embed_dim; j += 1) {
283
+ const int rot_offset = j;
284
+ const int x_index = 2 * rot_offset;
285
+ const int y_index = 2 * rot_offset + 1;
286
+
287
+ const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
288
+ const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
289
+
290
+ const float x = static_cast<float>(head_query[x_index]);
291
+ const float y = static_cast<float>(head_query[y_index]);
292
+
293
+ head_query[x_index] = static_cast<c10::Half>(x * cos - y * sin);
294
+ head_query[y_index] = static_cast<c10::Half>(y * cos + x * sin);
295
+ }
296
+ }
297
+ }
298
+
299
+ if (key == nullptr) {
300
+ return;
301
+ }
302
+
303
+#pragma omp parallel for collapse(2)
304
+ for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
305
+ for (int i = 0; i < num_kv_heads; ++i) {
306
+ int64_t pos = positions[token_idx];
307
+ const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
308
+ const c10::Half* cos_cache_ptr = cache_ptr;
309
+ const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
310
+ const int head_idx = i;
311
+ const int64_t token_head = token_idx * key_stride + head_idx * head_size;
312
+ c10::Half* head_key = key + token_head;
313
+ for (int j = 0; j < embed_dim; j += 1) {
314
+ const int rot_offset = j;
315
+ const int x_index = 2 * rot_offset;
316
+ const int y_index = 2 * rot_offset + 1;
317
+
318
+ const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
319
+ const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
320
+
321
+ const float x = static_cast<float>(head_key[x_index]);
322
+ const float y = static_cast<float>(head_key[y_index]);
323
+
324
+ head_key[x_index] = static_cast<c10::Half>(x * cos - y * sin);
325
+ head_key[y_index] = static_cast<c10::Half>(y * cos + x * sin);
326
+ }
327
+ }
328
+ }
329
+}
330
}; // namespace
331
332
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
vllm/platforms/cpu.py
+1
-1
@@ -50,7 +50,7 @@ class CpuPlatform(Platform):
50
@property
51
def supported_dtypes(self) -> list[torch.dtype]:
52
if self.get_cpu_architecture() == CpuArchEnum.POWERPC:
53
- return [torch.bfloat16, torch.float32]
53
+ return [torch.bfloat16, torch.float32, torch.float16]
54
elif self.get_cpu_architecture() == CpuArchEnum.ARM and sys.platform.startswith(
55
"darwin"
56
):