😇

uv で `pyproject.toml` になるべく情報を残しつつ flash-attention をインストールする

2024/11/14に公開

最近 uv の公式ドキュメントに flash-attention をインストールするための方法が追加されました。
https://docs.astral.sh/uv/concepts/projects/#build-isolation

具体的には以下のような pyproject.toml を用意します。

[project]
name = "uv-flash-attn"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = []

[project.optional-dependencies]
build = ["torch", "setuptools"]
compile = ["flash-attn"]

[tool.uv]
no-build-isolation-package = ["flash-attn"]

あとは

uv sync --extra build
uv sync --extra build --extra compile

と2段階に分けて uv sync するというものです。

できることなら uv sync 一発で完結してほしいものですが、インストールするパッケージの順番を指定できない以上それは難しそうです。
とは言っても、愚直に uv add で先に pytorch をインストールした後で uv sync を実行するよりは、依存パッケージのすべてを pyproject.toml に載せられるのはメリットのように思います。

Discussion