iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CPU][s390x] Optimize inference perf and add oneDNN INT8 GEMM for s390x (#50219)

Browse Source 
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Rehan Khan
2026-07-31 17:49:31 +08:00
committed by GitHub
co-authored by Li, Jiang <jiang1.li@intel.com>
parent 482cfc20f0
commit 88bc8fb56f
6 changed files with 182 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/cpu_extension.cmake
+18 -2
View File
@@ -257,7 +257,7 @@ endif()
# Build oneDNN for GEMM kernels
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)
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)
# Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
# TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
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
set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
FetchContent_MakeAvailable(oneDNN)
if(S390_FOUND)
FetchContent_GetProperties(oneDNN)
if(NOT onednn_POPULATED)
FetchContent_Populate(oneDNN)
# Patch s390x helpers.h: ALWAYS_INLINE on operator+= breaks C++20/GCC14
file(READ "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" _helpers_content)
string(REPLACE
"vec_type_t<T> &ALWAYS_INLINE operator+="
"ALWAYS_INLINE vec_type_t<T> &operator+="
_helpers_content "${_helpers_content}")
file(WRITE "${onednn_SOURCE_DIR}/src/cpu/s390x/helpers.h" "${_helpers_content}")
add_subdirectory("${onednn_SOURCE_DIR}" "${onednn_BINARY_DIR}")
endif()
else()
FetchContent_MakeAvailable(oneDNN)
endif()
set(CMAKE_BUILD_TYPE ${VLLM_BUILD_TYPE})
add_library(dnnl_ext OBJECT "csrc/cpu/dnnl_helper.cpp")
target_include_directories(
csrc/cpu/cpu_arch_macros.h
+11
View File
@@ -183,4 +183,15 @@
#endif // __powerpc__
// IBM Z (s390x) VXE
#ifdef __s390x__
// FP32Vec16::exp() in cpu_types_vxe.hpp delegates to FP32Vec8::exp(), which
// implements a vectorised 5-term minimax polynomial using VXE intrinsics.
#define DEFINE_FAST_EXP \
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
__attribute__((always_inline)) { return vec.exp(); }; \
auto fast_exp_f16 = fast_exp;
#endif // __s390x__
#endif
csrc/cpu/cpu_types_vxe.hpp
+150 -22
View File
@@ -140,10 +140,21 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
// Save 256 bits in two parts
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst(reg.val[1], 16, (signed short*)ptr);
}
void save(void* ptr, const int elem_num) const {
auto* dst = reinterpret_cast<c10::BFloat16*>(ptr);
union {
ss16x8x2_t r;
c10::BFloat16 values[16];
} ar;
ar.r = reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
dst[i] = ar.values[i];
}
}
};
const static __vector signed short zero = vec_splats((signed short)0);
@@ -322,13 +333,12 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
}
float reduce_sum() const {
AliasReg ar;
ar.reg = reg;
float result = 0;
unroll_loop<int, VEC_ELEM_NUM>(
[&result, &ar](int i) { result += ar.values[i]; });
return result;
__vector float sum = vec_add(reg.val[0], reg.val[1]);
__vector float hi = vec_sld(sum, sum, 8);
sum = vec_add(sum, hi);
__vector float lo = vec_sld(sum, sum, 4);
sum = vec_add(sum, lo);
return vec_extract(sum, 0);
}
FP32Vec8 exp() const {
@@ -672,6 +682,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = (__vector float)vec_mergel(v.reg.val[1], zero);
}
explicit FP32Vec16(const c10::Half* ptr) : FP32Vec16(FP16Vec16(ptr)) {}
explicit FP32Vec16(const FP16Vec16& v) {
__vector unsigned int raw_hi_0 =
(__vector unsigned int)vec_unpackh(v.reg.val[0]);
@@ -723,13 +735,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
}
float reduce_sum() const {
AliasReg ar;
ar.reg = reg;
float result = 0;
unroll_loop<int, VEC_ELEM_NUM>(
[&result, &ar](int i) { result += ar.values[i]; });
return result;
__vector float sum = vec_add(vec_add(reg.val[0], reg.val[1]),
vec_add(reg.val[2], reg.val[3]));
__vector float hi = vec_sld(sum, sum, 8);
sum = vec_add(sum, hi);
__vector float lo = vec_sld(sum, sum, 4);
sum = vec_add(sum, lo);
return vec_extract(sum, 0);
}
template <int group_size>
@@ -754,13 +766,68 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
}
float reduce_max() const {
AliasReg ar;
ar.reg = reg;
float result = ar.values[0];
unroll_loop<int, VEC_ELEM_NUM>([&result, &ar](int i) {
if (ar.values[i] > result) result = ar.values[i];
});
return result;
__vector float m = vec_max(vec_max(reg.val[0], reg.val[1]),
vec_max(reg.val[2], reg.val[3]));
__vector float hi = vec_sld(m, m, 8);
m = vec_max(m, hi);
__vector float lo = vec_sld(m, m, 4);
m = vec_max(m, lo);
return vec_extract(m, 0);
}
FP32Vec16 exp() const {
FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
auto lo_e = lo.exp();
auto hi_e = hi.exp();
return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
hi_e.reg.val[0], hi_e.reg.val[1]});
}
FP32Vec16 abs() const {
return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]),
vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
vec_min(reg.val[1], b.reg.val[1]),
vec_min(reg.val[2], b.reg.val[2]),
vec_min(reg.val[3], b.reg.val[3])}));
}
FP32Vec16 clamp(const FP32Vec16& lo, const FP32Vec16& hi) const {
return this->max(lo).min(hi);
}
float reduce_min() const {
__vector float m = vec_min(vec_min(reg.val[0], reg.val[1]),
vec_min(reg.val[2], reg.val[3]));
__vector float h = vec_sld(m, m, 8);
m = vec_min(m, h);
__vector float l = vec_sld(m, m, 4);
m = vec_min(m, l);
return vec_extract(m, 0);
}
FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
AliasReg ar_this, ar_b;
ar_this.reg = reg;
ar_b.reg = b.reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
ar_this.values[i] = std::min(ar_this.values[i], ar_b.values[i]);
}
return FP32Vec16(ar_this.reg);
}
FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
AliasReg ar_this, ar_b;
ar_this.reg = reg;
ar_b.reg = b.reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
ar_this.values[i] = std::max(ar_this.values[i], ar_b.values[i]);
}
return FP32Vec16(ar_this.reg);
}
void save(float* ptr) const {
@@ -769,6 +836,67 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
vec_xst(reg.val[2], 32, ptr);
vec_xst(reg.val[3], 48, ptr);
}
void save(float* ptr, const int elem_num) const {
AliasReg ar;
ar.reg = reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
ptr[i] = ar.values[i];
}
}
void save(c10::Half* ptr) const {
FP16Vec16 fp16(*this);
fp16.save(ptr);
}
void save(c10::Half* ptr, const int elem_num) const {
FP16Vec16 fp16(*this);
union {
ss16x8x2_t r;
c10::Half values[16];
} ar;
ar.r = fp16.reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
ptr[i] = ar.values[i];
}
}
};
struct INT8Vec16 : public Vec<INT8Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
union AliasReg {
__vector signed char reg;
int8_t values[VEC_ELEM_NUM];
};
__vector signed char reg;
explicit INT8Vec16(const FP32Vec16& vec) {
__vector signed int ret[4];
ret[0] = vec_signed(vec.reg.val[0]);
ret[1] = vec_signed(vec.reg.val[1]);
ret[2] = vec_signed(vec.reg.val[2]);
ret[3] = vec_signed(vec.reg.val[3]);
__vector signed short packed1 = vec_packs(ret[0], ret[1]);
__vector signed short packed2 = vec_packs(ret[2], ret[3]);
reg = vec_packs(packed1, packed2);
}
void save(void* ptr) const {
*reinterpret_cast<__vector signed char*>(ptr) = reg;
}
void save(int8_t* ptr, const int elem_num) const {
AliasReg ar;
ar.reg = reg;
for (int i = 0; i < elem_num && i < VEC_ELEM_NUM; ++i) {
ptr[i] = ar.values[i];
}
}
};
template <typename T>
csrc/cpu/torch_bindings.cpp
+1 -1
View File
@@ -398,7 +398,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Quantization
#if defined(__AVX512F__) || defined(__AVX2__) || \
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
defined(__riscv_v)
defined(__riscv_v) || defined(__s390x__)
// Helper function to release oneDNN handlers
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
&release_dnnl_matmul_handler);
docs/getting_started/installation/cpu.md
+1 -1
View File
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
- vLLM CPU supports quantizations:
- AWQ (x86 only)
- GPTQ (x86 only)
- compressed-tensor INT8 W8A8 (x86 only)
- compressed-tensor INT8 W8A8 (x86, s390x only)
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
docs/getting_started/installation/cpu.s390x.inc.md
+1 -1
View File
@@ -10,7 +10,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
- OS: `Linux`
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
- Instruction Set Architecture (ISA): VXE support is required. Works with Z15 and above.
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
--8<-- [end:requirements]