# syntax=docker/dockerfile:1.7
# ==========================================================================
# GLM-5.2 NF3-Hybrid vLLM image for DGX Spark (aarch64/sm_121a, GB10)
#
# 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
#
# Usage:
#   docker build -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
# ==========================================================================

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

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

# 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
RUN apt-get update && apt-get install -y --no-install-recommends \
    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

WORKDIR /tmp/flashinfer
RUN pip install --no-cache-dir -e . \
    --no-build-isolation \
    && python3 -c 'import flashinfer; print(f"FlashInfer OK: {flashinfer.__version__}")'

# ───── 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.
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")'

# 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__}")'

# vLLM entrypoint
ENTRYPOINT ["vllm", "serve"]