aws-neuron

Community maintained hardware plugin for vLLM on AWS Neuron

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

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

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

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.

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

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.

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.

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.

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.

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.

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

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