Download
MSA, or Memory Sparse Attention, is a research framework for scaling language-model memory to extremely long contexts. It replaces full attention over all tokens with sparse selection of compressed latent memory states. Document-wise rotary position encoding and top-k routing keep training and inference close to linear complexity. A tiered KV-cache design stores routing keys on GPU while larger content states can remain on CPU. Its Memory Parallel engine distributes scoring and transfers only selected memory back to the accelerator. Memory Interleave alternates retrieval, context expansion, and generation to improve multi-hop reasoning across distant segments. The project reports experiments extending from 16K to 100M tokens, including inference on two A800 GPUs.
Features
- Memory Sparse Attention architecture
- Document-wise rotary position encoding
- Top-k latent memory routing
- GPU and CPU tiered KV-cache compression
- Memory Parallel distributed inference
- Memory Interleave for multi-hop reasoning
Project Samples
