flash attn

2025-10-15 00:18:34 +09:00
parent 3b03453e5d
commit e51017897d
5 changed files with 69 additions and 52 deletions
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,4 +22,11 @@ dependencies = [
    "tqdm>=4.67.1",
    "transformers>=4.56.2",
    "wandb[media]>=0.22.0",
+    "flash-attn"
 ]
+
+[tool.uv.extra-build-dependencies]
+flash-attn = [{ requirement = "torch", match-runtime = true }]
+
+[tool.uv.extra-build-variables]
+flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }