v0.3.0
natten-mps v0.3.0
Neighborhood Attention for Apple Silicon — PyTorch MPS backend.
Highlights
- 1D, 2D, and 3D neighborhood attention with fused and split-QK/AV paths
- Metal compute shaders via
torch.mps.compile_shader— no native extension build step - Variable-length (varlen) attention support
- GQA / MQA, causal masking, dilation, additional KV
- Automatic backend selection (metal → pure fallback)
- CPU + MPS support; float16 / bfloat16 / float32
Install
pip install natten-mpsRequires PyTorch ≥ 2.8.0 and Apple Silicon.