Revert "Merge branch 'state-spaces:main' into feat/add-cu_seqlens"

This reverts commit aeb4d3b, reversing
changes made to cda4b5a.

Author: zigzagcai

.github/workflows/publish.yaml

@@ -44,7 +44,7 @@ jobs:
           # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
           os: [ubuntu-20.04]
           python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
-          torch-version: ['2.0.1', '2.1.2', '2.2.2', '2.3.1', '2.4.0']
+          torch-version: ['2.0.1', '2.1.2', '2.2.2', '2.3.1', '2.4.0.dev20240505']
           cuda-version: ['11.8.0', '12.2.2']
           # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
           # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
@@ -117,7 +117,7 @@ jobs:
           # This code is ugly, maybe there's a better way to do this.
           export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \
             minv = {'2.0': 117, '2.1': 118, '2.2': 118, '2.3': 118, '2.4': 118}[env['MATRIX_TORCH_VERSION']]; \
-            maxv = {'2.0': 118, '2.1': 121, '2.2': 121, '2.3': 121, '2.4': 124}[env['MATRIX_TORCH_VERSION']]; \
+            maxv = {'2.0': 118, '2.1': 121, '2.2': 121, '2.3': 121, '2.4': 121}[env['MATRIX_TORCH_VERSION']]; \
             print(max(min(int(env['MATRIX_CUDA_VERSION']), maxv), minv))" \
           )
           if [[ ${{ matrix.torch-version }} == *"dev"* ]]; then

README.md

@@ -21,8 +21,6 @@ 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.
 

mamba_ssm/modules/mha.py

@@ -180,8 +180,7 @@ def _update_kvcache_attention(self, q, kv, inference_params):
             ).transpose(1, 2)
         else:
             batch = q.shape[0]
-            kv_cache, _ = inference_params.key_value_memory_dict[self.layer_idx]
-            kv_cache = kv_cache[:batch]
+            kv_cache = inference_params.key_value_memory_dict[self.layer_idx][:batch]
             cache_seqlens = (
                 inference_params.lengths_per_sample[:batch]
                 if inference_params.lengths_per_sample is not None

pyproject.toml

@@ -19,6 +19,7 @@
 import torch
 from torch.utils.cpp_extension import (
     BuildExtension,
+    CppExtension,
     CUDAExtension,
     CUDA_HOME,
     HIP_HOME
@@ -348,13 +349,31 @@ 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",
+    ],
 )