PantheonSim / AWS Trainium and Inferentia simulator

A simulated AWS Trainium or Inferentia in your browser

PantheonSim, the GPU simulator for CUDA and ROCm, also hosts AWS Trainium and Inferentia machines in its playground. You get a Linux terminal with AWS’s NKI package installed, write a kernel in NKI, the Neuron Kernel Interface, and run it on nki.simulate, AWS’s own CPU simulator for NKI kernels. It is a functional simulation: kernels run and give answers, but there is no performance data of any kind.

What you get

Choose an AWS Neuron machine in the playground and start it. The machine is a container with Ubuntu 24.04, Python 3.12 and NumPy, with AWS’s nki package installed. It has no Neuron chip in it, and no driver, runtime or compiler for one. Before the shell opens, the machine sets NKI_SIMULATOR=1, so an @nki.jit kernel you call with NumPy arrays runs on the simulator instead of looking for a device.

Four machines are offered. Each is one simulated NeuronCore target, so the count is always one. Inferentia2 uses the same NeuronCore generation as Trainium1, and the two behave the same.

The target is the value of NEURON_PLATFORM_TARGET_OVERRIDE that the simulator reads. The machines show no memory size, because the simulator allocates each tensor on its own in the session’s memory and does not model the size of the on-chip buffers.

What runs it

The simulator is nki.simulate, part of the AWS Neuron nki package, which AWS publishes under the Apache-2.0 licence on its own package index. It is not PantheonSim’s simulator and PantheonSim did not change it: the playground installs AWS’s release, pinned by version and checksum, and keeps its licence and notice in the image. PantheonSim contributes the machine around it: the container, the editor and terminal in the browser, the limits, and the clean-up when the session ends. PantheonSim is not affiliated with or endorsed by Amazon Web Services.

That division matters when something looks wrong. A message about an NKI instruction, a tile shape or a data type comes from the simulator and is AWS’s to answer; a machine that will not start or a file that will not save is the playground’s.

Start a machine and run a kernel

1. Open the playground and choose AWS Trainium (trn1). The operating system is fixed at Ubuntu 24.04, because that is what the image is built on.
2. Press Start machine. The editor opens with three kernels: vector_add.py, matmul.py and exp.py.
3. Press Run. The editor saves the file and types python3 vector_add.py into the terminal.

This is vector_add.py, the first kernel:

"""Add two tensors on a simulated NeuronCore.

An NKI kernel is a Python function decorated with @nki.jit. On this machine it
runs on AWS's CPU simulator (NKI_SIMULATOR=1 is already set); the same file runs
unchanged on a Trainium or Inferentia instance, where the kernel is compiled and
runs on the device.
"""
import numpy as np
import nki
import nki.language as nl


@nki.jit
def vector_add(a_ptr, b_ptr):
    a = nl.load(a_ptr)                    # HBM -> SBUF; the first axis is the partition axis (at most 128)
    b = nl.load(b_ptr)
    out = nl.ndarray(a.shape, a.dtype, buffer=nl.shared_hbm)   # a kernel's output lives in shared HBM
    nl.store(out, nl.add(a, b))           # SBUF -> HBM
    return out


a = np.random.rand(128, 512).astype(np.float32)
b = np.random.rand(128, 512).astype(np.float32)
result = vector_add(a, b)
error = float(np.abs(result - (a + b)).max())
print(f"vector_add: max abs error against NumPy {error:.2e}")
print("PASS" if error == 0.0 else "FAIL")

In the terminal, nki-info prints the machine, the target and the versions the way nvidia-smi prints a GPU. The self-test, /opt/nki/selftest <target>, runs every example and checks the answers. The Tools menu in the playground has both, and a list of the examples in /work/examples.

A kernel is ordinary Python, so there is nothing to build. The same file, called with NumPy arrays on a real Trainium or Inferentia instance, is compiled and runs on the chip.

Checked against a real chip

A simulator can be wrong about the chip it stands for, so the same kernel files were run both ways on 24 October 2026, on a trn1.2xlarge and an inf2.xlarge, with the simulator and with the chip. Both instances gave the same results.

Two differences matter in practice. A kernel’s output has to be in shared HBM (buffer=nl.shared_hbm) to compile for the chip; the simulator also accepts plain HBM, so a kernel that passes here with plain HBM fails to compile on the device, and the examples use shared HBM for that reason. And a floating point matrix multiply on the chip is not bit for bit the simulator’s result, so compare with a tolerance, not with equality. The Trainium2 and Trainium3 targets were not checked against a real chip.

What it cannot tell you

• Speed. The simulator models no engine timing and no scheduling. A kernel that takes a second here says nothing about how long it takes on a chip, and neither does anything else a simulated run prints.
• Whether it compiles for the chip. The simulator runs kernels as Python with no compile step, so it accepts things the NKI compiler rejects. It does enforce some chip rules, such as at most 128 partitions, no copy straight from PSUM and an fp32 matrix multiply output on Trainium1 and Inferentia2, but AWS lists others as not checked yet. Compile on a real instance before you rely on a kernel.
• Memory limits. The on-chip SBUF and PSUM sizes are not enforced, so a kernel that needs more than the chip has still runs.
• The rest of the Neuron software. The machine has NKI and NumPy only. There is no PyTorch, no torch-neuronx and no compiler for whole models, and no neuron-ls or neuron-monitor. Only part of the collectives are supported by the simulator.
• The network. By default the machine has none, so bring files in with the editor’s Open button.

Like every playground machine it is shut down after a time limit and after a period without input, and its files are deleted with it.

Questions

For NVIDIA, AMD and Tenstorrent machines in the same playground, see running CUDA without a GPU, running ROCm and HIP without an AMD GPU, the Tenstorrent simulator, or the GPU index.