Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Introduction

⚠️ Warning: teenygrad is still under active development and not yet ready for production use. This book, and the crates it documents, will change frequently.

teenygrad is a high-performance, memory-safe Rust ML training and inference library, in the spirit of PyTorch and tinygrad. It targets devices from microcontrollers to data centers with statically-typed GPU kernels and full async support.

Vision

We envision a world where machine learning is truly ubiquitous: from the tiniest embedded sensors to the largest distributed clusters, every device can harness the power of AI efficiently, safely, and without constraints.

Goals:

  • Highly concurrent, memory-safe — designed for scalability and safety.
  • Low memory footprint — optimized to run even on the smallest devices.
  • Statically typed — GPU kernels and ML algorithms are statically typed.
  • No performance compromises — hardware-accelerated wherever possible.
  • Broad hardware support — not just NVIDIA and AMD.
  • Extensible — build high-performance accelerators for your own hardware.
  • Embedded-friendly — core training/inference components are no_std compatible (teeny-core builds no_std by default).
  • Full async support, multi-threaded by default.

How this book is organized

  • Getting Started walks through installing the toolchain and building your first model.
  • Core Concepts covers the tensor/graph types, the compilation pipeline, the dtype system, and name scoping — the ideas every other chapter builds on.
  • Neural Network Layers covers the nn layer API.
  • Kernels & Backends covers how kernels are written (the Triton-like DSL) and compiled to device code.
  • Compiler Internals goes under the hood of the FXGraph compiler.
  • PyTorch Interop covers teeny-torch, the PyO3-based Python bridge.
  • CLI & AOT Compilation covers teeny-cli and teeny-llm.

This book is a guide and mental model; it deliberately doesn’t duplicate API reference material — for that, see the generated rustdoc at docs.teenygrad.org/api (or docs.rs for the crates that build there).

Community

Join our Discord to discuss the project or get help.

Installation & Toolchain

Rust

Most teeny-* crates build with a plain stable (or nightly) Rust toolchain and no system dependencies — see each crate’s own README (linked from its API docs) for specifics. Two crates are exceptions:

  • teeny-cuda requires the CUDA toolkit headers/libraries on the host at build time (its build.rs runs bindgen against them unconditionally). See its README for CUDA_INCLUDE_DIR/CUDA_LIB_DIR.
  • teeny-onnx vendors the upstream ONNX proto schema as a git submodule — run git submodule update --init support/teeny-onnx/onnx in a workspace checkout (not needed when depending on the published crate).

The teenyc compiler (runtime only)

Compiling/running kernels — through teeny-compiler’s LLVM/MLIR backend, teeny-triton’s DSL, or teeny-cuda’s AOT/JIT path — shells out at runtime to a custom compiler binary, teenyc (a fork of rustc with an MLIR codegen backend). This is not needed to cargo build/cargo doc any teeny-* crate, only to actually compile and run kernels.

export TEENYC_PATH=/path/to/teenyc   # falls back to `teenyc` on $PATH if unset

The supported way to obtain it is via cargo-teeny:

cargo install --git https://github.com/spinorml/cargo-teeny
cargo teeny install-toolchain

This mirrors the toolchain setup used by downstream projects such as vision-rs.

System setup (Ubuntu)

For working on the teenygrad workspace itself (not just depending on its published crates):

sudo apt-get install build-essential z3 libz3-dev lld

or run ./setup_ubuntu.sh from the repository root.

TODO: expand this chapter with GPU compute-capability selection, PTX ISA version overrides (TEENYC_PTX_VERSION), and cross-compilation notes once they’re documented per-crate.

Your First Model

The canonical end-to-end example in the workspace today is teeny-vision’s LeNet-5 over MNIST:

use teeny_vision::mnist::mnist_lenet5;

let model = mnist_lenet5::<f32>();

