tokamax

What it is and what it does

Tokamax is a JAX library that wraps hand-tuned accelerator kernels for common deep-learning operations, built on top of JAX's Pallas framework. It provides optimized implementations of dot-product attention (FlashAttention), gated linear units, layer and RMS normalization, mixture-of-experts routing, and linear softmax cross-entropy loss, with support for both NVIDIA GPUs and Google TPUs. The library also exposes an autotuning system to discover optimal kernel configurations for your hardware and input shapes, and utilities for benchmarking kernel execution time separately from Python overhead.

Tokamax is designed for researchers and practitioners building large-scale models who need fine-grained control over kernel selection and performance. It lets you choose between multiple implementations (e.g., Triton, Mosaic, XLA) for each operation, or allow automatic selection. The package is young—still in active development with API changes expected—but offers a path to both use pre-optimized kernels and build custom ones by inheriting from the Op class.

Use it for:

• Accelerate transformer attention layers on H100 or TPU hardware by swapping standard JAX attention with tokamax.dot_product_attention
• Optimize layer normalization and RMS norm in deep networks using tokamax.layer_norm with hardware-specific implementations
• Build mixture-of-experts models with efficient ragged tensor routing via tokamax.ragged_dot on GPU or TPU
• Autotune kernel configurations for your specific hardware and input shapes to maximize throughput
• Benchmark actual kernel execution time (excluding Python overhead) using tokamax.benchmark with CUPTI or default profiling
• Export JAX functions containing custom kernels to StableHLO with device-specific guarantees via tokamax.DISABLE_JAX_EXPORT_CHECKS

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Tokamax provides custom accelerator kernels for JAX, including optimized implementations of attention, normalization, and mixture-of-experts operations for NVIDIA GPUs and Google TPUs, plus tooling to build and autotune custom kernels.

Yes, if you are training or serving large models on NVIDIA GPUs or TPUs and want to use pre-optimized kernels for attention, normalization, and mixture-of-experts without rewriting core operations. The low install friction and permissive license make adoption straightforward. However, expect API changes and incomplete features—suitable for research and experimentation, less so for production systems requiring stability guarantees. No known vulnerabilities.

Install

tokamax on PyPI

pip install tokamax

uv add tokamax

poetry add tokamax

Installing tokamax

Before you install

Low install friction with a pure Python wheel. Active maintenance status as of 148 days since latest release. However, the package declares itself still heavily under development with incomplete features and expected API changes.

License in practice

Apache License 2.0 is permissive and allows commercial use, modification, and distribution with minimal restrictions—primarily requiring license and copyright notice preservation.

Quickstart

pip install tokamax

import jax
import jax.numpy as jnp
import tokamax

def loss(x, scale):
  x = tokamax.layer_norm(x, scale=scale, offset=None)
  x = tokamax.dot_product_attention(x, x, x)
  return jnp.sum(x)

f_grad = jax.jit(jax.grad(loss))

Requires Python >=3.11 and JAX with jaxlib installed; GPU or TPU hardware needed for actual acceleration.

Verify before relying

• Stability and API compatibility guarantees beyond the stated 'heavily under development' warning
• Performance benchmarks comparing implementations against standard JAX operations
• Supported GPU models and compute capabilities beyond the H100 example

Package facts

Evidence: tokamax-0.0.12-py3-none-any.whl

Tags

More Artificial Intelligence packages

Further reading