GithubStar/ktransformers
1
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-10-16 21:01:14 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
9b0da97c917cc33345096777ad25e1fd5d8dddc7
ktransformers/docker/Dockerfile
T

767 lines
28 KiB
Docker
Raw Blame History

ARG CUDA_VERSION=12.8.1
ARG DSV4_BASE_IMAGE=docker.m.daocloud.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04
FROM docker.1ms.run/nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 AS base
ARG TARGETARCH
ARG GRACE_BLACKWELL=0
ARG HOPPER_SBO=0
ARG CPU_VARIANT=x86-intel-multi
ARG BUILD_ALL_CPU_VARIANTS=1
# Proxy settings for build-time network access
ARG HTTP_PROXY
ARG HTTPS_PROXY
ARG http_proxy
ARG https_proxy
ENV HTTP_PROXY=${HTTP_PROXY} \
HTTPS_PROXY=${HTTPS_PROXY} \
http_proxy=${http_proxy} \
https_proxy=${https_proxy}
ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2
ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6
ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
ARG BUILD_AND_DOWNLOAD_PARALLEL=8
ARG SGL_KERNEL_VERSION=0.3.19
ARG SGL_VERSION=0.5.6.post1
ARG USE_LATEST_SGLANG=0
ARG GDRCOPY_VERSION=2.5.1
ARG UBUNTU_MIRROR
ARG GITHUB_ARTIFACTORY=github.com
ARG FLASHINFER_VERSION=0.5.3
# ktransformers wheel version (cu128torch28 for CUDA 12.8 + PyTorch 2.8)
ARG KTRANSFORMERS_VERSION=0.5.3
ARG KTRANSFORMERS_WHEEL=ktransformers-0.5.3+cu128torch28fancy-cp312-cp312-linux_x86_64.whl
# flash_attn wheel for fine-tune env
ARG FLASH_ATTN_WHEEL=flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl
ENV DEBIAN_FRONTEND=noninteractive \
CUDA_HOME=/usr/local/cuda \
GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \
FLASHINFER_VERSION=${FLASHINFER_VERSION}
# Add GKE default lib and bin locations
ENV PATH="${PATH}:/usr/local/nvidia/bin" \
LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/nvidia/lib:/usr/local/nvidia/lib64"
# Replace Ubuntu sources with Tsinghua mirror for Ubuntu 24.04 (noble)
RUN if [ -n "$UBUNTU_MIRROR" ]; then \
echo "deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble main restricted universe multiverse" > /etc/apt/sources.list && \
echo "deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble-updates main restricted universe multiverse" >> /etc/apt/sources.list && \
echo "deb https://mirrors.tuna.tsinghua.edu.cn/ubuntu/ noble-backports main restricted universe multiverse" >> /etc/apt/sources.list && \
echo "deb http://security.ubuntu.com/ubuntu/ noble-security main restricted universe multiverse" >> /etc/apt/sources.list && \
rm -f /etc/apt/sources.list.d/ubuntu.sources; \
fi
# Install system dependencies (organized by category for better caching)
RUN --mount=type=cache,target=/var/cache/apt,id=base-apt \
echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
&& apt-get update && apt-get install -y --no-install-recommends --allow-change-held-packages \
# Core system utilities
tzdata \
ca-certificates \
software-properties-common \
netcat-openbsd \
kmod \
unzip \
openssh-server \
curl \
wget \
lsof \
locales \
# Build essentials
build-essential \
cmake \
perl \
patchelf \
ccache \
git \
git-lfs \
# MPI and NUMA
libopenmpi-dev \
libnuma1 \
libnuma-dev \
numactl \
# transformers multimodal VLM
ffmpeg \
# InfiniBand/RDMA
libibverbs-dev \
libibverbs1 \
libibumad3 \
librdmacm1 \
libnl-3-200 \
libnl-route-3-200 \
libnl-route-3-dev \
libnl-3-dev \
ibverbs-providers \
infiniband-diags \
perftest \
# Development libraries
libgoogle-glog-dev \
libgtest-dev \
libjsoncpp-dev \
libunwind-dev \
libboost-all-dev \
libssl-dev \
libgrpc-dev \
libgrpc++-dev \
libprotobuf-dev \
protobuf-compiler \
protobuf-compiler-grpc \
pybind11-dev \
libhiredis-dev \
libcurl4-openssl-dev \
libczmq4 \
libczmq-dev \
libfabric-dev \
# Package building tools
devscripts \
debhelper \
fakeroot \
dkms \
check \
libsubunit0 \
libsubunit-dev \
# Development tools
gdb \
ninja-build \
vim \
tmux \
htop \
zsh \
tree \
less \
rdma-core \
# NCCL
libnccl2 \
libnccl-dev \
&& rm -rf /var/lib/apt/lists/* \
&& apt-get clean
# GDRCopy installation
RUN mkdir -p /tmp/gdrcopy && cd /tmp \
&& curl --retry 3 --retry-delay 2 -fsSL -o v${GDRCOPY_VERSION}.tar.gz \
https://${GITHUB_ARTIFACTORY}/NVIDIA/gdrcopy/archive/refs/tags/v${GDRCOPY_VERSION}.tar.gz \
&& tar -xzf v${GDRCOPY_VERSION}.tar.gz && rm v${GDRCOPY_VERSION}.tar.gz \
&& cd gdrcopy-${GDRCOPY_VERSION}/packages \
&& CUDA=/usr/local/cuda ./build-deb-packages.sh \
&& dpkg -i gdrdrv-dkms_*.deb libgdrapi_*.deb gdrcopy-tests_*.deb gdrcopy_*.deb \
&& cd / && rm -rf /tmp/gdrcopy
# Fix DeepEP IBGDA symlink
RUN ln -sf /usr/lib/$(uname -m)-linux-gnu/libmlx5.so.1 /usr/lib/$(uname -m)-linux-gnu/libmlx5.so
# Set up locale
RUN locale-gen en_US.UTF-8
ENV LANG=en_US.UTF-8 \
LANGUAGE=en_US:en \
LC_ALL=en_US.UTF-8
########################################################
########## Install Miniconda ###########################
########################################################
RUN mkdir -p /opt/miniconda3 \
&& wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /opt/miniconda3/miniconda.sh \
&& bash /opt/miniconda3/miniconda.sh -b -u -p /opt/miniconda3 \
&& rm /opt/miniconda3/miniconda.sh
# Add conda to PATH
ENV PATH="/opt/miniconda3/bin:${PATH}"
# Accept conda TOS
RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main \
&& conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
# Configure conda to use Tsinghua mirror
RUN conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main \
&& conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free \
&& conda config --set show_channel_urls yes
########################################################
########## DeepSeek-V4-Flash one-click image ############
########################################################
# This target is self-contained so the DSV4 image and its merged SGLang API
# support can be reviewed as a Dockerfile-only change.
FROM ${DSV4_BASE_IMAGE} AS dsv4
SHELL ["/bin/bash", "-o", "pipefail", "-c"]
ARG IMAGE_VERSION=dev
ARG BUILD_DATE=unknown
ARG KTRANSFORMERS_REPOSITORY=https://github.com/kvcache-ai/ktransformers.git
ARG KTRANSFORMERS_REF=6d460cc10780f2e0ad541b5d58ad28086dd32ef0
ARG SGLANG_REPOSITORY=https://github.com/kvcache-ai/sglang.git
ARG SGLANG_REF=0980b6da258192896a9d884fbe564fd1f9ec07e3
ARG PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
ARG TORCH_VERSION=2.9.1
ARG SGL_KERNEL_VERSION=0.3.21
ARG FLASHINFER_VERSION=0.6.15.post1
ARG TRANSFORMERS_VERSION=4.57.1
ARG TILELANG_VERSION=0.1.10
ARG TVM_FFI_VERSION=0.1.11
ARG CPUINFER_CUDA_ARCHS=80;86;89;90;100;120
ARG BUILD_JOBS=16
ENV CUDA_HOME=/usr/local/cuda \
DEBIAN_FRONTEND=noninteractive \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
PIP_NO_CACHE_DIR=1 \
PYTHONUNBUFFERED=1 \
SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=false \
MODEL_PATH=/model \
PORT=30000 \
DSV4_REASONING_EFFORT_PROFILE=official \
REASONING_PARSER=deepseek-v4 \
TOOL_CALL_PARSER=deepseekv4 \
PATH=/opt/venv/bin:/usr/local/bin:${PATH}
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
build-essential \
ca-certificates \
curl \
git \
libhwloc-dev \
libnuma-dev \
libopenmpi-dev \
numactl \
pkg-config \
software-properties-common \
&& add-apt-repository -y ppa:deadsnakes/ppa \
&& apt-get update \
&& apt-get install -y --no-install-recommends \
python3.11 \
python3.11-dev \
python3.11-venv \
&& rm -rf /var/lib/apt/lists/* \
&& python3.11 -m venv /opt/venv
RUN python -m pip config set global.index-url "${PIP_INDEX_URL}" \
&& python -m pip install --upgrade setuptools wheel \
&& python -m pip install "cmake>=3.31,<4" ninja
RUN python -m pip install \
"torch==${TORCH_VERSION}" \
"torchvision==0.24.1" \
"torchaudio==${TORCH_VERSION}" \
--index-url https://download.pytorch.org/whl/cu128 \
--extra-index-url "${PIP_INDEX_URL}" \
--retries 10 \
--timeout 120
ARG FLASHINFER_CUBIN_INDEX_URL=https://flashinfer.ai/whl
ARG FLASHINFER_CUBIN_TRUSTED_HOST=flashinfer.ai
RUN python -m pip install \
"sgl-kernel==${SGL_KERNEL_VERSION}" \
"flashinfer-python==${FLASHINFER_VERSION}" \
"transformers==${TRANSFORMERS_VERSION}" \
"tilelang==${TILELANG_VERSION}" \
"apache-tvm-ffi==${TVM_FFI_VERSION}" \
"cuda-python==13.2.0" \
"nvidia-cutlass-dsl==4.6.1" \
IPython \
aiohttp \
'anthropic>=0.20.0' \
blobfile==3.0.0 \
build \
compressed-tensors \
datasets \
decord2 \
einops \
fastapi \
gguf \
'grpcio>=1.78.0' \
'grpcio-health-checking>=1.78.0' \
'grpcio-reflection>=1.78.0' \
hf_transfer \
'huggingface_hub<1.0' \
interegular \
'llguidance>=0.7.11,<0.8.0' \
modelscope \
msgspec \
numpy \
nvidia-ml-py \
openai==2.6.1 \
openai-harmony==0.0.4 \
orjson \
outlines==0.1.11 \
packaging \
partial_json_parser \
pillow \
'prometheus-client>=0.20.0' \
psutil \
py-spy \
pybase64 \
pydantic \
python-multipart \
'pyzmq>=25.1.2' \
requests \
scipy \
sentencepiece \
setproctitle \
'smg-grpc-proto>=0.3.3' \
soundfile==0.13.1 \
tiktoken \
timm==1.0.16 \
torch_memory_saver==0.0.9 \
torchao==0.9.0 \
torchcodec==0.8.0 \
tqdm \
uvicorn \
uvloop \
xgrammar==0.1.27 \
&& python -m pip install \
--index-url "${FLASHINFER_CUBIN_INDEX_URL}" \
--trusted-host "${FLASHINFER_CUBIN_TRUSTED_HOST}" \
--no-deps \
"flashinfer-cubin==${FLASHINFER_VERSION}" \
&& python -m pip install --no-deps quack-kernels==0.2.4
RUN git clone --filter=blob:none "${KTRANSFORMERS_REPOSITORY}" /workspace/ktransformers \
&& git -C /workspace/ktransformers checkout --detach "${KTRANSFORMERS_REF}" \
&& git -C /workspace/ktransformers submodule update --init --recursive \
&& git -C /workspace/ktransformers/third_party/sglang fetch \
--depth 1 "${SGLANG_REPOSITORY}" "${SGLANG_REF}" \
&& git -C /workspace/ktransformers/third_party/sglang checkout --detach FETCH_HEAD
RUN export SGLANG_KT_VERSION="$(python -c "exec(open('/workspace/ktransformers/version.py').read()); print(__version__)")" \
&& python -m pip install --no-deps \
/workspace/ktransformers/third_party/sglang/python \
&& cd /workspace/ktransformers/kt-kernel \
&& CPUINFER_BUILD_ALL_VARIANTS=1 \
CPUINFER_USE_CUDA=1 \
CPUINFER_CUDA_ARCHS="${CPUINFER_CUDA_ARCHS}" \
CPUINFER_PARALLEL="${BUILD_JOBS}" \
./install.sh build \
&& python -m pip install --no-deps /workspace/ktransformers \
&& rm -rf /workspace/ktransformers/kt-kernel/build
RUN python - <<'PY'
import importlib.metadata
from sglang.srt.entrypoints.openai import encoding_dsv4
from sglang.srt.function_call.function_call_parser import FunctionCallParser
expected = {
"ktransformers": "0.7.0.post1",
"kt-kernel": "0.7.0.post1",
"sglang-kt": "0.7.0.post1",
"torch": "2.9.1+cu128",
"flashinfer-python": "0.6.15.post1",
"transformers": "4.57.1",
}
for distribution, version in expected.items():
actual = importlib.metadata.version(distribution)
if actual != version:
raise RuntimeError(f"{distribution}: expected {version}, got {actual}")
if set(encoding_dsv4.REASONING_EFFORT_PROFILES["official"]) != {
"low",
"high",
"max",
}:
raise RuntimeError("missing official DeepSeek-V4 reasoning effort profile")
if "deepseekv4" not in FunctionCallParser.ToolCallParserEnum:
raise RuntimeError("missing DeepSeek-V4 tool-call parser")
PY
COPY --chmod=0755 <<'ENTRYPOINT' /usr/local/bin/entrypoint-dsv4
#!/usr/bin/env bash
set -euo pipefail
log() {
printf '[dsv4-entrypoint] %s\n' "$*" >&2
}
die() {
log "ERROR: $*"
exit 1
}
if [[ $# -gt 0 && "$1" != --* ]]; then
exec "$@"
fi
MODEL_PATH="${MODEL_PATH:-/model}"
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-30000}"
TP="${TP:-1}"
CONTEXT_LENGTH="${CONTEXT_LENGTH:-16384}"
MEM_FRACTION="${MEM_FRACTION:-0.90}"
KT_GPU_EXPERTS="${KT_GPU_EXPERTS:-0}"
KT_GPU_PREFILL_TOKEN_THRESHOLD="${KT_GPU_PREFILL_TOKEN_THRESHOLD:-2048}"
CHUNKED_PREFILL_SIZE="${CHUNKED_PREFILL_SIZE:-4096}"
MAX_PREFILL_TOKENS="${MAX_PREFILL_TOKENS:-$CHUNKED_PREFILL_SIZE}"
MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-2}"
SWA_FULL_TOKENS_RATIO="${SWA_FULL_TOKENS_RATIO:-0.4}"
WATCHDOG_TIMEOUT="${WATCHDOG_TIMEOUT:-1200}"
DSV4_REASONING_EFFORT_PROFILE="${DSV4_REASONING_EFFORT_PROFILE:-official}"
REASONING_PARSER="${REASONING_PARSER:-deepseek-v4}"
TOOL_CALL_PARSER="${TOOL_CALL_PARSER:-deepseekv4}"
[[ "$TP" =~ ^[1-9][0-9]*$ ]] || die "TP must be a positive integer"
[[ -r "$MODEL_PATH/config.json" ]] || die "missing $MODEL_PATH/config.json"
find "$MODEL_PATH" -maxdepth 1 -type f \
\( -name '*.safetensors' -o -name '*.gguf' \) -print -quit | grep -q . ||
die "no model weight files found in $MODEL_PATH"
GPU_INFO="$(TP="$TP" python - <<'PY'
import os
import torch
if not torch.cuda.is_available():
raise SystemExit("PyTorch cannot access CUDA")
tp = int(os.environ["TP"])
if torch.cuda.device_count() < tp:
raise SystemExit(f"TP={tp} exceeds the visible GPU count")
supported = {
(8, 0): ("8.0", "8.0"),
(8, 6): ("8.6", "8.6"),
(8, 9): ("8.9", "8.9"),
(9, 0): ("9.0+PTX", "9.0a"),
(10, 0): ("10.0+PTX", "10.0a"),
(12, 0): ("12.0+PTX", "12.0a"),
}
torch_arches = []
flashinfer_arches = []
for index in range(tp):
capability = torch.cuda.get_device_capability(index)
if capability not in supported:
raise SystemExit(f"unsupported GPU capability: SM{capability[0]}{capability[1]}")
torch_arch, flashinfer_arch = supported[capability]
if torch_arch not in torch_arches:
torch_arches.append(torch_arch)
if flashinfer_arch not in flashinfer_arches:
flashinfer_arches.append(flashinfer_arch)
print("|".join((";".join(torch_arches), ";".join(flashinfer_arches))))
PY
)" || die "GPU validation failed"
IFS='|' read -r DEFAULT_TORCH_ARCH DEFAULT_FLASHINFER_ARCH <<< "$GPU_INFO"
export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-$DEFAULT_TORCH_ARCH}"
export FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST:-$DEFAULT_FLASHINFER_ARCH}"
export SGLANG_DSV4_MODE="${SGLANG_DSV4_MODE:-2604}"
export SGLANG_DSV4_2604_SUBMODE="${SGLANG_DSV4_2604_SUBMODE:-2604B}"
if [[ -z "${KT_CPUINFER_THREADS:-}" ]]; then
PHYSICAL_CORES="$(lscpu -p=CORE,SOCKET | awk -F, '!/^#/ {print $1 "," $2}' | sort -u | wc -l)"
(( PHYSICAL_CORES > 4 )) && KT_CPUINFER_THREADS=$((PHYSICAL_CORES - 4)) || KT_CPUINFER_THREADS="$PHYSICAL_CORES"
fi
if [[ -z "${KT_THREADPOOL_COUNT:-}" ]]; then
KT_THREADPOOL_COUNT="$(lscpu -p=NODE | awk -F, '!/^#/ && $1 >= 0 {print $1}' | sort -u | wc -l)"
(( KT_THREADPOOL_COUNT > 0 )) || KT_THREADPOOL_COUNT=1
fi
cmd=(
python -m sglang.launch_server
--host "$HOST"
--port "$PORT"
--model "$MODEL_PATH"
--kt-weight-path "$MODEL_PATH"
--kt-method MXFP4
--kt-num-gpu-experts "$KT_GPU_EXPERTS"
--kt-cpuinfer "$KT_CPUINFER_THREADS"
--kt-threadpool-count "$KT_THREADPOOL_COUNT"
--kt-gpu-prefill-token-threshold "$KT_GPU_PREFILL_TOKEN_THRESHOLD"
--kt-enable-dynamic-expert-update
--tensor-parallel-size "$TP"
--context-length "$CONTEXT_LENGTH"
--attention-backend flashinfer
--mem-fraction-static "$MEM_FRACTION"
--chunked-prefill-size "$CHUNKED_PREFILL_SIZE"
--max-prefill-tokens "$MAX_PREFILL_TOKENS"
--max-running-requests "$MAX_RUNNING_REQUESTS"
--swa-full-tokens-ratio "$SWA_FULL_TOKENS_RATIO"
--watchdog-timeout "$WATCHDOG_TIMEOUT"
--disable-shared-experts-fusion
--trust-remote-code
--cuda-graph-bs 1
--cuda-graph-max-bs 1
--disable-radix-cache
--skip-server-warmup
--reasoning-parser "$REASONING_PARSER"
--tool-call-parser "$TOOL_CALL_PARSER"
--json-model-override-args "{\"dsv4_reasoning_effort_profile\":\"$DSV4_REASONING_EFFORT_PROFILE\"}"
)
cmd+=("$@")
printf '[dsv4-entrypoint] launching:' >&2
printf ' %q' "${cmd[@]}" >&2
printf '\n' >&2
exec numactl --interleave=all "${cmd[@]}"
ENTRYPOINT
LABEL org.opencontainers.image.title="KTransformers DeepSeek-V4-Flash" \
org.opencontainers.image.version="${IMAGE_VERSION}" \
org.opencontainers.image.created="${BUILD_DATE}" \
io.ktransformers.source-revision="${KTRANSFORMERS_REF}" \
io.ktransformers.sglang-revision="${SGLANG_REF}" \
io.ktransformers.reasoning-effort-profile="official"
WORKDIR /workspace
EXPOSE 30000
HEALTHCHECK --interval=30s --timeout=10s --start-period=15m --retries=20 \
CMD curl --fail --silent --show-error "http://127.0.0.1:$PORT/health" >/dev/null || exit 1
STOPSIGNAL SIGTERM
ENTRYPOINT ["/usr/local/bin/entrypoint-dsv4"]
CMD []
########################################################
########## Dual Conda Environment Setup ################
########################################################
FROM base AS framework
ARG CUDA_VERSION
ARG BUILD_AND_DOWNLOAD_PARALLEL
ARG SGL_KERNEL_VERSION
ARG SGL_VERSION
ARG USE_LATEST_SGLANG
ARG FLASHINFER_VERSION
ARG GRACE_BLACKWELL
ARG GRACE_BLACKWELL_DEEPEP_BRANCH
ARG HOPPER_SBO
ARG HOPPER_SBO_DEEPEP_COMMIT
ARG DEEPEP_COMMIT
ARG GITHUB_ARTIFACTORY
ARG KTRANSFORMERS_VERSION
ARG KTRANSFORMERS_WHEEL
ARG FLASH_ATTN_WHEEL
ARG FUNCTIONALITY=sft
WORKDIR /workspace
# Create conda environments (fine-tune only needed for sft mode)
RUN conda create -n serve python=3.12 -y \
&& if [ "$FUNCTIONALITY" = "sft" ]; then conda create -n fine-tune python=3.12 -y; fi
# Set pip mirror for conda envs
RUN /opt/miniconda3/envs/serve/bin/pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple \
&& if [ "$FUNCTIONALITY" = "sft" ]; then \
/opt/miniconda3/envs/fine-tune/bin/pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple; \
fi
# Clone repositories (sglang is included as a submodule in ktransformers)
RUN git clone --depth 1 https://${GITHUB_ARTIFACTORY}/kvcache-ai/ktransformers.git /workspace/ktransformers \
&& cd /workspace/ktransformers && git submodule update --init --recursive \
&& ln -s /workspace/ktransformers/third_party/sglang /workspace/sglang \
&& if [ "$FUNCTIONALITY" = "sft" ]; then \
git clone --depth 1 https://${GITHUB_ARTIFACTORY}/hiyouga/LLaMA-Factory.git /workspace/LLaMA-Factory; \
fi
# Download ktransformers wheel and flash_attn wheel for fine-tune env (sft mode only)
RUN if [ "$FUNCTIONALITY" = "sft" ]; then \
curl --retry 3 --retry-delay 2 -fsSL -o /workspace/${KTRANSFORMERS_WHEEL} \
https://${GITHUB_ARTIFACTORY}/kvcache-ai/ktransformers/releases/download/v${KTRANSFORMERS_VERSION}/${KTRANSFORMERS_WHEEL} \
&& curl --retry 3 --retry-delay 2 -fsSL -o /workspace/${FLASH_ATTN_WHEEL} \
https://${GITHUB_ARTIFACTORY}/Dao-AILab/flash-attention/releases/download/v2.8.3/${FLASH_ATTN_WHEEL}; \
fi
########################################################
# Environment 1: serve (sglang + kt-kernel)
########################################################
# Upgrade pip and install basic tools in serve env
RUN --mount=type=cache,target=/root/.cache/pip \
/opt/miniconda3/envs/serve/bin/pip install --upgrade pip setuptools wheel html5lib six
# Install sgl-kernel
RUN --mount=type=cache,target=/root/.cache/pip \
case "$CUDA_VERSION" in \
12.6.1) CUINDEX=126 ;; \
12.8.1) CUINDEX=128 ;; \
12.9.1) CUINDEX=129 ;; \
13.0.1) CUINDEX=130 ;; \
*) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
esac \
&& if [ "$CUDA_VERSION" = "12.6.1" ]; then \
/opt/miniconda3/envs/serve/bin/pip install https://${GITHUB_ARTIFACTORY}/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu124-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps \
; \
elif [ "$CUDA_VERSION" = "12.8.1" ] || [ "$CUDA_VERSION" = "12.9.1" ]; then \
/opt/miniconda3/envs/serve/bin/pip install sgl-kernel==${SGL_KERNEL_VERSION} \
; \
elif [ "$CUDA_VERSION" = "13.0.1" ]; then \
/opt/miniconda3/envs/serve/bin/pip install https://github.com/sgl-project/whl/releases/download/v${SGL_KERNEL_VERSION}/sgl_kernel-${SGL_KERNEL_VERSION}+cu130-cp310-abi3-manylinux2014_$(uname -m).whl --force-reinstall --no-deps \
; \
fi
# Install SGLang in serve env (version aligned with ktransformers)
RUN --mount=type=cache,target=/root/.cache/pip \
case "$CUDA_VERSION" in \
12.6.1) CUINDEX=126 ;; \
12.8.1) CUINDEX=128 ;; \
12.9.1) CUINDEX=129 ;; \
13.0.1) CUINDEX=130 ;; \
esac \
&& export SGLANG_KT_VERSION=$(python3 -c "exec(open('/workspace/ktransformers/version.py').read()); print(__version__)") \
&& echo "Installing sglang-kt v${SGLANG_KT_VERSION}" \
&& cd /workspace/sglang \
&& /opt/miniconda3/envs/serve/bin/pip install -e "python[all]" --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX}
# Download FlashInfer cubin for serve env
RUN --mount=type=cache,target=/root/.cache/pip \
FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning \
/opt/miniconda3/envs/serve/bin/python -m flashinfer --download-cubin
# Install DeepEP in serve env
RUN set -eux; \
if [ "$GRACE_BLACKWELL" = "1" ]; then \
git clone https://github.com/fzyzcjy/DeepEP.git /workspace/DeepEP && \
cd /workspace/DeepEP && \
git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh; \
elif [ "$HOPPER_SBO" = "1" ]; then \
git clone https://github.com/deepseek-ai/DeepEP.git -b antgroup-opt /workspace/DeepEP && \
cd /workspace/DeepEP && \
git checkout ${HOPPER_SBO_DEEPEP_COMMIT} && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh; \
else \
curl --retry 3 --retry-delay 2 -fsSL -o /tmp/${DEEPEP_COMMIT}.zip \
https://${GITHUB_ARTIFACTORY}/deepseek-ai/DeepEP/archive/${DEEPEP_COMMIT}.zip && \
unzip -q /tmp/${DEEPEP_COMMIT}.zip -d /tmp && rm /tmp/${DEEPEP_COMMIT}.zip && \
mv /tmp/DeepEP-${DEEPEP_COMMIT} /workspace/DeepEP && \
cd /workspace/DeepEP && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh; \
fi
RUN --mount=type=cache,target=/root/.cache/pip \
cd /workspace/DeepEP && \
case "$CUDA_VERSION" in \
12.6.1) CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' ;; \
12.8.1) CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0' ;; \
12.9.1|13.0.1) CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0;10.3' ;; \
*) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
esac && \
. /opt/miniconda3/etc/profile.d/conda.sh && conda activate serve && \
TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} \
pip install --no-build-isolation .
# Install NCCL for serve env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
/opt/miniconda3/envs/serve/bin/pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
/opt/miniconda3/envs/serve/bin/pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
fi
# Install kt-kernel in serve env with all CPU variants
RUN . /opt/miniconda3/etc/profile.d/conda.sh && conda activate serve \
&& cd /workspace/ktransformers/kt-kernel \
&& CPUINFER_BUILD_ALL_VARIANTS=1 ./install.sh build
########################################################
# Environment 2: fine-tune (LLaMA-Factory + ktransformers) - sft mode only
########################################################
# Install dependency libraries for ktransformers (CUDA 11.8 runtime required)
RUN if [ "$FUNCTIONALITY" = "sft" ]; then \
conda install -n fine-tune -y -c conda-forge libstdcxx-ng gcc_impl_linux-64 \
&& conda install -n fine-tune -y -c nvidia/label/cuda-11.8.0 cuda-runtime; \
fi
# Install PyTorch 2.8 in fine-tune env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$FUNCTIONALITY" = "sft" ]; then \
case "$CUDA_VERSION" in \
12.6.1) CUINDEX=126 ;; \
12.8.1) CUINDEX=128 ;; \
12.9.1) CUINDEX=129 ;; \
13.0.1) CUINDEX=130 ;; \
esac \
&& /opt/miniconda3/envs/fine-tune/bin/pip install --upgrade pip setuptools wheel hatchling \
&& /opt/miniconda3/envs/fine-tune/bin/pip install \
torch==2.8.0 \
torchvision \
torchaudio \
--extra-index-url https://download.pytorch.org/whl/cu${CUINDEX}; \
fi
# Install LLaMA-Factory in fine-tune env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$FUNCTIONALITY" = "sft" ]; then \
cd /workspace/LLaMA-Factory \
&& /opt/miniconda3/envs/fine-tune/bin/pip install -e ".[torch,metrics]" --no-build-isolation; \
fi
# Install ktransformers wheel in fine-tune env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$FUNCTIONALITY" = "sft" ]; then \
/opt/miniconda3/envs/fine-tune/bin/pip install /workspace/${KTRANSFORMERS_WHEEL}; \
fi
# Install flash_attn wheel in fine-tune env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$FUNCTIONALITY" = "sft" ]; then \
/opt/miniconda3/envs/fine-tune/bin/pip install /workspace/${FLASH_ATTN_WHEEL}; \
fi
# Install NCCL for fine-tune env
RUN --mount=type=cache,target=/root/.cache/pip \
if [ "$FUNCTIONALITY" = "sft" ]; then \
if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
/opt/miniconda3/envs/fine-tune/bin/pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
/opt/miniconda3/envs/fine-tune/bin/pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
fi; \
fi
########################################################
# Cleanup and final setup
########################################################
# Clean up downloaded wheels
RUN if [ "$FUNCTIONALITY" = "sft" ]; then \
rm -f /workspace/${KTRANSFORMERS_WHEEL} /workspace/${FLASH_ATTN_WHEEL}; \
fi
# Initialize conda for bash
RUN /opt/miniconda3/bin/conda init bash
# Create shell aliases for convenience
RUN echo '\n# Conda environment aliases\nalias serve="conda activate serve"' >> /root/.bashrc \
&& if [ "$FUNCTIONALITY" = "sft" ]; then \
echo 'alias finetune="conda activate fine-tune"' >> /root/.bashrc; \
fi
########################################################
# Extract version information for image naming
########################################################
# Extract versions from each component and save to versions.env
RUN set -x && \
# KTransformers version (single source of truth for both kt-kernel and sglang-kt)
cd /workspace/ktransformers && \
KTRANSFORMERS_VERSION=$(python3 -c "exec(open('version.py').read()); print(__version__)" 2>/dev/null || echo "unknown") && \
echo "KTRANSFORMERS_VERSION=$KTRANSFORMERS_VERSION" > /workspace/versions.env && \
echo "Extracted KTransformers version: $KTRANSFORMERS_VERSION" && \
\
# sglang-kt version = ktransformers version (aligned)
echo "SGLANG_KT_VERSION=$KTRANSFORMERS_VERSION" >> /workspace/versions.env && \
echo "sglang-kt version (aligned): $KTRANSFORMERS_VERSION" && \
\
# LLaMA-Factory version (from fine-tune environment, sft mode only)
if [ "$FUNCTIONALITY" = "sft" ]; then \
. /opt/miniconda3/etc/profile.d/conda.sh && conda activate fine-tune && \
cd /workspace/LLaMA-Factory && \
LLAMAFACTORY_VERSION=$(python -c "import sys; sys.path.insert(0, 'src'); from llamafactory import __version__; print(__version__)" 2>/dev/null || echo "unknown") && \
echo "LLAMAFACTORY_VERSION=$LLAMAFACTORY_VERSION" >> /workspace/versions.env && \
echo "Extracted LLaMA-Factory version: $LLAMAFACTORY_VERSION"; \
else \
echo "LLAMAFACTORY_VERSION=none" >> /workspace/versions.env && \
echo "LLaMA-Factory not installed (infer mode)"; \
fi && \
\
# Display all versions
echo "=== Version Summary ===" && \
cat /workspace/versions.env
WORKDIR /workspace
CMD ["/bin/bash"]
Reference in New Issue View Git Blame Copy Permalink