That model is used in two different ways elsewhere in the workspace, both worth reading as reference implementations:

  • Training/inference example: models/teeny-vision/examples/mnist.rs — a full, runnable example (cargo run --example mnist, with TEENYC_PATH set — see Installation).
  • Ahead-of-time kernel compilation: utilities/teeny-cli/src/main.rs — traces the model with a symbolic input and AOT-compiles its kernels via teeny_cli::aot_compile.

The shape of a model

Models are built from teeny-core’s nn layer types, implementing the Layer trait. A traced model produces a symbolic computational graph (SymTensor in, Graph out) — see Tensors & the Computational Graph and Compilation Flow for what happens to that graph next.

TODO: walk through building a small model from scratch (not just pointing at the existing LeNet-5 example) once the nn layer API stabilizes enough to commit to in a book.

Tensors & the Computational Graph

teeny-core::graph (API docs) defines the types every other crate builds on:

  • Shape — Vec<Option<usize>>; None entries are symbolic/dynamic dimensions (e.g. a batch dimension left unbound until trace time).
  • DtypeRepr — the element dtype of a tensor (F32, etc. — see The Dtype System).
  • Op — the set of graph operations (elementwise, reductions, layout, etc.).
  • GraphNode / Graph — the traced computational graph: nodes are Op applications, edges are data dependencies.
  • SymTensor — a symbolic tensor handle used while tracing a model: calling layers with a SymTensor input records Ops into a Graph rather than executing eagerly.

Tracing

Building a Graph from a model is “tracing”: construct a SymTensor::input(dtype, shape), call your model with it (via the nn::Layer trait), and read back the Graph that was recorded as a side effect. See teeny-cli’s aot_compile for a concrete example of this pattern.

What happens to the Graph next is covered in Compilation Flow.

Compilation Flow

At a high level, a model goes from traced graph to running device code like this:

graph LR
    A["Model (teeny-core nn::Layer)"] -->|trace| B["Graph (teeny-core::graph)"]
    B -->|lower| C["FXGraph"]
    C -->|compile| D["teeny-compiler backends"]
    D -->|LLVM/MLIR backend, via teenyc| E["Device object code"]
    D -->|ndarray backend| F["CPU (ndarray)"]
    E --> G["Device drivers (teeny-cuda, ...)"]
  • Trace: your model, built from teeny-core::nn layers, is called with a symbolic input, recording a teeny-core::graph::Graph — see Tensors & the Computational Graph.
  • Lower: the graph is lowered to FXGraph, teeny-compiler’s intermediate representation.
  • Compile: teeny-compiler (API docs) compiles FXGraph to a target backend:
    • The LLVM/MLIR backend shells out to the custom teenyc compiler at runtime (-Zcodegen-backend=mlir) — see The LLVM/MLIR Backend.
    • The ndarray backend (feature-gated, default-on) runs on CPU without teenyc.
  • Device drivers: compiled device code is loaded and run through a driver crate — today, teeny-cuda for NVIDIA GPUs. teeny-cpu and teeny-vulkan are planned but not yet implemented (only drivers/teeny-cuda exists in the workspace today).

Kernels themselves (the actual per-op device code, e.g. attention, matmul, elementwise ops) are defined separately — see Kernels & Backends.

Ahead-of-time vs. just-in-time

The same compilation pipeline can run either eagerly (JIT, at model-run time) or ahead of time (AOT, producing cached artifacts you deploy without a compiler on the target device) — see teeny-cli and Ahead-of-Time Compilation.

The Dtype System

teeny-core::dtype (API docs) defines the element-type traits used throughout the workspace — Dtype, Float, and tensor traits (Tensor, RankedTensor, EagerTensor, Comparison) that generic kernel/layer code is written against (e.g. teeny-kernels’ attention kernels are generic over D: Float).

Host-only functionality that doesn’t belong in the kernel DSL context (like converting a float to its little-endian byte representation) lives in a separate bytes submodule (dtype::bytes::FloatBytes) rather than on the core Float trait itself — this split exists because teeny-triton’s build.rs embeds dtype/mod.rs’s source directly into the kernel DSL (see Writing a Triton Kernel), so anything in that file needs to make sense in a no_core DSL context, not just as normal host-side Rust.

