feat: rewrite GLM-5.2 NF3-Hybrid Dockerfile (single-stage, local repos, real overlay kernels)

recipes/glm52-nf3-hybrid/Dockerfile

@@ -1,214 +1,105 @@
 # syntax=docker/dockerfile:1.7
 # ==========================================================================
-# GLM-5.2 NF3-Hybrid vLLM image for DGX Spark (aarch64/sm_121a, GB10)
+# GLM-5.2 NF3-Hybrid — DGX Spark (aarch64/sm_121a)
 #
-# Builds the eldritch-enlightenment vLLM + voipmonitor/b12x NF3-hybrid stack
-# from source, then overlays the NF3 Triton kernels + hybrid loader.
-#
-# Image: vllm-nf3-hybrid:probe
-# Build time: ~1h20m on a single DGX Spark
+# Builds vLLM (local-inference-lab/eldritch-enlightenment) + b12x
+# + FlashInfer from source, then overlays the NF3 Triton kernels.
 #
 # Usage:
-#   docker build -t vllm-nf3-hybrid:probe .
+#   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 .
 #
-# Upstream:
-#   vLLM: local-inference-lab/vllm @ 45c1582e9b80ba83e71c3a6458e71da4736fbdc4
-#         (dev/eldritch-enlightenment — NF3 hybrid loader, DCP ag_rs dispatch,
-#          low-bit MLA KV-cache dtypes, MTP/spec-decode cherries)
-#   b12x: voipmonitor/b12x @ f3686b555d639823b276c2080f173145eed7f007
-#         (NF3 W4A16 kernel, NVFP4 packed weights, sm_121 register table)
-#   FlashInfer: 0.6.14 (aarch64 wheels from flashinfer.ai)
-#   CUDA: 13.2 / PyTorch 2.12.0+cu132
-#   Arch: TORCH_CUDA_ARCH_LIST=12.1a, CUTE_DSL_ARCH=sm_121a
+# 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 PYTHON_VERSION=3.12
-ARG MAX_JOBS=16
 ARG GPU_ARCH=12.1a
 
