diff --git a/build/torch210-cxx11-cu126-x86_64-linux/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index 15f1a1e591cc8ba101f4d8e2beb5cdde593c9f91..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:ed029d363b3bb2149788b3577e7b0a7c5783d7ce40d5ad0de45679b5cdfe5484 -size 444257200 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_muomt3evauf7u.abi3.so b/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_muomt3evauf7u.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..0128c8886b22bde2eed9fbab1226d4039400fd69 --- /dev/null +++ b/build/torch210-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_muomt3evauf7u.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8eb110259b4313446e301660d897825a591775c7f9ed9853970b2fd6358e1eb9 +size 444257232 diff --git a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py b/build/torch210-cxx11-cu126-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..3eeddce9d97b6bcf003d246542f788262bbde456 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_muomt3evauf7u +ops = torch.ops._mamba_ssm_cuda_muomt3evauf7u def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_muomt3evauf7u::{op_name}" diff --git a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json b/build/torch210-cxx11-cu126-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu126-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/models/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/models/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/modules/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/modules/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu126-x86_64-linux/utils/generation.py b/build/torch210-cxx11-cu126-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch210-cxx11-cu126-x86_64-linux/utils/generation.py +++ b/build/torch210-cxx11-cu126-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: diff --git a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index 98aa7268367924fbb72ab8163437d2a79534c9f2..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b0997257c9217211de7446898e89756ab9a6237989ea6817f3bb018abde9df3c -size 610662240 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_p3tnkmr4mmukm.abi3.so b/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_p3tnkmr4mmukm.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..8162d110f8d0b876a47ff580c61a48c7c71e9d32 --- /dev/null +++ b/build/torch210-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_p3tnkmr4mmukm.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81bd32914c56686dc50c48ed195db1540a9a90658c9dbcb0355f995e90dfe180 +size 610662272 diff --git a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..60c1bdf884f37ec7ac9c2198695be12644792b95 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_p3tnkmr4mmukm +ops = torch.ops._mamba_ssm_cuda_p3tnkmr4mmukm def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_p3tnkmr4mmukm::{op_name}" diff --git a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu128-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/models/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/models/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/modules/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/modules/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu128-x86_64-linux/utils/generation.py b/build/torch210-cxx11-cu128-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch210-cxx11-cu128-x86_64-linux/utils/generation.py +++ b/build/torch210-cxx11-cu128-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: diff --git a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index 2683c1fb7f2e7c7007a9b9e6821acef8a179fb40..0000000000000000000000000000000000000000 --- a/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:722c9a77eedb9a593885db984a9fc3a2e8bcea0e1c769a1b79c6f1312a68dddb -size 446637448 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_yndcxrj64jqjy.abi3.so b/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_yndcxrj64jqjy.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..5dd4c138c3c3ca52a9aa0c7773797947a048c300 --- /dev/null +++ b/build/torch210-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_yndcxrj64jqjy.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f1416b3f9ce1d7c2a33874baa9fca4de03305f693aff3aa952f990943002d3d7 +size 446637480 diff --git a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..4ff87b2273e3ec260e5cf17dd246f9efb9ca27ed 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_yndcxrj64jqjy +ops = torch.ops._mamba_ssm_cuda_yndcxrj64jqjy def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_yndcxrj64jqjy::{op_name}" diff --git a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch210-cxx11-cu130-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/models/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/models/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/modules/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/modules/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch210-cxx11-cu130-x86_64-linux/utils/generation.py b/build/torch210-cxx11-cu130-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch210-cxx11-cu130-x86_64-linux/utils/generation.py +++ b/build/torch210-cxx11-cu130-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: diff --git a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/__init__.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index 487a1cc82964c6c129c40fcfc2b0e05dcaa64c6f..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5a243449941f1718d8fb655cdeb3a41dafc4984d7717ac774a79f6dbcc215d20 -size 444242136 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_lknozciqovp5m.abi3.so b/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_lknozciqovp5m.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..3016a4b6c962e5b79c630feb7a18240a29a8f185 --- /dev/null +++ b/build/torch29-cxx11-cu126-x86_64-linux/_mamba_ssm_cuda_lknozciqovp5m.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81da74fdca330af5de6fdbfed5ee56c7acb25cd8799f2db37fbd58690f80c810 +size 444242168 diff --git a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..805ec77446e2a26ac33d6afbb2d2a1be54d877d9 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_lknozciqovp5m +ops = torch.ops._mamba_ssm_cuda_lknozciqovp5m def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_lknozciqovp5m::{op_name}" diff --git a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu126-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/models/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/models/__init__.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/modules/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/modules/__init__.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/ops/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/ops/__init__.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu126-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/ops/triton/__init__.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu126-x86_64-linux/utils/generation.py b/build/torch29-cxx11-cu126-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch29-cxx11-cu126-x86_64-linux/utils/generation.py +++ b/build/torch29-cxx11-cu126-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: diff --git a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/__init__.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index 83928a1de3528800100c07bf70da70055cc271bb..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:32a752606508348bbe26d3095a81c6b26d8bf1d89528e5661d563a5efdc119de -size 610651464 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_eomlxkg5rvjoa.abi3.so b/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_eomlxkg5rvjoa.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..ad0181a99c5c36bed54942b45fa08c85d9651be1 --- /dev/null +++ b/build/torch29-cxx11-cu128-x86_64-linux/_mamba_ssm_cuda_eomlxkg5rvjoa.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fb2c0f3bad94842e6d43cf37fc651e9b31c5240d48898b598501f0dc27fe9930 +size 610651496 diff --git a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..9f030b8cccd9a0f086f71d828949c84daaaa2812 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_eomlxkg5rvjoa +ops = torch.ops._mamba_ssm_cuda_eomlxkg5rvjoa def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_eomlxkg5rvjoa::{op_name}" diff --git a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu128-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/models/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/models/__init__.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/modules/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/modules/__init__.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/ops/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/ops/__init__.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu128-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/ops/triton/__init__.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu128-x86_64-linux/utils/generation.py b/build/torch29-cxx11-cu128-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch29-cxx11-cu128-x86_64-linux/utils/generation.py +++ b/build/torch29-cxx11-cu128-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: diff --git a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py index a767386abca1cf69349f8777984014ba0c00ae9d..c209429f2f64022865f5ebea75a87a5eda2625cd 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/__init__.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/__init__.py @@ -1,14 +1,12 @@ __version__ = "2.2.4" +from .ops import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined from .ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn -from .modules.mamba_simple import Mamba -from .modules.mamba2 import Mamba2 -from .models.mixer_seq_simple import MambaLMHeadModel +from .modules import Mamba, Mamba2 +from .models import MambaLMHeadModel +falcon_mamba_inner_fn = mamba_inner_fn __all__ = [ - "selective_scan_fn", - "mamba_inner_fn", - "Mamba", - "Mamba2", - "MambaLMHeadModel", + "selective_scan_fn", "mamba_inner_fn", "falcon_mamba_inner_fn", "selective_state_update", "mamba_chunk_scan_combined", "mamba_split_conv1d_scan_combined", + "Mamba", "Mamba2", "MambaLMHeadModel", ] diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so deleted file mode 100644 index bca0d4b1afda2479a2808466bcfcc311f34bd8a3..0000000000000000000000000000000000000000 --- a/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_1a63054.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:85c1aec6878d012e5de30321bdc4bb72b3c41caff53198c0fe65d0f7442212fd -size 446626672 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_szbxxtsgftczo.abi3.so b/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_szbxxtsgftczo.abi3.so new file mode 100644 index 0000000000000000000000000000000000000000..840e6cf3e9abcdd44edc3b453476925303ce78f1 --- /dev/null +++ b/build/torch29-cxx11-cu130-x86_64-linux/_mamba_ssm_cuda_szbxxtsgftczo.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4b5cf96ebc7f14251217da815b24e44a271503a86b1ad4a96ff016a62dad96b0 +size 446626712 diff --git a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py index f980cce99cdf1b9a9d08fed6e7bb8bd4e09d1a7a..45033ee0dc3ecb80bdaf41e170ca58cb70b71ac8 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/_ops.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _mamba_ssm_1a63054 -ops = torch.ops._mamba_ssm_1a63054 +from . import _mamba_ssm_cuda_szbxxtsgftczo +ops = torch.ops._mamba_ssm_cuda_szbxxtsgftczo def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_mamba_ssm_1a63054::{op_name}" \ No newline at end of file + return f"_mamba_ssm_cuda_szbxxtsgftczo::{op_name}" diff --git a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json index 9cf5deed9898dce769f4cc73913d3530b92a0bd8..ad205a907d250bdf72167cb8e31fb02ea92dd12d 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/metadata.json +++ b/build/torch29-cxx11-cu130-x86_64-linux/metadata.json @@ -1,4 +1,6 @@ { "version": 1, - "python-depends": [] + "python-depends": [ + "einops" + ] } \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/models/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/models/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..c64f5ff98f8579a185159d3b2cc63ebb5ebca7b2 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/models/__init__.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/models/__init__.py @@ -0,0 +1,5 @@ +from .mixer_seq_simple import MambaLMHeadModel + +__all__ = [ + "MambaLMHeadModel", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/modules/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/modules/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..ad51e45c79edd194726a6e515987377feea5f650 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/modules/__init__.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/modules/__init__.py @@ -0,0 +1,7 @@ +from .mamba_simple import Mamba +from .mamba2 import Mamba2 + +__all__ = [ + "Mamba", + "Mamba2", +] \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/ops/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/ops/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1e994c1801667556fe746e2081fab2b082a4b2cd 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/ops/__init__.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/ops/__init__.py @@ -0,0 +1 @@ +from .triton import selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/ops/triton/__init__.py b/build/torch29-cxx11-cu130-x86_64-linux/ops/triton/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b210bb0fb69b3592633f08c353f36fea4ec1dd74 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/ops/triton/__init__.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/ops/triton/__init__.py @@ -0,0 +1,2 @@ +from .selective_state_update import selective_state_update +from .ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined \ No newline at end of file diff --git a/build/torch29-cxx11-cu130-x86_64-linux/utils/generation.py b/build/torch29-cxx11-cu130-x86_64-linux/utils/generation.py index 330672afe2f0021d09c8913fb42421996714a8d3..873210ce4b1db4fdde3dc3f794b4e86efc3e8741 100644 --- a/build/torch29-cxx11-cu130-x86_64-linux/utils/generation.py +++ b/build/torch29-cxx11-cu130-x86_64-linux/utils/generation.py @@ -11,8 +11,7 @@ import torch.nn.functional as F from einops import rearrange, repeat from torch import Tensor from torch.profiler import ProfilerActivity, profile, record_function -from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput, TextStreamer - +from transformers.generation import GenerateDecoderOnlyOutput, TextStreamer @dataclass class InferenceParams: