Tokamax is a library of custom accelerator kernels. It provides highly optimized implementations of key operators, including mixture-of-experts building blocks like grouped matrix multiplication and token gathers/scatters, attention variants like paged attention, splash attention, gated deltanet, Kimi delta attention and multiheaded latent attention, and TopK computation. The operator portfolio will evolve over time to support the latest OSS models running on the latest hardware accelerators, including Google TPUs and NVIDIA GPUs. The library also provides foundational infrastructure to simplify kernel authoring, maintenance, and deployment; for example, features like autotuning, low-overhead benchmarking, and so on.
Tokamax kernels are typically built on top of JAX and Pallas, but also expose native PyTorch APIs on TPUs for developer convenience.
Detailed documentation can be found at openxla.org/tokamax, including a list of supported ops on different hardware accelerators, benchmarking infrastructure, and autotuning.
The latest Tokamax PyPI release:
pip install -U tokamaxThe latest bleeding edge version from Github, with no stability guarantees:
pip install git+https://github.com/openxla/tokamax.gitPlease see the Colab notebooks here for usage examples.
Please see documentation on how to contribute and our code style guide.