Bump FlashAttn version and add deterministic option for FAv2 (#585)

* Deterministic FA, bump minimum supported version Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * fixes Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Fix MQA/GQA Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> * Address review comments Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> --------- Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
NVIDIA · Jan 6, 2024 · f2bd53c · f2bd53c
1 parent e2a7531
commit f2bd53c
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 92 deletions.
diff --git a/setup.py b/setup.py
@@ -284,7 +284,7 @@ def add_unique(l: List[str], vals: Union[str, List[str]]) -> None:
 
     # Framework-specific requirements
     if "pytorch" in frameworks():
-        add_unique(install_reqs, ["torch", "flash-attn>=1.0.6,<=2.3.3,!=2.0.9,!=2.1.0"])
+        add_unique(install_reqs, ["torch", "flash-attn>=2.0.6,<=2.4.2,!=2.0.9,!=2.1.0"])
         add_unique(test_reqs, ["numpy", "onnxruntime", "torchvision"])
     if "jax" in frameworks():
         if not found_pybind11():