ringattention

Why we included this project

Ring Attention sounds simple on paper but is fiddly to get right, and this repo is the JAX implementation behind the Ring Attention and Blockwise Parallel Transformer papers. It splits attention and feedforward work across devices and overlaps communication with computation, so the context length you can train on grows with the number of accelerators instead of stopping at whatever fits in one chip's memory. The package installs with pip and exposes `ringattention` and `blockwise_feedforward`, with shard_map examples showing how to slot them into a model. It is research code, so plan on adapting it to your own training loop rather than treating it as a drop-in library. If you are pushing long-context models on TPUs or multi-GPU JAX setups, this is the most direct place to start.