{"categories":[{"label":"Artificial Intelligence","url":"https://skillfed.io/packages/category/scientific-engineering-artificial-intelligence/7"}],"enrichment":{"capability":"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.","skillfed_tags":["jax-ecosystem","gpu-optimization","kernel-tuning"],"use_cases":["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"],"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.\n\nTokamax 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\u2014still in active development with API changes expected\u2014but offers a path to both use pre-optimized kernels and build custom ones by inheriting from the Op class.","worth_installing":"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\u2014suitable for research and experimentation, less so for production systems requiring stability guarantees. No known vulnerabilities."},"id":"tokamax","links":{"html":"https://skillfed.io/packages/tokamax","md":"https://skillfed.io/packages/tokamax.md","pypi":"https://pypi.org/project/tokamax/"},"maintenance":{"status":"active"},"meta":{"latest_release":"2026-03-19","license_spdx":null,"license_treatment":"permissive","name":"tokamax","python_support":"supports_current","summary":"A Pallas Custom Kernel Library."},"popularity":{"monthly_downloads":241972,"position":8865,"tier":"top_15000"},"security":{"n_vulnerabilities":0},"version":"0.0.12"}