|
@@ -48,19 +48,17 @@ RUN rm -f /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources \
|
|
|
FROM apt-base AS torch-base
|
|
FROM apt-base AS torch-base
|
|
|
ARG PYPI_INDEX_URL
|
|
ARG PYPI_INDEX_URL
|
|
|
ARG PYTORCH_INDEX_URL
|
|
ARG PYTORCH_INDEX_URL
|
|
|
-ARG FLASHINFER_WHL_INDEX
|
|
|
|
|
ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION
|
|
ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION
|
|
|
ARG FLASHINFER_VERSION
|
|
ARG FLASHINFER_VERSION
|
|
|
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
|
|
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
|
|
|
uv pip install \
|
|
uv pip install \
|
|
|
--index-url ${PYPI_INDEX_URL} \
|
|
--index-url ${PYPI_INDEX_URL} \
|
|
|
--extra-index-url ${PYTORCH_INDEX_URL} \
|
|
--extra-index-url ${PYTORCH_INDEX_URL} \
|
|
|
- --extra-index-url ${FLASHINFER_WHL_INDEX} \
|
|
|
|
|
--index-strategy unsafe-best-match \
|
|
--index-strategy unsafe-best-match \
|
|
|
"torch==${TORCH_VERSION}" \
|
|
"torch==${TORCH_VERSION}" \
|
|
|
"torchvision==${TORCHVISION_VERSION}" \
|
|
"torchvision==${TORCHVISION_VERSION}" \
|
|
|
"torchaudio==${TORCHAUDIO_VERSION}" \
|
|
"torchaudio==${TORCHAUDIO_VERSION}" \
|
|
|
- "flashinfer==${FLASHINFER_VERSION}"
|
|
|
|
|
|
|
+ "flashinfer-python==${FLASHINFER_VERSION}"
|
|
|
|
|
|
|
|
# The fork's r12 composition: base branch + sequential merge of the open
|
|
# The fork's r12 composition: base branch + sequential merge of the open
|
|
|
# PR heads (ascending). The PR head pseudo-refs are GitHub's stable refs.
|
|
# PR heads (ascending). The PR head pseudo-refs are GitHub's stable refs.
|