# syntax=docker/dockerfile:1.7
# karmic-kraken (local-inference-lab vLLM fork, dev/karmic-kraken) — GB10/Spark build.
#
# Clone-at-pin build, NO staged tarballs: the Dockerfile clones every source
# at the pins in versions.env and gates each one (HEAD commit + tree hash).
# A fresh clone of this repo + `bash build.sh` reproduces this image end to
# end. Current tip: vLLM 636afde9 (2026-10-08), which takes B12X from
# FlashInfer 0.7.1 (the standalone b12x repo is archived — no separate wheel).
#
# Runtime: CUDA 13.4.2 devel aarch64 base + torch 2.13.0 cu130 + flashinfer
# 0.7.1 (carries B12X as flashinfer/experimental/b12x + the top-level `b12x`
# import package). GB10 sm_121 fixes baked in (patch_cuda_sm12x.py, eugr
# patch set).
#
# Build: cd docker-images/karmic-kraken && TAG=... bash build.sh

ARG CUDA_BASE_IMAGE
ARG CUDA_BASE_DIGEST
ARG UV_VERSION

FROM ${CUDA_BASE_IMAGE}@${CUDA_BASE_DIGEST} AS apt-base
ARG PYPI_INDEX_URL
ARG APT_SOURCE
ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    UV_SYSTEM_PYTHON=1 UV_BREAK_SYSTEM_PACKAGES=1 UV_LINK_MODE=copy
COPY locks/apt-sources.list /tmp/apt-sources-snapshot.list
COPY locks/apt-packages.txt /tmp/apt-packages.txt
COPY locks/python-bootstrap.txt /tmp/python-bootstrap.txt
COPY locks/deps.lock /tmp/deps.lock
# APT_SOURCE (versions.env):
#   snapshot (default, historical behavior): the pinned point-in-time mirror
#       + exact package versions — strictly reproducible, but
#       snapshot.ubuntu.com is slow (~30 kB/s) and intermittently 503s
#       (12-min apt stall, 2026-09-08).
#   current: the live archive.ubuntu.com — fast and reliable; installs the
#       pinned package names unversioned. Acceptable: the engine stack is
#       pinned by the vLLM/B12X/LMCache source pins, not by apt.
RUN set -eux; \
    if [ "${APT_SOURCE:-snapshot}" = "current" ]; then \
        printf '%s\n' \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble main restricted universe multiverse" \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble-updates main restricted universe multiverse" \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble-security main restricted universe multiverse" \
            > /etc/apt/sources.list; \
        PKGS="$(grep -vE '^\s*(#|$)' /tmp/apt-packages.txt | cut -d= -f1)"; \
    else \
        cp /tmp/apt-sources-snapshot.list /etc/apt/sources.list; \
        PKGS="$(grep -vE '^\s*(#|$)' /tmp/apt-packages.txt)"; \
    fi; \
    rm -f /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources; \
    apt-get update; \
    apt-get install -y --no-install-recommends ${PKGS}; \
    rm -rf /var/lib/apt/lists/*; \
    python3 -m pip install --no-cache-dir --no-deps \
        -r /tmp/python-bootstrap.txt --index-url "${PYPI_INDEX_URL}"

FROM apt-base AS torch-base
ARG PYPI_INDEX_URL
ARG PYTORCH_INDEX_URL
ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION
ARG FLASHINFER_VERSION
ARG B12X_VERSION
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv pip install \
      --index-url ${PYPI_INDEX_URL} \
      --extra-index-url ${PYTORCH_INDEX_URL} \
      --index-strategy unsafe-best-match \
      "torch==${TORCH_VERSION}" \
      "torchvision==${TORCHVISION_VERSION}" \
      "torchaudio==${TORCHAUDIO_VERSION}" \
      "flashinfer-python==${FLASHINFER_VERSION}" \
      "b12x==${B12X_VERSION}"

