├── .github └── workflows │ ├── python-publish.yml │ └── test.yml ├── .gitignore ├── LICENSE ├── README.md ├── data ├── README.md └── enwik8.gz ├── fig2.png ├── native_sparse_attention_pytorch ├── __init__.py ├── compress_networks.py ├── native_sparse_attention.py ├── tensor_typing.py ├── transformer.py └── triton_native_sparse_attention.py ├── pyproject.toml ├── test_flex_masks.py ├── test_triton_nsa.py ├── tests ├── test_custom_compress_mlp.py └── test_sparse_attn.py └── train.py /.github/workflows/python-publish.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/.github/workflows/python-publish.yml -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/.github/workflows/test.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/README.md -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/data/README.md -------------------------------------------------------------------------------- /data/enwik8.gz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/data/enwik8.gz -------------------------------------------------------------------------------- /fig2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/fig2.png -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/__init__.py -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/compress_networks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/compress_networks.py -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/native_sparse_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/native_sparse_attention.py -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/tensor_typing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/tensor_typing.py -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/transformer.py -------------------------------------------------------------------------------- /native_sparse_attention_pytorch/triton_native_sparse_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/native_sparse_attention_pytorch/triton_native_sparse_attention.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/pyproject.toml -------------------------------------------------------------------------------- /test_flex_masks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/test_flex_masks.py -------------------------------------------------------------------------------- /test_triton_nsa.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/test_triton_nsa.py -------------------------------------------------------------------------------- /tests/test_custom_compress_mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/tests/test_custom_compress_mlp.py -------------------------------------------------------------------------------- /tests/test_sparse_attn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/tests/test_sparse_attn.py -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lucidrains/native-sparse-attention-pytorch/HEAD/train.py --------------------------------------------------------------------------------