make properly pip installable without needing build isolation (uv pip install now works too) (#402)

Co-authored-by: Tri Dao <tridao@users.noreply.github.com>

Author: kszlim

README.md

@@ -21,6 +21,8 @@ with an efficient hardware-aware design and implementation in the spirit of [Fla
 
 - [Option] `pip install causal-conv1d>=1.4.0`: an efficient implementation of a simple causal Conv1d layer used inside the Mamba block.
 - `pip install mamba-ssm`: the core Mamba package.
+- `pip install mamba-ssm[causal-conv1d]`: To install core Mamba package and causal-conv1d.
+- `pip install mamba-ssm[dev]`: To install core Mamba package and dev depdencies.
 
 It can also be built from source with `pip install .` from this repository.
 

pyproject.toml

@@ -0,0 +1,46 @@
+[project]
+name = "mamba_ssm"
+description = "Mamba state-space model"
+readme = "README.md"
+authors = [
+    { name = "Tri Dao", email = "tri@tridao.me" },
+    { name = "Albert Gu", email = "agu@cs.cmu.edu" }
+]
+requires-python = ">= 3.7"
+dynamic = ["version"]
+license = { file = "LICENSE" }  # Include a LICENSE file in your repo
+keywords = ["cuda", "pytorch", "state-space model"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: BSD License",
+    "Operating System :: Unix"
+]
+dependencies = [
+    "torch",
+    "ninja",
+    "einops",
+    "triton",
+    "transformers",
+    "packaging",
+    "setuptools>=61.0.0",
+]
+urls = { name = "Repository", url = "https://github.com/state-spaces/mamba"}
+
+[project.optional-dependencies]
+causal-conv1d = [
+    "causal-conv1d>=1.2.0"
+]
+dev = [
+    "pytest"
+]
+
+
+[build-system]
+requires = [
+    "setuptools>=61.0.0",
+    "wheel",
+    "torch",
+    "packaging",
+    "ninja",
+]
+build-backend = "setuptools.build_meta"

setup.py

@@ -19,7 +19,6 @@
 import torch
 from torch.utils.cpp_extension import (
     BuildExtension,
-    CppExtension,
     CUDAExtension,
     CUDA_HOME,
     HIP_HOME
@@ -349,31 +348,13 @@ def run(self):
             "mamba_ssm.egg-info",
         )
     ),
-    author="Tri Dao, Albert Gu",
-    author_email="tri@tridao.me, agu@cs.cmu.edu",
-    description="Mamba state-space model",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    url="https://github.com/state-spaces/mamba",
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: BSD License",
-        "Operating System :: Unix",
-    ],
+    
     ext_modules=ext_modules,
     cmdclass={"bdist_wheel": CachedWheelsCommand, "build_ext": BuildExtension}
     if ext_modules
     else {
         "bdist_wheel": CachedWheelsCommand,
-    },
-    python_requires=">=3.8",
-    install_requires=[
-        "torch",
-        "packaging",
-        "ninja",
-        "einops",
-        "triton",
-        "transformers",
-        # "causal_conv1d>=1.4.0",
-    ],
+    }
 )