native-sparse-attention-pytorch

收录理由

Transformer 的训练开销会随上下文长度呈平方级增长，这个库直接给出了一个可 pip 安装的解法：它复现了 DeepSeek 在《Native Sparse Attention》论文中提出的稀疏注意力机制。使用时只需实例化 SparseAttention 模块，为滑动窗口、压缩块和选择块设置几个超参数，就能替换掉原本的全注意力。内核基于 Triton 和 PyTorch 的 Flex Attention 构建，而非手写 CUDA，因此代码本身也是一份可读性很好的参考，方便你快速原型验证硬件对齐的稀疏注意力。仓库还附带了一个 Enwik8 语言建模示例，能端到端训练该模块，团队在投入生产内核之前，可以先借此观察长上下文场景下的实际效果。