feat: execute FlashKDA on RTX 5090

2026-07-29 13:36:49 +08:00
parent 1146208b5e
commit 2ef846f751
16 changed files with 1189 additions and 51 deletions
@@ -0,0 +1,30 @@
+FROM nvidia/cuda:13.0.2-devel-ubuntu24.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+      build-essential \
+      git \
+      ninja-build \
+      python3 \
+      python3-dev \
+      python3-pip \
+      python3-venv \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN python3 -m venv /opt/flashkda \
+    && /opt/flashkda/bin/python -m pip install --upgrade \
+      pip \
+      setuptools \
+      wheel \
+      ninja \
+      numpy \
+      pytest \
+      --index-url https://pypi.tuna.tsinghua.edu.cn/simple \
+    && /opt/flashkda/bin/python -m pip install \
+      torch==2.11.0 \
+      --index-url https://download.pytorch.org/whl/cu130
+
+ENV PATH=/opt/flashkda/bin:${PATH}
+ENV CUDA_HOME=/usr/local/cuda