TT-Forge: Compile Anything

vLLM is a curated serving runtime. It knows exactly which models it supports, it has them tuned and tested, and it presents a clean HTTP API for inference. Tremendous for what it does. But it covers a specific list.

TT-Forge is the other gate. You bring the model — any PyTorch nn.Module, any JAX function, any ONNX export — and the compiler traces it, lowers it to Tensix operations, and hands back something that runs on your QB2 hardware. One call. Hardware execution. No server, no model list to consult.

If vLLM is the highway, TT-Forge is the ability to go anywhere.

Before You Begin — Install Forge

Forge is not part of a default tt-installer run. tt-installer sets up the base — driver, firmware, hugepages, and the ~/.tenstorrent-venv environment that holds tt-smi and tt-flash. Forge itself you install as a pip wheel from Tenstorrent’s package index. That’s how the TT-Forge docs want you to do it — not a container wrapper, not a 45-minute source build.

(tt-installer can give you Forge as a container instead, via --install-forge-container, which writes a tt-forge wrapper to ~/.local/bin/. It’s off unless you ask for it, and it’s a separate path from the wheel below — pick one.)

First confirm the base is ready (Ubuntu 24.04, Python 3.12):

tt-smi   # should show the System Management Interface

Then make Forge a venv of its own. Don’t install it into ~/.tenstorrent-venv: Forge pins its own torch / torch-xla stack, and resolving that on top of your hardware tooling means a bad Forge install is also a broken tt-smi.

python3 -m venv ~/.venvs/forge
source ~/.venvs/forge/bin/activate

Then install the frontend for your framework:

PyTorch & JAX — TT-XLA (the primary frontend):

pip install pjrt-plugin-tt --extra-index-url https://pypi.eng.aws.tenstorrent.com/
tt-forge-install     # pulls in any missing system dependencies

pip install tt-forge is the convenience meta-package that wraps the same thing.

pip install --no-deps torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cpu

ONNX / TensorFlow / PaddlePaddle — TT-Forge-ONNX (single-chip only):

sudo apt-get install -y libgomp1 libmpc3
uv pip install tt_forge_onnx tt_tvm --extra-index-url https://pypi.eng.aws.tenstorrent.com/

The Compilation Paths

Two frontends cover every framework. Both lower to the same TT-MLIR compiler and the same Tensix backend — the framework you start from doesn’t change where you land.

TT-XLA is the primary frontend and the one to reach for first: it takes both PyTorch (through torch-xla) and JAX (through jax.jit), and it’s the only path that scales across multiple chips. TT-Forge-ONNX is the TVM-based route for models that arrive as ONNX, TensorFlow, or PaddlePaddle graphs, and it’s single-chip only.

Your First Compile

ResNet-50 is the right first target — well-understood architecture, small enough to compile fast. This is the canonical PyTorch quickstart from the TT-Forge docs:

import torch
import torch_xla.core.xla_model as xm
import torch_xla.runtime as xr
import tt_torch  # registers "tt" as a torch.compile backend
from torchvision.models import resnet50, ResNet50_Weights

# Point PyTorch/XLA at the Tenstorrent device
xr.set_device_type("TT")
device = xm.xla_device()

# Load ResNet-50 in bfloat16 — Blackhole's native float format
model = resnet50(weights=ResNet50_Weights.DEFAULT).to(torch.bfloat16).eval()

# Compile for Tensix and move the compiled model onto the device
compiled_model = torch.compile(model, backend="tt").to(device)

# Run inference on hardware
input_tensor = torch.randn(1, 3, 224, 224, dtype=torch.bfloat16).to(device)
with torch.no_grad():
    output = compiled_model(input_tensor)

print(output.cpu().argmax(dim=-1).item())   # predicted ImageNet class

What torch.compile(model, backend="tt") does: torch-xla traces the model into a StableHLO graph, the TT-MLIR pipeline lowers that graph to Tensix kernels, and you get back a callable that dispatches to hardware. The first compilation is slow (tens of seconds for ResNet, longer for large models). Subsequent calls with the same input shapes hit a compiled cache and run fast.

Loading in torch.bfloat16 matters: Blackhole is bfloat16-native, so it gives you full hardware throughput. Float32 works, but leaves performance on the table.

Here is the chip view during compilation and inference:

The compile step dispatches weight loads from DRAM then fans work across the Tensix grid.

The tt-forge-models Zoo

Writing model-loading boilerplate for hundreds of architectures is tedious. Somebody already did it. tt-forge-models is the standardized model zoo for TT-Forge — 800+ model variants tested in CI, every one exposing the same ModelLoader interface and loadable in two lines.

The repo lives at ~/code/tt-forge-models and on GitHub at tenstorrent/tt-forge-models.