# DeepGEMM: cloned in-build (proven to work on the nodes via git; pins are
# gated again in fetch-deps.sh against the vLLM tree's own declaration).
FROM torch-base AS deepgemm-builder
ARG DEEPGEMM_REPO DEEPGEMM_BRANCH DEEPGEMM_REF
RUN set -eux; \
    git config --global protocol.version 0; \
    git clone --branch ${DEEPGEMM_BRANCH} ${DEEPGEMM_REPO} /deps/deepgemm; \
    cd /deps/deepgemm; \
    git checkout ${DEEPGEMM_REF}; \
    git submodule update --init --depth 1 -- third-party/cutlass third-party/fmt

# ---- vLLM @ r27 integration branch, gated by commit + tree hash ----
FROM torch-base AS vllm-clone
ARG VLLM_REPO VLLM_PIN VLLM_BRANCH
RUN set -eux; \
    git config --global protocol.version 0; \
    git init -q /workspace/vllm; \
    git -C /workspace/vllm fetch -q --depth 1 ${VLLM_REPO} ${VLLM_PIN}; \
    git -C /workspace/vllm checkout -q FETCH_HEAD; \
    cd /workspace/vllm; \
    test "$(git rev-parse HEAD)" = "${VLLM_PIN}"; \
    echo "vllm at $(git rev-parse HEAD) (branch ${VLLM_BRANCH})"; \
    git config --global --add safe.directory /workspace/vllm

# ---- All build sub-deps, cloned at pins the r27 tree itself declares ----
# REGEN_LOCK=1: fetch-deps.sh rewrites /tmp/deps.lock from the staged state
# (pin cuts) instead of verifying against it — export the stage to harvest
# the new lock, then commit it.
FROM vllm-clone AS deps-fetched
ARG REGEN_LOCK=""
COPY fetch-deps.sh /tmp/fetch-deps.sh
RUN if [ -n "$REGEN_LOCK" ]; then bash /tmp/fetch-deps.sh --regen-lock /tmp/deps.lock; else bash /tmp/fetch-deps.sh; fi

FROM deps-fetched AS vllm-builder
WORKDIR /workspace/vllm
# Hard-declared (do NOT rely on stage-inherited ENVs): the buildkit build
# has no GPU, so torch's device-capability autodetect is unavailable and
# the wheel must be told the target arch explicitly (the r27 slip: kernels
# came out sm120-only -> cudaErrorNoKernelImageForDevice at boot).
ENV TORCH_CUDA_ARCH_LIST=12.1a GIT_PROTOCOL=0
ARG DEEPGEMM_REF
ARG VLLM_PRETEND_VERSION
ARG SOURCE_DATE_EPOCH
ARG BUILD_JOBS
# Empty default overrides (profile `FLASH_KDA_SRC_DIR=`) make the vLLM tree
# fall back to its own pinned FlashKDA + in-tree packed-checkpoints patch
# (required by the r28.1+ flashkda.cmake API gate); the default keeps the
# r27-style shared /deps2 checkout for the older trees.
ARG FLASH_KDA_SRC_DIR=
# eugr/spark-vllm-docker build-time patches (main, 2026-09-02): all are
# source-aware/idempotent (skip when the vulnerable pattern is absent), so
# running the full set is safe on the r27 tree. Env gates are opt-in and ON
# for this build (sm_121a + b12x):
#   - preserve_sm12x_target: keep 12.1a in CUDA_SUPPORTED_ARCHS (their
#     standard CUDA build collapses 12.1a -> 12.0)
#   - sm120_cooperative_topk: keep SM90-only (Spark fix)
#   - mrv2_speculator_cudagraph_pool / spark_kv_cache_cleanup (UMA snapshot
#     flow) / b12x c128a alignment / flashinfer b12x swigluoai / routed
#     expert weight shape / topk_softplus_sqrt / others: see each script.
COPY patches-eugr/patch_vllm_*.py /tmp/eugr-patches/
COPY --from=deepgemm-builder /deps/deepgemm /deps/deepgemm
RUN VLLM_PRESERVE_SM12X_TARGET=1 \
    VLLM_PATCH_B12X_C128A_ALIGNMENT=1 \
    python3 /tmp/eugr-patches/patch_vllm_mrv2_speculator_cudagraph_pool.py . \
    && python3 /tmp/eugr-patches/patch_vllm_b12x_c128a_topk_alignment.py . \
    && python3 /tmp/eugr-patches/patch_vllm_flashinfer_b12x_swigluoai.py . \
    && python3 /tmp/eugr-patches/patch_vllm_topk_softplus_sqrt_control_flow.py . \
    && python3 /tmp/eugr-patches/patch_vllm_preserve_sm12x_target.py . \
    && python3 /tmp/eugr-patches/patch_vllm_diffusion_tensor_causal.py . \
    && python3 /tmp/eugr-patches/patch_vllm_gemma4_mtp_embedding_share.py . \
    && python3 /tmp/eugr-patches/patch_vllm_sm120_cooperative_topk.py . \
    && python3 /tmp/eugr-patches/patch_vllm_autogptq_symmetric_moe_qzeros.py . \
    && python3 /tmp/eugr-patches/patch_vllm_disable_minimax_qk_rmsnorm_ipc.py . \
    && python3 /tmp/eugr-patches/patch_vllm_routed_experts_weight_shape.py . \
    && python3 /tmp/eugr-patches/patch_vllm_spark_kv_cache_cleanup.py .
