English | 中文

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, KDA (Kimi Delta Attention), and friends — are written in the DSL itself; the compiler contains only tile-level vocabulary primitives and general passes, no application-specific intrinsics. Design decisions and measured history live in the git log and the issue tracker.

Built by Kimi K3 (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 on compiler design, not an official project from Moonshot AI.

For stable reproduction use the stable branch (advanced in lockstep with main). Every figure below is regenerated by a re-runnable script under benchmarks/; the scripts and their CSVs are the source of truth for any number.

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 machine's measured limit (calibrated here), 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 third figure trains the same small GPT in default fp32 on both stacks — the curves overlap.

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-python-bindings mlir
./.conda-env/bin/pip install cuda-python numpy pytest        # torch/triton: benchmark baselines only
PY=./.conda-env/bin/python

$PY examples/vecadd.py                  # DSL → MLIR → PTX → GPU end to end
$PY -m pytest tests/ -q                 # 1619 passed (numpy/fp64 oracles only, no torch import)

Or install as a git package (MLIR toolchain via the pip mlir-wheels index — note the wheel is ~1.3 GB; on slow links prefer the conda-forge route above, which mirrors the same 22.1.0 toolchain):

pip install "mlir>=22" -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest
pip install cuda-python numpy
pip install "git+https://dev.msh.team/wangzhiyuan/minitriton.git"
python -c "import minitriton; print(minitriton.__version__)"

Examples and usage guide

Everyday tasks, one command each (torch/triton are comparison baselines only, never runtime dependencies):

# 1. write a kernel: @tl.kernel DSL → MLIR → PTX (examples/vecadd.py, matmul.py)
$PY examples/vecadd.py

# 2. fuse eager code: @tl.compile + the fusion rule engine (rules R1-R8/R2')
$PY -m pytest tests/test_fusion.py -q          # 20+ runnable fusion examples

# 3. train a model end to end (corpus = the repo's own source, zero external data)
CUDA_VISIBLE_DEVICES=0 $PY examples/train_gpt.py --steps 200   # ~50M char-GPT

# 4. mixed precision: bf16 training path
#    (bf16 params + fp32 master weights + dynamic loss scaling, CUDA-graphed step)
$PY benchmarks/training/train_gpt_precision_bench.py --steps 60 --impls tl_bf16_graph

# 5. streams: overlap independent kernels on side streams (opt-in)
MT_STREAM_OVERLAP=1 $PY benchmarks/training/train_gpt_precision_bench.py --steps 60 --impls tl_bf16

# 6. benchmarks + CI gate (core suite vs torch eager and torch.compile)
CUDA_VISIBLE_DEVICES=0 $PY benchmarks/ops/ce_report.py
CUDA_VISIBLE_DEVICES=0 $PY benchmarks/ci_report.py   # rc=0 means PASS

# 7. scheduling helpers: canonical pipeline skeletons and composed
#    row-op reference forms live in minitriton/sched/ (import and reuse)

Distributed (NCCL data-parallel, gradients all-reduced once per step):

CUDA_VISIBLE_DEVICES=0,1 $PY examples/train_gpt.py --ddp --steps 120 --loss-csv /tmp/ddp.csv

Graphics / physics (minitriton.viz, window or headless PNG frames):

$PY examples/rigid_demo.py --steps 240 --frames /tmp/rigid_frames   # rigid-body particles
$PY examples/sph_demo.py   --steps 240 --frames /tmp/sph_frames     # SPH fluid
$PY examples/mpm_demo.py   --stats 50                               # MPM continuum (console stats)

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. Static contact sheet · 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; no fp16/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).
• Performance ceilings are documented, not hidden: some kernel families sit below torch/cuBLAS (see the rooflines; evidence chains are in the git log and issues). Cold-start compilation is serial per process (the test gate parallelizes at process level with pytest-xdist).
• Distributed: NCCL data-parallel only, tested at 2 GPUs single-node. sparse/ and distributed/ are the youngest packages.
• Ecosystem & maturity: torch-like API, not torch-compatible; no HuggingFace/ONNX bridges. Teaching-grade but measurement-honest — every claim has a re-runnable script; known gaps are recorded in the git log and the issue tracker.
• Input validation is deliberately thin: embedding and cross-entropy indices are not bounds-checked (out-of-range indices are undefined behavior, including out-of-bounds writes in the backward scatter); tensor-tensor pow with a negative base returns NaN (integer scalar exponents like x ** 3 work).
• No middle-layer optimizer: the DSL lowers to hand-scheduled MLIR (pipelining, smem swizzles, ldmatrix/mma orchestration and ~25 inline PTX hotspots are hand-written per kernel family in compiler/builder.py), with only three small general W-IR passes on top. This is a teaching-grade hand-scheduled tile DSL, not an optimizing compiler — a new kernel family costs hand-written schedule code, not a new pass.

Repository layout

├── minitriton/          # the package
│   ├── frontend/        # @tl.kernel AST frontend (types/code_generator/semantics)
│   ├── compiler/        # builder (tile→MLIR) layout (layout algebra)
│   │                    #   wir/passes (general passes) lowering (→PTX→cubin)
│   ├── fusion.py        # graph-level fusion rule engine (registry + arbitration)
│   ├── compile.py       # @tl.compile (trace + executor)
│   ├── sched/           # scheduling choreography skeletons + row-op reference forms
│   ├── runtime/         # cuda_driver / buffer / cache / autotune / allocator / streams
│   ├── device/
│   │   ├── cuda/        # ops.py (routing shells) elementwise.py (generator + expr vocab)
│   │   │   └── kernels/ # the single kernel library (matmul/attention/norm/ce/misc/kda, pure DSL)
│   │   └── cpu/         # numpy oracle
│   ├── ops/             # eager op shared semantics (broadcast/promotion/dispatch)
│   ├── nn/  autograd/  distributed/  sparse/  precision.py (precision switches)
├── examples/            # vecadd → matmul → train_gpt (~50M char-GPT, self-contained corpus)
├── benchmarks/
│   ├── ci_report.py     # CI performance gate (geomean thresholds + nonzero exit)
│   ├── roofline/        # plot_style (unified style) + dual rooflines + solve_tril
│   ├── matmul/  attention/  kda/  ops/  training/  probes/
├── tests/               # pytest (numpy/fp64 oracles; no torch imports)
├── docs/                # logos (architecture_map.py is a dev-time IR map)
└── build/               # compile intermediates (gitignored, local debugging)

Engineering rules

• Every performance comparison ships with a figure + CSV + baseline implementation (torch eager / torch.compile / triton, same methodology); figures go through the benchmarks/roofline/plot_style.py style family ((dark, light) × (png, svg, pdf); light PDFs for print/LaTeX)
• Negative results and unmet targets are recorded, not hidden (git log + issue tracker)
• DSL first: application kernels are written in the Python DSL and lowered through the generic pipeline; compiler intrinsics are limited to the tile-vocabulary set plus case-by-case justified expressibility exceptions (AGENTS.md §1 three-zone rule)
• Every kernel has a numpy/fp64 oracle test; benchmarks/training/grad_check.py (vs torch autograd, atol 1e-4) must not regress