ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04

FROM ${BASE_IMAGE} AS base-image

ARG BASE_IMAGE

ARG UV_VERSION=0.11.27

ENV DEBIAN_FRONTEND=noninteractive

ENV TRANSFORMERS_OFFLINE=0

ENV HYDRA_FULL_ERROR=1

ENV PYTHONUNBUFFERED=1

ENV UV_PYTHON=3.13

ENV UV_PROJECT_ENVIRONMENT=/opt/venv

ENV UV_LINK_MODE=copy

ENV UV_NO_CACHE=1

ENV UV_NO_PROGRESS=1

ENV VIRTUAL_ENV=/opt/venv

ENV PATH="/root/.local/bin:${VIRTUAL_ENV}/bin:${PATH}"

RUN <<"EOF" bash -ex

apt-get update

apt-get install -y --only-upgrade --no-install-recommends \

libssl3t64 \

openssl

apt-get install -y \

curl \

git

apt-get clean

EOF

RUN <<"EOF" bash -ex

curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" | sh

uv --version

EOF

RUN <<"EOF" bash -euxo pipefail

case "${BASE_IMAGE}" in

*cuda12*) cuda_flavor=cu12 ;;

*cuda13*) cuda_flavor=cu13 ;;

*)

echo "Cannot derive CUDA flavor from BASE_IMAGE='${BASE_IMAGE}'. Expected image tag containing 'cuda12' or 'cuda13'."

exit 1

;;

esac

cuda_major_minor="$(sed -n 's/.*cuda\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' <<<"${BASE_IMAGE}")"

if [[ -z "${cuda_major_minor}" ]]; then

echo "Cannot derive CUDA major.minor from BASE_IMAGE='${BASE_IMAGE}'. Expected image tag containing e.g. 'cuda12.9' or 'cuda13.2'."

exit 1

fi

cat >/usr/local/bin/nemo-cuda-flavor <<SCRIPT

set -euo pipefail

echo "${cuda_flavor}"

SCRIPT

cat >/usr/local/bin/nemo-install-cuda-python <<SCRIPT

set -euo pipefail

cuda_major_minor="${cuda_major_minor}"

cuda_major="\${cuda_major_minor%%.*}"

cuda_minor="\${cuda_major_minor#*.}"

cuda_next_minor="\$((cuda_minor + 1))"

uv pip install \

"cuda-bindings>=\${cuda_major_minor}.0,<\${cuda_major}.\${cuda_next_minor}" \

"cuda-python>=\${cuda_major_minor}.0,<\${cuda_major}.\${cuda_next_minor}"

SCRIPT

chmod +x /usr/local/bin/nemo-cuda-flavor /usr/local/bin/nemo-install-cuda-python

echo "Derived CUDA flavor: ${cuda_flavor}"

echo "Derived CUDA Python major.minor: ${cuda_major_minor}"

EOF

WORKDIR /workspace

COPY pyproject.toml uv.lock /workspace/

COPY nemo/__init__.py nemo/package_info.py /workspace/nemo/

RUN <<"EOF" bash -ex

cuda_flavor="$(nemo-cuda-flavor)"

uv sync --link-mode copy --locked --extra all --extra "${cuda_flavor}" --group test

nemo-install-cuda-python

EOF

RUN <<"EOF" bash -ex

uv pip install \

dill \

orjson

case "$(nemo-cuda-flavor)" in

cu12) torchcodec_index=https://download.pytorch.org/whl/cu129 ;;

cu13) torchcodec_index=https://download.pytorch.org/whl/cu132 ;;

esac

uv pip install --index-url "${torchcodec_index}" torchcodec

EOF

FROM base-image AS automodel-deps

ARG GPU_TARGET=h100plus

ARG DEEPEP_REPO=https://github.com/deepseek-ai/DeepEP.git

ARG DEEPEP_REF=7febc6e25660af0f54d95dd781ecdcd62265ecca

ARG FLASH_ATTN_MAX_JOBS=8

ARG FLASH_ATTN_NVCC_THREADS=1

ENV PIP_NO_CACHE_DIR=1

COPY external/patches/deepep_v1_a100_disable_nvshmem.patch /tmp/patches/

RUN <<"EOF" bash -euxo pipefail

apt-get update

apt-get install -y \

build-essential \

ca-certificates \

cmake \

ninja-build \

patch

apt-get clean

rm -rf /var/lib/apt/lists/*

EOF

RUN <<"EOF" bash -euxo pipefail

cat >/usr/local/bin/automodel-build-env <<'SCRIPT'

case "${GPU_TARGET}" in

h100plus)

export DEEPEP_TORCH_CUDA_ARCH_LIST="9.0+PTX"

export TORCH_CUDA_ARCH_LIST="9.0;10.0;12.0"

export NVTE_CUDA_ARCHS="90;100;120"

export FLASH_ATTN_CUDA_ARCHS="90;100;120"

export CMAKE_CUDA_ARCHITECTURES="90;100;120"

;;

a100)

export DEEPEP_TORCH_CUDA_ARCH_LIST="8.0"

export TORCH_CUDA_ARCH_LIST="8.0"

export NVTE_CUDA_ARCHS="80"

export FLASH_ATTN_CUDA_ARCHS="80"

export CMAKE_CUDA_ARCHITECTURES="80"

export DISABLE_SM90_FEATURES=1

;;

*)

echo "Unsupported GPU_TARGET='${GPU_TARGET}'. Expected 'h100plus' or 'a100'."

exit 1

;;

esac

CUDA_FLAVOR="$(nemo-cuda-flavor)"

AUTOMODEL_CCCL_INCLUDES="/usr/local/cuda/include/cccl"

PYTHON_CCCL_INCLUDE="${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages/nvidia/${CUDA_FLAVOR}/include/cccl"

if [[ -d "${PYTHON_CCCL_INCLUDE}" ]]; then

AUTOMODEL_CCCL_INCLUDES="${AUTOMODEL_CCCL_INCLUDES}:${PYTHON_CCCL_INCLUDE}"

fi

export CPATH="${AUTOMODEL_CCCL_INCLUDES}${CPATH:+:${CPATH}}"

export CPLUS_INCLUDE_PATH="${AUTOMODEL_CCCL_INCLUDES}${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}"

export NVTE_FRAMEWORK=pytorch

export MAX_JOBS="${FLASH_ATTN_MAX_JOBS}"

export NVCC_THREADS="${FLASH_ATTN_NVCC_THREADS}"

SCRIPT

mkdir -p /opt/automodel-src

EOF

COPY external/patches/flash_attn_4_sm120_disable_tma_o.patch /tmp/patches/

RUN <<"EOF" bash -euxo pipefail

source /usr/local/bin/automodel-build-env

uv pip install --upgrade \

cmake \

ninja \

nvidia-mathdx \

packaging \

pip \

psutil \

pybind11 \

setuptools \

setuptools-scm \

wheel

automodel_extra=compiled

if [[ "${GPU_TARGET}" == "h100plus" ]]; then

automodel_extra=compiled

elif [[ "${GPU_TARGET}" == "a100" ]]; then

automodel_extra=compiled-a100

fi

cuda_flavor="$(nemo-cuda-flavor)"

uv sync \

--inexact \

--link-mode copy \

--locked \

--extra all \

--extra "${cuda_flavor}" \

--extra "${automodel_extra}" \

--group test

nemo-install-cuda-python

automodel_site_packages="$(

if [[ "${GPU_TARGET}" == "a100" ]]; then

git clone "${DEEPEP_REPO}" /opt/automodel-src/DeepEP

git -C /opt/automodel-src/DeepEP checkout -f "${DEEPEP_REF}"

patch -d /opt/automodel-src/DeepEP -p1 < /tmp/patches/deepep_v1_a100_disable_nvshmem.patch

TORCH_CUDA_ARCH_LIST="${DEEPEP_TORCH_CUDA_ARCH_LIST}" \

uv pip install \

--no-build-isolation \

/opt/automodel-src/DeepEP

fi

if [[ "${GPU_TARGET}" == "h100plus" ]]; then

# flash-attn-4 requires apache-tvm-ffi 0.1.11, while mamba-ssm

# currently constrains the solved environment to apache-tvm-ffi<=0.1.9.

# Keep CUTLASS and Quack at the versions resolved by uv.lock. Automodel

# main now requires quack-kernels 0.6.1, which in turn requires CUTLASS

# 4.6.0; installing the old 4.5.2 override leaves the runtime inconsistent.

uv pip install \

--no-deps \

"apache-tvm-ffi==0.1.11" \

"torch-c-dlpack-ext==0.1.5"

mkdir -p /opt/automodel-src/flash-attn-4 /tmp/flash-attn-4

python -m pip download \

--no-binary=flash-attn-4 \

--no-deps \

--dest /tmp/flash-attn-4 \

"flash-attn-4==4.0.0b15"

tar -xf /tmp/flash-attn-4/flash_attn_4-*.tar.gz -C /opt/automodel-src/flash-attn-4 --strip-components=1

patch -d /opt/automodel-src/flash-attn-4 -p1 < /tmp/patches/flash_attn_4_sm120_disable_tma_o.patch

MAX_JOBS="${FLASH_ATTN_MAX_JOBS}" \

NVCC_THREADS="${FLASH_ATTN_NVCC_THREADS}" \

uv pip install \

--no-build-isolation \

--no-deps \

/opt/automodel-src/flash-attn-4

fi

EOF

COPY external/patches/flash_attn_4_cutlass46_compat.patch /tmp/patches/

RUN <<"EOF" bash -euxo pipefail

if [[ "${GPU_TARGET}" == "h100plus" ]]; then

patch -d "${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages/flash_attn/cute" -p1 \

< /tmp/patches/flash_attn_4_cutlass46_compat.patch

fi

EOF

FROM base-image

COPY --from=automodel-deps /opt/venv /opt/venv

ENV PYTHON_SITE_PACKAGES="${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages" \

CUDNN_HOME="${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages/nvidia/cudnn" \

CURAND_HOME="${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages/nvidia/curand" \

NVRTC_HOME="${VIRTUAL_ENV}/lib/python${UV_PYTHON}/site-packages/nvidia/cuda_nvrtc" \

NVTE_FRAMEWORK=pytorch

ENV LD_LIBRARY_PATH="${PYTHON_SITE_PACKAGES}/torch/lib:${PYTHON_SITE_PACKAGES}/nvidia/cublas/lib:${PYTHON_SITE_PACKAGES}/nvidia/cuda_cupti/lib:${PYTHON_SITE_PACKAGES}/nvidia/cuda_nvrtc/lib:${PYTHON_SITE_PACKAGES}/nvidia/cuda_runtime/lib:${PYTHON_SITE_PACKAGES}/nvidia/cudnn/lib:${PYTHON_SITE_PACKAGES}/nvidia/cufft/lib:${PYTHON_SITE_PACKAGES}/nvidia/cufile/lib:${PYTHON_SITE_PACKAGES}/nvidia/curand/lib:${PYTHON_SITE_PACKAGES}/nvidia/cusolver/lib:${PYTHON_SITE_PACKAGES}/nvidia/cusparse/lib:${PYTHON_SITE_PACKAGES}/nvidia/cusparselt/lib:${PYTHON_SITE_PACKAGES}/nvidia/nccl/lib:${PYTHON_SITE_PACKAGES}/nvidia/nvjitlink/lib:${PYTHON_SITE_PACKAGES}/nvidia/nvshmem/lib:${PYTHON_SITE_PACKAGES}/nvidia/nvtx/lib:${LD_LIBRARY_PATH}"

ARG NVIDIA_BUILD_ID

ENV NVIDIA_BUILD_ID=${NVIDIA_BUILD_ID:-<unknown>}

LABEL com.nvidia.build.id="${NVIDIA_BUILD_ID}"

ARG NVIDIA_BUILD_REF

LABEL com.nvidia.build.ref="${NVIDIA_BUILD_REF}"

ARG RC_DATE=00.00

ARG TARGETARCH

ARG INSTALL_FFMPEG=false

RUN <<"EOF" bash -ex

if [ "${INSTALL_FFMPEG}" = "true" ]; then

apt-get update

apt-get install -y ffmpeg

apt-get clean

rm -rf /var/lib/apt/lists/*

fi

EOF

COPY nemo /workspace/nemo

RUN echo "This distribution includes open source which is archived at the following URL: https://opensource.nvidia.com/oss/teams/nvidia/nemo/${RC_DATE}:linux-${TARGETARCH}/index.html" > NOTICES.txt && \

echo "For further inquiries or assistance, contact us at oss-requests@nvidia.com" >> NOTICES.txt