Skip to content

v0.3.0

Choose a tag to compare

@ssmall256 ssmall256 released this 27 Feb 14:45
· 6 commits to main since this release

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-mps

Requires PyTorch ≥ 2.8.0 and Apple Silicon.