diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 64df94e947..cdfc6cb6f8 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -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 &ALWAYS_INLINE operator+=" + "ALWAYS_INLINE vec_type_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( diff --git a/csrc/cpu/cpu_arch_macros.h b/csrc/cpu/cpu_arch_macros.h index cf3dd4b961..ccec4c2e87 100644 --- a/csrc/cpu/cpu_arch_macros.h +++ b/csrc/cpu/cpu_arch_macros.h @@ -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 diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index 26a14e8ea6..9b28300ced 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -140,10 +140,21 @@ struct BF16Vec16 : public Vec { 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(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 { } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&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 { 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 { } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&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 @@ -754,13 +766,68 @@ struct FP32Vec16 : public Vec { } float reduce_max() const { - AliasReg ar; - ar.reg = reg; - float result = ar.values[0]; - unroll_loop([&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 { 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 { + 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 diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 45dfce7e01..e6130473e3 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -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); diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 1534d4d9e6..925f54dcec 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -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? diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 407e13baf7..0c07c255e2 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -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]