diff --git a/.devcontainer/rocm/Dockerfile b/.devcontainer/rocm/Dockerfile index 0f0fbe84cd..392a1fc7e0 100644 --- a/.devcontainer/rocm/Dockerfile +++ b/.devcontainer/rocm/Dockerfile @@ -80,10 +80,17 @@ RUN \ echo 'eval "$(micromamba shell hook --shell bash)"' >> /home/$USERNAME/.bashrc && \ echo "micromamba activate ${MAMBA_ENV_NAME}" >> /home/$USERNAME/.bashrc -# Create a new micromamba env and install needed packages for flashinfer development +# Create a new micromamba env and install needed packages for flashinfer development. +# torch uses -f/--find-links (not --index-url) because the radeon repo is a flat +# wheel listing, not a PEP 503 simple index — with --index-url pip requests the +# per-package path (.../torch/), which 404s, and the install fails with "No +# matching distribution found". After installing, assert torch is a ROCm/HIP build +# (torch.version.hip is not None) so the build fails fast if -f ever resolves a +# PyPI CPU/CUDA wheel instead of the radeon ROCm wheel. RUN /bin/micromamba create -n ${MAMBA_ENV_NAME} python=${PY_VERSION} gtest gmock bash-completion -c conda-forge && \ /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --no-cache-dir cmake ninja "scikit-build-core>=0.4.3" "setuptools-scm>=9.2" pre-commit numpy pytest pytest-cov pytest-xdist pytest-rerunfailures pybind11 ruff && \ - /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/ && \ + /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install --no-cache-dir torch==${TORCH_VERSION} -f https://repo.radeon.com/rocm/manylinux/rocm-rel-${ROCM_VERSION}/ && \ + /bin/micromamba run -n ${MAMBA_ENV_NAME} python -c "import torch, sys; hip = torch.version.hip; print(f'torch {torch.__version__} (hip={hip})'); sys.exit('ERROR: installed torch has no ROCm/HIP build (hip is None)' if hip is None else 0)" && \ /bin/micromamba run -n ${MAMBA_ENV_NAME} pip install amd_aiter==${AITER_VERSION} --extra-index-url https://pypi.amd.com/rocm-${AITER_ROCM_VERSION}/simple SHELL ["/usr/local/bin/_dockerfile_shell.sh"] diff --git a/README.md b/README.md index 1201942203..5b06a92117 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ Install the matching ROCm-enabled PyTorch wheel from : ```bash -pip install torch==2.9.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ +pip install torch==2.9.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ ``` Other versions may work but have not been tested. Replace `7.2` with the @@ -141,11 +141,14 @@ pip install amd-flashinfer --index-url https://pypi.amd.com/simple/ Install the matching ROCm-enabled torch package from : ```bash -pip install torch==2.9.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ +pip install torch==2.9.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/ ``` -**NOTE:** Use `--index-url` (not `-f`) so pip cannot silently fall back -to a CPU-only PyPI wheel. +**NOTE:** The radeon repo is a flat wheel listing, not a PEP 503 index, so +`--index-url` fails with "No matching distribution found" — use `-f` +(`--find-links`). pip still prefers the ROCm wheel over a same-version PyPI +wheel because its `+rocm` local version ranks higher. Confirm the result +with `python -c "import torch; assert torch.version.hip, 'not a ROCm build'"`. ### Trying the Examples diff --git a/flashinfer/hip_utils.py b/flashinfer/hip_utils.py index c28aa4f11f..7868377816 100644 --- a/flashinfer/hip_utils.py +++ b/flashinfer/hip_utils.py @@ -536,7 +536,7 @@ def check_torch_rocm_compatibility() -> None: " 1. Uninstall current PyTorch:\n" " pip uninstall torch\n\n" " 2. Install PyTorch for ROCm:\n" - " pip install torch==2.7.1 --index-url https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/\n\n" + " pip install torch==2.7.1 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/\n\n" "See https://github.com/rocm/flashinfer for detailed installation instructions.\n" + "=" * 70 @@ -559,11 +559,11 @@ def check_torch_rocm_compatibility() -> None: f" PyTorch ROCm version: {torch_rocm_major_minor}\n\n" f"This may cause runtime errors or crashes.\n\n" f"To fix, reinstall PyTorch for your ROCm version:\n" - f" pip install torch==2.7.1 --index-url " + f" pip install torch==2.7.1 -f " f"https://repo.radeon.com/rocm/manylinux/rocm-rel-{system_rocm}/\n\n" f"Or if using uv:\n" f" export FLASHINFER_ROCM_VERSION={system_rocm}\n" - f" uv pip install torch==2.7.1 --index-url " + f" uv pip install torch==2.7.1 -f " f"https://repo.radeon.com/rocm/manylinux/rocm-rel-{system_rocm}/\n" f"{'=' * 70}", RuntimeWarning,