FROM ghcr.io/anemll/dspark-vllm-gx10:0.1.1

ARG EXLLAMAV3_COMMIT=c5d9c657966ffeeaa9353f0cc899f18629da4a13
ARG EXLLAMAV3_ARCHIVE_SHA256=87cf1ed1abad0fd2eb84744418e6a2ea3ec09a874552c1a3e0f226eac03cb8f9
ARG FLASHINFER_VERSION=0.6.13

ENV DEBIAN_FRONTEND=noninteractive \
    TORCH_CUDA_ARCH_LIST=12.1 \
    MAX_JOBS=4 \
    LIBRARY_PATH=/usr/local/lib/python3.12/dist-packages/nvidia/cu13/lib

COPY patch_exllamav3_arm64.py /opt/patch_exllamav3_arm64.py
COPY vendor/exllamav3-c5d9c657966ffeeaa9353f0cc899f18629da4a13.tar.gz /opt/exllamav3.tar.gz

RUN echo "${EXLLAMAV3_ARCHIVE_SHA256}  /opt/exllamav3.tar.gz" | sha256sum -c - \
    && for header in \
      /usr/local/lib/python3.12/dist-packages/nvidia/cu13/include/cusparse*.h \
      /usr/local/lib/python3.12/dist-packages/nvidia/cu13/include/cusolver*.h; \
    do \
      ln -sf "$header" "/usr/local/cuda/include/$(basename "$header")"; \
    done \
    && EXLLAMA_NOCOMPILE=1 python3 -m pip install \
      --no-cache-dir --no-deps --no-build-isolation \
      /opt/exllamav3.tar.gz \
    && python3 /opt/patch_exllamav3_arm64.py \
      /usr/local/lib/python3.12/dist-packages/exllamav3/exllamav3_ext \
    && python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
      "flashinfer-python==${FLASHINFER_VERSION}" \
    && python3 - <<'PY'
import importlib.metadata as im
import flashinfer

expected = "0.6.13"
actual = im.version("flashinfer-python")
cubin = im.version("flashinfer-cubin")
if actual != expected or cubin != expected or flashinfer.__version__ != expected:
    raise SystemExit(
        f"FlashInfer ABI mismatch: python={actual} cubin={cubin} "
        f"module={flashinfer.__version__} expected={expected}"
    )
print(f"FlashInfer ABI coherent: python={actual} cubin={cubin}")
PY

COPY sitecustomize.py /usr/local/lib/python3.12/dist-packages/sitecustomize.py

LABEL org.opencontainers.image.source="https://github.com/turboderp-org/exllamav3" \
      org.opencontainers.image.revision="${EXLLAMAV3_COMMIT}" \
      ai.laguna.exllamav3_archive_sha256="${EXLLAMAV3_ARCHIVE_SHA256}" \
      ai.laguna.flashinfer_python="${FLASHINFER_VERSION}" \
      ai.laguna.flashinfer_cubin="${FLASHINFER_VERSION}" \
      ai.laguna.runtime_fix="graph-safe hot expert remap and additive TP1 reconstruction of TP2 EXL3 slices" \
      ai.laguna.cuda_graphs="required" \
      ai.laguna.purpose="Laguna TP1 or TP2 saliency hybrid NVFP4 plus EXL3 tails"