Directory structure is not uniformly flat — verified against the real repo. Models with one obvious task (ResNet, CLIP, DeiT) go straight to the framework directory; models that serve several downstream tasks (BERT, DINOv2, Llama) add a task-type directory in between:

tt-forge-models/
  resnet/
    pytorch/
      loader.py       # ModelLoader class
  bert/
    masked_lm/
      pytorch/
        loader.py
      jax/
        loader.py
    question_answering/
      pytorch/
        loader.py
    sequence_classification/
      pytorch/
        loader.py
  clip/
    pytorch/
      loader.py
  dinov2/
    image_classification/
      pytorch/
        loader.py
      jax/
        loader.py       # Flax variant
  llama/
    causal_lm/
      pytorch/
        loader.py

If an import at the depth shown above 404s, that’s why — check the task-type directories for the model you want rather than assuming the flat pattern.

Every loader.py exports a ModelLoader class with two instance methods (not static — ModelLoader.load_model(...) raises TypeError: missing 1 required positional argument: 'self', confirmed live; instantiate first). load_model() returns a standard PyTorch nn.Module and load_inputs() returns matching sample tensors — so you compile them exactly like any other model:

import torch, tt_torch
from third_party.tt_forge_models.bert.masked_lm.pytorch import ModelLoader

# Load the pretrained model and representative inputs
loader = ModelLoader()
model = loader.load_model(dtype_override=torch.bfloat16)
inputs = loader.load_inputs(dtype_override=torch.bfloat16)

# compile for Tensix and run — same torch.compile path as before
compiled = torch.compile(model, backend="tt").to(device)
output = compiled(inputs.to(device))

Five models worth knowing immediately:

Models not on this table: BLOOM, GPT-2, LLaMA, YOLOv4, BEiT, and 190+ more. Browse the full zoo in the forge-models-zoo lesson.

JAX and TT-XLA

For JAX and Flax models, the compilation path uses TT-XLA. Import pjrt_plugin_tt and the TT hardware backend registers automatically:

import jax
import jax.numpy as jnp
import pjrt_plugin_tt  # registers TT hardware as the XLA backend

# Any JAX function — jax.jit traces it and compiles to TT hardware
@jax.jit
def predict(params, x):
    return model.apply(params, x)

output = predict(params, batch)

The pjrt_plugin_tt import is the entire setup. After that, jax.jit compiles to Tensix cores instead of CPU or GPU. Flax transformer models slot directly into this pattern — load the model, load weights, wrap model.apply in jax.jit, run inference.

Full walkthrough: TT-XLA / JAX lesson.

Compiletron: The Expedition Game

Someone at Tenstorrent decided the best way to stress-test the compiler stack across hundreds of model architectures was to make it a roguelike game. They were right.

tt-forge-compiletron is a model-compilation expedition game (it lives at tenstorrent/tt-forge-compiletron). You launch expeditions into the model zoo. Each expedition compiles a different model. The chip runs the compilation. You score points. You build a bestiary.

Set it up, then start it:

cd ~/code/tt-forge-compiletron
bash scripts/install.sh          # installs forge venv, XLA venv, clones tt-forge-models
python3 expedition.py run --tui

A three-screen Textual TUI opens. The countdown is four seconds — then the expedition begins automatically.

The bestiary (data/bestiary.json) is a persistent record of every model you’ve successfully compiled. Base score per compile: 200 points. First time you compile a model, ever: multiplier of 5, making it 1,000 points. Freshness and rarity bonuses stack on top. The scoring structure incentivizes breadth: you gain more by compiling 10 new models than by recompiling the same model 10 times.

Compiletron supports both compiler backends from a single interface — verified against the actual --backend argument, which has four choices, not two:

Side quests activate when the mesh is busy with a large model compilation. Idle chips get assigned fast curated models to compile in parallel, keeping hardware utilization high and points accumulating while you wait. The game manages chip allocation automatically.

For unattended recording (VHS demos, overnight compilation runs), use --auto-quit SECS — verified against the actual argument: it’s a countdown in seconds after the summary screen appears, not a count of models compiled:

python3 expedition.py run --tui --auto-quit 30

The expedition runs its normal course; once the summary screen shows, the game waits 30 seconds and then quits on its own — bestiary saved, score written to disk — instead of waiting on a keypress.

All four chips busy — main expedition on Chip 0, three side quests running simultaneously.

Compiletron was built to find compiler bugs. It works through that bestiary systematically, surfacing edge cases in graph lowering and kernel generation that sequential targeted testing would miss. Every expedition you run contributes data to that effort. Points are real. The bestiary is real. And the compiler gets better.

← Performance Tuning | Going Deeper →