Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions builder/manywheel/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
## Building lmdeploy builder images

To build all lmdeploy builder images, such as
`lmdeploy-builder:manylinux_2_28-cuda12.8`, execute:
`lmdeploy-builder:manylinux_2_28-cuda12.8` and
`lmdeploy-builder:manylinux_2_28-cuda13.0`, execute:

```bash
./build_all_lmdeploy_builders.sh
Expand All @@ -15,7 +16,7 @@ WITH_PUSH=true ./build_all_lmdeploy_builders.sh
For custom builds with specific versions:

```bash
MANY_LINUX_VERSION=2_28 GPU_ARCH_VERSION=12.8 ./build_lmdeploy_builder.sh
MANY_LINUX_VERSION=2_28 GPU_ARCH_VERSION=13.0 ./build_lmdeploy_builder.sh
```

## Build lmdeploy wheels
Expand Down
45 changes: 32 additions & 13 deletions docker/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,15 +1,20 @@
# Base images
ARG IMAGE_TYPE=final
ARG CUDA_VERSION=cu12
ARG CUDA_VERSION=cu128

FROM nvidia/cuda:13.0.2-devel-ubuntu22.04 AS cu13
FROM nvidia/cuda:13.0.3-devel-ubuntu24.04 AS cu130
ENV CUDA_VERSION_SHORT=cu130
ENV PYTORCH_CUDA_VERSION=cu130

FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 AS cu12.8
FROM nvidia/cuda:12.8.1-devel-ubuntu24.04 AS cu128
ENV CUDA_VERSION_SHORT=cu128
# PyTorch 2.12 no longer publishes cu128 wheels; cu126 is its CUDA 12 build.
ENV PYTORCH_CUDA_VERSION=cu126

FROM nvidia/cuda:12.6.3-devel-ubuntu22.04 AS cu12
ENV CUDA_VERSION_SHORT=cu126
# Backward-compatible stage aliases.
FROM cu130 AS cu13
FROM cu128 AS cu12.8
FROM cu128 AS cu12

# Builder image
FROM ${CUDA_VERSION} AS dev
Expand All @@ -20,7 +25,13 @@ ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

RUN --mount=type=cache,target=/root/.cache \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list && \
if [ -f /etc/apt/sources.list ]; then \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list; \
fi && \
if [ -d /etc/apt/sources.list.d ]; then \
find /etc/apt/sources.list.d -type f -name '*.sources' -exec \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} +; \
fi && \
apt-get update -y && \
apt-get install -y --no-install-recommends \
tzdata wget curl openssh-server ssh sudo git-core \
Expand All @@ -40,7 +51,9 @@ ENV PYTHON_VERSION=${PYTHON_VERSION}
COPY requirements/common.txt requirements/runtime_cuda.txt /tmp/requirements/

RUN --mount=type=cache,target=/root/.cache \
pip install -r /tmp/requirements/runtime_cuda.txt --extra-index-url https://download.pytorch.org/whl/${CUDA_VERSION_SHORT}
grep -E '^torch(vision)?([<>=]|$)' /tmp/requirements/runtime_cuda.txt > /tmp/requirements/torch.txt && \
pip install -r /tmp/requirements/torch.txt --index-url https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION} && \
pip install -r /tmp/requirements/runtime_cuda.txt
Comment on lines +54 to +56

RUN mkdir -p /wheels

Expand All @@ -67,24 +80,30 @@ COPY --from=thirdparty-wheels /wheels /wheels
COPY --from=lmdeploy-wheel /wheels /wheels

# Runtime image
FROM nvidia/cuda:13.0.2-base-ubuntu22.04 AS cu13-base
FROM nvidia/cuda:13.0.3-base-ubuntu24.04 AS cu130-base
ENV CUDA_VERSION_SHORT=cu130
ENV PYTORCH_CUDA_VERSION=cu130

FROM nvidia/cuda:12.8.1-base-ubuntu22.04 AS cu12.8-base
FROM nvidia/cuda:12.8.1-base-ubuntu24.04 AS cu128-base
ENV CUDA_VERSION_SHORT=cu128
ENV PYTORCH_CUDA_VERSION=cu126

FROM nvidia/cuda:12.6.3-base-ubuntu22.04 AS cu12-base
ENV CUDA_VERSION_SHORT=cu126
# Backward-compatible stage aliases.
FROM cu130-base AS cu13-base
FROM cu128-base AS cu12.8-base
FROM cu128-base AS cu12-base

FROM ${CUDA_VERSION}-base AS final
ARG PYTHON_VERSION=3.12
ENV PYTHON_VERSION=${PYTHON_VERSION}

# Install the CUDA runtime requirements explicitly so torch is selected from
# the matching PyTorch CUDA index before local wheels resolve dependencies.
# Some dependencies such as timm (required by InternVL models) are listed in requirements/serve.txt.
COPY requirements/serve.txt /tmp/requirements/serve.txt
COPY requirements/common.txt requirements/runtime_cuda.txt requirements/serve.txt /tmp/requirements/
COPY docker/install.sh /tmp/install.sh
RUN --mount=type=cache,target=/root/.cache \
--mount=type=cache,target=/wheels,from=wheels,source=/wheels \
--mount=type=bind,target=/wheels,from=wheels,source=/wheels \
/tmp/install.sh

# explicitly set ptxas path for triton
Expand Down
17 changes: 12 additions & 5 deletions docker/Dockerfile.jetson
Original file line number Diff line number Diff line change
Expand Up @@ -27,12 +27,19 @@ RUN --mount=type=cache,target=/root/.cache \
COPY . /opt/lmdeploy
WORKDIR /opt/lmdeploy

# The Jetson AI Lab index may lag xgrammar releases. Force xgrammar from PyPI,
# while keeping the LMDeploy wheel and its CUDA/Torch dependencies on the Jetson index.
# Install the Jetson-specific Torch stack from the Jetson AI Lab index first.
# Resolve generic LMDeploy dependencies from PyPI so packages missing from the
# Jetson index, such as pydantic, do not fail dependency resolution.
RUN --mount=type=cache,target=/root/.cache \
--mount=type=cache,target=/opt/pytorch \
pip install --upgrade pip && \
pip install build change-wheel-version && \
python -m build -w -o /wheels -v . && \
change_wheel_version --local-version cu126 --delete-old-wheel /wheels/lmdeploy*.whl && \
pip install -v --no-deps --only-binary=:all: xgrammar==0.1.33 --index-url https://pypi.org/simple/ && \
pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/
grep -E '^torch(vision)?([<>=]|$)' requirements/runtime_cuda.txt > /tmp/requirements_torch.txt && \
mkdir -p /tmp/pytorch-wheels && \
pip download --no-deps --only-binary=:all: -r /tmp/requirements_torch.txt -d /tmp/pytorch-wheels \
--index-url https://pypi.jetson-ai-lab.io/jp6/cu126/+simple/ && \
pip install -v /tmp/pytorch-wheels/*.whl --index-url https://pypi.org/simple/ && \
pip install -v --no-deps --only-binary=:all: xgrammar==0.1.33 \
--index-url https://pypi.org/simple/ && \
pip install -v /wheels/lmdeploy*.whl --index-url https://pypi.org/simple/
17 changes: 13 additions & 4 deletions docker/Dockerfile_dev
Original file line number Diff line number Diff line change
@@ -1,16 +1,23 @@
FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 AS cu12.8
FROM nvidia/cuda:12.8.1-devel-ubuntu24.04 AS cu128

# environment variables
ENV DEBIAN_FRONTEND=noninteractive \
TZ=Etc/UTC \
PATH=/opt/py3/bin:/root/.local/bin:${PATH} \
TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \
CUDA_VERSION_SHORT=cu128
CUDA_VERSION_SHORT=cu128 \
PYTORCH_CUDA_VERSION=cu126

# Install dependencies and create python virtual environment
RUN --mount=type=cache,target=/var/cache/apt \
--mount=type=cache,target=/root/.cache \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list && \
if [ -f /etc/apt/sources.list ]; then \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list; \
fi && \
if [ -d /etc/apt/sources.list.d ]; then \
find /etc/apt/sources.list.d -type f -name '*.sources' -exec \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} +; \
fi && \
apt-get update -y && \
apt-get install -y --no-install-recommends \
tzdata wget curl openssh-server ssh sudo git-core \
Expand All @@ -33,7 +40,9 @@ RUN mkdir -p /wheels

# install lmdeploy and its dependencies
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install -r requirements_cuda.txt --extra-index-url https://download.pytorch.org/whl/cu128 && \
grep -E '^torch(vision)?([<>=]|$)' requirements/runtime_cuda.txt > /tmp/requirements_torch.txt && \
uv pip install -r /tmp/requirements_torch.txt --index-url https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION} && \
uv pip install -r requirements_cuda.txt && \
Comment on lines +43 to +45
uv pip install -e .

RUN --mount=type=cache,target=/root/.cache/uv \
Expand Down
2 changes: 1 addition & 1 deletion docker/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

mkdir -p /wheels

if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then
if [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then
pip install nvidia-nccl-cu13
else
pip install nvidia-nccl-cu12
Expand Down
99 changes: 90 additions & 9 deletions docker/install.sh
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,40 @@
if [ ! -f "/opt/py3/bin/python" ]; then
# install system packages
export DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list
if [ -f /etc/apt/sources.list ]; then
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' /etc/apt/sources.list
fi
if [ -d /etc/apt/sources.list.d ]; then
find /etc/apt/sources.list.d -type f -name '*.sources' -exec \
sed -i 's|http://archive.ubuntu.com|http://azure.archive.ubuntu.com|g' {} +
fi
apt-get update -y
apt-get install -y --no-install-recommends \
tzdata wget curl ssh sudo git-core vim libibverbs1 ibverbs-providers ibverbs-utils librdmacm1 libibverbs-dev rdma-core libmlx5-1

if [[ ${PYTHON_VERSION} != "3.10" ]]; then
# Ubuntu may expose the interpreter package without the matching -dev and
# -venv packages. Check the complete set before deciding whether the PPA
# is needed.
python_packages=(
"python${PYTHON_VERSION}"
"python${PYTHON_VERSION}-dev"
"python${PYTHON_VERSION}-venv"
)
need_python_ppa=0
for package in "${python_packages[@]}"; do
if ! apt-cache show "${package}" >/dev/null 2>&1; then
need_python_ppa=1
break
fi
done
if ((need_python_ppa)); then
apt-get install -y --no-install-recommends software-properties-common
add-apt-repository -y ppa:deadsnakes/ppa
apt-get update -y
fi

# install python, create virtual env
apt-get install -y --no-install-recommends \
python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv
apt-get install -y --no-install-recommends "${python_packages[@]}"

pushd /opt >/dev/null
python${PYTHON_VERSION} -m venv py3
Expand Down Expand Up @@ -46,13 +66,35 @@ export PATH=/opt/py3/bin:$PATH

pip install -U pip wheel setuptools

if [[ "${CUDA_VERSION_SHORT}" = "cu130" ]]; then
# In the production image, install torch/torchvision from the matching CUDA
# index first. Installing all local wheels directly would otherwise let pip
# resolve torch from the default PyPI index.
if [ -f /tmp/requirements/runtime_cuda.txt ]; then
grep -E '^torch(vision)?([<>=]|$)' \
/tmp/requirements/runtime_cuda.txt > /tmp/requirements/torch.txt
pip install -r /tmp/requirements/torch.txt \
--index-url "https://download.pytorch.org/whl/${PYTORCH_CUDA_VERSION}"
pip install -r /tmp/requirements/runtime_cuda.txt
Comment on lines +73 to +77
fi

if [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then
pip install nvidia-nvshmem-cu13==3.4.5
else
pip install nvidia-nvshmem-cu12==3.4.5
fi

for package in deep_ep deep_gemm flash_attn_3; do
if ! compgen -G "/wheels/${package}-*.whl" >/dev/null; then
echo "${package} wheel is missing from /wheels" >&2
exit 1
fi
done
pip install /wheels/*.whl
python - <<'PY'
import deep_ep
import deep_gemm
import lmdeploy.pytorch.third_party.flash_attn_interface
PY
pip install dlslime==0.0.2.post1

pip install ninja einops packaging
Expand All @@ -62,9 +104,48 @@ if [ -f /tmp/requirements/serve.txt ]; then
pip install -r /tmp/requirements/serve.txt
fi

# PyTorch cu130 pins an older NCCL release in its wheel metadata. Check the
# declared dependencies before intentionally upgrading NCCL for GIN support.
pip check

if [[ "${CUDA_VERSION_SHORT}" = "cu128" ]]; then
# As described in https://github.com/InternLM/lmdeploy/pull/4313,
# window registration may cause memory leaks in NCCL 2.27, NCCL 2.28+ resolves the issue,
# but turbomind engine will use nccl GIN for EP in future, which is brought in since 2.29
pip install "nvidia-nccl-cu12>2.29"
pip install "nvidia-nccl-cu12>=2.30.4"
elif [[ "${CUDA_VERSION_SHORT}" == cu13* ]]; then
pip install "nvidia-nccl-cu13>=2.30.4"
fi

python - <<'PY'
import os
from importlib.metadata import distributions, version

import torch
from packaging.version import Version

expected = os.environ['PYTORCH_CUDA_VERSION']
actual = torch.version.cuda
installed = {
dist.metadata['Name'].lower().replace('_', '-')
for dist in distributions()
if dist.metadata['Name']
}
if expected == 'cu126':
assert actual is not None and actual.startswith('12.6'), actual
nccl_package = 'nvidia-nccl-cu12'
nccl_minimum = Version('2.29')
unexpected = sorted(name for name in installed if name.endswith('-cu13'))
elif expected == 'cu130':
assert actual is not None and actual.startswith('13.0'), actual
nccl_package = 'nvidia-nccl-cu13'
nccl_minimum = Version('2.30.4')
unexpected = sorted(name for name in installed if name.endswith('-cu12'))
else:
raise AssertionError(f'unsupported PYTORCH_CUDA_VERSION: {expected}')
assert not unexpected, f'unexpected CUDA packages: {unexpected}'
nccl_version = version(nccl_package)
if expected == 'cu126':
assert Version(nccl_version) > nccl_minimum, nccl_version
else:
assert Version(nccl_version) >= nccl_minimum, nccl_version
print(f'Validated torch {torch.__version__} with CUDA {actual} and '
f'{nccl_package} {nccl_version}')
PY
14 changes: 10 additions & 4 deletions docker/prepare_3rdparty_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,16 @@ pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/d
# sm100 compilation for Flash MLA requires NVCC 12.9 or higher
FLASH_MLA_DISABLE_SM100=1 pip wheel -v --no-build-isolation --no-deps -w /wheels "git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_VERSION}"

# flash_attn_3 (prebuilt wheels; CUDA + torch must match this image)
TORCH_VER=$(python3 -c "import torch; print(''.join(torch.__version__.split('+')[0].split('.')))")
FA3_WHEELS_URL="https://windreamer.github.io/flash-attention3-wheels/${CUDA_VERSION_SHORT}_torch${TORCH_VER}"
pip download --no-deps -d /wheels "flash_attn_3" --find-links "${FA3_WHEELS_URL}"
# flash_attn_3 (official, LibTorch ABI-stable wheels for torch >= 2.9).
# Select the FA3 channel matching the PyTorch binary. CUDA 13 FA3 wheels are
# published in the cu130 channel.
FA3_CUDA_VERSION="${PYTORCH_CUDA_VERSION:-${CUDA_VERSION_SHORT}}"
if [[ "${FA3_CUDA_VERSION}" == cu13* ]]; then
FA3_CUDA_VERSION=cu130
fi
pip download --no-deps --only-binary=:all: -d /wheels \
"flash-attn-3==3.0.0" \
--index-url "https://download.pytorch.org/whl/${FA3_CUDA_VERSION}"

# GDRCopy debs
apt-get update -y \
Expand Down
2 changes: 1 addition & 1 deletion lmdeploy/pytorch/check_env/triton.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

from .base import BaseChecker

MAX_TRITON_VERSION = '3.6.0'
MAX_TRITON_VERSION = '3.7.1'
MIN_TRITON_VERSION = '3.0.0'


Expand Down
6 changes: 3 additions & 3 deletions requirements/runtime_cuda.txt
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,6 @@ opencv-python-headless
peft<=0.14.0
prometheus_client
tilelang==0.1.11; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
torch<=2.10.0,>=2.0.0
torchvision<=0.25.0,>=0.15.0
triton<=3.6.0,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
torch<=2.12.1,>=2.0.0
torchvision<=0.27.1,>=0.15.0
triton<=3.7.1,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
Loading