72 lines
3.0 KiB
Docker
72 lines
3.0 KiB
Docker
ARG BASE_IMAGE=nvcr.io/nvidia/nemo:26.06
|
|
FROM ${BASE_IMAGE}
|
|
|
|
ARG INSTALL_FLASH_ATTN4=1
|
|
ARG PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
|
|
ARG PIP_EXTRA_INDEX_URL=https://pypi.org/simple
|
|
ARG CUTLASS_DSL_LIBS_BASE_URL=https://mirrors.aliyun.com/pypi/packages/97/68/c1247ab848f26c4ab56e562eea0e3f31fc14c9aaf0d883afaa92d8f05592/nvidia_cutlass_dsl_libs_base-4.5.2-cp312-cp312-manylinux_2_28_x86_64.whl
|
|
ARG CUTLASS_DSL_LIBS_CU13_URL=https://mirrors.aliyun.com/pypi/packages/03/60/443e559139da15ab544761ac14f4206dffb981af48cc9856cd5b5b7cf0e7/nvidia_cutlass_dsl_libs_cu13-4.5.2-cp312-cp312-manylinux_2_28_x86_64.whl
|
|
|
|
ENV PIP_INDEX_URL=${PIP_INDEX_URL}
|
|
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
|
ENV PIP_DISABLE_PIP_VERSION_CHECK=1
|
|
|
|
RUN if [ "${INSTALL_FLASH_ATTN4}" = "1" ]; then \
|
|
/usr/bin/python3 -m pip uninstall -y flash-attn flash-attn-3 flash-attn-4 || true; \
|
|
python3 -m pip uninstall -y flash-attn flash-attn-3 flash-attn-4 || true; \
|
|
MAX_JOBS="${MAX_JOBS:-16}" python3 -m pip install --pre --no-cache-dir --no-build-isolation \
|
|
"${CUTLASS_DSL_LIBS_BASE_URL}" "${CUTLASS_DSL_LIBS_CU13_URL}" \
|
|
"nvidia-cutlass-dsl==4.5.2" "flash-attn-4==4.0.0b11"; \
|
|
fi
|
|
|
|
RUN set -eux; \
|
|
python3 - <<'PY'
|
|
from pathlib import Path
|
|
import shutil
|
|
|
|
site = Path("/usr/local/lib/python3.12/dist-packages")
|
|
src = site / "vllm/third_party/triton_kernels"
|
|
dst = site / "triton_kernels"
|
|
if not (src / "matmul_ogs.py").exists() or not (src / "tensor.py").exists():
|
|
raise SystemExit(f"missing vLLM triton_kernels source: {src}")
|
|
if not dst.exists():
|
|
raise SystemExit(f"missing top-level triton_kernels package: {dst}")
|
|
backup = site / "triton_kernels_nv26_backup"
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
shutil.copytree(dst, backup, ignore=shutil.ignore_patterns("__pycache__"))
|
|
shutil.rmtree(dst)
|
|
shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__"))
|
|
PY
|
|
|
|
RUN python3 - <<'PY'
|
|
import importlib.metadata as metadata
|
|
import importlib.util
|
|
import flash_attn.cute.interface
|
|
required = ["torch", "transformer-engine", "megatron-core", "megatron-bridge"]
|
|
for package in required:
|
|
print(package, metadata.version(package))
|
|
expected = {
|
|
"flash-attn-4": "4.0.0b11",
|
|
"nvidia-cutlass-dsl": "4.5.2",
|
|
"nvidia-cutlass-dsl-libs-base": "4.5.2",
|
|
"nvidia-cutlass-dsl-libs-cu13": "4.5.2",
|
|
}
|
|
for package, expected_version in expected.items():
|
|
actual_version = metadata.version(package)
|
|
if actual_version != expected_version:
|
|
raise SystemExit(f"{package} version mismatch: {actual_version} != {expected_version}")
|
|
print(package, actual_version)
|
|
for legacy_package in ["flash-attn", "flash-attn-3"]:
|
|
try:
|
|
version = metadata.version(legacy_package)
|
|
except metadata.PackageNotFoundError:
|
|
continue
|
|
raise SystemExit(f"legacy package still installed: {legacy_package} {version}")
|
|
spec = importlib.util.find_spec("triton_kernels.matmul_ogs")
|
|
if spec is None:
|
|
raise SystemExit("triton_kernels.matmul_ogs MISSING")
|
|
print("triton_kernels.matmul_ogs", spec.origin)
|
|
print("flash_attn.cute.interface import OK")
|
|
PY
|