# Large Scale Distributed Linear Algebra With Tensor Processing Units

arXiv: https://arxiv.org/abs/2112.09017 · HTML: https://mlsanity.com/paper/2112.09017/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 16 Dec 2021.

## Built on by (most cited first, 4 of 4)

- [Solving the Discretised Neutron Diffusion Equations using Neural Networks](https://mlsanity.com/paper/2301.09939/) (2023)
- [Neural Physics: Using AI Libraries to Develop Physics-Based Solvers for Incompressible Computational Fluid Dynamics](https://mlsanity.com/paper/2402.17913/) (2024)
- [Distributed computing for physics-based data-driven reduced modeling at scale: Application to a rotating detonation rocket engine](https://mlsanity.com/paper/2407.09994/) (2024)
- [Comprehensive Performance Modeling and System Design Insights for Foundation Models](https://mlsanity.com/paper/2410.00273/) (2024)

## Abstract (arXiv)

We have repurposed Google Tensor Processing Units (TPUs), application-specific chips developed for machine learning, into large-scale dense linear algebra supercomputers. The TPUs' fast inter-core interconnects (ICI)s, physically two-dimensional network topology, and high-bandwidth memory (HBM) permit distributed matrix multiplication algorithms to rapidly become computationally bound. In this regime, the matrix-multiply units (MXU)s dominate the runtime, yielding impressive scaling, performance, and raw size: operating in float32 precision, a full 2048-core pod of third generation TPUs can multiply two matrices with linear size N= 220= 1 048 576 in about 2 minutes. Via curated algorithms emphasizing large, single-core matrix multiplications, other tasks in dense linear algebra can similarly scale. As examples, we present (i) QR decomposition; (ii) resolution of linear systems; and (iii) the computation of matrix functions by polynomial iteration, demonstrated by the matrix polar factorization.