Views
No views yet
1# vLLM + LMCache Multi-Stage Dockerfile
2# + Docker cache
3# + Ccache
4# + uv cache
5# Version: 2026-01-26
6# Build: TMPDIR=vllm-dockercache podman build -v ./vllm-ccache:/root/.ccache -t vllm-202601-cu129 -f vllm-lmcache-Dockerfile
7
8#################### ARGUMENTS ####################
9ARG CUDA_VERSION=12.9.1
10ARG LMCACHE_GIT_REF=dev
11ARG VLLM_GIT_REF=main
12ARG FLASHINFER_VERSION=0.6.1
13ARG WHEELS_DIR=/tmp/wheels
14
15#################### BUILD STAGE ####################
16# Full build environment with all development tools
17FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04 AS build
18
19ARG CUDA_VERSION
20ARG LMCACHE_GIT_REF
21ARG VLLM_GIT_REF
22ARG WHEELS_DIR
23
24# Build environment
25ENV CUDA_VERSION=${CUDA_VERSION}
26ENV WHEELS_DIR=${WHEELS_DIR}
27
28# Build config
29ENV UV_LINK_MODE=copy
30ENV UV_HTTP_TIMEOUT=500
31ENV UV_INDEX_STRATEGY="unsafe-best-match"
32ENV MAX_JOBS=128
33ENV NVCC_THREADS=8
34ENV CMAKE_BUILD_TYPE=Release
35ENV USE_CUDA=1
36ENV CCACHE_DIR=/root/.ccache
37ENV CUDA_HOME=/usr/local/cuda
38ENV NVCC_GENCODE="-gencode=arch=compute_120,code=sm_120"
39ENV TORCH_CUDA_ARCH_LIST='12.0'
40ENV FLASH_ATTN_CUDA_ARCHS=120
41ENV VLLM_FLASH_ATTN_VERSION=2
42# Note: flashinfer is now installed as pre-compiled wheel in runtime
43# ENV FLASHINFER_ENABLE_AOT=1
44ENV VLLM_TARGET_DEVICE=cuda
45ENV LMCACHE_NVCC_THREADS=8
46ENV LMCACHE_MAX_JOBS=32
47ENV LMCACHE_CUDA_VERSION=${CUDA_VERSION}
48ENV LMCACHE_CUDA_ARCHS=12.0
49ENV LMCACHE_TORCH_CUDA_ARCH_LIST=12.0
50ENV LMCACHE_VLLM_FA_CMAKE_GPU_ARCHES=120
51ENV VLLM_DOCKER_BUILD_CONTEXT=1
52ENV PATH="/opt/venv/bin:$PATH"
53
54# System packages
55RUN apt-get update && apt-get install -y --no-install-recommends \
56 build-essential \
57 curl \
58 ca-certificates \
59 python3.12 \
60 python3.12-venv \
61 python3.12-dev \
62 python3-pip \
63 git \
64 ccache \
65 && rm -rf /var/lib/apt/lists/*
66
67# Create venv
68RUN python3 -m venv /opt/venv
69RUN /opt/venv/bin/pip install --no-cache-dir --upgrade pip
70RUN /opt/venv/bin/pip install --no-cache-dir uv
71
72# Create wheel output directory
73RUN mkdir -p ${WHEELS_DIR}
74
75# Build tools
76RUN --mount=type=cache,target=/root/.cache/uv \
77 /opt/venv/bin/uv pip install ninja setuptools setuptools_scm
78
79# App
80# ---------------------------------------------------------------
81# PyTorch
82RUN --mount=type=cache,target=/root/.cache/uv \
83 /opt/venv/bin/uv pip install --pre torch>=2.9.0 torchvision torchaudio \
84 --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION%.*}
85
86# Clone vLLM
87WORKDIR /workspace
88RUN git clone --branch ${VLLM_GIT_REF} https://github.com/vllm-project/vllm
89
90# vLLM: Specialize for SM120 (RTX 5090, RTX Pro 6000) to save hours of compilation time
91WORKDIR /workspace/vllm
92RUN sed -i \
93 -e 's/ALLSPARK_ARCHS "8.0;8.6;8.7;8.9"/ALLSPARK_ARCHS "12.0"/g' \
94 -e 's/MARLIN_ARCHS "8.0+PTX"/MARLIN_ARCHS "12.0"/g' \
95 -e 's/MARLIN_FP8_ARCHS "8.9;12.0"/MARLIN_FP8_ARCHS "12.0"/g' \
96 -e 's/MARLIN_OTHER_ARCHS "7.5;8.0+PTX"/MARLIN_OTHER_ARCHS "12.0"/g' \
97 -e 's/MARLIN_MOE_ARCHS "8.0+PTX"/MARLIN_MOE_ARCHS "12.0"/g' \
98 -e 's/MARLIN_MOE_FP8_ARCHS "8.9;12.0"/MARLIN_MOE_FP8_ARCHS "12.0"/g' \
99 -e 's/MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX"/MARLIN_MOE_OTHER_ARCHS "12.0"/g' \
100 -e 's/HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}"/HADACORE_ARCHS "12.0" "${CUDA_ARCHS}"/g' \
101 -e 's/"7.5;8.0;8.7;8.9+PTX" "${CUDA_ARCHS}"/"12.0" "${CUDA_ARCHS}"/g' \
102 CMakeLists.txt
103
104# vLLM build requirements
105RUN --mount=type=cache,target=/root/.cache/uv \
106 /opt/venv/bin/uv pip install -r requirements/build.txt \
107 --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION%.*}
108
109# Build vLLM wheel
110RUN --mount=type=cache,target=/root/.cache/uv \
111 --mount=type=cache,target=/root/.ccache \
112 CCACHE_NOHASHDIR="true" \
113 /opt/venv/bin/python3 setup.py bdist_wheel --dist-dir ${WHEELS_DIR} \
114 | grep -vE "^copying|^creating|^writing|^adding"
115
116# Clone LMCache
117WORKDIR /workspace
118RUN git clone --branch ${LMCACHE_GIT_REF} https://github.com/LMCache/LMCache
119
120# Build LMCache wheel
121WORKDIR /workspace/LMCache
122RUN --mount=type=cache,target=/root/.cache/uv \
123 --mount=type=cache,target=/root/.ccache \
124 CCACHE_NOHASHDIR="true" \
125 /opt/venv/bin/python3 setup.py bdist_wheel --dist-dir ${WHEELS_DIR} \
126 | grep -vE "^copying|^creating|^writing|^adding"
127
128# ccache stats
129WORKDIR /workspace
130RUN --mount=type=cache,target=/root/.ccache,sharing=locked \
131 ccache -s
132
133#################### RUNTIME STAGE ####################
134# Lean production image without build tools
135FROM nvcr.io/nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu24.04 AS runtime
136
137ARG CUDA_VERSION
138ARG FLASHINFER_VERSION
139ARG WHEELS_DIR
140
141ENV UV_LINK_MODE=copy
142ENV UV_HTTP_TIMEOUT=500
143ENV UV_INDEX_STRATEGY="unsafe-best-match"
144ENV CUDA_VERSION=${CUDA_VERSION}
145ENV FLASHINFER_VERSION=${FLASHINFER_VERSION}
146ENV FLASHINFER_CUDA_ARCH_LIST="12.0"
147ENV WHEELS_DIR=${WHEELS_DIR}
148ENV DEBIAN_FRONTEND=noninteractive
149ENV VLLM_TARGET_DEVICE=cuda
150ENV PATH="/opt/venv/bin:$PATH"
151
152# Distro setup
153RUN CUDA_VERSION_DASH=$(echo ${CUDA_VERSION} | cut -d. -f1,2 | tr '.' '-') && \
154 apt-get update -y && \
155 apt-get install -y --no-install-recommends \
156 # Runtime packages
157 kmod \
158 # Install CUDA development tools for runtime JIT compilation
159 # (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
160 build-essential \
161 cuda-nvcc-${CUDA_VERSION_DASH} \
162 # Python
163 python3.12 \
164 python3.12-venv \
165 python3.12-dev \
166 python3-pip \
167 && rm -rf /var/lib/apt/lists/*
168
169# Create venv
170RUN python3 -m venv /opt/venv
171RUN /opt/venv/bin/pip install --no-cache-dir --upgrade pip
172RUN /opt/venv/bin/pip install --no-cache-dir uv
173
174# Install packages in venv
175WORKDIR /tmp
176
177# PyTorch (use uv cache for fast install)
178RUN --mount=type=cache,target=/root/.cache/uv \
179 /opt/venv/bin/uv pip install --pre torch>=2.9.0 torchvision torchaudio \
180 --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION%.*}
181
182RUN --mount=type=cache,target=/root/.cache/uv \
183 /opt/venv/bin/uv pip install torch-c-dlpack-ext \
184 --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION%.*}
185
186# Copy all wheels from build stage & install them
187COPY --from=build ${WHEELS_DIR} /tmp/wheels
188RUN /opt/venv/bin/uv pip install /tmp/wheels/*.whl
189
190# Clean up
191RUN rm -rf /tmp/wheels
192
193# Install FlashInfer pre-compiled kernel cache and binaries
194# https://docs.flashinfer.ai/installation.html
195RUN --mount=type=cache,target=/root/.cache/uv \
196 /opt/venv/bin/uv pip install flashinfer-python flashinfer-cubin==${FLASHINFER_VERSION} \
197 && /opt/venv/bin/uv pip install flashinfer-jit-cache==${FLASHINFER_VERSION} \
198 --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
199 && /opt/venv/bin/flashinfer show-config
200
201# Allow z.ai GLM-4.6V and GLM-4.7-Flash models
202RUN --mount=type=cache,target=/root/.cache/uv \
203 apt-get update -y && \
204 apt-get install -y --no-install-recommends git && \
205 /opt/venv/bin/uv pip install git+https://github.com/huggingface/transformers.git && \
206 apt-get purge -y --auto-remove git && rm -rf /var/lib/apt/lists/*
207
208# TODO: Unsure why this is needed - remove it ASAP. Pulled by OpenCV for vLLM image processing
209RUN apt-get update -y && \
210 apt-get install -y --no-install-recommends libxcb1 \
211 && rm -rf /var/lib/apt/lists/*
212
213WORKDIR /workspace
214
215CMD ["bash"]1# Model configuration (Mandatory)
2MODEL="mratsim/GLM-4.7-Flash-FP8"
3MODELNAME="GLM-4.7-Flash"
4GPU_UTIL=0.90
5CONTEXT_SIZE=202752
6
7# Prevent memory fragmentation
8export PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb:512
9
10# Prevent vLLM from using 100% CPU when idle (Very Recommended)
11export VLLM_SLEEP_WHEN_IDLE=1
12
13vllm serve "${MODEL}" \
14 --served-model-name "${MODELNAME}" \
15 --gpu-memory-utilization ${GPU_UTIL} \
16 --max-model-len "${CONTEXT_SIZE}" \
17 --tool-call-parser glm47 \
18 --reasoning-parser glm45 \
19 --enable-auto-tool-choice1import os
2
3from llmcompressor import model_free_ptq
4
5os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
6os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:512")
7
8MODEL_ID = "zai-org/GLM-4.7-Flash"
9MODEL_OUT = MODEL_ID.split("/")[1] + "-FP8"
10
11model_free_ptq(
12 model_stub=MODEL_ID,
13 save_directory=MODEL_OUT,
14 scheme="FP8_BLOCK",
15 ignore=[
16 "lm_head",
17 "re:.*mlp\\.gate$", # MoE router
18 "re:.*kv_a_proj_with_mqa$",
19 "re:.*q_a_proj$",
20 "model.embed_tokens",
21 ],
22 max_workers=16,
23 device="cuda:0",
24)
25
26print(f"SUCCESS: files saved in {MODEL_OUT}")LayerNorm in Quantization. Kovaleva et al. (2021); Wei et al. (2022) find that outliers in the LayerNorm parameters of BERT (Devlin et al., 2019) cause difficulties in model compression. Given the importance of LayerNorm, all the quantization methods we discuss above leave LayerNorm unquantized.
Fig. 3: Maximum absolute value over layers for a LLaMA3-8B. Each color represent a different projection and we clearly see that down_proj has the biggest spikes in input and output. We also observe that RMSNorm propagate spikes through the entire model According to [5] Figure 5(a) illustrates the extremal ratio across layers and modules in LLaMA2-7B, highlighting that weight outliers are concentrated in the down-projection matrices Wdown ℓ of the second layer and the last two layers. Figures 5(b) and 5(c) provide detailed visualizations of these outliers in the last two layers.
k blocks have a significantly higher impact on model quality than for the same last k blocks.