Skip to content

Packaged Ring Attention with Blockwise Transformers for Near-Infinite Context implemented in Jax + Flax.

License

Notifications You must be signed in to change notification settings

dtunai/Ring-Attention-Jax

Repository files navigation

Ring Attention Sharded Jax

Packaged version of standard ring attention mechanism using JAX used in LWM.

About

Packaged Ring Attention with Blockwise Transformers for Near-Infinite Context implemented in Jax + Flax.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published