[CPU][s390x] Optimize inference perf and add oneDNN INT8 GEMM for s390x (#50219)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rehan Khan committed
Jul 31, 2026 at 15:19 UTC
88bc8fb56f69667bcb67b482ecf729385d7ae972
6 files changed
+182
-27
cmake/cpu_extension.cmake
+18
-2
@@ -257,7 +257,7 @@ endif()
257
258
259
# Build oneDNN for GEMM kernels
260
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND)
260
+if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND OR S390_FOUND)
261
# Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
262
# TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
263
set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -368,7 +368,23 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
368
369
set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
370
set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
371
- FetchContent_MakeAvailable(oneDNN)
371
+
372
+ if(S390_FOUND)
373
+ FetchContent_GetProperties(oneDNN)
374
+ if(NOT onednn_POPULATED)
375
+ FetchContent_Populate(oneDNN)
376
+ # Patch s390x helpers.h: ALWAYS_INLINE on operator+= breaks C++20/GCC14
377
+ file(READ "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" _helpers_content)
378
+ string(REPLACE
379
+ "vec_type_t<T> &ALWAYS_INLINE operator+="
380
+ "ALWAYS_INLINE vec_type_t<T> &operator+="
381
+ _helpers_content "${_helpers_content}")
382
+ file(WRITE "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" "${_helpers_content}")
383
+ add_subdirectory("${onednn_SOURCE_DIR}" "${onednn_BINARY_DIR}")
384
+ endif()
385
+ else()
386
+ FetchContent_MakeAvailable(oneDNN)
387
+ endif()
388
set(CMAKE_BUILD_TYPE ${VLLM_BUILD_TYPE})
389
add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp")
390
target_include_directories(
csrc/cpu/cpu_arch_macros.h
+11
@@ -183,4 +183,15 @@
183
184
#endif // __powerpc__
185
186
+// IBM Z (s390x) VXE
187
+#ifdef __s390x__
188
+ // FP32Vec16::exp() in cpu_types_vxe.hpp delegates to FP32Vec8::exp(), which
189
+ // implements a vectorised 5-term minimax polynomial using VXE intrinsics.
190
+ #define DEFINE_FAST_EXP \
191
+ auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
192
+ __attribute__((always_inline)) { return vec.exp(); }; \
193
+ auto fast_exp_f16 = fast_exp;
194
+
195
+#endif // __s390x__
196
+
197
#endif
csrc/cpu/cpu_types_vxe.hpp
+150
-22
@@ -140,10 +140,21 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
140
explicit BF16Vec16(const FP32Vec16&);
141
142
void save(void* ptr) const {
143
- // Save 256 bits in two parts
143
vec_xst(reg.val[0], 0, (signed short*)ptr);
144
vec_xst(reg.val[1], 16, (signed short*)ptr);
145
}
146
+
147
+ void save(void* ptr, const int elem_num) const {
148
+ auto* dst = reinterpret_cast<c10::BFloat16*>(ptr);
149
+ union {
150
+ ss16x8x2_t r;
151
+ c10::BFloat16 values[16];
152
+ } ar;
153
+ ar.r = reg;
154
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
155
+ dst[i] = ar.values[i];
156
+ }
157
+ }
158
};
159
160
const static __vector signed short zero = vec_splats((signed short)0);
@@ -322,13 +333,12 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
333
}
334
335
float reduce_sum() const {
325
- AliasReg ar;
326
- ar.reg = reg;
327
- float result = 0;
328
- unroll_loop<int, VEC_ELEM_NUM>(
329
- [&result, &ar](int i) { result += ar.values[i]; });
330
-
331
- return result;
336
+ __vector float sum = vec_add(reg.val[0], reg.val[1]);
337
+ __vector float hi = vec_sld(sum, sum, 8);
338
+ sum = vec_add(sum, hi);
339
+ __vector float lo = vec_sld(sum, sum, 4);
340
+ sum = vec_add(sum, lo);
341
+ return vec_extract(sum, 0);
342
}
343
344
FP32Vec8 exp() const {
@@ -672,6 +682,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
682
reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
683
}
684
685
+ explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {}
686
+
687
explicit FP32Vec16(const FP16Vec16& v) {
688
__vector unsigned int raw_hi_0 =
689
(__vector unsigned int)vec_unpackh(v.reg.val[0]);
@@ -723,13 +735,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
735
}
736
737
float reduce_sum() const {
726
- AliasReg ar;
727
- ar.reg = reg;
728
- float result = 0;
729
- unroll_loop<int, VEC_ELEM_NUM>(
730
- [&result, &ar](int i) { result += ar.values[i]; });
731
-
732
- return result;
738
+ __vector float sum = vec_add(vec_add(reg.val[0], reg.val[1]),
739
+ vec_add(reg.val[2], reg.val[3]));
740
+ __vector float hi = vec_sld(sum, sum, 8);
741
+ sum = vec_add(sum, hi);
742
+ __vector float lo = vec_sld(sum, sum, 4);
743
+ sum = vec_add(sum, lo);
744
+ return vec_extract(sum, 0);
745
}
746
747
template <int group_size>
@@ -754,13 +766,68 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
766
}
767
768
float reduce_max() const {
757
- AliasReg ar;
758
- ar.reg = reg;
759
- float result = ar.values[0];
760
- unroll_loop<int, VEC_ELEM_NUM>([&result, &ar](int i) {
761
- if (ar.values[i] > result) result = ar.values[i];
762
- });
763
- return result;
769
+ __vector float m = vec_max(vec_max(reg.val[0], reg.val[1]),
770
+ vec_max(reg.val[2], reg.val[3]));
771
+ __vector float hi = vec_sld(m, m, 8);
772
+ m = vec_max(m, hi);
773
+ __vector float lo = vec_sld(m, m, 4);
774
+ m = vec_max(m, lo);
775
+ return vec_extract(m, 0);
776
+ }
777
+
778
+ FP32Vec16 exp() const {
779
+ FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
780
+ FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
781
+ auto lo_e = lo.exp();
782
+ auto hi_e = hi.exp();
783
+ return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
784
+ hi_e.reg.val[0], hi_e.reg.val[1]});
785
+ }
786
+
787
+ FP32Vec16 abs() const {
788
+ return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]),
789
+ vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
790
+ }
791
+
792
+ FP32Vec16 min(const FP32Vec16& b) const {
793
+ return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
794
+ vec_min(reg.val[1], b.reg.val[1]),
795
+ vec_min(reg.val[2], b.reg.val[2]),
796
+ vec_min(reg.val[3], b.reg.val[3])}));
797
+ }
798
+
799
+ FP32Vec16 clamp(const FP32Vec16& lo, const FP32Vec16& hi) const {
800
+ return this->max(lo).min(hi);
801
+ }
802
+
803
+ float reduce_min() const {
804
+ __vector float m = vec_min(vec_min(reg.val[0], reg.val[1]),
805
+ vec_min(reg.val[2], reg.val[3]));
806
+ __vector float h = vec_sld(m, m, 8);
807
+ m = vec_min(m, h);
808
+ __vector float l = vec_sld(m, m, 4);
809
+ m = vec_min(m, l);
810
+ return vec_extract(m, 0);
811
+ }
812
+
813
+ FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
814
+ AliasReg ar_this, ar_b;
815
+ ar_this.reg = reg;
816
+ ar_b.reg = b.reg;
817
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
818
+ ar_this.values[i] = std::min(ar_this.values[i], ar_b.values[i]);
819
+ }
820
+ return FP32Vec16(ar_this.reg);
821
+ }
822
+
823
+ FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
824
+ AliasReg ar_this, ar_b;
825
+ ar_this.reg = reg;
826
+ ar_b.reg = b.reg;
827
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
828
+ ar_this.values[i] = std::max(ar_this.values[i], ar_b.values[i]);
829
+ }
830
+ return FP32Vec16(ar_this.reg);
831
}
832
833
void save(float* ptr) const {
@@ -769,6 +836,67 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
836
vec_xst(reg.val[2], 32, ptr);
837
vec_xst(reg.val[3], 48, ptr);
838
}
839
+
840
+ void save(float* ptr, const int elem_num) const {
841
+ AliasReg ar;
842
+ ar.reg = reg;
843
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
844
+ ptr[i] = ar.values[i];
845
+ }
846
+ }
847
+
848
+ void save(c10::Half* ptr) const {
849
+ FP16Vec16 fp16(*this);
850
+ fp16.save(ptr);
851
+ }
852
+
853
+ void save(c10::Half* ptr, const int elem_num) const {
854
+ FP16Vec16 fp16(*this);
855
+ union {
856
+ ss16x8x2_t r;
857
+ c10::Half values[16];
858
+ } ar;
859
+ ar.r = fp16.reg;
860
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
861
+ ptr[i] = ar.values[i];
862
+ }
863
+ }
864
+};
865
+
866
+struct INT8Vec16 : public Vec<INT8Vec16> {
867
+ constexpr static int VEC_ELEM_NUM = 16;
868
+
869
+ union AliasReg {
870
+ __vector signed char reg;
871
+ int8_t values[VEC_ELEM_NUM];
872
+ };
873
+
874
+ __vector signed char reg;
875
+
876
+ explicit INT8Vec16(const FP32Vec16& vec) {
877
+ __vector signed int ret[4];
878
+ ret[0] = vec_signed(vec.reg.val[0]);
879
+ ret[1] = vec_signed(vec.reg.val[1]);
880
+ ret[2] = vec_signed(vec.reg.val[2]);
881
+ ret[3] = vec_signed(vec.reg.val[3]);
882
+
883
+ __vector signed short packed1 = vec_packs(ret[0], ret[1]);
884
+ __vector signed short packed2 = vec_packs(ret[2], ret[3]);
885
+
886
+ reg = vec_packs(packed1, packed2);
887
+ }
888
+
889
+ void save(void* ptr) const {
890
+ *reinterpret_cast<__vector signed char*>(ptr) = reg;
891
+ }
892
+
893
+ void save(int8_t* ptr, const int elem_num) const {
894
+ AliasReg ar;
895
+ ar.reg = reg;
896
+ for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
897
+ ptr[i] = ar.values[i];
898
+ }
899
+ }
900
};
901
902
template <typename T>
csrc/cpu/torch_bindings.cpp
+1
-1
@@ -398,7 +398,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
398
// Quantization
399
#if defined(__AVX512F__) || defined(__AVX2__) || \
400
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
401
- defined(__riscv_v)
401
+ defined(__riscv_v) || defined(__s390x__)
402
// Helper function to release oneDNN handlers
403
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
404
&release_dnnl_matmul_handler);
docs/getting_started/installation/cpu.md
+1
-1
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
315
- vLLM CPU supports quantizations:
316
- AWQ (x86 only)
317
- GPTQ (x86 only)
318
- - compressed-tensor INT8 W8A8 (x86 only)
318
+ - compressed-tensor INT8 W8A8 (x86, s390x only)
319
320
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
321
docs/getting_started/installation/cpu.s390x.inc.md
+1
-1
@@ -10,7 +10,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
10
11
- OS: `Linux`
12
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
13
-- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
13
+- Instruction Set Architecture (ISA): VXE support is required. Works with Z15 and above.
14
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
15
16
--8<-- [end:requirements]