feat: vllm-gb10 based Dockerfile for GLM-5.2 NF3-Hybrid

recipes/glm52-nf3-hybrid/Dockerfile

@@ -1,105 +1,140 @@
 # syntax=docker/dockerfile:1.7
-# ==========================================================================
-# GLM-5.2 NF3-Hybrid — DGX Spark (aarch64/sm_121a)
-#
-# Builds vLLM (local-inference-lab/eldritch-enlightenment) + b12x
-# + FlashInfer from source, then overlays the NF3 Triton kernels.
-#
-# Usage:
-#   docker build \
-#     --build-arg VLLM_SRC=/path/to/vllm \
-#     --build-arg B12X_SRC=/path/to/b12x \
-#     --build-arg NCCL_SRC=/path/to/nccl \
-#     -t vllm-nf3-hybrid:probe .
-#
-# On dgx4.lan:
-#   VLLM_SRC=/home/kilork/projects/github.com/vllm-eldritch
-#   B12X_SRC=/home/kilork/projects/github.com/b12x-nf3
-#   NCCL_SRC=/home/kilork/projects/github.com/nccl-canonical
-# ==========================================================================
-
-ARG CUDA_VERSION=13.2.0
-ARG UBUNTU_VERSION=24.04
-ARG GPU_ARCH=12.1a
-
-FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}
-
-ARG GPU_ARCH
+# GLM-5.2 NF3-Hybrid for DGX Spark (aarch64/sm_121a)
+# Adapted from timothystewart6/vllm-gb10
+
+ARG CUDA_BASE_IMAGE
+
+# ───── Stage 1: system-base ──────────────────────────────────────────
+FROM ${CUDA_BASE_IMAGE} AS system-base
+
+ARG UV_VERSION
+ARG PYTORCH_INDEX_URL
 ARG DEBIAN_FRONTEND=noninteractive
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_BREAK_SYSTEM_PACKAGES=1 \
-    TORCH_CUDA_ARCH_LIST=${GPU_ARCH} \
-    CUTE_DSL_ARCH=sm_121a \
-    FLASHINFER_CUDA_ARCH_LIST=${GPU_ARCH} \
-    MAX_JOBS=$(nproc) \
-    CMAKE_BUILD_PARALLEL_LEVEL=$(nproc)
+    UV_SYSTEM_PYTHON=1 UV_LINK_MODE=copy
 
