TrnRT

// Overview

TrnRT is an open-source inference engine for large language models on AWS Trainium, built by RunAnywhere and released under the Apache-2.0 license. It serves an OpenAI-compatible API from a single `trn2.48xlarge` instance (16 Trainium2 chips). GLM-5.3-Flash — Z.ai's 320B-parameter Mixture-of-Experts model with 18B active parameters, run in FP8 with a 1M-token context — is the first and so far only supported model; the README says the engine, router, caches and compile pipeline are built to carry more model families.

The engine's distinguishing idea is the decode megakernel. Instead of a chain of small kernels, each layer half runs as one fused kernel written in NKI, the Neuron Kernel Interface, covering the KDA attention, the fused DSA layer, mHC and the MoE feed-forward block. The project's reasoning is that decode on Trainium is dominated by small instructions, kernel boundaries and exposed collectives rather than memory bandwidth, so each kernel keeps its working set in on-chip SBUF memory, uses both physical cores symmetrically and moves the collectives inside the kernel. RunAnywhere's launch post reports that this fusion cut a batch-8 decode step from 32.2 ms to 24.2 ms.

Around the kernels sits a serving stack: two engines per box, each tensor-parallel over 32 logical NeuronCores, behind a session-affine router on port 8000; a radix-tree prefix cache over FP8 KV pages with a host-memory tier for evicted pages and state checkpoints; and an API that handles chat and completions, streaming, tool calls and reasoning output. Compilation is done offline on any x86_64 host with no NeuronCore — a five-step pipeline (shapes, decisions, capture, compile, strip) that takes hours — so the serving box loads a prebuilt cache and compiles nothing by default.

On SemiAnalysis's AgentX harness, which replays real coding-agent sessions, the README reports 3.57M tokens per minute (28.9 requests per minute) on one `trn2.48xlarge` at decode batch 12 with 32 concurrent agents and a 1M-token context, which it prices at $0.040 per million tokens at the spot rate and $0.167 on a capacity block. The launch post notes these runs are self-run, that about 95% of counted tokens are cached prompt input, and that there is no GPU baseline on the same traces.

// What it does

• NKI decode megakernels — one fused kernel per layer half for KDA attention, the fused DSA layer, mHC and the MoE FFN
• 1M-token context through sparse attention (top 2,048) and linear attention with fixed-size state
• A prefix cache built as a radix tree over FP8 KV pages, with a host-memory tier for evicted pages and state checkpoints
• Two tensor-parallel engines per `trn2.48xlarge` behind a session-affine, prefix-aware router
• OpenAI-compatible chat and completions with streaming, tool calls and reasoning output
• Offline compile on any x86_64 host, so the serving instance loads a ready compile cache instead of compiling

// Getting started

You need Docker, a `trn2.48xlarge` to serve on and any x86_64 host to compile on; the project's docs/AWS.md covers Trainium setup for newcomers. These steps follow the README's quickstart.

git clone https://github.com/RunanywhereAI/trnrt.git && cd trnrt
docker build -f docker/Dockerfile \
  --build-arg TRNRT_VERSION=$(python3 -c 'import trnrt; print(trnrt.__version__)') \
  --build-arg TRNRT_REVISION=$(git rev-parse HEAD) -t trnrt .

OUT=/data/trnrt REV=eb9eb208eb0d988989d07a6a12d0fdeb5f52574a
mkdir -p $OUT/nki $OUT/work
trnrt() { docker run --rm -e HF_TOKEN -v $OUT:$OUT \
  -v $OUT/nki:/var/tmp/nki-intermediate-cache trnrt "$@"; }
trnrt compile shapes zai-org/GLM-5.3-Flash --revision $REV --out $OUT
trnrt compile decisions --model zai-org/GLM-5.3-Flash --revision $REV --out $OUT
trnrt compile capture --out $OUT -- --model $OUT/shapes
trnrt compile compile --out $OUT --work-root $OUT/work
trnrt compile strip --out $OUT

hf download zai-org/GLM-5.3-Flash --revision eb9eb208eb0d988989d07a6a12d0fdeb5f52574a \
  --local-dir /models/GLM-5.3-Flash
docker run --rm --privileged --network host --ipc=host \
  -v /data/trnrt/cache:/cache \
  -v /data/trnrt/nki:/var/tmp/nki-intermediate-cache \
  -v /models/GLM-5.3-Flash:/model \
  -e NEURON_LIBTORCH_CACHE_ROOT=/cache \
  trnrt serve --model /model --port 8000

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused")
reply = client.chat.completions.create(
    model="glm",
    messages=[{"role": "user", "content": "Write a haiku about Trainium."}],
    extra_body={"session_id": "my-session"},
)
print(reply.choices[0].message.content)

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// When to use it

• Reach for it when you want to serve GLM-5.3-Flash on AWS Trainium2 capacity instead of GPUs
• Reach for it for long-running coding-agent workloads with long, heavily reused prompts, where the session-affine router and prefix cache pay off
• Reach for it as a worked reference for writing fused NKI megakernels on Trainium — docs/MEGAKERNELS.md walks through every kernel

// How TrnRT compares

TrnRT alongside other open-source serving & deployment tools AI/TLDR tracks, ranked by GitHub stars.