This commit is contained in:
hailin 2025-09-16 19:38:25 +08:00
parent 4e38d1f0d2
commit 45c70451e2
1 changed files with 5 additions and 2 deletions

View File

@ -25,6 +25,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
libjpeg-dev libpng-dev ca-certificates && \
python3 -m pip install --no-cache-dir --upgrade pip wheel setuptools sympy pyyaml typing-extensions numpy
RUN python3 -m pip install --no-cache-dir numpy requests packaging build
WORKDIR /opt
RUN git clone --recursive -b v2.7.1 https://github.com/pytorch/pytorch.git
@ -39,6 +41,7 @@ RUN echo "Building PyTorch with USE_DISTRIBUTED=$USE_DISTRIBUTED" && \
ARG CUDA_VERSION=12.6.1
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS builder-extras
ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0"
ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8
RUN apt-get update && apt-get install -y --no-install-recommends \
@ -74,8 +77,8 @@ WORKDIR /opt/flashinfer
ENV FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0a"
# 先做 AOT 预编译,再直接打 wheel不隔离使用同一份自编 torch
RUN python3 -m flashinfer.aot && \
python3 -m pip install --no-cache-dir build && \
RUN python3 -m pip install --no-cache-dir numpy requests build && \
python3 -m flashinfer.aot && \
python3 -m build --no-isolation --wheel && \
ls -lh dist/