DtypeRepr (in teeny-core::graph, not dtype) is the runtime/graph-level enum representation of a dtype (e.g. DtypeRepr::F32), distinct from the compile-time Dtype/Float traits — used when tracing a graph (SymTensor::input(dtype, shape)) where the dtype isn’t a type parameter.

TODO: expand with a full trait reference and guidance on writing dtype-generic kernel code, once the trait hierarchy stabilizes.

Name Scopes

teeny-core::name_scope (API docs), available under the std feature, provides scoped naming for graph nodes/parameters — useful for giving generated graphs (and, downstream, debug output, checkpoints, profiler traces) human-readable, hierarchical names instead of anonymous IDs.

use teeny_core::name_scope::name_scope;

let _guard = name_scope("encoder/layer0");
// anything built while `_guard` is alive is implicitly namespaced under "encoder/layer0"

current_scope() returns the active scope name, if any (None outside of any name_scope guard).

Because this depends on std (thread-local or similar scope-tracking machinery), it’s gated behind teeny-core’s std feature — unavailable in the default no_std build.

TODO: document nested scopes and how scope names interact with parameter naming once that’s wired into the nn layer API.

Building Models with nn

teeny-core::nn (API docs) provides the Layer trait and a set of standard layers:

  • Convolution: conv1d, conv2d, conv3d
  • Normalization: batchnorm, groupnorm, instancenorm, layernorm, rmsnorm
  • Other: activation, flatten, linear, pad, pool

Every layer implements Layer<I>, with an associated Output type and a call(&self, input: I) -> Self::Output method — this is what gets invoked during tracing to build up a Graph.

Composite models (e.g. teeny-vision::mnist::mnist_lenet5) are just structs composing these layers and implementing Layer themselves by chaining calls to their fields’ call methods — see models/teeny-vision/src/mnist/ for a complete example.

TODO: document the parameter-management/initialization story once it’s stable enough to commit to in the book (constructing layers with random vs. loaded weights, interaction with name scopes).

CPU, CUDA, and Vulkan Backends

Device backends live under drivers/ in the workspace. Today, only teeny-cuda (NVIDIA, via bindgen-generated driver bindings) is implemented. teeny-cpu and teeny-vulkan are on the roadmap but don’t exist yet — the ndarray-backed CPU path in teeny-compiler (the ndarray feature, on by default) is the current CPU story, distinct from a dedicated teeny-cpu driver crate.

teeny-kernels’ cuda feature (on by default) enables the teeny-cuda backend for its kernel implementations. Building it requires the CUDA toolkit — see Installation & Toolchain.

Kernel definitions (the actual per-op logic — attention, matmul, elementwise ops, etc.) are backend-agnostic, written once against the teeny-triton DSL and compiled per-backend — see Writing a Triton Kernel.

TODO: document the teeny-cuda device/runtime API directly (device selection, memory management, profiling via cuda_profiler_start/cuda_profiler_stop) once a driver-level chapter is warranted.

Writing a Triton Kernel

teeny-triton provides a Triton-like DSL for writing kernels in Rust. A kernel is a function generic over the DSL’s Triton/Dtype/Float traits, marked with the #[kernel] attribute macro:

use teeny_macros::kernel;
use teeny_core::dtype::Float;

#[kernel]
pub fn my_kernel<T: Triton, D: Float, const HEAD_DIM: i32>(
    // pointers, strides, block-size const-generics, etc.
) {
    // kernel body, written against T's tensor/pointer operations
}

See kernels/teeny-kernels/src/nn/attention/flash_attn2.rs for a real, well-documented example (Flash Attention 2 forward/backward) — a good template to copy for new kernels.

How it compiles

At teeny-triton’s own build time, build.rs reads the DSL source under src/triton/ (plus teeny-core’s dtype definitions) and embeds it as a string constant (teeny_triton::triton_lang::TRITON) — pure text processing, no compiler invocation.

