# syntax=docker/dockerfile:1.7
# GLM-5.2 NF3-Hybrid — rapid build atop community DGX Spark image
#   Base: ghcr.io/timothystewart6/vllm-gb10:v0.23.0-gb10.1
#     (torch 2.11.0+cu130, vLLM, flashinfer, aarch64/sm_121a)
#   Adds: b12x (voipmonitor NF3-hybrid fork) + Triton NF3 overlay

ARG BASE_IMAGE=ghcr.io/timothystewart6/vllm-gb10:v0.24.0-gb10.5
FROM ${BASE_IMAGE}

ARG TORCH_CUDA_ARCH_LIST=12.1a
ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    CUTE_DSL_ARCH=sm_121a

# Install build tools + b12x deps
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential ninja-build cmake git curl wget python3-dev ccache \
    libibverbs-dev rdma-core \
    && rm -rf /var/lib/apt/lists/*

# Build b12x from source (NVFP4 + NF3 CuTe-DSL kernels)
RUN pip install --no-cache-dir uv setuptools-rust
RUN git clone --recursive https://github.com/voipmonitor/b12x.git /opt/b12x \
    && cd /opt/b12x \
    && git checkout f3686b555d639823b276c2080f173145eed7f007
WORKDIR /opt/b12x
RUN pip install --no-cache-dir -e . --no-build-isolation \
    -C setup.args="--cuda-arch=${TORCH_CUDA_ARCH_LIST}"

# Copy 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
RUN python3 -c "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('OK'); import nf3_kernel; print('OK'); import nvfp4_kernel; print('OK'); import nf3_replan; print('OK'); import b12x; print(f'OK b12x'); import torch; print(f'OK torch {torch.__version__} cuda {torch.version.cuda}'); import vllm; print(f'OK vllm {vllm.__version__}')"

ENTRYPOINT ["vllm", "serve"]
