---
id: tokamax
version: "0.0.12"
license: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# tokamax — A Pallas Custom Kernel Library.

License: permissive · Maintenance: active · Downloads: 242.0K/mo

## 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 above — 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

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

- License: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 242.0K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

JAX custom kernels GPU TPU, flash attention implementation, accelerator kernel optimization, JAX Pallas kernels, layer norm gated linear unit, mixture of experts kernel, kernel autotuning JAX, jax-ecosystem, gpu-optimization, kernel-tuning

[View on SkillFed](https://skillfed.io/packages/tokamax) · [View on PyPI](https://pypi.org/project/tokamax/)