At your kernel’s compile time (via teeny-compiler’s LLVM/MLIR backend), that DSL text plus your kernel’s source is written out and compiled by the custom teenyc compiler — see The LLVM/MLIR Backend.

TODO: document the available DSL operations (tensor loads/stores, arithmetic, reductions, control flow) as a proper reference once the DSL surface stabilizes.

The #[kernel] Macro

teeny-macros provides the #[kernel] attribute macro (teeny_macros::kernel), which marks a function as a kernel definition consumed by teeny-triton/teeny-kernels:

use teeny_macros::kernel;

#[kernel]
pub fn my_kernel(/* ... */) {
    // kernel body
}

See Writing a Triton Kernel for how a #[kernel]-annotated function fits into the broader compilation pipeline.

TODO: document exactly what code transformation #[kernel] performs (macros/teeny-macros/src/macros/kernel.rs) once it’s stable enough to describe without immediately going stale.

The egg/z3 Optimizer

teeny-compiler depends on egg (e-graph based equality saturation) for graph-level optimization/rewriting as part of the FXGraph compilation pipeline (see Compilation Flow). teeny-fxgraph and teeny-torch also depend on egg.

z3 (SMT solving) is referenced in project documentation as a planned complement to egg for optimization, but is not currently wired up as a dependency of any workspace crate — treat it as roadmap, not shipped functionality, until a z3 dependency actually appears in a Cargo.toml.

TODO: document the specific e-graph rewrite rules and what they optimize for, once this layer has enough surface area to be worth a dedicated walkthrough.

The LLVM/MLIR Backend

teeny-compiler::compiler::backend::llvm compiles a kernel by shelling out to the custom teenyc compiler — a fork of rustc with an MLIR codegen backend — as a subprocess:

teenyc <kernel-source>.rs \
  -Copt-level=3 \
  -Zcodegen-backend=mlir \
  --emit=obj \
  -o <output>.o \
  --target=nvptx64-nvidia-cuda \
  --crate-type=lib \
  -C overflow-checks=off \
  --frontend=triton

A few details worth knowing:

  • TEENYC_PATH (env var) points at the teenyc binary; see Installation & Toolchain. This is a runtime dependency of teeny-compiler, not a build-time one — teeny-compiler itself builds with a plain toolchain.
  • -Zcodegen-backend is an unstable rustc flag. teenyc is distributed on the stable channel (real version numbers, normal feature-gating) — RUSTC_BOOTSTRAP=1 is set on the subprocess call to permit this specific unstable flag against a stable-channel compiler, the same narrowly-scoped mechanism rustc’s own bootstrap, bindgen, and miri’s installer rely on.
  • --frontend=triton tells teenyc to parse the input as teeny-triton’s DSL rather than plain Rust — the kernel source file written out for compilation is the DSL text (teeny_triton::triton_lang::TRITON) concatenated with your kernel’s own source.
  • Caching: compiled kernels are cached by a hash of the kernel ID plus target CPU/PTX-version overrides, so repeated compiles of the same kernel/target are skipped.
  • TEENYC_PTX_VERSION: overrides teenyc’s otherwise-conservative default PTX ISA floor for a given -Ctarget-cpu — see teeny-compiler’s README/rustdoc for when you need this (e.g. targeting a GPU architecture newer than teenyc’s default assumption).

TODO: document the ndarray CPU backend as a parallel section once it has comparable surface area documented here.

Using teeny-torch

teeny-torch is a PyO3-based Python extension module (built as a cdylib, Python module name teenygrad) bridging teenygrad into PyTorch-style Python workflows.

Not published to crates.io. teeny-torch (and teeny-fxgraph, which it depends on) are excluded from the crates.io publish set (publish = false) — they’re the PyTorch compatibility layer, not part of the public Rust SDK. teeny-torch is distributed as a Python package (PyPI), not a Rust crate.

It depends on teeny-core, teeny-fxgraph, and teeny-compiler, and uses egg (see The egg/z3 Optimizer) and flatbuffers for its own serialization needs (built via flatc-rust in build.rs).

