flash

server/pdm.lock

@@ -5,7 +5,7 @@
 groups = ["default"]
 strategy = ["cross_platform"]
 lock_version = "4.4"
-content_hash = "sha256:4c9d99229e31396fab6ad1bc7aaf336366423f7e144888cb4edb1a4d5dd6a7dc"
+content_hash = "sha256:d655d8ba97dd513f7d0f51e54a640d0045fc33d784c63ddc0c6cc8d36d981bbc"
 
 [[package]]
 name = "accelerate"
@@ -123,6 +123,16 @@ files = [
     {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"},
 ]
 
+[[package]]
+name = "einops"
+version = "0.7.0"
+requires_python = ">=3.8"
+summary = "A new flavour of deep learning operations"
+files = [
+    {file = "einops-0.7.0-py3-none-any.whl", hash = "sha256:0f3096f26b914f465f6ff3c66f5478f9a5e380bb367ffc6493a68143fbbf1fd1"},
+    {file = "einops-0.7.0.tar.gz", hash = "sha256:b2b04ad6081a3b227080c9bf5e3ace7160357ff03043cd66cc5b2319eb7031d1"},
+]
+
 [[package]]
 name = "fastapi"
 version = "0.109.0"
@@ -148,6 +158,21 @@ files = [
     {file = "filelock-3.13.1.tar.gz", hash = "sha256:521f5f56c50f8426f5e03ad3b281b490a87ef15bc6c526f168290f0c7148d44e"},
 ]
 
+[[package]]
+name = "flash-attn"
+version = "2.4.2"
+requires_python = ">=3.7"
+summary = ""
+dependencies = [
+    "einops",
+    "ninja",
+    "packaging",
+    "torch",
+]
+files = [
+    {file = "flash_attn-2.4.2.tar.gz", hash = "sha256:eb822a8c4219b610e9d734cbc8cd9ee4547f27433815a2b90dc1462766feefc1"},
+]
+
 [[package]]
 name = "fsspec"
 version = "2023.12.2"

server/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "transformers[torch]>=4.36.2",
     "ninja>=1.11.1.1",
     "packaging>=23.2",
+    "flash-attn>=2.4.2",
 ]
 requires-python = ">=3.11"
 readme = "README.md"