build: Ninja as a dependency in JIT mode (#659)

flashinfer-ai · Dec 13, 2024 · 16549dc · 16549dc
1 parent ed89d24
commit 16549dc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/custom_backend.py b/custom_backend.py
@@ -6,11 +6,10 @@
 
 
 def _get_requires_for_build():
-    return (
-        ["torch", "ninja"]
-        if os.environ.get("FLASHINFER_ENABLE_AOT", "0") == "1"
-        else ["ninja"]
-    )
+    requires = []
+    if os.environ.get("FLASHINFER_ENABLE_AOT", "0") == "1":
+        requires += ["torch", "ninja"]
+    return requires
 
 
 def get_requires_for_build_wheel(config_settings=None):

diff --git a/setup.py b/setup.py
@@ -83,7 +83,7 @@ def generate_cuda() -> None:
 
 ext_modules = []
 cmdclass = {}
-install_requires = ["torch"]
+install_requires = ["torch", "ninja"]
 generate_build_meta({})
 generate_cuda()