trainium

Foundation model benchmarking tool. Run any model on any AWS platform and benchmark for performance across instance type and serving stack options.

Training Object Permanence in World Models — the codebase

Community maintained hardware plugin for vLLM on AWS Neuron

AI/ML and HPC at planetary scale. One API. Every Accelerator. Any Region. MIT-0 licensed

Predict AWS Trainium val_bpb before spending chip hours: a calibrated simulator + GPU proxy from the Trainium Frontier challenge

vLLM Omni backend plugin for diffusion and multimodal generation on AWS Trainium

Production LLM pipeline on AWS Trainium and Inferentia: LoRA fine-tune Llama 3.1 8B on a trn1.2xlarge, ship the adapter through S3, serve it with vLLM on an inf2.xlarge, and measure everything (TTFT/TPOT percentiles, tokens/s, MFU, goodput at SLO) with compile costs included and failures recorded as receipts.

AWS TorchNeuron Deep Learning Projects using Trainium1 Instances

Open-source calculator for AI accelerator compatibility across TPUs, Trainium/Inferentia, Gaudi, and NVIDIA GPUs

Streaming video diffusion on one AWS Trainium2 chip. StreamDiffusionV2’s causal Wan2.1 1.3B model ported to Neuron with a resident KV cache and a 4-NeuronCore layer pipeline: 5.6 → 38.4 FPS end to end, matching the CPU reference (cosine 0.99996). Plus a live SD-Turbo webcam restyle at 43.8 FPS.

A hybrid testbed for evaluating top open-source LLMs (like gpt-oss-20b and Llama 3.3) on local, cloud GPUs, and AWS Inferentia2/Trainium instances, focusing on vLLM optimization, capacity management, kernel bypass, hardware-software co-design, as well as supporting infrastructure such as NCCL, RDMA, NVMeoF.

DFlash speculative decoding for Qwen3-8B on AWS Trainium2 with vLLM-Neuron; developed at NYU × Annapurna Labs Hack the Chip.

Pseudo-spectral direct numerical simulation (DNS) of the Taylor-Green vortex on AWS Neuron: NKI kernels on Inferentia2 and Trainium1, 3-D FFTs as matmuls, all-to-all collectives inside the kernel and a libnrt C driver, in fp32 up to 512^3, checked against an fp64 oracle and Trainium1 references. Sample code for HPC and CFD engineers.