A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 12:30 UTC
[CPU] Enable FP16 (Half dtype) support for s390x (#34116)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
R3hankhan committed
Feb 11, 2026 at 12:11 UTC
d1b837f0ae6a0152d820194a181e809ffaef6864
3 files changed
+244
-9
csrc/cpu/cpu_attn_impl.hpp
+1
-1
@@ -821,7 +821,7 @@ struct VecTypeTrait<c10::BFloat16> {
821
using vec_t = vec_op::BF16Vec16;
822
};
823
824
-#if !defined(__powerpc__) && !defined(__s390x__)
824
+#if !defined(__powerpc__)
825
template <>
826
struct VecTypeTrait<c10::Half> {
827
using vec_t = vec_op::FP16Vec16;
csrc/cpu/cpu_types_vxe.hpp
+241
-6
@@ -16,10 +16,12 @@ namespace vec_op {
16
#define vec_sr(a, b) ((a) >> (b)) // Vector Shift Right Algebraic
17
#define vec_sl(a, b) ((a) << (b)) // Vector Shift Left
18
19
-// FIXME: FP16 is not fully supported in Torch-CPU
20
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
21
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
22
- AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
19
+// NOTE: FP16 (Half) is supported on s390x via custom bit-manipulation
20
+// conversion. PyTorch itself lacks native s390x FP16 support.
21
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
22
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
23
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
24
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
25
26
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
27
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -86,6 +88,39 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
88
}
89
};
90
91
+struct FP16Vec8 : public Vec<FP16Vec8> {
92
+ constexpr static int VEC_ELEM_NUM = 8;
93
+
94
+ __vector signed short reg;
95
+
96
+ explicit FP16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
97
+ explicit FP16Vec8(const FP32Vec8&);
98
+
99
+ void save(void* ptr) const {
100
+ *reinterpret_cast<__vector signed short*>(ptr) = reg;
101
+ }
102
+};
103
+
104
+struct FP16Vec16 : public Vec<FP16Vec16> {
105
+ constexpr static int VEC_ELEM_NUM = 16;
106
+
107
+ ss16x8x2_t reg;
108
+
109
+ explicit FP16Vec16(const void* ptr) {
110
+ // Load 256 bits (16 FP16 values) in two parts
111
+ reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
112
+ reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
113
+ }
114
+
115
+ explicit FP16Vec16(const FP32Vec16&);
116
+
117
+ void save(void* ptr) const {
118
+ // Save 256 bits in two parts
119
+ vec_xst(reg.val[0], 0, (signed short*)ptr);
120
+ vec_xst(reg.val[1], 16, (signed short*)ptr);
121
+ }
122
+};
123
+
124
struct BF16Vec16 : public Vec<BF16Vec16> {
125
constexpr static int VEC_ELEM_NUM = 16;
126
@@ -108,6 +143,92 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
143
144
const static __vector signed short zero = vec_splats((signed short)0);
145
146
+FORCE_INLINE __vector float fp16_to_fp32_bits(__vector unsigned int x) {
147
+ const __vector unsigned int mask_sign = {0x8000, 0x8000, 0x8000, 0x8000};
148
+ const __vector unsigned int mask_exp = {0x7C00, 0x7C00, 0x7C00, 0x7C00};
149
+ const __vector unsigned int mask_mant = {0x03FF, 0x03FF, 0x03FF, 0x03FF};
150
+ const __vector unsigned int bias_adj = {112, 112, 112, 112};
151
+ const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F,
152
+ 0x1F}; // FP16 NaN/Inf exponent
153
+ const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF,
154
+ 0xFF}; // FP32 NaN/Inf exponent
155
+
156
+ __vector unsigned int s = (x & mask_sign) << 16;
157
+ __vector unsigned int e = (x & mask_exp) >> 10;
158
+ __vector unsigned int m = (x & mask_mant) << 13;
159
+
160
+ // Check for NaN/Inf: exponent = 0x1F in FP16
161
+ __vector __bool int is_nan_inf = vec_cmpeq(e, exp_max_fp16);
162
+
163
+ // Normal: adjust bias; NaN/Inf: set to 0xFF
164
+ __vector unsigned int e_normal = e + bias_adj;
165
+ e = vec_sel(e_normal, exp_max_fp32, is_nan_inf);
166
+
167
+ return (__vector float)(s | (e << 23) | m);
168
+}
169
+
170
+FORCE_INLINE __vector unsigned int fp32_to_fp16_bits(__vector float f_in) {
171
+ __vector unsigned int in = (__vector unsigned int)f_in;
172
+
173
+ const __vector unsigned int mask_sign_32 = {0x80000000, 0x80000000,
174
+ 0x80000000, 0x80000000};
175
+ const __vector unsigned int mask_exp_32 = {0x7F800000, 0x7F800000, 0x7F800000,
176
+ 0x7F800000};
177
+ const __vector unsigned int mask_mant_32 = {0x007FFFFF, 0x007FFFFF,
178
+ 0x007FFFFF, 0x007FFFFF};
179
+
180
+ // Use SIGNED integers for exponent math to handle underflow check
181
+ const __vector signed int bias_adj = {112, 112, 112, 112};
182
+ const __vector signed int zero = {0, 0, 0, 0};
183
+ const __vector signed int max_exp = {31, 31, 31, 31}; // Max FP16 exp
184
+ const __vector unsigned int exp_max_fp32 = {0xFF, 0xFF, 0xFF, 0xFF};
185
+ const __vector unsigned int exp_max_fp16 = {0x1F, 0x1F, 0x1F, 0x1F};
186
+
187
+ __vector unsigned int s = (in & mask_sign_32) >> 16;
188
+ __vector unsigned int e_u = (in & mask_exp_32) >> 23;
189
+
190
+ // Check for NaN/Inf: exponent = 0xFF in FP32
191
+ __vector __bool int is_nan_inf = vec_cmpeq(e_u, exp_max_fp32);
192
+
193
+ __vector signed int e_s = (__vector signed int)e_u;
194
+ e_s = vec_sub(e_s, bias_adj);
195
+ e_s = vec_max(e_s, zero);
196
+ e_s = vec_min(e_s, max_exp);
197
+ __vector unsigned int e_normal = (__vector unsigned int)e_s;
198
+
199
+ __vector unsigned int e_final = vec_sel(e_normal, exp_max_fp16, is_nan_inf);
200
+
201
+ const __vector unsigned int one_v = {1, 1, 1, 1};
202
+ const __vector unsigned int mask_sticky = {0xFFF, 0xFFF, 0xFFF, 0xFFF};
203
+
204
+ __vector unsigned int round_bit = (in >> 12) & one_v;
205
+ __vector unsigned int sticky = in & mask_sticky;
206
+ __vector unsigned int m = (in & mask_mant_32) >> 13;
207
+ __vector unsigned int lsb = m & one_v; // LSB of mantissa for tie-breaking
208
+
209
+ // Round up if: round_bit && (sticky || lsb)
210
+ __vector __bool int sticky_nonzero =
211
+ vec_cmpgt(sticky, (__vector unsigned int){0, 0, 0, 0});
212
+ __vector __bool int lsb_set = vec_cmpeq(lsb, one_v);
213
+ __vector __bool int round_up =
214
+ vec_and(vec_cmpeq(round_bit, one_v), vec_or(sticky_nonzero, lsb_set));
215
+
216
+ m = vec_sel(m, m + one_v, round_up);
217
+
218
+ const __vector unsigned int mant_mask = {0x3FF, 0x3FF, 0x3FF, 0x3FF};
219
+ const __vector unsigned int max_normal_exp = {0x1E, 0x1E, 0x1E, 0x1E};
220
+ __vector __bool int mant_overflows = vec_cmpgt(m, mant_mask);
221
+ __vector __bool int would_overflow_to_inf =
222
+ vec_and(mant_overflows, vec_cmpeq(e_final, max_normal_exp));
223
+ __vector unsigned int e_inc = vec_min(e_final + one_v, exp_max_fp16);
224
+ e_final = vec_sel(e_final, e_inc, mant_overflows);
225
+ m = vec_and(m, mant_mask);
226
+ e_final = vec_sel(e_final, max_normal_exp, would_overflow_to_inf);
227
+ m = vec_sel(m, mant_mask, would_overflow_to_inf);
228
+
229
+ return s | (e_final << 10) | m;
230
+}
231
+
232
struct BF16Vec32 : public Vec<BF16Vec32> {
233
constexpr static int VEC_ELEM_NUM = 32;
234
@@ -180,6 +301,18 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
301
reg.val[1] = (__vector float)vec_mergel(v.reg, zero);
302
}
303
304
+ explicit FP32Vec8(const FP16Vec8& v) {
305
+ // Cast to UNSIGNED short vector to prevent sign-extension during unpack
306
+ __vector unsigned short raw_u = (__vector unsigned short)v.reg;
307
+
308
+ // Unpack 8x16-bit to two 4x32-bit vectors (Zero extended)
309
+ __vector unsigned int raw_hi = (__vector unsigned int)vec_unpackh(raw_u);
310
+ __vector unsigned int raw_lo = (__vector unsigned int)vec_unpackl(raw_u);
311
+
312
+ reg.val[0] = fp16_to_fp32_bits(raw_hi);
313
+ reg.val[1] = fp16_to_fp32_bits(raw_lo);
314
+ }
315
+
316
float reduce_sum() const {
317
AliasReg ar;
318
ar.reg = reg;
@@ -531,6 +664,22 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
664
reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
665
}
666
667
+ explicit FP32Vec16(const FP16Vec16& v) {
668
+ __vector unsigned int raw_hi_0 =
669
+ (__vector unsigned int)vec_unpackh(v.reg.val[0]);
670
+ __vector unsigned int raw_lo_0 =
671
+ (__vector unsigned int)vec_unpackl(v.reg.val[0]);
672
+ reg.val[0] = fp16_to_fp32_bits(raw_hi_0);
673
+ reg.val[1] = fp16_to_fp32_bits(raw_lo_0);
674
+
675
+ __vector unsigned int raw_hi_1 =
676
+ (__vector unsigned int)vec_unpackh(v.reg.val[1]);
677
+ __vector unsigned int raw_lo_1 =
678
+ (__vector unsigned int)vec_unpackl(v.reg.val[1]);
679
+ reg.val[2] = fp16_to_fp32_bits(raw_hi_1);
680
+ reg.val[3] = fp16_to_fp32_bits(raw_lo_1);
681
+ }
682
+
683
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
684
685
FP32Vec16 operator*(const FP32Vec16& b) const {
@@ -628,8 +777,10 @@ struct VecType<c10::BFloat16> {
777
using vec_type = BF16Vec8;
778
};
779
631
-// On s390x, FP16 (Half) is not natively supported, use FP32 vectors instead
632
-using FP16Vec16 = FP32Vec16;
780
+template <>
781
+struct VecType<c10::Half> {
782
+ using vec_type = FP16Vec8;
783
+};
784
785
template <typename T>
786
void storeFP32(float v, T* ptr) {
@@ -650,6 +801,52 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
801
*ptr = *(v_ptr + 1);
802
}
803
804
+template <>
805
+inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
806
+ // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
807
+ // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
808
+ // produce incorrect results for some inputs. Process each of the 4 vectors
809
+ // separately.
810
+ uint32_t in;
811
+ std::memcpy(&in, &v, sizeof(in));
812
+
813
+ uint32_t s = (in & 0x80000000) >> 16; // Sign
814
+ uint32_t e = (in & 0x7F800000) >> 23; // Exponent
815
+ uint32_t round_bit = (in >> 12) & 1;
816
+ uint32_t sticky = (in & 0xFFF) != 0; // Any bits in [11..0]
817
+ uint32_t m = (in & 0x007FFFFF) >> 13;
818
+ uint32_t lsb = m & 1; // LSB of mantissa for tie-breaking
819
+
820
+ // Check for NaN/Inf before rounding
821
+ bool is_nan_inf = (e == 0xFF);
822
+
823
+ if (round_bit && (sticky || lsb)) {
824
+ m++;
825
+ // Handle mantissa overflow: if m overflows 10 bits, increment exponent
826
+ if (m > 0x3FF) {
827
+ m = 0;
828
+ e++;
829
+ }
830
+ }
831
+
832
+ if (is_nan_inf) {
833
+ // NaN/Inf: preserve it
834
+ e = 0x1F;
835
+ } else {
836
+ // Normal: adjust bias (127 - 15), flush subnormals to zero
837
+ e = (e >= 112) ? (e - 112) : 0;
838
+ // If exponent overflows to Inf range, saturate to max normal FP16 value
839
+ if (e > 0x1E) {
840
+ e = 0x1E; // Max normal exponent
841
+ m = 0x3FF; // Max mantissa
842
+ }
843
+ }
844
+
845
+ uint16_t fp16 = (uint16_t)(s | (e << 10) | m);
846
+
847
+ *reinterpret_cast<uint16_t*>(ptr) = fp16;
848
+}
849
+
850
#ifndef __VEC_CLASS_FP_NAN
851
#define __VEC_CLASS_FP_NAN (1 << 6)
852
#endif
@@ -803,6 +1000,44 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
1000
reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
1001
}
1002
1003
+inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
1004
+ // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
1005
+ // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
1006
+ // produce incorrect results for some inputs. Process each of the 4 vectors
1007
+ // separately.
1008
+ __vector unsigned int res_hi = fp32_to_fp16_bits(v.reg.val[0]);
1009
+ __vector unsigned int res_lo = fp32_to_fp16_bits(v.reg.val[1]);
1010
+
1011
+ const __vector unsigned char perm_pack = {
1012
+ 2, 3, 6, 7, 10, 11, 14, 15, // Select lower 2 bytes from res_hi
1013
+ 18, 19, 22, 23, 26, 27, 30, 31 // Select lower 2 bytes from res_lo
1014
+ };
1015
+
1016
+ reg = vec_perm((__vector signed short)res_hi, (__vector signed short)res_lo,
1017
+ perm_pack);
1018
+}
1019
+
1020
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
1021
+ // Use bit-manipulation for IEEE FP32 to FP16 conversion since vector
1022
+ // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
1023
+ // produce incorrect results for some inputs. Process each of the 4 vectors
1024
+ // separately.
1025
+ __vector unsigned int res_0 = fp32_to_fp16_bits(v.reg.val[0]);
1026
+ __vector unsigned int res_1 = fp32_to_fp16_bits(v.reg.val[1]);
1027
+ __vector unsigned int res_2 = fp32_to_fp16_bits(v.reg.val[2]);
1028
+ __vector unsigned int res_3 = fp32_to_fp16_bits(v.reg.val[3]);
1029
+
1030
+ const __vector unsigned char perm_pack = {
1031
+ 2, 3, 6, 7, 10, 11, 14, 15, // Lower 2 bytes from first vector
1032
+ 18, 19, 22, 23, 26, 27, 30, 31 // Lower 2 bytes from second vector
1033
+ };
1034
+
1035
+ reg.val[0] = vec_perm((__vector signed short)res_0,
1036
+ (__vector signed short)res_1, perm_pack);
1037
+ reg.val[1] = vec_perm((__vector signed short)res_2,
1038
+ (__vector signed short)res_3, perm_pack);
1039
+}
1040
+
1041
// 1D softmax over `n` elements in `input`, writes result to `output`.
1042
// Uses FP32Vec8 for main body, scalar tail handling.
1043
// Requirement: n > 0
csrc/cpu/mla_decode.cpp
+2
-2
@@ -18,8 +18,8 @@ struct KernelVecType<float> {
18
19
template <>
20
struct KernelVecType<c10::Half> {
21
-#if defined(__powerpc64__) || defined(__s390x__)
22
- // Power and s390x architecture-specific vector types
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;