# syntax=docker/dockerfile:1
# Blackwell (sm_100) kernel stack on top of the cloud image: DeepEP for expert parallelism,
# flash-attn-4 and the quack / cutlass DSL kernels behind SonicMoE. linux/amd64 only.
ARG BASE_TAG=main-py3.12-cu130-2.14.1
FROM axolotlai/axolotl-cloud-uv:$BASE_TAG

ARG DEEP_EP_REF=v1.2.1
ARG TORCH_CUDA_ARCH_LIST="10.0"
ARG MAX_JOBS=16

# DeepEP main needs CUDA 13.1; v1.2.1 builds on this image's CUDA 13.0 once
# CUDA 13's relocated cuda/std headers are on the include path. It links torch's
# nvidia-nvshmem-cu13 wheel, which setup.py finds on its own. The tag's build backend fails
# under `pip install .`, so build the wheel with setuptools directly.
RUN git clone --depth 1 --branch $DEEP_EP_REF https://github.com/deepseek-ai/DeepEP.git /tmp/DeepEP && \
    cd /tmp/DeepEP && \
    TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST MAX_JOBS=$MAX_JOBS \
    CFLAGS=-I/usr/local/cuda/include/cccl CXXFLAGS=-I/usr/local/cuda/include/cccl \
    NVCC_PREPEND_FLAGS=-I/usr/local/cuda/include/cccl \
    python setup.py bdist_wheel -d /tmp/wheels && \
    uv pip install --no-deps /tmp/wheels/deep_ep-*.whl && \
    cd /workspace && rm -rf /tmp/DeepEP /tmp/wheels && \
    python -c "import deep_ep; print(deep_ep.Buffer)"

# flash-attn-4 and the quack kernels behind SonicMoE, pinned by axolotl's `blackwell` extra. Both
# require an unbounded torch, so pin the image's build (DeepEP above links against it).
RUN --mount=type=cache,target=/root/.cache/uv \
    uv pip install --no-build-isolation -e "/workspace/axolotl[blackwell]" \
        "torch==$(python -c 'import torch; print(torch.__version__)')" && \
    python -c "import importlib.metadata as m; [print(p, m.version(p)) for p in ('flash-attn-4', 'quack-kernels', 'nvidia-cutlass-dsl', 'deep-ep')]"
