Chroma is a small ML compiler and runtime that turns fixed-shape PyTorch models into native CPU, CUDA, and Metal programs that run without PyTorch. It contains a graph importer, operator lowering, codegen, memory planning, and a Python API.
From a model to a native program
Chroma starts from a fixed-shape torch.export graph. Because every dimension
and dtype is known ahead of time, the compiler can fold constant work, choose
kernels, and plan all temporary storage before inference begins.
The output is a Python extension, a weight file, the generated source, and a
manifest. A compiled program takes NumPy arrays and loads without PyTorch. A
jit wrapper also compiles PyTorch functions or modules on first call, caching
one program per input shape and dtype.
Fusion and dispatch
Elementwise operations stay as deferred expressions until something consumes them, so chains fuse. An expression is written to a buffer only when two or more operations read it, or when it nests too deeply (32 levels on CPU, 24 on GPU).
The CPU backend emits C++20 expression templates with shapes and strides as
template params, so an expression like (x + 1).sin().sum(-1) can become
a single loop per row with no intermediates. Compatible FP32 matrix products call
Accelerate on macOS or CBLAS on Linux.
The CUDA backend calls cuBLAS for FP32 linear layers, with TF32 disabled. Everything else, including fused arithmetic, expert and attention contractions, reductions, indexing, and top-k, runs as kernels from a GPU generator shared with Metal.
The Metal backend pairs generated shaders with three hand-written kernels: a tiled FP32 matmul, plus a BF16-weight matmul and an RMSNorm kernel adapted from gpt-oss. The two gpt-oss kernels run only for BF16-representable weights at larger widths, so the FP32 dev model below uses just the tiled matmul. Wide reductions and top-k use generated kernels with one 32-lane SIMD group per row, and each call encodes a single command buffer.
Memory
Every temporary has a lifetime, from the kernel that produces it to its last reader, and a planner packs those lifetimes into one aligned arena. For the model below, 1.90 MiB of temporaries fit in a 74 KiB arena, 26× less than giving each its own space.
Inference
The dev model uses gpt-oss’s architecture at a small size. It has 59.08 million parameters: 12 layers, hidden and intermediate widths of 128, 12 experts with four active per token, and a 201,088-token vocabulary. Its weights are seeded random values, so it exercises the compiler and numerics but produces no useful language.
Each call runs eight tokens in FP32 and returns all eight rows of logits. Chroma and PyTorch use identical weights and inputs. Medians pool two runs of 100 calls, each after 10 warmups. CPU runs use one thread for both PyTorch and BLAS, and GPU timings on both sides include host-device copies and synchronization. All results were measured on October 8, 2026.
The gap is this large because eager PyTorch dispatches each of the model’s 1000+ ATen operations from Python, and on a GPU most of them become separate kernel launches. Chroma fuses them into 289 CPU kernels, 361 CUDA kernels, or 325 Metal dispatches, all issued from native code.
The full method and raw reports are in benchmarks.md.
Decoding with a KV cache
Prefill computes the prompt’s keys and values once, and each later step processes one new token against a rolling cache. When the cache fills, the oldest rows are overwritten while positions keep advancing.
The comparison below uses a 256-token cache and a 32-token prompt, timed over 32 decode steps.
The uncached baseline is Chroma’s own graph recomputing the full 256-token window at every step, and the prefill graph also pads the prompt to 256 rows.
Caching saves less on the A100 for two reasons. The uncached step is already cheap there (7.39 ms, against 37.58 ms on the M3 CPU), and a one-token step still issues hundreds of small kernel launches, so fixed launch overhead likely sets a floor under the cached step. The next step is to move the cache and shared weights into device memory; today the prefill and decode programs each hold their own copy of the weights.
Closing
Chroma was very interesting to build from an ML perspective, and it gave me real insight into how systems like JAX + XLA and PyTorch + ExecuTorch work.
I hope to return to this at some point and work on topics such as autograd and more complex fusion boundaries.
Thank you for reading!