Tlaloc documentation

Compile-time automatic differentiation and GPU serving for Kotlin

GitHub · Maven Central · API reference · Changelog

Tlaloc is a Kotlin library for differentiable array programs. A K2 compiler plugin turns grad { } into gradient code when your code compiles, and tensors carry their rank, dtype and axis names in the Kotlin type. Programs lower to StableHLO and run on NVIDIA GPUs through PJRT, on Apple-silicon GPUs through MetalHLO, or on the CPU in a reference interpreter.

val g = grad { a: DTensor<Rank2<Sym, Sym>, F32> -> (a matmul a).sum().toFloat() }
val a = Tensors.f32Matrix<Sym, Sym>(2, 2, floatArrayOf(1f, 2f, 3f, 4f))
println(g(a).hostF32().toList())   // [7.0, 11.0, 9.0, 13.0]

Warning

This is an alpha, 0.1.0-alpha03. APIs can change between alphas without a deprecation cycle; COMPATIBILITY.md lists what may break.

What are you trying to do?

To	Read
take a first gradient	Getting started, then examples/quickstart
train a model on a GPU	examples/gpu-training, then the case study
fine-tune a Hugging Face model	examples/fine-tune, examples/lora-finetune
serve a model	the serving runbook, examples/triton-llm
read the derivative the compiler wrote	Readable reverse mode
understand why results differ between backends	Numerical precision
find out why a step is slow	Measuring and profiling
check whether something is supported	Capabilities, FAQ

Contents

• Get started

  ---

  Install, take a first gradient, set up a GPU.

  • Getting started: install, first gradient, GPU setup, configuration, troubleshooting
  • Examples: eighteen standalone programs, each its own Gradle build
  • FAQ: short answers, with links

• How it works

  ---

  Where derivatives are taken, when programs compile, and what each backend supports.

  • How Tlaloc runs a program: grad { } and capture, compiles, buffers, transfers, device memory
  • Readable reverse mode: gradients printed as Kotlin that compiles
  • Dynamic shapes: bounded axes, buckets, and what recompiles
  • Distributed: meshes, shardings and collectives: what is emitted, what has run
  • Custom kernels: custom derivative rules, KPTX GPU kernels, Mosaic TPU kernels

• Analyze your workload

  ---

  What each backend computes, and where a step's time goes.

  • Numerical precision: f32, f64 and bf16 by backend, and the tests' tolerances
  • Measuring and profiling: timing a step, seeing what XLA compiled, profiling a served model
  • A training step on the device: 25.78 ms to 2.22 ms a step by keeping parameters and optimizer state on the GPU

• Serve

  ---

  Compile a model into an artifact and serve it without a JVM.

  • Serving architecture: how serving works, and where the time goes
  • Serving runbook: export an artifact, run it, plug it into vLLM
  • TPU bring-up: the steps for a first TPU session

• Reference

  ---

  • Capabilities: every capability, its status, and the test behind it
  • Compatibility: what may change between alphas
  • The emit contract: the StableHLO Tlaloc emits, for backend implementers
  • The Maestro descriptor: the workflow descriptor format
  • API reference: per module; ./gradlew apiDocs builds it locally

• For maintainers

  ---

  • Releasing: cutting a release
  • Vendoring: the vendored Maestro tree
  • Open issues: what is known to be wrong or missing
  • Doc snippets: the setup and checks behind the docs' Kotlin snippets (scripts/check-doc-snippets.py)
  • This site: scripts/docs.sh serve previews it, scripts/docs.sh build builds it

Records

Plans, audits and design notes, kept as written: plans/, audits/, design/, work-log/, papers/.