py-flash-attn: add v2.6.3 (#46487)

* [py-flash-attn] Add version 2.6.3

* Update dependencies according to the latest version

* Add max_jobs environmental variable to avoid oom error

---------

Co-authored-by: aurianer <8trash-can8@protonmail.ch>
This commit is contained in:
Auriane R. 2024-10-02 16:11:50 +02:00 committed by GitHub
parent b1707c2e3c
commit f3faeb0f77
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -20,6 +20,7 @@ class PyFlashAttn(PythonPackage):
license("BSD")
version("main", branch="main")
version("2.6.3", sha256="5bfae9500ad8e7d2937ebccb4906f3bc464d1bf66eedd0e4adabd520811c7b52")
version(
"2.5.9.post1", sha256="a92db1683a5b141a0f4371d251ae9f73e9aef629b3a58a50d0ef430266c68782"
)
@ -33,18 +34,21 @@ class PyFlashAttn(PythonPackage):
depends_on("cxx", type="build") # generated
with default_args(type="build"):
depends_on("py-ninja")
depends_on("py-packaging")
depends_on("py-psutil")
depends_on("py-setuptools")
depends_on("ninja")
with default_args(type=("build", "run")):
depends_on("py-torch+cuda")
depends_on("ninja")
depends_on("py-einops")
depends_on("py-packaging")
with default_args(type=("build", "link", "run")):
depends_on("py-pybind11")
depends_on("python@3.7:", type=("build", "run"))
depends_on("python@3.7:", type=("build", "run"), when="@:2.5")
depends_on("python@3.8:", type=("build", "run"), when="@2.6:")
def setup_build_environment(self, env):
# If oom error, try lowering the number of jobs with `spack install -j`
env.set("MAX_JOBS", make_jobs)