{"id":31135991,"url":"https://github.com/sweep76/sparse-attention","last_synced_at":"2025-10-13T15:38:17.589Z","repository":{"id":300966395,"uuid":"1007729372","full_name":"Sweep76/sparse-attention","owner":"Sweep76","description":null,"archived":false,"fork":false,"pushed_at":"2025-07-06T11:58:43.000Z","size":80,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2025-09-18T07:54:15.357Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Sweep76.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-06-24T12:46:46.000Z","updated_at":"2025-07-06T11:58:46.000Z","dependencies_parsed_at":null,"dependency_job_id":"19605e1f-1e34-40c2-a452-7125ec0f82c6","html_url":"https://github.com/Sweep76/sparse-attention","commit_stats":null,"previous_names":["sweep76/sparse-attention"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/Sweep76/sparse-attention","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Sweep76%2Fsparse-attention","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Sweep76%2Fsparse-attention/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Sweep76%2Fsparse-attention/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Sweep76%2Fsparse-attention/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Sweep76","download_url":"https://codeload.github.com/Sweep76/sparse-attention/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Sweep76%2Fsparse-attention/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279015933,"owners_count":26085777,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-13T02:00:06.723Z","response_time":61,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2025-09-18T07:53:09.746Z","updated_at":"2025-10-13T15:38:17.540Z","avatar_url":"https://github.com/Sweep76.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n\n# Native Sparse Attention Triton\n\n\u003c/div\u003e\n\nThis repository implements the sparse attention mechanism introduced in the paper [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://arxiv.org/abs/2502.11089) and provides an efficient training implementation based on [Triton](https://github.com/triton-lang/triton).\n\n🎉 We now support both training and inference for Native Sparse Attention (variable-length version, including prefilling, decoding, and KV cache management). We have provided a toy model at `model.ToyNSALlama`, which supports `forward` function for training and `generate` function for inference. Welcome to try it out!\n\n## Requirements\nEnsure the following dependencies are installed:\n- PyTorch \u003e= 2.1.0\n- triton \u003e= 3.0.0\n- einops \u003e= 0.7.0\n- flash_attn \u003e= 2.6.3\n\n## Usage\n\n### Notes\n1. PyTorch implementations (`ops.torch`) are intended for debugging only.\n2. For production use, prefer Triton operators (`ops.triton`).\n3. All implementations are based on the varlen approach similiar to flash_attn_func_varlen. Please concatenate the inputs of a batch before use.\n4. Only support attention head dimension less than 128 for now.\n\n### Install\n\nYou can install `native_sparse_attention` using pip:\n\n```shell\npip install git+https://github.com/Sweep76/sparse-attention.git\n```\n\n### Functions\n\n\nYou can import those functions from the `ops` module:\n\n```python\nimport torch\nfrom native_sparse_attention.ops import linear_compress, compressed_attention, topk_sparse_attention\n\n# input example\nnum_q_heads = 64\nnum_kv_heads = 4\nhead_dim = 128\nkernel_size = 32\nkernel_stride = 16\nblock_size = 64\ntopk = 16\ncu_seqlens = torch.Tensor([0, 1024, 8192, 16384]).to(torch.int32).cuda()\nquery = torch.randn(16384, num_q_heads, head_dim).to(torch.bfloat16).cuda()\nkey = torch.randn(16384, num_kv_heads, head_dim).to(torch.bfloat16).cuda()\nvalue = torch.randn(16384, num_kv_heads, head_dim).to(torch.bfloat16).cuda()\n\n# weight example\nw = (\n    torch.randn(num_kv_heads, kernel_size * head_dim, head_dim)\n    .to(torch.bfloat16)\n    .cuda()\n)\npe = torch.randn(num_kv_heads, kernel_size, head_dim).to(torch.bfloat16).cuda()\n\n# 1. key value compression\ncompressed_key, compressed_cu_seqlens = linear_compress(\n    key, w, cu_seqlens, kernel_size, kernel_stride, pe\n)\ncompressed_value, _ = linear_compress(\n    value, w, cu_seqlens, kernel_size, kernel_stride, None\n)\n\n# 2. attention between query and compressed key value\ncompressed_attn_output, topk_idx = compressed_attention(\n    query,\n    compressed_key,\n    compressed_value,\n    kernel_size,\n    kernel_stride,\n    block_size,\n    topk,\n    cu_seqlens,\n    compressed_cu_seqlens,\n    init_blocks=1,\n    local_blocks=2,\n)\n\n# 3. topk sparse attention\nsparse_attn_output = topk_sparse_attention(\n    query,\n    key,\n    value,\n    topk_idx,\n    block_size,\n    cu_seqlens,\n)\n```\n\n### Module\n\nThe `modules` directory also provides implementations based on `torch.nn.module` for easy integration into models.\n\n```python\nfrom native_sparse_attention.modules import NativeSparseAttention, RopeConfig\n\nNSA_Layer = NativeSparseAttention(\n    compress_type=\"linear\",\n    hidden_size=4096,\n    num_q_heads=64,\n    num_kv_heads=4,\n    head_dim=128,\n    kernel_size=32,\n    kernel_stride=16,\n    block_size=64,\n    topk=8,\n    init_blocks=1,\n    local_blocks=2,\n    window_size=512,\n    rope_config=RopeConfig(\n        max_position_embeddings=32768,\n        head_dim=128,\n        rope_theta=500000,\n        rope_scaling={\n            \"factor\": 4.0,\n            \"high_freq_factor\": 4.0,\n            \"low_freq_factor\": 1.0,\n            \"original_max_position_embeddings\": 8192,\n            \"rope_type\": \"llama3\",\n        },\n    ),\n)\n```\n\n### Model\n\nWe offer two simplified LLaMA models in the `model` directory, featuring self-attention and native sparse attention. For more details on how to use these models, please refer to [this link](https://github.com/XunhaoLai/native-sparse-attention-triton/tree/main/native_sparse_attention/model#readme).\n\n\n```python\nfrom native_sparse_attention.model import ToyNSALlamaConfig, InferenceConfig, ToyNSALlama\n\nconfig = ToyNSALlamaConfig(\n    hidden_size=4096,\n    intermediate_size=14336,\n    num_hidden_layers=8,\n    num_attention_heads=32,\n    num_key_value_heads=2,\n    head_dim=128,\n    rope_theta=500000.0,\n    rope_scaling={\n        \"factor\": 8.0,\n        \"high_freq_factor\": 4.0,\n        \"low_freq_factor\": 1.0,\n        \"original_max_position_embeddings\": 8192,\n        \"rope_type\": \"llama3\",\n    },\n    compress_type=\"weightedpool\",\n    kernel_size=32,\n    kernel_stride=16,\n    block_size=64,\n    topk=8,\n    init_blocks=1,\n    local_blocks=2,\n    window_size=512,\n)\ninference_config = InferenceConfig(\n    max_batch_size=4,\n    max_length=8192,\n    max_new_tokens=128,\n)\nmodel = ToyNSALlama(config, inference_config).cuda().bfloat16()\n```\n\n## Testing\n\nSome test scripts are available in the `test` folder and can be run directly for unit testing. For example:\n\n```bash\npython test/test_topk_sparse_attention.py\npython test/test_nsa_module.py\npython test/test_nsa_model.py\n```\n\n### Benchmarks\n\nHere are the speed benchmarks conducted on a single NVIDIA A100 GPU or H100 GPU for the `topk_sparse_attention` function: \n\nA100 GPU speed benchmarks:\n```sh\n** forward with block size 64 **:\n          N       Flash  Triton-Flash  Triton-Top8  Triton-Top16\n0    2048.0     0.414144      0.635648     0.633440      1.009184\n1    4096.0     1.400304      2.267552     1.179808      1.916736\n2    8192.0     5.223776      8.528160     2.266816      3.723168\n3   16384.0    20.225697     32.745537     4.468128      7.359168\n4   32768.0    79.587715    128.951065     8.517440     14.142848\n5   65536.0   321.240479    511.652100    17.249599     30.991360\n6  131072.0  1349.810425   2063.245605    36.400482     67.884544\n\n** backward with block size 64 **:\n          N        Flash  Triton-Flash  Triton-Top8  Triton-Top16\n0    2048.0     1.315440      2.348560     1.941568      2.691040\n1    4096.0     4.271584      8.553184     3.647744      5.032160\n2    8192.0    15.323984     32.665440     5.650144      9.066112\n3   16384.0    58.753281    127.675964    11.160832     17.113279\n4   32768.0   227.770462    504.572693    21.723392     34.715614\n5   65536.0   899.181274   2059.718506    44.517181     76.309441\n6  131072.0  3587.918701   8530.726562   105.344734    182.970169\n```\n\nH100 GPU benchmarks:\n```sh\n** forward with block size 64 **:\n          N       Flash  Triton-Flash  Triton-Top8  Triton-Top16\n0    2048.0    0.259552      0.293888     0.584544      0.917664\n1    4096.0    0.846848      1.029904     1.094976      1.745136\n2    8192.0    3.043744      3.843392     2.128256      3.396880\n3   16384.0   11.743568     14.791360     4.190528      6.704192\n4   32768.0   45.968513     57.532478     7.614496     12.417440\n5   65536.0  187.234375    228.093948    14.840048     24.511856\n6  131072.0  810.890381    914.693970    29.470400     48.990192\n\n** backward with block size 64 **:\n          N        Flash  Triton-Flash  Triton-Top8  Triton-Top16\n0    2048.0     0.798976      1.096096     1.117312      1.380016\n1    4096.0     2.545680      3.826336     1.669760      2.214880\n2    8192.0     9.029760     14.411633     2.772096      3.947456\n3   16384.0    34.144016     58.945698     5.201344      7.538912\n4   32768.0   135.718369    233.369247     9.968864     15.154192\n5   65536.0   541.053894    929.337646    21.089870     33.818878\n6  131072.0  2139.974854   3785.540527    54.918144     93.750717\n```\n\nHere comes another speed benchmark result for testing `compressed_attention` function on a single NVIDIA A100 GPU or H100 GPU:\n\nA100 GPU speed benchmarks:\n```sh\n** forward with kernel 32 and stride 16 **:\n          N       Flash  Triton-Flash  Compressed  Compressed-wo-Score\n0    2048.0     0.413664      0.635488    0.655024             0.170816\n1    4096.0     1.396416      2.247648    1.132304             0.377152\n2    8192.0     5.234656      8.526400    2.879200             0.977952\n3   16384.0    19.988865     32.755199    9.426448             2.943024\n4   32768.0    79.419907    128.955170   30.284096             9.901120\n5   65536.0   321.590210    511.615509  112.260544            36.001602\n6  131072.0  1346.996338   2069.837891  423.099518           136.820038\n\n** backward with kernel 32 and stride 16 **:\n          N        Flash  Triton-Flash  Compressed\n0    2048.0     1.322560      2.352000    0.486784\n1    4096.0     4.270832      8.552608    0.971392\n2    8192.0    15.515680     32.671329    2.603744\n3   16384.0    59.345055    128.377472    8.499456\n4   32768.0   230.626144    506.581238   30.064833\n5   65536.0   919.260498   2068.642578  113.466560\n6  131072.0  3646.603760   8498.374023  439.623444\n```\n\nH100 GPU speed benchmarks:\n```sh\n** forward with kernel 32 and stride 16 **:\n          N       Flash  Triton-Flash  Compressed  Compressed-wo-Score\n0    2048.0    0.259488      0.297152    0.485920             0.103232\n1    4096.0    0.847376      1.030400    0.710208             0.217760\n2    8192.0    3.044016      3.875840    1.607360             0.516016\n3   16384.0   11.823104     14.829360    4.970272             1.440288\n4   32768.0   46.204750     57.527809   15.004992             4.584736\n5   65536.0  187.324249    227.909958   53.009087            16.134224\n6  131072.0  810.707214    910.106873  191.245728            60.154270\n\n** backward with kernel 32 and stride 16 **:\n          N        Flash  Triton-Flash  Compressed\n0    2048.0     0.797728      1.090640    0.283104\n1    4096.0     2.547088      3.834592    0.550464\n2    8192.0     9.021520     14.421088    1.249184\n3   16384.0    34.159508     58.793377    3.743440\n4   32768.0   136.844070    233.447708   12.640032\n5   65536.0   537.559814    929.360229   46.054817\n6  131072.0  2135.629883   3782.351562  175.587296\n```\n\nAll the speed benchmarks above were tested with 64 query heads, 4 key/value heads, and a head dimension of 128.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsweep76%2Fsparse-attention","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fsweep76%2Fsparse-attention","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fsweep76%2Fsparse-attention/lists"}