Paper Feed

Issue 40 · Project 01 GitHub Tooling ✓ read

FrederikRogalski/tinycoral

tinygrad on Google's Coral Edge TPU: reverse-engineered ISA, our own code generator and USB driver. A whole transformer runs on the stick. Written by AI agents.

TL;DR: tinycoral is a from-scratch software stack for Google's Coral USB Edge TPU — reverse-engineered ISA, custom code generator, pure-Python USB driver — wired in as a tinygrad backend (DEV=CORAL). It runs things Google's own toolchain cannot (a complete TinyStories-15M transformer as one on-chip program, 176.7 tok/s at batch 1), validates its generated programs byte-for-byte against edgetpu_compiler as an offline oracle, and was written almost entirely by AI agents with a human directing. Interesting twice over: as a compiler/reverse-engineering feat on a closed accelerator, and as a concrete data point on what agentic engineering can now ship.

What it is

The Coral Edge TPU has always been a black box: you compile a TFLite model with Google's closed edgetpu_compiler and run it through libedgetpu, with a fixed op set and no general compute story. tinycoral replaces the entire stack. The ISA was decoded (building on geohot's edgetpuxray), a code generator emits DarwiNN programs directly, and a libusb-based Python driver handles DFU firmware upload, registers, the single-endpoint DMA protocol, and recovery of a hung chip. Google's compiler is used only offline, as a correctness oracle: wherever it has an equivalent program, tinycoral's output is checked byte-identical.

The result is a real tinygrad device backend, not a demo harness: any tinygrad model compiles through the normal scheduler, and kernels that are provably quantized matmuls/convs get routed to the TPU while the rest fall back to clang on the host.

How it works

tinygrad model float, plain nn coral.quantize uint8, BN folding kernel select AST proof: matmul/conv chain fusion one TPU call codegen: ISA FC, conv, attention Python USB driver DFU, DMA, recovery Edge TPU stick weights resident edgetpu_compiler (offline) byte-for-byte oracle
The whole path runs inside tinygrad's compile step; the runtime only moves bytes. Google's compiler appears only as an offline byte-exactness check, never at run time.

Key moves: post-training quantization rewrites nn.Conv2d/nn.Linear into uint8 tinygrad ops checked on calibration data; a selection pass proves from the kernel AST that it is exactly a quantized matmul/conv; chain fusion merges runs of TPU kernels (joined by exact uint8 max pools) into single programs so activations never leave the chip; the runner keeps weights resident in tile memory. Beyond what Google's compiler supports, the codegen has softmax on the scalar core, fused FFN blocks, attention, and all six TinyStories layers as one program per token.

Evidence

The numbers are specific and honestly caveated. MobileNet v1 runs Google's own program bit-exact and 2% faster (3.89 vs 3.96 ms). Single FC layers beat Google's programs (0.176 vs 0.215 ms). The MNIST CNN fused into one 32-image program hits 22,994 img/s vs 2,615 through Google's stack — though the README admits Google's stack reaches 26,127 img/s given the same batching trick. TinyStories-15M, which edgetpu_compiler simply cannot compile, runs at 1,815 tok/s at batch 256 and 176.7 tok/s with the whole transformer on-chip at batch 1.

TinyStories-15M, batch 256, tokens/stok/s05,00010,00015,00020,00025,00030,000640tinygrad CPU (1 core)1,815tinycoral (CORAL)5,435numpy + Accelerate (all cores)28,000tinygrad METAL (M1 Pro)from bench/RESULTS.md; the Coral only makes sense on weak hosts, as the README says itself

The repo is complete: driver, decoded ISA docs (docs/isa/*.md), codegen, backend, offline byte-comparison tests against the compiler in Docker, hardware tests, a MOCKCORAL=1 numpy bit-model for stick-free development, a web demo, and STORY.md documenting the agentic workflow (human set goals, made calls, replugged the stick when it hung).

Caveats

This is a 3-day-old, 16-star, one-device project. Coverage through the scheduler is matmuls, convs, and fused FFN/classifier blocks; activations, attention, norms, and RoPE still run on the host in the standard path, and the whole-transformer program bypasses tinygrad's scheduler entirely. Whole-network fusion only covers VALID stride-1 convs (odd kernels ≤5×5) + 2×2 maxpool + final linear. Batch-1 matmuls are USB-bound (~55 MB/s device→host), so performance only wins when data stays on-chip. Tested only on macOS with one stick; the setup pins a specific tinygrad commit. And practically: a modern laptop CPU beats the Coral anyway — the value is the reverse-engineering, the compiler, and the generality unlocked on a $60 closed accelerator, not raw throughput.

Try it

git clone https://github.com/FrederikRogalski/tinycoral && cd tinycoral
git clone https://github.com/tinygrad/tinygrad tinygrad-src && \
  git -C tinygrad-src checkout d2892164e20e9c53350abaa7e7a8cc7b7376f8d4
python3.12 -m venv .venv && .venv/bin/pip install -e tinygrad-src -e .
.venv/bin/python -m coral.install
git clone https://github.com/google-coral/libedgetpu ref/libedgetpu  # firmware + register map

# no stick needed — numpy bit-model:
MOCKCORAL=1 DEV=CORAL python examples/stories.py
# with a Coral on USB 3, whole transformer on chip:
DEV=CORAL CHIP=1 python examples/server.py

Models (stories15M.bin, tokenizer.bin) come from karpathy/llama2.c. Read STORY.md even if you never plug in a stick — the agentic-engineering account may be the most durable part.