-# ───── Stage 0: System base ─────────────────────────────────────────────
-FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} AS system-base
-
-ARG MAX_JOBS
-ARG PYTHON_VERSION
-ARG DEBIAN_FRONTEND=noninteractive
-
-ENV MAX_JOBS=${MAX_JOBS}
-ENV CMAKE_BUILD_PARALLEL_LEVEL=${MAX_JOBS}
-ENV NINJAFLAGS="-j${MAX_JOBS}"
-ENV MAKEFLAGS="-j${MAX_JOBS}"
-ENV PIP_BREAK_SYSTEM_PACKAGES=1
-ENV PIP_CACHE_DIR=/root/.cache/pip
-ENV UV_SYSTEM_PYTHON=1
-ENV UV_LINK_MODE=copy
-
-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 \
-    && rm -rf /var/lib/apt/lists/* \
-    && pip install uv
-
-# Install PyTorch 2.12.0+cu132 (aarch64)
-RUN pip install --no-cache-dir torch==2.12.0
-RUN pip install --no-cache-dir torchvision==0.22.0
-
-# ───── Stage 1: Build NCCL (patched, for multi-node) ───────────────────
-FROM system-base AS nccl-build
-
-ARG GPU_ARCH
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV CUTE_DSL_ARCH=sm_121a
-
-RUN git clone --depth=1 \
-    https://github.com/local-inference-lab/nccl-canonical.git \
-    /tmp/nccl
-
-WORKDIR /tmp/nccl
-RUN make -j${MAX_JOBS} src.build \
-    CUDA_HOME=/usr/local/cuda \
-    NVCC_GENCODE="-gencode=arch=compute_${GPU_ARCH},code=sm_${GPU_ARCH}" \
-    && make install
-
-# ───── Stage 2: Build b12x ─────────────────────────────────────────────
-FROM system-base AS b12x-build
-
-ARG GPU_ARCH
-ARG MAX_JOBS
-
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV CUTE_DSL_ARCH=sm_121a
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-
-# b12x NF3-hybrid fork
-RUN git clone https://github.com/voipmonitor/b12x.git /tmp/b12x \
-    && cd /tmp/b12x \
-    && git checkout f3686b555d639823b276c2080f173145eed7f007 \
-    && pip install --no-cache-dir -e . \
-        --no-build-isolation \
-        -C setup.args="--cuda-arch=${GPU_ARCH}" \
-    && python3 -c 'import b12x; print(f"b12x OK: {b12x.__version__}")'
-
-# ───── Stage 3: Build vLLM (eldritch-enlightenment) ────────────────────
-FROM system-base AS vllm-build
+FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}
 
 ARG GPU_ARCH
-ARG MAX_JOBS
 ARG DEBIAN_FRONTEND=noninteractive
 
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV CUTE_DSL_ARCH=sm_121a
-ENV VLLM_TARGET_DEVICE=cuda
-ENV VLLM_BUILD_WHEEL=1
+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)
 
-# Copy b12x from previous stage
-COPY --from=b12x-build /usr/local/lib/python3.12/dist-packages/b12x* \
-    /usr/local/lib/python3.12/dist-packages/
-COPY --from=b12x-build /tmp/b12x /tmp/b12x
-
-# Install build deps
+# ── System packages ─────────────────────────────────────────────────
 RUN apt-get update && apt-get install -y --no-install-recommends \
-    rapidjson-dev \
-    libxml2-dev \
+    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/*
 
-# Clone vLLM at the eldritch-enlightenment pin
-RUN git clone https://github.com/local-inference-lab/vllm.git /tmp/vllm \
-    && cd /tmp/vllm \
-    && git checkout 45c1582e9b80ba83e71c3a6458e71da4736fbdc4
-
-WORKDIR /tmp/vllm
-
-# Install Python deps
-RUN pip install --no-cache-dir \
-    -r requirements/common.txt \
-    -r requirements/cuda.txt \
-    pybind11 \
-    setuptools-scm
-
-# Build vLLM
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 setup.py bdist_wheel --dist-dir=/tmp/vllm-dist \
-    && pip install --no-cache-dir /tmp/vllm-dist/*.whl \
-    && python3 -c 'import vllm; print(f"vLLM OK: {vllm.__version__}")'
-
-# ───── Stage 4: Install FlashInfer ──────────────────────────────────────
-FROM system-base AS flashinfer-build
-
-ARG GPU_ARCH
-ARG MAX_JOBS
-
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV MAX_JOBS=${MAX_JOBS}
-
-# Clone flashinfer at the eldritch-compatible pin
-RUN git clone https://github.com/flashinfer-ai/flashinfer.git /tmp/flashinfer \
-    && cd /tmp/flashinfer \
-    && git checkout 25dd814e03791e370f96c3148242f0dc8de504ac
+# ── Python build essentials ─────────────────────────────────────────
+RUN pip install --no-cache-dir uv pybind11 setuptools-scm wheel
 
-WORKDIR /tmp/flashinfer
-RUN pip install --no-cache-dir -e . \
-    --no-build-isolation \
-    && python3 -c 'import flashinfer; print(f"FlashInfer OK: {flashinfer.__version__}")'
+# ── 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
 
-# ───── Stage 5: Final runtime image ─────────────────────────────────────
-FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${UBUNTU_VERSION} AS final
-
-ARG GPU_ARCH
-
-ENV TORCH_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV CUTE_DSL_ARCH=sm_121a
-ENV FLASHINFER_CUDA_ARCH_LIST=${GPU_ARCH}
-ENV DEBIAN_FRONTEND=noninteractive
-ENV PIP_BREAK_SYSTEM_PACKAGES=1
-
-# Runtime libs
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3 \
-    python3-pip \
-    libibverbs1 \
-    libibverbs-dev \
-    rdma-core \
-    libcudnn9-cuda-13 \
-    && rm -rf /var/lib/apt/lists/*
-
-# Copy all built artifacts
-COPY --from=vllm-build /usr/local/lib/python3.12/dist-packages/ \
-    /usr/local/lib/python3.12/dist-packages/
-COPY --from=vllm-build /tmp/vllm/vllm /usr/local/lib/python3.12/dist-packages/vllm
-COPY --from=flashinfer-build /usr/local/lib/python3.12/dist-packages/flashinfer* \
-    /usr/local/lib/python3.12/dist-packages/
-COPY --from=nccl-build /usr/local/lib/libnccl* /usr/local/lib/
-RUN ldconfig
-
-# ───── NF3 hybrid overlay ──────────────────────────────────────────────
-# These are placed in the Python path; the hybrid_loader's import hook
-# patches vllm.model_executor.layers.quantization.modelopt at import time.
+# ── NCCL ────────────────────────────────────────────────────────────
+COPY nccl /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
+WORKDIR /opt/b12x
+RUN pip install --no-cache-dir -e . --no-build-isolation \
+    -C setup.args="--cuda-arch=${GPU_ARCH}"
+
+# ── FlashInfer ──────────────────────────────────────────────────────
+RUN git clone --recursive https://github.com/flashinfer-ai/flashinfer.git /opt/flashinfer \
+    && cd /opt/flashinfer \
+    && git checkout 25dd814e03791e370f96c3148242f0dc8de504ac \
+    && git submodule update --init --recursive
+WORKDIR /opt/flashinfer
+RUN 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
+
+# Build vLLM wheel
+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 ─────────────────────────────────────────────
 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 so hybrid_loader fires before vLLM modelopt imports
-RUN python3 -c 'import sys; open("/usr/local/lib/python3.12/dist-packages/hybrid_loader.pth", "w").write("import hybrid_loader\n"); print("hybrid_loader import hook installed")'
+# .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 overlay imports
-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__}"); import b12x; print(f"b12x OK: {b12x.__version__}"); import vllm; print(f"vLLM OK: {vllm.__version__}")'
+# ── 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__}')"
 
-# vLLM entrypoint
 ENTRYPOINT ["vllm", "serve"]