aws-neuron

A native PyTorch backend for AWS Trainium with support for eager execution, torch.compile, distributed training, and custom NKI kernels

Community maintained hardware plugin for vLLM on AWS Neuron

Open Telco AI (OTel): telecom LLMs, embeddings, rerankers, datasets, and benchmarks. Includes 27 training recipes across multiple software stacks and 10 inference stacks, spanning AMD, NVIDIA, Apple, Intel, and AWS Trainium hardware. Platform support varies by stack.

AWS Neuron's extension of TorchTitan that adds AWS Trainium support with TorchNeuron for training generative AI models

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

Terraform blueprints for running AI/ML workloads on Amazon ECS - vLLM, Triton, and TEI on GPU, Inferentia2, and Fargate

FFT and complex-valued tensor operations for AWS Trainium via NKI (cuFFT-equivalent) — Cooley-Tukey, Bluestein, STFT, ComplexTensor and complex NN layers.

Random number generation for AWS Trainium via NKI (cuRAND-equivalent) — Philox counter-based RNG, standard distributions, Sobol / Halton / Latin-hypercube quasi-random sequences for Monte Carlo and QMC.

Sentence Transformers on EC2 Inf1

Scientific computing library suite for AWS Trainium via NKI — the cuFFT/cuBLAS/cuRAND/cuSOLVER/cuSPARSE/cuTENSOR equivalents for Neuron. Python-first, PyTorch fallback everywhere, Apache-2.0.

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.

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.

BLAS Levels 1–3 for AWS Trainium via NKI (cuBLAS-equivalent) — GEMM with stationary-tile reuse, batched GEMM, TRSM, validated DF-MP2 for quantum chemistry.

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.

Sparse matrix operations for AWS Trainium via NKI (cuSPARSE-equivalent) — CSR/COO formats, SpMV and SpMM via gather-matmul-scatter, Schwarz integral screening for quantum chemistry.

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

Clean-room port of NVIDIA Aerial 5G L1 to AWS Trainium via NKI. Agentic, hardware-in-the-loop. Views are the author's own.

Linear solvers and eigendecomposition for AWS Trainium via NKI (cuSOLVER-equivalent) — Jacobi eigh, Cholesky/LU/QR factorizations, CG/GMRES iterative solvers, Newton-Schulz inverse square root.

Tensor contractions for AWS Trainium via NKI (cuTENSOR-equivalent) — einsum with contraction planning, CP/PARAFAC and Tucker decompositions, density-fitted post-Hartree-Fock patterns.