@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ada19210a3
Fork 0

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 15, 21:00 UTC

Adding cpu inference with VXE ISA for s390x architecture (#12613)

Signed-off-by: Dilip Gowda Bhagavan <dilip.bhagavan@ibm.com> Signed-off-by: Rishika Kedia <rishika.kedia@in.ibm.com> Co-authored-by: Rishika Kedia <rishika.kedia@in.ibm.com>

Dilip Gowda Bhagavan committed Mar 6, 2025 at 22:10 UTC ada19210a356378e274c5a2c847096f5d27abf2b
7 files changed +653 -8
Dockerfile.s390x new
+152
@@ -0,0 +1,152 @@
1 +# Base UBI image for s390x architecture
2 +ARG BASE_UBI_IMAGE_TAG=9.5-1736404155
3 +ARG PYTHON_VERSION=3.12
4 +FROM registry.access.redhat.com/ubi9/ubi-minimal:${BASE_UBI_IMAGE_TAG} AS base
5 +
6 +# Install basic dependencies
7 +ARG PYTHON_VERSION
8 +ENV PYTHON_VERSION=${PYTHON_VERSION}
9 +
10 +WORKDIR /workspace
11 +
12 +ENV LANG=C.UTF-8 \
13 + LC_ALL=C.UTF-8
14 +
15 +# Install development utilities
16 +RUN microdnf install -y \
17 + which procps findutils tar vim git gcc gcc-gfortran g++ make patch zlib-devel \
18 + libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
19 + openssl-devel openblas openblas-devel autoconf automake libtool cmake && \
20 + microdnf clean all
21 +
22 +# Python Installation
23 +FROM base AS python-install
24 +ARG PYTHON_VERSION
25 +
26 +ENV VIRTUAL_ENV=/opt/vllm
27 +ENV PATH="$VIRTUAL_ENV/bin:$PATH"
28 +ENV PYTHON_VERSION=${PYTHON_VERSION}
29 +RUN microdnf install -y \
30 + python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip python${PYTHON_VERSION}-wheel && \
31 + python${PYTHON_VERSION} -m venv $VIRTUAL_ENV && pip install --no-cache -U pip wheel uv && microdnf clean all
32 +
33 +FROM python-install AS pyarrow
34 +
35 +# Build Apache Arrow
36 +WORKDIR /tmp
37 +RUN --mount=type=cache,target=/root/.cache/uv \
38 + git clone https://github.com/apache/arrow.git && \
39 + cd arrow/cpp && \
40 + mkdir release && cd release && \
41 + cmake -DCMAKE_BUILD_TYPE=Release \
42 + -DCMAKE_INSTALL_PREFIX=/usr/local \
43 + -DARROW_PYTHON=ON \
44 + -DARROW_PARQUET=ON \
45 + -DARROW_ORC=ON \
46 + -DARROW_FILESYSTEM=ON \
47 + -DARROW_WITH_LZ4=ON \
48 + -DARROW_WITH_ZSTD=ON \
49 + -DARROW_WITH_SNAPPY=ON \
50 + -DARROW_JSON=ON \
51 + -DARROW_CSV=ON \
52 + -DARROW_DATASET=ON \
53 + -DPROTOBUF_PROTOC_EXECUTABLE=/usr/bin/protoc \
54 + -DARROW_DEPENDENCY_SOURCE=BUNDLED \
55 + .. && \
56 + make -j$(nproc) && \
57 + make install && \
58 + cd ../../python && \
59 + export PYARROW_PARALLEL=4 && \
60 + export ARROW_BUILD_TYPE=release && \
61 + uv pip install -r requirements-build.txt && \
62 + python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
63 +
64 +FROM python-install AS numa-build
65 +# Install numactl (needed for numa.h dependency)
66 +WORKDIR /tmp
67 +RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.16.tar.gz && \
68 + tar -xvzf v2.0.16.tar.gz && \
69 + cd numactl-2.0.16 && \
70 + ./autogen.sh && \
71 + ./configure && \
72 + make
73 +
74 +# Set include path
75 +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
76 +
77 +FROM python-install AS rust
78 +ENV CARGO_HOME=/root/.cargo
79 +ENV RUSTUP_HOME=/root/.rustup
80 +ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
81 +
82 +RUN curl https://sh.rustup.rs -sSf | sh -s -- -y && \
83 + . "$CARGO_HOME/env" && \
84 + rustup default stable && \
85 + rustup show
86 +
87 +FROM python-install AS torch-vision
88 +# Install torchvision
89 +ARG TORCH_VERSION=2.7.0.dev20250304
90 +ARG TORCH_VISION_VERSION=v0.20.1
91 +WORKDIR /tmp
92 +RUN --mount=type=cache,target=/root/.cache/uv \
93 + git clone https://github.com/pytorch/vision.git && \
94 + cd vision && \
95 + git checkout $TORCH_VISION_VERSION && \
96 + uv pip install -v torch==${TORCH_VERSION} --extra-index-url https://download.pytorch.org/whl/nightly/cpu && \
97 + python setup.py bdist_wheel
98 +
99 +# Final build stage
100 +FROM python-install AS vllm-cpu
101 +ARG PYTHON_VERSION
102 +
103 +# Set correct library path for torch and numactl
104 +ENV LD_LIBRARY_PATH="/opt/vllm/lib64/python${PYTHON_VERSION}/site-packages/torch/lib:/usr/local/lib:$LD_LIBRARY_PATH"
105 +ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
106 +ENV UV_LINK_MODE=copy
107 +ENV CARGO_HOME=/root/.cargo
108 +ENV RUSTUP_HOME=/root/.rustup
109 +ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
110 +
111 +COPY . /workspace/vllm
112 +WORKDIR /workspace/vllm
113 +
114 +RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.16,target=/numactl \
115 + make -C /numactl install
116 +
117 +# Install dependencies, including PyTorch and Apache Arrow
118 +RUN --mount=type=cache,target=/root/.cache/uv \
119 + --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
120 + --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
121 + --mount=type=bind,from=pyarrow,source=/tmp/arrow/python/dist,target=/tmp/arrow-wheels \
122 + --mount=type=bind,from=torch-vision,source=/tmp/vision/dist,target=/tmp/vision-wheels/ \
123 + sed -i '/^torch/d' requirements-build.txt && \
124 + ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/pyarrow-*.whl | head -n 1) && \
125 + VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl | head -n 1) && \
126 + uv pip install -v \
127 + $ARROW_WHL_FILE \
128 + $VISION_WHL_FILE \
129 + --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
130 + --index-strategy unsafe-best-match \
131 + -r requirements-build.txt \
132 + -r requirements-cpu.txt
133 +
134 +# Build and install vllm
135 +RUN --mount=type=cache,target=/root/.cache/uv \
136 + VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
137 + uv pip install "$(echo dist/*.whl)[tensorizer]"
138 +
139 +# setup non-root user for vllm
140 +RUN umask 002 && \
141 + useradd --uid 2000 --gid 0 vllm && \
142 + mkdir -p /home/vllm && \
143 + chmod g+rwx /home/vllm
144 +
145 +COPY LICENSE /licenses/vllm.md
146 +COPY examples/*.jinja /app/data/template/
147 +
148 +USER 2000
149 +WORKDIR /home/vllm
150 +
151 +# Set the default entrypoint
152 +ENTRYPOINT ["python", "-m", "vllm.entrypoints.openai.api_server"]
\ No newline at end of file
cmake/cpu_extension.cmake
+10 -1
@@ -81,6 +81,7 @@ else()
81 find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
82 find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
83 find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
84 + find_isa(${CPUINFO} "S390" S390_FOUND)
85 endif()
86
87
@@ -129,8 +130,16 @@ elseif (ASIMD_FOUND)
130 elseif(APPLE_SILICON_FOUND)
131 message(STATUS "Apple Silicon Detected")
132 set(ENABLE_NUMA OFF)
133 +elseif (S390_FOUND)
134 + message(STATUS "S390 detected")
135 + # Check for S390 VXE support
136 + list(APPEND CXX_COMPILE_FLAGS
137 + "-mvx"
138 + "-mzvector"
139 + "-march=native"
140 + "-mtune=native")
141 else()
133 - message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA or ARMv8 support.")
142 + message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA, S390X ISA or ARMv8 support.")
143 endif()
144
145 #
csrc/cpu/attention.cpp
+2 -2
@@ -24,8 +24,8 @@ struct KernelVecType<float> {
24
25 template <>
26 struct KernelVecType<c10::Half> {
27 -#ifdef __powerpc64__
28 - // Power architecture-specific vector types
27 +#if defined(__powerpc64__) || defined(__s390x__)
28 + // Power and s390x architecture-specific vector types
29 using q_load_vec_type = vec_op::FP32Vec8;
30 using k_load_vec_type = vec_op::FP32Vec16;
31 using v_load_vec_type = vec_op::FP32Vec16;
csrc/cpu/cpu_types.hpp
+3
@@ -7,6 +7,9 @@
7 #elif defined(__POWER9_VECTOR__)
8 // ppc implementation
9 #include "cpu_types_vsx.hpp"
10 +#elif defined(__s390x__)
11 + // s390 implementation
12 + #include "cpu_types_vxe.hpp"
13 #elif defined(__aarch64__)
14 // arm implementation
15 #include "cpu_types_arm.hpp"
csrc/cpu/cpu_types_vxe.hpp new
+480
@@ -0,0 +1,480 @@
1 +
2 +#ifndef CPU_TYPES_VXE_HPP
3 +#define CPU_TYPES_VXE_HPP
4 +
5 +#include <vecintrin.h>
6 +#include <cmath>
7 +#include <torch/all.h>
8 +namespace vec_op {
9 +
10 +#define vec_neg(a) (-(a))
11 +#define vec_add(a, b) ((a) + (b))
12 +#define vec_sub(a, b) ((a) - (b))
13 +#define vec_mul(a, b) ((a) * (b))
14 +#define vec_div(a, b) ((a) / (b))
15 +#define vec_sr(a, b) ((a) >> (b)) // Vector Shift Right Algebaic
16 +#define vec_sl(a, b) ((a) << (b)) // Vector Shift Left
17 +
18 +// FIXME: FP16 is not fully supported in Torch-CPU
19 +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
20 + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
21 + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
22 +
23 +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
24 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
25 +
26 +#ifndef CPU_OP_GUARD
27 + #define CPU_KERNEL_GUARD_IN(NAME)
28 + #define CPU_KERNEL_GUARD_OUT(NAME)
29 +#else
30 + #define CPU_KERNEL_GUARD_IN(NAME) \
31 + std::cout << #NAME << " invoked." << std::endl;
32 + #define CPU_KERNEL_GUARD_OUT(NAME) \
33 + std::cout << #NAME << " exit." << std::endl;
34 +#endif
35 +
36 +#define FORCE_INLINE __attribute__((always_inline)) inline
37 +
38 +namespace {
39 +template <typename T, T... indexes, typename F>
40 +constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
41 + (f(std::integral_constant<T, indexes>{}), ...);
42 +}
43 +}; // namespace
44 +
45 +template <typename T, T count, typename F,
46 + typename = std::enable_if_t<std::is_invocable_v<F, T>>>
47 +constexpr void unroll_loop(F&& f) {
48 + unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
49 +}
50 +
51 +template <typename T>
52 +struct Vec {
53 + constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }
54 +};
55 +
56 +typedef struct ss16x8x2_t {
57 + __vector signed short val[2];
58 +} ss16x8x2_t;
59 +
60 +typedef struct ss16x8x4_t {
61 + __vector signed short val[4];
62 +} ss16x8x4_t;
63 +
64 +typedef struct f32x4x2_t {
65 + __vector float val[2];
66 +} f32x4x2_t;
67 +
68 +typedef struct f32x4x4_t {
69 + __vector float val[4];
70 +} f32x4x4_t;
71 +
72 +struct FP32Vec8;
73 +struct FP32Vec16;
74 +
75 +struct BF16Vec8 : public Vec<BF16Vec8> {
76 + constexpr static int VEC_ELEM_NUM = 8;
77 +
78 + __vector signed short reg;
79 +
80 + explicit BF16Vec8(const void* ptr) : reg(*(__vector signed short*)ptr) {}
81 + explicit BF16Vec8(const FP32Vec8&);
82 +
83 + void save(void* ptr) const {
84 + *reinterpret_cast<__vector signed short*>(ptr) = reg;
85 + }
86 +};
87 +
88 +struct BF16Vec16 : public Vec<BF16Vec16> {
89 + constexpr static int VEC_ELEM_NUM = 16;
90 +
91 + ss16x8x2_t reg;
92 +
93 + explicit BF16Vec16(const void* ptr) {
94 + // Load 256 bits in two parts
95 + reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
96 + reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
97 + }
98 +
99 + explicit BF16Vec16(const FP32Vec16&);
100 +
101 + void save(void* ptr) const {
102 + // Save 256 bits in two parts
103 + vec_xst(reg.val[0], 0, (signed short*)ptr);
104 + vec_xst(reg.val[1], 16, (signed short*)ptr);
105 + }
106 +};
107 +
108 +const static __vector signed short zero = vec_splats((signed short)0);
109 +
110 +struct BF16Vec32 : public Vec<BF16Vec32> {
111 + constexpr static int VEC_ELEM_NUM = 32;
112 +
113 + ss16x8x4_t reg;
114 + explicit BF16Vec32(const void* ptr)
115 + : reg(*reinterpret_cast<const ss16x8x4_t*>(ptr)) {}
116 +
117 + explicit BF16Vec32(ss16x8x4_t data) : reg(data) {}
118 +
119 + explicit BF16Vec32(const BF16Vec8& vec8_data)
120 + : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
121 +
122 + void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
123 +};
124 +
125 +struct FP32Vec4 : public Vec<FP32Vec4> {
126 + constexpr static int VEC_ELEM_NUM = 4;
127 + union AliasReg {
128 + __vector float reg;
129 + float values[VEC_ELEM_NUM];
130 + };
131 +
132 + __vector float reg;
133 +
134 + explicit FP32Vec4(float v) : reg(vec_splats(v)) {}
135 +
136 + explicit FP32Vec4() : reg(vec_splats(0.0f)) {}
137 +
138 + explicit FP32Vec4(const float* ptr) : reg(vec_xl(0, ptr)) {}
139 +
140 + explicit FP32Vec4(__vector float data) : reg(data) {}
141 +
142 + explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {}
143 +};
144 +
145 +struct FP32Vec8 : public Vec<FP32Vec8> {
146 + constexpr static int VEC_ELEM_NUM = 8;
147 + union AliasReg {
148 + f32x4x2_t reg;
149 + float values[VEC_ELEM_NUM];
150 + };
151 +
152 + f32x4x2_t reg;
153 +
154 + explicit FP32Vec8(float v) {
155 + reg.val[0] = vec_splats(v);
156 + reg.val[1] = vec_splats(v);
157 + }
158 +
159 + explicit FP32Vec8() {
160 + reg.val[0] = vec_splats(0.0f);
161 + reg.val[1] = vec_splats(0.0f);
162 + }
163 +
164 + explicit FP32Vec8(const float* ptr) {
165 + reg.val[0] = vec_xl(0, ptr);
166 + reg.val[1] = vec_xl(16, ptr);
167 + }
168 +
169 + explicit FP32Vec8(f32x4x2_t data) : reg(data) {}
170 +
171 + explicit FP32Vec8(const FP32Vec8& data) {
172 + reg.val[0] = data.reg.val[0];
173 + reg.val[1] = data.reg.val[1];
174 + }
175 +
176 + explicit FP32Vec8(const BF16Vec8& v) {
177 + reg.val[0] = (__vector float)vec_mergeh(zero, v.reg);
178 + reg.val[1] = (__vector float)vec_mergel(zero, v.reg);
179 + }
180 +
181 + float reduce_sum() const {
182 + AliasReg ar;
183 + ar.reg = reg;
184 + float result = 0;
185 + unroll_loop<int, VEC_ELEM_NUM>(
186 + [&result, &ar](int i) { result += ar.values[i]; });
187 +
188 + return result;
189 + }
190 +
191 + FP32Vec8 exp() const {
192 + // TODO: Vectorize this
193 + AliasReg ar;
194 + ar.reg = reg;
195 + f32x4x4_t ret;
196 + ret.val[0][0] = std::exp(ar.values[0]);
197 + ret.val[0][1] = std::exp(ar.values[1]);
198 + ret.val[0][2] = std::exp(ar.values[2]);
199 + ret.val[0][3] = std::exp(ar.values[3]);
200 + ret.val[1][0] = std::exp(ar.values[4]);
201 + ret.val[1][1] = std::exp(ar.values[5]);
202 + ret.val[1][2] = std::exp(ar.values[6]);
203 + ret.val[1][3] = std::exp(ar.values[7]);
204 + return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
205 + }
206 +
207 + FP32Vec8 tanh() const {
208 + // TODO: Vectorize this
209 + AliasReg ar;
210 + ar.reg = reg;
211 + f32x4x4_t ret;
212 + ret.val[0][0] = std::tanh(ar.values[0]);
213 + ret.val[0][1] = std::tanh(ar.values[1]);
214 + ret.val[0][2] = std::tanh(ar.values[2]);
215 + ret.val[0][3] = std::tanh(ar.values[3]);
216 + ret.val[1][0] = std::tanh(ar.values[4]);
217 + ret.val[1][1] = std::tanh(ar.values[5]);
218 + ret.val[1][2] = std::tanh(ar.values[6]);
219 + ret.val[1][3] = std::tanh(ar.values[7]);
220 + return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
221 + }
222 +
223 + FP32Vec8 er() const {
224 + // TODO: Vectorize this
225 + AliasReg ar;
226 + ar.reg = reg;
227 + f32x4x4_t ret;
228 + ret.val[0][0] = std::erf(ar.values[0]);
229 + ret.val[0][1] = std::erf(ar.values[1]);
230 + ret.val[0][2] = std::erf(ar.values[2]);
231 + ret.val[0][3] = std::erf(ar.values[3]);
232 + ret.val[1][0] = std::erf(ar.values[4]);
233 + ret.val[1][1] = std::erf(ar.values[5]);
234 + ret.val[1][2] = std::erf(ar.values[6]);
235 + ret.val[1][3] = std::erf(ar.values[7]);
236 + return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
237 + }
238 +
239 + FP32Vec8 operator*(const FP32Vec8& b) const {
240 + return FP32Vec8(
241 + {vec_mul(reg.val[0], b.reg.val[0]), vec_mul(reg.val[1], b.reg.val[1])});
242 + }
243 +
244 + FP32Vec8 operator+(const FP32Vec8& b) const {
245 + return FP32Vec8(
246 + {vec_add(reg.val[0], b.reg.val[0]), vec_add(reg.val[1], b.reg.val[1])});
247 + }
248 +
249 + FP32Vec8 operator-(const FP32Vec8& b) const {
250 + return FP32Vec8(
251 + {vec_sub(reg.val[0], b.reg.val[0]), vec_sub(reg.val[1], b.reg.val[1])});
252 + }
253 +
254 + FP32Vec8 operator/(const FP32Vec8& b) const {
255 + return FP32Vec8(
256 + {vec_div(reg.val[0], b.reg.val[0]), vec_div(reg.val[1], b.reg.val[1])});
257 + }
258 +
259 + void save(float* ptr) const {
260 + vec_xst(reg.val[0], 0, ptr);
261 + vec_xst(reg.val[1], 16, ptr);
262 + }
263 +};
264 +
265 +struct FP32Vec16 : public Vec<FP32Vec16> {
266 + constexpr static int VEC_ELEM_NUM = 16;
267 + union AliasReg {
268 + f32x4x4_t reg;
269 + float values[VEC_ELEM_NUM];
270 + };
271 +
272 + f32x4x4_t reg;
273 +
274 + explicit FP32Vec16(float v) {
275 + reg.val[0] = vec_splats(v);
276 + reg.val[1] = vec_splats(v);
277 + reg.val[2] = vec_splats(v);
278 + reg.val[3] = vec_splats(v);
279 + }
280 +
281 + explicit FP32Vec16() {
282 + reg.val[0] = vec_splats(0.0f);
283 + reg.val[1] = vec_splats(0.0f);
284 + reg.val[2] = vec_splats(0.0f);
285 + reg.val[3] = vec_splats(0.0f);
286 + }
287 +
288 + explicit FP32Vec16(const float* ptr) {
289 + reg.val[0] = vec_xl(0, ptr);
290 + reg.val[1] = vec_xl(16, ptr);
291 + reg.val[2] = vec_xl(32, ptr);
292 + reg.val[3] = vec_xl(48, ptr);
293 + }
294 +
295 + explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
296 +
297 + explicit FP32Vec16(const FP32Vec16& data) {
298 + reg.val[0] = data.reg.val[0];
299 + reg.val[1] = data.reg.val[1];
300 + reg.val[2] = data.reg.val[2];
301 + reg.val[3] = data.reg.val[3];
302 + }
303 +
304 + explicit FP32Vec16(const FP32Vec4& data) {
305 + reg.val[0] = data.reg;
306 + reg.val[1] = data.reg;
307 + reg.val[2] = data.reg;
308 + reg.val[3] = data.reg;
309 + }
310 +
311 + explicit FP32Vec16(const FP32Vec8& data) {
312 + reg.val[0] = data.reg.val[0];
313 + reg.val[1] = data.reg.val[1];
314 + reg.val[2] = data.reg.val[0];
315 + reg.val[3] = data.reg.val[1];
316 + }
317 +
318 + explicit FP32Vec16(const BF16Vec16& v) {
319 + reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
320 + reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
321 + reg.val[2] = (__vector float)vec_mergeh(zero, v.reg.val[1]);
322 + reg.val[3] = (__vector float)vec_mergel(zero, v.reg.val[1]);
323 + }
324 +
325 + explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
326 +
327 + FP32Vec16 operator*(const FP32Vec16& b) const {
328 + return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
329 + vec_mul(reg.val[1], b.reg.val[1]),
330 + vec_mul(reg.val[2], b.reg.val[2]),
331 + vec_mul(reg.val[3], b.reg.val[3])}));
332 + }
333 +
334 + FP32Vec16 operator+(const FP32Vec16& b) const {
335 + return FP32Vec16(f32x4x4_t({vec_add(reg.val[0], b.reg.val[0]),
336 + vec_add(reg.val[1], b.reg.val[1]),
337 + vec_add(reg.val[2], b.reg.val[2]),
338 + vec_add(reg.val[3], b.reg.val[3])}));
339 + }
340 +
341 + FP32Vec16 operator-(const FP32Vec16& b) const {
342 + return FP32Vec16(f32x4x4_t({vec_sub(reg.val[0], b.reg.val[0]),
343 + vec_sub(reg.val[1], b.reg.val[1]),
344 + vec_sub(reg.val[2], b.reg.val[2]),
345 + vec_sub(reg.val[3], b.reg.val[3])}));
346 + }
347 +
348 + FP32Vec16 operator/(const FP32Vec16& b) const {
349 + return FP32Vec16(f32x4x4_t({vec_div(reg.val[0], b.reg.val[0]),
350 + vec_div(reg.val[1], b.reg.val[1]),
351 + vec_div(reg.val[2], b.reg.val[2]),
352 + vec_div(reg.val[3], b.reg.val[3])}));
353 + }
354 +
355 + float reduce_sum() const {
356 + AliasReg ar;
357 + ar.reg = reg;
358 + float result = 0;
359 + unroll_loop<int, VEC_ELEM_NUM>(
360 + [&result, &ar](int i) { result += ar.values[i]; });
361 +
362 + return result;
363 + }
364 +
365 + template <int group_size>
366 + float reduce_sub_sum(int idx) {
367 + static_assert(VEC_ELEM_NUM % group_size == 0);
368 +
369 + AliasReg ar;
370 + ar.reg = reg;
371 + float result = 0;
372 + const int start = idx * group_size;
373 + unroll_loop<int, group_size>(
374 + [&result, &start, ar](int i) { result += ar.values[start + i]; });
375 +
376 + return result;
377 + }
378 +
379 + void save(float* ptr) const {
380 + vec_xst(reg.val[0], 0, ptr);
381 + vec_xst(reg.val[1], 16, ptr);
382 + vec_xst(reg.val[2], 32, ptr);
383 + vec_xst(reg.val[3], 48, ptr);
384 + }
385 +};
386 +
387 +template <typename T>
388 +struct VecType {
389 + using vec_type = void;
390 +};
391 +
392 +template <typename T>
393 +using vec_t = typename VecType<T>::vec_type;
394 +
395 +template <>
396 +struct VecType<float> {
397 + using vec_type = FP32Vec8;
398 +};
399 +
400 +template <>
401 +struct VecType<c10::BFloat16> {
402 + using vec_type = BF16Vec8;
403 +};
404 +
405 +template <typename T>
406 +void storeFP32(float v, T* ptr) {
407 + *ptr = v;
408 +}
409 +
410 +inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
411 + acc = acc + a * b;
412 +}
413 +
414 +namespace c10 {
415 +struct BFloat16 {
416 + uint16_t value; // Assume BFloat16 is defined as a struct containing a 16-bit
417 + // value.
418 +};
419 +} // namespace c10
420 +
421 +template <>
422 +inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
423 + c10::BFloat16 __attribute__((__may_alias__))* v_ptr =
424 + reinterpret_cast<c10::BFloat16*>(&v);
425 + *ptr = *(v_ptr + 1);
426 +}
427 +
428 +#ifndef __VEC_CLASS_FP_NAN
429 + #define __VEC_CLASS_FP_NAN (1 << 6)
430 +#endif
431 +
432 +const static __vector unsigned char omask = {2, 3, 6, 7, 10, 11, 14, 15,
433 + 18, 19, 22, 23, 26, 27, 30, 31};
434 +const static __vector unsigned int bias = {0x00007fff, 0x00007fff, 0x00007fff,
435 + 0x00007fff};
436 +const static __vector unsigned int nan = {0x7fc00000, 0x7fc00000, 0x7fc00000,
437 + 0x7fc00000};
438 +const static __vector unsigned int sh16 = {16, 16, 16, 16};
439 +const static __vector unsigned int one = {1, 1, 1, 1};
440 +
441 +inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
442 + __vector unsigned int inp0 = (__vector unsigned int)(v.reg.val[0]);
443 + __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
444 + int cc;
445 + __vector __bool int sel0 =
446 + vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
447 + __vector __bool int sel1 =
448 + vec_fp_test_data_class(v.reg.val[1], __VEC_CLASS_FP_NAN, &cc);
449 + inp0 = vec_sel(inp0, nan, sel0) >> sh16;
450 + inp1 = vec_sel(inp1, nan, sel1) >> sh16;
451 + reg = (__vector signed short)vec_perm(inp0, inp1, omask);
452 +}
453 +
454 +inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
455 + __vector unsigned int inp0 = (__vector unsigned int)(v.reg.val[0]);
456 + __vector unsigned int inp1 = (__vector unsigned int)(v.reg.val[1]);
457 + __vector unsigned int inp2 = (__vector unsigned int)(v.reg.val[2]);
458 + __vector unsigned int inp3 = (__vector unsigned int)(v.reg.val[3]);
459 + int cc;
460 + __vector __bool int sel0 =
461 + vec_fp_test_data_class(v.reg.val[0], __VEC_CLASS_FP_NAN, &cc);
462 + __vector __bool int sel1 =
463 + vec_fp_test_data_class(v.reg.val[1], __VEC_CLASS_FP_NAN, &cc);
464 + __vector __bool int sel2 =
465 + vec_fp_test_data_class(v.reg.val[2], __VEC_CLASS_FP_NAN, &cc);
466 + __vector __bool int sel3 =
467 + vec_fp_test_data_class(v.reg.val[3], __VEC_CLASS_FP_NAN, &cc);
468 + inp0 = vec_sel(inp0, nan, sel0) >> sh16;
469 + inp1 = vec_sel(inp1, nan, sel1) >> sh16;
470 + inp2 = vec_sel(inp2, nan, sel2) >> sh16;
471 + inp3 = vec_sel(inp3, nan, sel3) >> sh16;
472 + reg.val[0] = (__vector signed short)vec_perm(inp0, inp1, omask);
473 + reg.val[1] = (__vector signed short)vec_perm(inp2, inp3, omask);
474 +}
475 +
476 +inline void prefetch(const void* addr) { void __dcbt(const void* addr); }
477 +
478 +}; // namespace vec_op
479 +
480 +#endif
\ No newline at end of file
csrc/cpu/quant.cpp
+1 -1
@@ -25,7 +25,7 @@ struct KernelVecType<c10::BFloat16> {
25
26 template <>
27 struct KernelVecType<c10::Half> {
28 -#ifdef __powerpc64__
28 +#if defined(__powerpc64__) || defined(__s390x__)
29 // Power architecture-specific vector type
30 using load_vec_type = vec_op::FP32Vec16;
31 #else
requirements-cpu.txt
+5 -4
@@ -2,14 +2,15 @@
2 -r requirements-common.txt
3
4 # Dependencies for CPUs
5 -torch==2.5.1+cpu; platform_machine != "ppc64le" and platform_machine != "aarch64" and platform_system != "Darwin"
6 -torch==2.5.1; platform_machine == "ppc64le" or platform_machine == "aarch64" or platform_system == "Darwin"
5 +torch==2.5.1+cpu; platform_machine != "ppc64le" and platform_machine != "aarch64" and platform_system != "Darwin" and platform_machine != "s390x"
6 +torch==2.5.1; platform_machine == "ppc64le" or platform_machine == "aarch64" or platform_system == "Darwin"
7 +torch==2.7.0.dev20250304; platform_machine == "s390x"
8
9 # required for the image processor of minicpm-o-2_6, this must be updated alongside torch
9 -torchaudio; platform_machine != "ppc64le"
10 +torchaudio; platform_machine != "ppc64le" and platform_machine != "s390x"
11 torchaudio==2.5.1; platform_machine == "ppc64le"
12
13 # required for the image processor of phi3v, this must be updated alongside torch
13 -torchvision; platform_machine != "ppc64le"
14 +torchvision; platform_machine != "ppc64le" and platform_machine != "s390x"
15 torchvision==0.20.1; platform_machine == "ppc64le"
16 datasets # for benchmark scripts