trnblas

BLAS operations for AWS Trainium via NKI (Neuron Kernel Interface).

Trainium ships no BLAS library. trnblas provides Level 1-3 BLAS operations with NKI kernel acceleration on the Tensor Engine, targeting scientific computing workloads that are GEMM-dominated.

Part of the trnsci scientific computing suite (github.com/trnsci).

Current phase

trnblas follows the trnsci 5-phase roadmap. Active work is tracked in phase-labeled GitHub issues:

• Phase 1 — correctness: complete as of v0.4.0 (GEMM, SYRK, MP2 energy reduction kernels hardware-validated on trn1; end-to-end DF-MP2 validated against PySCF at nanohartree tolerance).
• Phase 2 — precision (next): double-double FP64 GEMM for chemistry workloads. Unblocks trnsolver#27 and trntensor#28.
• Phase 3 — perf: tile sweeps, fused DF-MP2 kernels, true 3D batched GEMM, NEFF cache reuse.
• Phase 4 — multi-chip: tensor-parallel GEMM across NeuronCores.
• Phase 5 — generation: trn2 FP16-accumulate GEMM path.

Suite-wide tracker: trnsci/trnsci#1.

Why

NVIDIA has cuBLAS with 152 optimized routines. Trainium has torch.matmul. That's fine for ML training but insufficient for scientific computing codes that need TRSM, SYRK, SYMM, and batched GEMM with specific transpose/scaling semantics.

trnblas closes this gap — same BLAS API surface, NKI-accelerated GEMM on Trainium, PyTorch fallback everywhere else.

Install

pip install trnblas

# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
pip install "trnblas[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com

# Trn1 / PyTorch-XLA lineage
pip install "trnblas[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com

Kernels dispatch on whichever Neuron PyTorch lineage is installed:

• TorchNeuron (native): Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12, torch 2.14, torch-neuronx 2.14.*. Each kernel is a PyTorch custom op (torch.ops.trnblas.*) on the neuron device; no torch_xla.
• PyTorch/XLA: Trn1 (and any host on the torch-neuronx 2.9 lineage). Kernels dispatch through torch_xla, unchanged.

trnblas picks the native path when the native torch_neuronx and a neuron device are present, and the XLA path when torch_xla is installed; without nki it uses torch.matmul. Set TRNSCI_NKI_DISPATCH=native or xla (shared by all trnsci libraries) to force a path. With an editable install (pip install -e), set NKI_ENABLE_TRACE_CACHE=0: the NKI kernel cache only versions installed wheels and can serve stale kernels. See installation.

Usage

import torch
import trnblas

# Level 3 — Matrix multiply (the hot path)
C = trnblas.gemm(alpha=1.0, A=A, B=B, beta=0.5, C=C_init, transA=True)

# Batched GEMM (DF-MP2 tensor contractions)
C = trnblas.batched_gemm(1.0, A_batch, B_batch)

# Symmetric matrix multiply (Fock builds)
F = trnblas.symm(1.0, density, H_core, side="left")

# Triangular solve (Cholesky-based density fitting)
X = trnblas.trsm(1.0, L, B, uplo="lower")

# Symmetric rank-k update (metric construction)
J = trnblas.syrk(1.0, integrals, trans=True)

# Level 2 — Matrix-vector
y = trnblas.gemv(1.0, A, x, beta=1.0, y=y)

# Level 1 — Vector operations
y = trnblas.axpy(alpha, x, y)
d = trnblas.dot(x, y)
n = trnblas.nrm2(x)

DF-MP2 Example

# Run the density-fitted MP2 example
python examples/df_mp2.py --demo
python examples/df_mp2.py --nbasis 100 --nocc 20

The example demonstrates all core BLAS operations in a realistic quantum chemistry workflow: Cholesky factorization, triangular solve, half-transform GEMMs, metric contraction, and energy evaluation.

Real-molecule validation (via PySCF)

pip install trnblas[pyscf]
python examples/df_mp2_pyscf.py                       # H2O / STO-3G
python examples/df_mp2_pyscf.py --mol ch4 --basis cc-pvdz

Runs SCF + density fitting via PySCF, feeds the integrals through trnblas, and compares to PySCF's own DF-MP2 reference energy. Matches to < 10⁻⁷ Hartree on H2O, CH4, NH3 at cc-pvdz.

Operations

Status

• Level 1-3 BLAS with PyTorch backend
• GEMM with NKI dispatch stub
• DF-MP2 example
• NKI GEMM kernel validation on trn1/trn2
• NKI GEMM with stationary tile reuse
• Batched GEMM NKI kernel
• Double-double FP64 emulation
• Benchmarks vs cuBLAS

Related Projects

License

Apache 2.0 — Copyright 2026 Scott Friedman

Disclaimer

trnsci is an independent open-source project. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.

"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.

All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on trnsci.dev are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.

Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.

For official AWS guidance, see aws-neuron documentation and the AWS Trainium product page.