-# ── System packages ─────────────────────────────────────────────────
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential ninja-build cmake curl git wget \
-    python3-dev python3-pip python3-venv ccache \
-    libibverbs-dev libibverbs1 rdma-core \
-    libcudnn9-cuda-13 libcudnn9-dev-cuda-13 \
-    rapidjson-dev libxml2-dev \
-    && rm -rf /var/lib/apt/lists/*
+    python3-dev python3-pip ccache \
+    libibverbs-dev rdma-core libcudnn9-cuda-13 libcudnn9-dev-cuda-13 \
+    && rm -rf /var/lib/apt/lists/* \
+    && pip install --no-cache-dir uv==0.11.28
+
+# ───── Stage 2: torch-base ───────────────────────────────────────────
+FROM system-base AS torch-base
+
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+ARG TRITON_VERSION
+ARG PYTORCH_INDEX_URL
+
+RUN pip install --no-cache-dir --no-deps \
+    torch==2.12.0 torchvision==0.22.0 \
+    --index-url "${PYTORCH_INDEX_URL}"
 
-# ── Python build essentials ─────────────────────────────────────────
-RUN pip install --no-cache-dir uv pybind11 setuptools-scm wheel
+# ───── Stage 3: NCCL (nccl-canonical) ────────────────────────────────
+FROM torch-base AS nccl-build
 
-# ── PyTorch 2.12.0+cu132 (aarch64) ──────────────────────────────────
-RUN pip install --no-cache-dir --no-deps torch==2.12.0 \
-    && pip install --no-cache-dir --no-deps torchvision==0.22.0
+ARG NCCL_REPO
+ARG NCCL_REF
 
-# ── NCCL ────────────────────────────────────────────────────────────
-COPY nccl /opt/nccl
+RUN git clone --depth=1 -b "${NCCL_REF}" "${NCCL_REPO}" /opt/nccl
 WORKDIR /opt/nccl
 RUN make -j$(nproc) src.build \
     CUDA_HOME=/usr/local/cuda \
     NVCC_GENCODE="-gencode=arch=compute_121,code=sm_121" \
     && make install PREFIX=/usr/local && ldconfig
 
-# ── b12x ────────────────────────────────────────────────────────────
-COPY b12x /opt/b12x
+# ───── Stage 4: b12x ─────────────────────────────────────────────────
+FROM torch-base AS b12x-build
+
+ARG B12X_REPO
+ARG B12X_REF
+ARG TORCH_CUDA_ARCH_LIST=12.1a
+
+ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
+    CUTE_DSL_ARCH=sm_121a
+
+RUN git clone --recursive "${B12X_REPO}" /opt/b12x \
+    && cd /opt/b12x && git checkout "${B12X_REF}"
 WORKDIR /opt/b12x
 RUN pip install --no-cache-dir -e . --no-build-isolation \
-    -C setup.args="--cuda-arch=${GPU_ARCH}"
+    -C setup.args="--cuda-arch=${TORCH_CUDA_ARCH_LIST}"
+
+# ───── Stage 5: FlashInfer ───────────────────────────────────────────
+FROM torch-base AS flashinfer-build
+
+ARG FLASHINFER_REPO
+ARG FLASHINFER_REF
+ARG FLASHINFER_CUDA_ARCH_LIST=12.1a
+ARG TORCH_CUDA_ARCH_LIST=12.1a
+
+ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST} \
+    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
 
-# ── FlashInfer ──────────────────────────────────────────────────────
-RUN git clone --recursive https://github.com/flashinfer-ai/flashinfer.git /opt/flashinfer \
-    && cd /opt/flashinfer \
-    && git checkout 25dd814e03791e370f96c3148242f0dc8de504ac \
+RUN git clone --recursive "${FLASHINFER_REPO}" /opt/flashinfer \
+    && cd /opt/flashinfer && git checkout "${FLASHINFER_REF}" \
     && git submodule update --init --recursive
 WORKDIR /opt/flashinfer
-RUN uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
+RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
+    uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
     && cd flashinfer-cubin \
     && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
     && cd ../flashinfer-jit-cache \
     && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
     && pip install --no-cache-dir /wheels/*.whl
 
-# ── vLLM (eldritch-enlightenment) ───────────────────────────────────
-COPY vllm /opt/vllm
-WORKDIR /opt/vllm
-RUN git checkout 45c1582e9b80ba83e71c3a6458e71da4736fbdc4 2>/dev/null || true \
-    && git submodule update --init --recursive 2>/dev/null || true
+# ───── Stage 6: vLLM (eldritch-enlightenment) ────────────────────────
+FROM torch-base AS vllm-build
+
+ARG VLLM_REPO
+ARG VLLM_REF
+ARG TORCH_CUDA_ARCH_LIST=12.1a
+
+ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
+    CUTE_DSL_ARCH=sm_121a
+
+# Clone vLLM at eldritch-enlightenment pin
+RUN git clone --recursive "${VLLM_REPO}" /opt/vllm \
+    && cd /opt/vllm && git checkout "${VLLM_REF}" \
+    && git submodule update --init --recursive
 
-# Build vLLM wheel
+WORKDIR /opt/vllm
 RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
     uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
     && pip install --no-cache-dir /wheels/vllm*.whl
 
-# ── NF3 hybrid overlay ─────────────────────────────────────────────
+# ───── Stage 7: runner (final image) ─────────────────────────────────
+FROM torch-base AS runner
+
+# Copy built artifacts from each stage
+COPY --from=nccl-build     /usr/local/lib/libnccl*       /usr/local/lib/
+COPY --from=nccl-build     /usr/local/include/nccl.h     /usr/local/include/
+RUN ldconfig
+
+COPY --from=b12x-build     /usr/local/lib/python3.12/dist-packages/b12x* \
+    /usr/local/lib/python3.12/dist-packages/
+COPY --from=flashinfer-build /usr/local/lib/python3.12/dist-packages/flashinfer* \
+    /usr/local/lib/python3.12/dist-packages/
+COPY --from=vllm-build     /usr/local/lib/python3.12/dist-packages/ \
+    /usr/local/lib/python3.12/dist-packages/
+
+# ── NF3 hybrid overlay ──────────────────────────────────────────────
 COPY overlay/hybrid_loader.py  /usr/local/lib/python3.12/dist-packages/hybrid_loader.py
 COPY overlay/nf3_kernel.py     /usr/local/lib/python3.12/dist-packages/nf3_kernel.py
 COPY overlay/nvfp4_kernel.py   /usr/local/lib/python3.12/dist-packages/nvfp4_kernel.py
 COPY overlay/nf3_replan.py     /usr/local/lib/python3.12/dist-packages/nf3_replan.py
 COPY overlay/mxfp8_tier.json   /usr/local/lib/python3.12/dist-packages/mxfp8_tier.json
 
-# .pth hook: hybrid_loader fires when any Python process imports from vllm
 RUN python3 -c "import sys; open('/usr/local/lib/python3.12/dist-packages/hybrid_loader.pth','w').write('import hybrid_loader\n')"
 
-# ── Verify ──────────────────────────────────────────────────────────
-RUN python3 -c "import hybrid_loader; print('hybrid_loader OK'); import nf3_kernel; print('nf3_kernel OK'); import nvfp4_kernel; print('nvfp4_kernel OK'); import nf3_replan; print('nf3_replan OK'); import flashinfer; print(f'flashinfer OK: {flashinfer.__version__}'); import torch; print(f'PyTorch OK: {torch.__version__}, CUDA: {torch.version.cuda}'); import b12x; print(f'b12x OK: {b12x.__version__}'); import vllm; print(f'vLLM OK: {vllm.__version__}')"
+# Verify imports
+RUN python3 -c "import hybrid_loader; print('OK hybrid_loader'); import nf3_kernel; print('OK nf3_kernel'); import nvfp4_kernel; print('OK nvfp4_kernel'); import nf3_replan; print('OK nf3_replan'); import flashinfer; print(f'OK flashinfer {flashinfer.__version__}'); import torch; print(f'OK torch {torch.__version__} cuda {torch.version.cuda}'); import b12x; print(f'OK b12x {b12x.__version__}'); import vllm; print(f'OK vllm {vllm.__version__}')"
+
+ENV TORCH_CUDA_ARCH_LIST=12.1a \
+    CUTE_DSL_ARCH=sm_121a \
+    FLASHINFER_CUDA_ARCH_LIST=12.1a
 
 ENTRYPOINT ["vllm", "serve"]