PantheonSim / GPUs / AWS Neuron machines

Simulated AWS Trainium and Inferentia machines

PantheonSim, the GPU simulator for CUDA and ROCm, also offers 4 AWS Neuron machines in the playground: Trainium and Inferentia. They have no profile like the NVIDIA and AMD cards do. They run AWS’s own CPU simulator for NKI kernels, nki.simulate, which is a functional simulation with no performance data.

The machines

The last column says whether the machine’s figures were read from a real chip. HBM and the NeuronCore count that were not are shown as not measured, never guessed. Inferentia2 is the same NeuronCore generation as Trainium1, so kernels behave the same on both. The figures come from the machine catalog the playground itself reads.

How they work

A session starts a login shell on Ubuntu 24.04 with the NKI package installed. Kernels are written in Python and run on the CPU, with tensors held as NumPy arrays in the session’s own memory. PantheonSim’s own profiles, and commands such as vgpu run --gpu, do not apply to these machines. They are chosen in the playground, one machine per session.

The simulation is functional. It checks what a kernel computes, and says nothing about how fast a real chip would run it. The AWS Trainium and Inferentia simulator guide walks through a first kernel.

Also: NVIDIA GPUs, AMD GPUs, Tenstorrent machines, or all of them.

Questions