├── .github ├── scripts │ ├── build.sh │ ├── check_for_ngc_images.sh │ └── test.sh └── workflows │ ├── _build.yml │ ├── _build_in_container.yml │ ├── build.yml │ ├── build_in_container.yml │ └── publish.yaml ├── .gitignore ├── .gitmodules ├── AUTHORS ├── LICENSE ├── MANIFEST.in ├── README.md ├── assets ├── selection.png └── ssd_algorithm.png ├── benchmarks └── benchmark_generation_mamba_simple.py ├── csrc └── selective_scan │ ├── reverse_scan.cuh │ ├── selective_scan.cpp │ ├── selective_scan.h │ ├── selective_scan_bwd_bf16_complex.cu │ ├── selective_scan_bwd_bf16_real.cu │ ├── selective_scan_bwd_fp16_complex.cu │ ├── selective_scan_bwd_fp16_real.cu │ ├── selective_scan_bwd_fp32_complex.cu │ ├── selective_scan_bwd_fp32_real.cu │ ├── selective_scan_bwd_kernel.cuh │ ├── selective_scan_common.h │ ├── selective_scan_fwd_bf16.cu │ ├── selective_scan_fwd_fp16.cu │ ├── selective_scan_fwd_fp32.cu │ ├── selective_scan_fwd_kernel.cuh │ ├── static_switch.h │ └── uninitialized_copy.cuh ├── evals └── lm_harness_eval.py ├── mamba_ssm ├── __init__.py ├── distributed │ ├── __init__.py │ ├── distributed_utils.py │ └── tensor_parallel.py ├── models │ ├── __init__.py │ ├── config_mamba.py │ └── mixer_seq_simple.py ├── modules │ ├── __init__.py │ ├── block.py │ ├── mamba2.py │ ├── mamba2_simple.py │ ├── mamba_simple.py │ ├── mha.py │ ├── mlp.py │ └── ssd_minimal.py ├── ops │ ├── __init__.py │ ├── selective_scan_interface.py │ └── triton │ │ ├── __init__.py │ │ ├── k_activations.py │ │ ├── layer_norm.py │ │ ├── layernorm_gated.py │ │ ├── selective_state_update.py │ │ ├── softplus.py │ │ ├── ssd_bmm.py │ │ ├── ssd_chunk_scan.py │ │ ├── ssd_chunk_state.py │ │ ├── ssd_combined.py │ │ └── ssd_state_passing.py └── utils │ ├── __init__.py │ ├── generation.py │ ├── hf.py │ └── torch.py ├── pyproject.toml ├── rocm_patch └── rocm6_0.patch ├── setup.py ├── tests ├── ops │ ├── test_selective_scan.py │ └── triton │ │ ├── test_layernorm_gated.py │ │ ├── test_selective_state_update.py │ │ └── test_ssd.py └── test_generation.py └── usage.md /.github/scripts/build.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/scripts/build.sh -------------------------------------------------------------------------------- /.github/scripts/check_for_ngc_images.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/scripts/check_for_ngc_images.sh -------------------------------------------------------------------------------- /.github/scripts/test.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/scripts/test.sh -------------------------------------------------------------------------------- /.github/workflows/_build.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/workflows/_build.yml -------------------------------------------------------------------------------- /.github/workflows/_build_in_container.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/workflows/_build_in_container.yml -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/workflows/build.yml -------------------------------------------------------------------------------- /.github/workflows/build_in_container.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/workflows/build_in_container.yml -------------------------------------------------------------------------------- /.github/workflows/publish.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.github/workflows/publish.yaml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/.gitmodules -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | Tri Dao, tri@tridao.me 2 | Albert Gu, agu@andrew.cmu.edu 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/LICENSE -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/MANIFEST.in -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/README.md -------------------------------------------------------------------------------- /assets/selection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/assets/selection.png -------------------------------------------------------------------------------- /assets/ssd_algorithm.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/assets/ssd_algorithm.png -------------------------------------------------------------------------------- /benchmarks/benchmark_generation_mamba_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/benchmarks/benchmark_generation_mamba_simple.py -------------------------------------------------------------------------------- /csrc/selective_scan/reverse_scan.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/reverse_scan.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan.cpp -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan.h -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_bf16_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_bf16_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_bf16_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_bf16_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp16_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp16_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp16_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp16_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp32_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp32_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp32_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp32_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_kernel.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_bwd_kernel.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_common.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_common.h -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_bf16.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_fwd_bf16.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_fp16.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_fwd_fp16.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_fp32.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_fwd_fp32.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_kernel.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/selective_scan_fwd_kernel.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/static_switch.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/static_switch.h -------------------------------------------------------------------------------- /csrc/selective_scan/uninitialized_copy.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/csrc/selective_scan/uninitialized_copy.cuh -------------------------------------------------------------------------------- /evals/lm_harness_eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/evals/lm_harness_eval.py -------------------------------------------------------------------------------- /mamba_ssm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/__init__.py -------------------------------------------------------------------------------- /mamba_ssm/distributed/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/distributed/distributed_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/distributed/distributed_utils.py -------------------------------------------------------------------------------- /mamba_ssm/distributed/tensor_parallel.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/distributed/tensor_parallel.py -------------------------------------------------------------------------------- /mamba_ssm/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/models/config_mamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/models/config_mamba.py -------------------------------------------------------------------------------- /mamba_ssm/models/mixer_seq_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/models/mixer_seq_simple.py -------------------------------------------------------------------------------- /mamba_ssm/modules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/modules/block.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/block.py -------------------------------------------------------------------------------- /mamba_ssm/modules/mamba2.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/mamba2.py -------------------------------------------------------------------------------- /mamba_ssm/modules/mamba2_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/mamba2_simple.py -------------------------------------------------------------------------------- /mamba_ssm/modules/mamba_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/mamba_simple.py -------------------------------------------------------------------------------- /mamba_ssm/modules/mha.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/mha.py -------------------------------------------------------------------------------- /mamba_ssm/modules/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/mlp.py -------------------------------------------------------------------------------- /mamba_ssm/modules/ssd_minimal.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/modules/ssd_minimal.py -------------------------------------------------------------------------------- /mamba_ssm/ops/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/ops/selective_scan_interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/selective_scan_interface.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/k_activations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/k_activations.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/layer_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/layer_norm.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/layernorm_gated.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/layernorm_gated.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/selective_state_update.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/selective_state_update.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/softplus.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/softplus.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/ssd_bmm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/ssd_bmm.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/ssd_chunk_scan.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/ssd_chunk_scan.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/ssd_chunk_state.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/ssd_chunk_state.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/ssd_combined.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/ssd_combined.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/ssd_state_passing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/ops/triton/ssd_state_passing.py -------------------------------------------------------------------------------- /mamba_ssm/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/utils/generation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/utils/generation.py -------------------------------------------------------------------------------- /mamba_ssm/utils/hf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/utils/hf.py -------------------------------------------------------------------------------- /mamba_ssm/utils/torch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/mamba_ssm/utils/torch.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/pyproject.toml -------------------------------------------------------------------------------- /rocm_patch/rocm6_0.patch: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/rocm_patch/rocm6_0.patch -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/setup.py -------------------------------------------------------------------------------- /tests/ops/test_selective_scan.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/tests/ops/test_selective_scan.py -------------------------------------------------------------------------------- /tests/ops/triton/test_layernorm_gated.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/tests/ops/triton/test_layernorm_gated.py -------------------------------------------------------------------------------- /tests/ops/triton/test_selective_state_update.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/tests/ops/triton/test_selective_state_update.py -------------------------------------------------------------------------------- /tests/ops/triton/test_ssd.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/tests/ops/triton/test_ssd.py -------------------------------------------------------------------------------- /tests/test_generation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/tests/test_generation.py -------------------------------------------------------------------------------- /usage.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/state-spaces/mamba/HEAD/usage.md --------------------------------------------------------------------------------