chore: import upstream snapshot with attribution

docker/Dockerfile

@@ -0,0 +1,294 @@
+# syntax=docker/dockerfile:1-labs
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This CI Dockerfile supports CUDA 13 and CUDA 12 from a single BASE_IMAGE
+# build arg. The default is the recommended CUDA 13 image:
+#   nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04
+# The current recommended CUDA 12 image is:
+#   nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04
+#
+# The build derives CUDA_FLAVOR internally from BASE_IMAGE by matching
+# "cuda13" or "cuda12" in the image tag. That flavor selects the matching uv
+# extra (cu13 or cu12) and CUDA Python package include path. If BASE_IMAGE does
+# not contain either token, the build fails early.
+#
+# Example CUDA 13 H100+ build:
+#   docker buildx build -f docker/Dockerfile \
+#     --build-arg GPU_TARGET=h100plus .
+# Example CUDA 12 A100 build:
+#   docker buildx build -f docker/Dockerfile \
+#     --build-arg BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04 \
+#     --build-arg GPU_TARGET=a100 .
+#
+# GPU_TARGET controls compiled Automodel dependency tuning. "h100plus" builds
+# for SM90/SM100/SM120 and includes H100+ features such as DeepEP and
+# flash-attn-4. "a100" builds only SM80 and uses an A100-specific DeepEP patch
+# to avoid unsupported newer-GPU/NVSHMEM build paths. This keeps CI images
+# smaller and avoids compiling kernels for architectures a target image cannot
+# use.
+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 \
+    curl \
+    git
+apt-get clean
+EOF
+# NOTE: removed libsndfile1. Soundfile ships with its own version of libsndfile1
+
+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
+#!/usr/bin/env bash
+set -euo pipefail
+echo "${cuda_flavor}"
+SCRIPT
+cat >/usr/local/bin/nemo-install-cuda-python <<SCRIPT
+#!/usr/bin/env bash
+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
+# Container-only runtime utilities. Keep these out of pyproject.toml so they do
+# not become NeMo package dependencies.
+uv pip install \
+    dill \
+    orjson
+
+case "$(nemo-cuda-flavor)" in
+    cu12) torchcodec_index=https://download.pytorch.org/whl/cu126 ;;
+    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=v1.2.1
+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
+
+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 \
+        --no-deps \
+        /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.
+    cutlass_packages=(
+        "nvidia-cutlass-dsl==4.5.2"
+        "nvidia-cutlass-dsl-libs-base==4.5.2"
+    )
+    if [[ "$(nemo-cuda-flavor)" == "cu13" ]]; then
+        cutlass_packages+=("nvidia-cutlass-dsl-libs-cu13==4.5.2")
+    fi
+    uv pip install \
+        --no-deps \
+        "apache-tvm-ffi==0.1.11" \
+        "${cutlass_packages[@]}" \
+        "quack-kernels==0.5.0" \
+        "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
+
+FROM base-image
+COPY --from=automodel-deps /opt/venv /opt/venv
+
+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
+
+# NOTICES.txt file points to where the OSS source code is archived
+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