@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fd2f3fcfa5

[CPU] Add video inferencing via torchcodec on s390x (#58693)

Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Rehan Khan committed Sep 28, 2026 at 11:46 UTC fd2f3fcfa5654fcef8361fb8fcd0837382e09da8

3 files changed +50 -5

docker/Dockerfile.s390x

+42

@@ -141,6 +141,41 @@ RUN --mount=type=cache,target=/root/.cache/uv \

141 sed -i '/#include "internal\/pycore_atomic.h"/i\#include "dynamic_annotations.h"' numba/_dispatcher.cpp; \

142 fi && python setup.py bdist_wheel

143

144 +# Shared FFmpeg for torchcodec video decoding (UBI has no ffmpeg package).

145 +FROM base AS ffmpeg-build

146 +ARG FFMPEG_VERSION=n6.1.2

147 +WORKDIR /tmp

148 +RUN microdnf install -y pkgconf-pkg-config xz && \

149 + curl -LO https://github.com/FFmpeg/FFmpeg/archive/refs/tags/${FFMPEG_VERSION}.tar.gz && \

150 + tar -xzf ${FFMPEG_VERSION}.tar.gz && \

151 + cd FFmpeg-${FFMPEG_VERSION} && \

152 + ./configure --prefix=/usr/local --libdir=/usr/local/lib \

153 + --enable-shared --disable-static --enable-pic \

154 + --disable-doc --disable-debug --disable-ffplay --disable-ffprobe \

155 + --extra-cflags="-Wno-incompatible-pointer-types" && \

156 + make -j"$(nproc)" && \

157 + make install

158 +

159 +# Torchcodec video decoder (no s390x wheels; CPU / FFmpeg only s390x does not support audio decoding).

160 +FROM python-install AS torchcodec-builder

161 +ARG TORCHCODEC_VERSION=v0.14.0

162 +WORKDIR /tmp

163 +ENV LD_LIBRARY_PATH="/usr/local/lib:$LD_LIBRARY_PATH"

164 +ENV PKG_CONFIG_PATH="/usr/local/lib/pkgconfig:$PKG_CONFIG_PATH"

165 +RUN --mount=type=cache,target=/root/.cache/uv \

166 + --mount=type=bind,from=ffmpeg-build,source=/usr/local,target=/usr/local \

167 + git clone --depth 1 --branch ${TORCHCODEC_VERSION} \

168 + https://github.com/pytorch/torchcodec.git && \

169 + cd torchcodec && \

170 + uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cpu && \

171 + uv pip install pybind11 setuptools wheel && \

172 + pybind11_DIR="$(python -c 'import pybind11; print(pybind11.get_cmake_dir())')" && \

173 + export pybind11_DIR CMAKE_PREFIX_PATH="${pybind11_DIR}:${CMAKE_PREFIX_PATH}" && \

174 + export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1 && \

175 + export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1 && \

176 + export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build" && \

177 + python setup.py bdist_wheel

178 +

179 # Build OpenCV from source for s390x

180 FROM python-install AS opencv-builder

181 WORKDIR /tmp

@@ -179,6 +214,10 @@ RUN --mount=type=bind,from=gperftools-build,src=/tmp/gperftools-2.16,target=/gpe

214 make -C /gperftools install-libLTLIBRARIES && \

215 /sbin/ldconfig

216

217 +RUN --mount=type=bind,from=ffmpeg-build,source=/usr/local,target=/ffmpeg-prefix \

218 + cp -a /ffmpeg-prefix/lib/. /usr/local/lib/ && \

219 + /sbin/ldconfig

220 +

221 ENV LD_PRELOAD=/usr/local/lib/libtcmalloc_minimal.so.4

222

223 # Install dependencies, including PyTorch and Apache Arrow

@@ -190,17 +229,20 @@ RUN --mount=type=cache,target=/root/.cache/uv \

229 --mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \

230 --mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \

231 --mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \

232 + --mount=type=bind,from=torchcodec-builder,source=/tmp/torchcodec/dist,target=/tmp/torchcodec-wheels/ \

233 VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \

234 HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \

235 LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \

236 NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \

237 OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \

238 + TORCHCODEC_WHL_FILE=$(ls /tmp/torchcodec-wheels/*.whl) && \

239 uv pip install -v \

240 $VISION_WHL_FILE \

241 $HF_XET_WHL_FILE \

242 $LLVM_WHL_FILE \

243 $NUMBA_WHL_FILE \

244 $OPENCV_WHL_FILE \

245 + $TORCHCODEC_WHL_FILE \

246 --torch-backend cpu \

247 --index-strategy unsafe-best-match \

248 -r requirements/build/cpu.txt \

docs/getting_started/installation/cpu.s390x.inc.md

+5 -2

@@ -11,7 +11,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and

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 Z15 and above.

14 -- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`

14 +- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`, `torchcodec`

16 --8<-- [end:requirements]

17 --8<-- [start:set-up-using-python]

@@ -55,7 +55,8 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \

55 Execute the following commands to build and install vLLM from source.

57 !!! tip

58 - Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`.

58 + Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`, `torchcodec`.

59 + `torchcodec` also needs a shared FFmpeg (built in `Dockerfile.s390x`). Video decoding only; there is no s390x NVDEC path.

60 See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build.

62 !!! note "LLVM 20 required for llvmlite"

@@ -94,6 +95,7 @@ uv pip install -v \

95 /path/to/numba.whl \

96 /path/to/opencv_python_headless.whl \

97 /path/to/hf_xet.whl \

98 + /path/to/torchcodec.whl \

99 -r requirements/build/cpu.txt \

100 -r requirements/cpu.txt \

101 --torch-backend cpu \

@@ -111,6 +113,7 @@ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \

113 /path/to/numba.whl \

114 /path/to/opencv_python_headless.whl \

115 /path/to/hf_xet.whl \

116 + /path/to/torchcodec.whl \

117 -r requirements/build/cpu.txt \

118 -r requirements/cpu.txt && \

119 VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \

requirements/cpu.txt

+3 -3

@@ -3,7 +3,7 @@

4 setuptools==77.0.3 # this version can reuse CMake build dir

6 -numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding

6 +numba == 0.65.0 # Required for N-gram speculative decoding

8 # Dependencies for CPUs

9 torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"

@@ -13,10 +13,10 @@ torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or p

13 torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"

15 # required for the image processor of phi3v, this must be updated alongside torch

16 -torchvision; platform_machine != "s390x" and platform_machine != "riscv64"

16 +torchvision; platform_machine != "riscv64"

18 # required for the torchcodec video decoding backend

19 -torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64" and platform_machine != "ppc64le"

19 +torchcodec >= 0.14; platform_machine != "riscv64" and platform_machine != "ppc64le"

21 # Intel Extension for PyTorch, only for x86_64 CPUs

22 intel-openmp==2024.2.1; platform_machine == "x86_64"