TritonX

Ergonomic Triton kernels with autograd support and near speed-of-light performance.

Install

uv pip install tritonx

Usage

import tritonx as tx
import torch

# Matrix multiplication
out = tx.mm(a, b)

# Element-wise addition
out = tx.add(a, b)

# Einsum operations
out = tx.einsum("ijk,ikl->ijl", a, b)

Kernels

• tx.mm - Matrix multiplication
• tx.add - Element-wise addition
• tx.einsum - Einstein summation with autodiff

Development

git clone https://github.com/windsornguyen/tritonx
cd tritonx
uv sync

License

MIT

Citation

@software{nguyen2026tritonx,
  author = {Nguyen, Windsor},
  title = {TritonX: Ergonomic Triton Kernels},
  year = {2026},
  url = {https://github.com/windsornguyen/tritonx}
}