MoonshotAI/minitriton

MoonshotAI/minitriton

Language: Python

License: Apache-2.0

Stars: 29

Forks: 3

Open issues: 0

Created: 2026-07-23T07:39:34Z

Pushed: 2026-07-27T13:49:08Z

Default branch: main

Fork: no

Archived: no

README:

MiniTriton is a teaching-grade but production-minded tile compiler. You write kernels in a small Python-embedded tile DSL; they lower through upstream MLIR dialects to PTX and run on NVIDIA GPUs. On top of that sits an eager tensor library (eager execution and tile-level compilation, one API). All application kernels — [flash attention](https://arxiv.org/abs/2205.14135), [KDA (Kimi Delta Attention)](https://arxiv.org/abs/2510.26692), and friends — are written in the DSL itself; the compiler contains only tile-level vocabulary primitives and general passes — no application-specific intrinsics (the one expressibility exception is a generic linear-algebra primitive, tl.solve_tril; see Limitations). Known gaps are documented in Limitations below. The tile programming model deliberately follows Triton's — a familiar, well-specified target to validate an agent-built compiler against; upstream triton appears here only as a benchmark baseline.

> Built by [Kimi K3](https://www.moonshot.cn) (Moonshot AI). The > whole stack — the DSL frontend, the MLIR compiler, CUDA kernels, > autograd/nn, benchmarks, figures and documentation — was designed, > implemented, measured and written by Moonshot AI's K3 model, with > engineering direction and review by the maintainer. > > Disclaimer: this is a demonstration of K3's capability on compiler > design — not a Moonshot product, and not for production use.

Start here

• Use the tensor library: examples/vecadd.py → examples/matmul.py

→ examples/train_gpt.py — same API as the package (minitriton.tensor / minitriton.nn / minitriton.autograd)

• Write a GPU kernel: the tile DSL in examples/vecadd.py

(@tl.kernel); the kernel corpus to learn from lives in minitriton/device/cuda/kernels/ — pure DSL, no compiler intrinsics

• See the numbers: the figures below regenerate from

benchmarks/roofline/ (performance, incl. baselines) and benchmarks/training/ (convergence & correctness evidence)

• Hack the compiler: minitriton/frontend/ (DSL → AST) →

minitriton/compiler/ (builder → passes → lowering to PTX), with tests/ as the executable specification

Headline results

| CUDA-core roofline (fp32) | Tensor-core roofline (tf32/bf16) | train_gpt convergence (default fp32 path) | |---|---|---| |  |  |  |

How to read a roofline: the x-axis is arithmetic intensity (FLOPs per DRAM byte), the y-axis is measured GFLOP/s. The solid line is the measured limit, the dashed line is the vendor spec sheet. Red is minitriton, blue is torch (eager or cuBLAS), green is torch.compile, purple is triton. Implementations that share a shape sit at the same point; the baseline series are drawn larger and lighter, so an overlap shows as a faint halo around the red marker. Points below the roofline's rivals are published too, on purpose. The star marker is the full gpt50m train step measured end to end (steady-state ms/step, identical accounting on every stack; gpt50m is the benchmark's ~50M config, L10/D640/H10 — 49.4M at T128 to 50.0M at T1024 via the positional table; the torch baseline uses explicit-eager attention below block 512 and SDPA at block 512+) — read the red stars against the torch ones directly (eager and compile; the fp32 panel shows eager only): that is the current, honest state of the library, wins and losses alike. The third figure trains the same small GPT in default fp32 on both stacks — the curves overlap. It is an L20 measurement snapshot captured at vocab=204 (the figure's own ln(vocab) annotation is self-consistent with that); the train corpus is the source tree itself, so the runtime vocab moves as the code changes — 103 at this writing, and the scripts always use the runtime value.

Quick start

Requirements: NVIDIA GPU (sm_89 verified; sm_80+ in principle), system CUDA ≥ 12 (ptxas), Python ≥ 3.10. Toolchain via conda-forge:

conda create -p ./.conda-env -c conda-forge python=3.10 mlir=22.1.0 mlir-python-bindings=22.1.0 llvm-tools=22.1.0
./.conda-env/bin/pip install "cuda-python>=12,=22" -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest
pip install "cuda-python>=12,

*SPH dam-break, 50k particles: field-sampled water (navy background,
cyan surface, white compressed foam), 480 steps @ 48 fps; regenerate
with `examples/render_gallery.py` (extra deps: `pip install pillow`
plus a system `ffmpeg`).
[Static contact sheet](assets/gallery_water.png) ·
[PDF](assets/gallery_water.pdf) for print.*

## Limitations

- **Hardware**: verified on NVIDIA L20 (sm_89) only; sm_80+ in principle
but untested; no AMD/ROCm; the CPU backend is a numpy *oracle* for
tests, not a usable runtime.
- **Dtypes**: fp32 and bf16 are first-class, tf32 is opt-in; fp16 works
for elementwise, cast and norm/CE paths (16-bit norms compute in f32
internally and cast back on store); fp16 matmul runs through f32
compute with a cast back (numerically correct — there is no fp16
tensor-core path); no fp8/quantized paths. The default path is not
bit-identical to torch (different reduction order, `ex2.approx`).
- **Operator coverage**: the op set is what the benchmarks and examples
exercise — not a torch replacement (no conv/pool/recurrent ops; most
ops assume contiguous inputs; new shape/stride keys trigger a fresh
compile + autotune pass). `@tl.compile` fuses forward-only: gradients
do not flow through a compiled graph (backward on its output raises —
compute the loss/backward outside it, unlike torch.compile). The
layer_norm/rms_norm/softmax row kernels (forward and backward) cap
the reduced row width at C ≤ 8192 and raise `NotImplementedError`
above it; cross_entropy covers wider rows through its split route.
KDA is forward-only (no tape — backward is not implemented), and
in-place indexed assignment (`t[i] = v`) raises `NotImplementedError`.
- **Autograd**: the tape is single-use — `backward()` frees it and a
second backward through the same graph raises (there is no
`retain_graph`); with several losses, sum them first and backprop
the sum (`(l1 + l2).backward()`). There is no version...