diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 87dae3336..2ae094b7b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -155,14 +155,13 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install wheel pytest pytest-cov nvidia-pyindex
+          pip install -e "client/[test]"
+          pip install -e "server/[tensorrt]"
           {
-            python -m pip install torch==1.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
-            python -m pip install git+https://github.com/HazyResearch/flash-attention.git
+            pip install -e "server/[flash-attn]"
           } || {
             echo "flash attention was not installed."
           }
-          pip install -e "client/[test]"
-          pip install -e "server/[tensorrt]"
       - name: Test
         id: test
         run: |
diff --git a/server/setup.py b/server/setup.py
index 78275b0c1..faaff6cfb 100644
--- a/server/setup.py
+++ b/server/setup.py
@@ -59,6 +59,7 @@
         'tensorrt': ['nvidia-tensorrt'],
         'transformers': ['transformers>=4.16.2'],
         'search': ['annlite>=0.3.10'],
+        'flash-attn': ['flash-attn'],
     },
     classifiers=[
         'Development Status :: 5 - Production/Stable',