# The vllm tip pins flashinfer-python==0.7.1 (final) in requirements/cuda.txt,
# which does NOT exist on PyPI (only 0.7.1rc5). Relax to the installed rc5 so
# the wheel's dep metadata is satisfiable (flashinfer 0.7.1rc5 is in torch-base);
# otherwise `pip install` of the wheel fails to resolve the dep and vllm is absent.
RUN set -eux; \
    sed -i 's/flashinfer-python==0.7.1/flashinfer-python==0.7.1rc5/' /workspace/vllm/requirements/cuda.txt; \
    grep 'flashinfer-python' /workspace/vllm/requirements/cuda.txt
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    --mount=type=cache,id=cargo-registry,target=/root/.cargo/registry \
    --mount=type=cache,id=rust-target,target=/workspace/vllm/vllm/target \
    git config --global protocol.version 0 \
    && uv pip install setuptools-rust setuptools-scm \
    && (command -v cargo >/dev/null 2>&1 || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --no-modify-path) \
    && export PATH=/root/.cargo/bin:$PATH \
    && (command -v protoc >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y --no-install-recommends protobuf-compiler libprotobuf-dev >/dev/null && rm -rf /var/lib/apt/lists/*)) \
    && _j="${BUILD_JOBS:-$(nproc)}" \
    && export GIT_PROTOCOL=0 MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
        VLLM_REQUIRE_RUST_FRONTEND=1 CARGO_BUILD_JOBS="${_j}" \
        SETUPTOOLS_SCM_PRETEND_VERSION="${VLLM_PRETEND_VERSION}" \
        TRITON_KERNELS_SRC_DIR=/deps2/triton_kernels \
        FMHA_SM100_SRC_DIR=/deps2/MSA \
        FLASH_MLA_SRC_DIR=/deps2/FlashMLA \
        FLASH_KDA_SRC_DIR="${FLASH_KDA_SRC_DIR}" \
        VLLM_FLASH_ATTN_SRC_DIR=/deps2/flash-attention \
        QUTLASS_SRC_DIR=/deps2/qutlass \
        TML_FA4_SRC_DIR=/deps2/tml-fa4 \
        VLLM_CUTLASS_SRC_DIR=/deps2/cutlass \
        CUTLASS_INCLUDE_DIR_PRE=/deps2/cutlass/include \
    && python3 use_existing_torch.py \
    && uv build --no-build-isolation --wheel . --out-dir=/wheels -v

# ---- B12X: now ships inside flashinfer-python 0.7.1 (installed in torch-base) ----
# The standalone lukealonso/b12x repo is archived; its code continues in
# FlashInfer, whose wheel installs the top-level `b12x` package that vLLM
# imports. No separate b12x wheel is built or installed (a standalone b12x
# package must not sit next to flashinfer 0.7.1 — same namespace).

# ---- Runner: flattened GB10 serving image ----
FROM torch-base AS runner
ARG PYPI_INDEX_URL
COPY checksums/tiktoken.sha256 /tmp/tiktoken.sha256
COPY patch_cuda_sm12x.py /tmp/patch_cuda_sm12x.py
# GLM-5.3 chat template (from the kindling stack): the karmic recipe serves
# GLM-5.3-Flash with --chat-template /usr/local/share/glm53-chat-template.jinja;
# bake it in so the karmic image matches the kindling one (fair A/B).
COPY glm53-chat-template.jinja /usr/local/share/glm53-chat-template.jinja
WORKDIR /workspace
RUN mkdir -p tiktoken_encodings \
 && wget -O tiktoken_encodings/o200k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken \
 && wget -O tiktoken_encodings/cl100k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken \
 && cd tiktoken_encodings && sha256sum -c /tmp/tiktoken.sha256
ENV TIKTOKEN_ENCODINGS_BASE=/workspace/tiktoken_encodings
RUN --mount=type=bind,from=vllm-builder,source=/wheels,target=/vllm-wheels \
    --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    set -eux; \
    echo '=== /vllm-wheels contents ==='; ls -la /vllm-wheels/; \
    pip install --no-cache-dir --no-build-isolation /vllm-wheels/*.whl; \
    python3 -c "import vllm; print('vllm installed OK:', vllm.__version__)" \
 && python3 /tmp/patch_cuda_sm12x.py \
 && find /usr/local/lib/python3.12/dist-packages/vllm \
      -name "__pycache__" -type d -prune -exec rm -rf {} + 2>/dev/null || true
# b12x loader allocation patch (discord, 2026-09-07): pinned_wc -> managed in
# b12x/integration/vllm/loader.py. Legacy-only: with FlashInfer 0.7.1's bundled
# b12x (auto-allocation, no pinned_wc) the file is ABSENT, so the patch is N/A.
# Tolerant of all three cases: pinned_wc present (patch it), file present but
# clean (skip), file absent (new flashinfer-bundled b12x — N/A).
RUN set -eux; \
    F=/usr/local/lib/python3.12/dist-packages/b12x/integration/vllm/loader.py; \
    if [ -f "$F" ] && grep -q 'pinned_wc' "$F" 2>/dev/null; then \
      sed -i 's/allocation="pinned_wc"/allocation="managed"/; \
              s/"allocation": "pinned_wc"/"allocation": "managed"/' "$F"; \
      test "$(grep -c 'allocation="managed"' "$F")" = "1"; \
      grep -Fq '"allocation": "managed"' "$F"; \
      ! grep -q 'pinned_wc' "$F"; \
    elif [ -f "$F" ]; then \
      echo "b12x loader carries no pinned_wc (new auto-allocation b12x) — patch skipped"; \
    else \
      echo "b12x loader file absent (b12x ships inside flashinfer 0.7.1 with auto-allocation) — patch N/A"; \
    fi
# Build gate: the r27 stack must import with the GB10 patches in place and
# carry the qualified package trees (marker spot-checks from the source lock).
RUN set -eux; \
    python3 -c "import vllm; print('vllm', vllm.__version__)"; \
    python3 -c "import b12x; print('b12x', getattr(b12x,'__version__','(bundled in flashinfer 0.7.1)'))"; \
    python3 -c "import vllm._flashkda_C; print('flashkda_C ok')"; \
    python3 -c "from vllm.config.attention import AttentionBackendEnum as E; \
assert E.B12X.value and 'b12x' in E.B12X.value.lower(); print('B12X enum ok')"; \
    grep -q "AttentionBackendEnum.B12X," /usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py; \
    grep -Fq 'def apply_top_k_top_p_probs' /usr/local/lib/python3.12/dist-packages/vllm/v1/sample/ops/topk_topp_sampler.py; \
    if [ -f /usr/local/lib/python3.12/dist-packages/b12x/comm/pcie/_oneshot_cute.py ]; then \
      grep -Fq '_PREPARED_ONESHOT_LAUNCHERS' /usr/local/lib/python3.12/dist-packages/b12x/comm/pcie/_oneshot_cute.py; \
    else \
      echo "b12x oneshot marker file absent (new flashinfer-bundled layout) — spot-check skipped"; \
    fi
ENTRYPOINT []
CMD ["bash"]