TODO: document the actual Python-facing API and installation instructions once this crate has a stable public interface and a packaging story for PyPI.

teeny-cli and Ahead-of-Time Compilation

teeny-cli provides ahead-of-time (AOT) kernel compilation as both a reusable library and a runnable reference binary.

As a library

Flatten teeny_cli::AotArgs (a clap-derived struct) into your own binary’s CLI, and call teeny_cli::aot_compile:

#[derive(clap::Parser)]
struct Cli {
    #[command(flatten)]
    aot: teeny_cli::AotArgs,
}

This is the pattern vision-rs uses in its own binaries — see teeny-cli’s own main.rs for the full reference implementation (AOT-compiling teeny-vision’s LeNet-5).

aot_compile traces your model with a symbolic input (see Tensors & the Computational Graph), then compiles every kernel the traced graph references for the backend/config selected by --device. Currently only --device cuda is implemented; adding a new backend means adding a match arm in aot_compile, not changing AotArgs.

Requires teenyc at runtime

Like the rest of the compilation pipeline, this needs TEENYC_PATH set — see Installation & Toolchain. It’s typically driven via cargo teeny aot (from cargo-teeny) rather than run directly.

cargo teeny aot --bin teeny-cli --device cuda --options "capability=sm_90"

teeny-llm

teeny-llm is the workspace’s LLM inference application — vLLM-style serving, aimed at 2x the throughput (see the roadmap).

⚠️ Early scaffolding. The teeny-llm binary currently only prints a placeholder banner. The planned direction is a single binary offering both an interactive mode and an OpenAPI-compatible HTTP server mode — earlier separate agent/console binaries were dropped in favor of this single-binary design.

cargo install teeny-llm
teeny-llm

TODO: fill this chapter in as teeny-llm gains real functionality — request/response API shape, model loading, batching/scheduling behavior.

Contributing to Teenygrad

Thank you for your interest in contributing! The full, authoritative guide lives in the repository’s CONTRIBUTING.md — this chapter summarizes the essentials.

We follow the contributor model used by Ubuntu, and require a signed CLA before accepting contributions. Reach out on Discord with any questions.

All contributions require:

  • An issue first — file a feature request or bug report to get a maintainer’s attention and find the right project for it.
  • A pull request, approved by the appropriate reviewers.

Engineering standards

See the repository’s AGENTS.md and CLAUDE.md for the full Rust engineering standards (idiomatic Rust, error handling, unsafe documentation requirements, testing expectations). In short: prefer clarity over cleverness, minimize cloning/allocation, avoid unwrap()/expect() in non-test code, and add tests for behavior changes.

Documentation standards

  • Every publishable crate should have a crate-level //! doc comment (see Introduction for how this book fits alongside per-crate rustdoc) and /// doc comments on public items.
  • New crates need a crate-specific README.md (not the shared workspace one) — see any existing crate’s README.md for the expected shape (a short description, a “Prerequisites” section covering any system dependencies/toolchain requirements, and a minimal usage example).

FAQ & Roadmap

Roadmap

2026 Q1 — Torch Inductor in Rust

  • 2026 and beyond
    • Q3: teeny-llm — vLLM-style inference, 2x faster
    • Q4: performance optimization
  • 2027 and beyond
    • Q2: embedded support
    • Q3: sparsity/quantization support
    • Q4: observability/metrics

(See the repository README for the current, authoritative roadmap — this chapter may lag behind.)

FAQ

Why create this project when PyTorch, TensorFlow, and tinygrad already exist?

Those frameworks excel at development and deployment on large-scale infrastructure. We believe the future of AI lies in devices of all sizes — from edge devices to massive clusters. Existing frameworks are relatively heavy; TensorFlow Lite comes closest to this vision, but a modern, memory-safe language was preferred over C.

At the same time, we didn’t want to sacrifice distributed training or other advanced features offered by larger projects — hence teenygrad: a lightweight but powerful alternative.

I’m interested in contributing — how do I get started?

See Contributing to Teenygrad.