Tensor cores and matrix math: the operation at the heart of AI chips

Updated

Tensor cores are specialized processor units that perform the matrix multiplications at the core of neural networks. They trade some precision for much higher throughput, using formats like TF32 and FP16. This article explains the operation, its role in deep learning, and how to use it efficiently.

Matrix multiplication in deep learning

Neural networks, including convolutional networks and transformers, rely heavily on matrix multiplications. A fully connected layer computes a product of a weight matrix and an input vector, while a convolutional layer performs a series of dot products between filters and local image regions, as described in the Stanford CS231n course notes. These operations are often expressed as general matrix multiplications (GEMMs), the same building blocks used in linear algebra libraries.

GEMMs are so fundamental that their performance often determines the speed of training and inference. The NVIDIA matrix multiplication guide notes that fully connected, recurrent, and convolutional layers all rely on GEMM operations. Because these operations repeat millions of times during training, even small improvements in matrix multiplication efficiency translate into large overall gains.

What tensor cores do

Tensor cores are hardware units inside modern GPUs that perform matrix multiply-accumulate operations in a single step. They take two small matrices, multiply them, and add the result to an accumulator, as the NVIDIA Tensor Cores page explains. This differs from general-purpose CUDA cores, which handle one arithmetic operation at a time.

Tensor cores are designed to accelerate the mixed-precision computing used in AI workloads. They dynamically adapt calculations to increase throughput while preserving accuracy, according to the same NVIDIA page. This makes them central to training large models and running inference efficiently.

Precision and throughput

Tensor cores operate on reduced-precision numbers to achieve higher throughput. The TF32 format, introduced with the Ampere architecture, uses 8 exponent bits and 10 mantissa bits, matching the range of FP32 while offering the precision of FP16, as detailed in the NVIDIA TF32 blog post. TF32 is the default mode for single-precision training on Ampere GPUs and requires no script changes.

Mixed precision training, described in the NVIDIA mixed precision guide, stores weights and activations in FP16 while accumulating in FP32. This reduces memory usage and bandwidth, and on GPUs with tensor cores it can be several times faster than FP32 training. The guide explains that loss scaling preserves small gradient values that would otherwise underflow in FP16.

The NVIDIA Ampere architecture blog lists additional supported precisions on third-generation tensor cores, including BF16, FP64, INT8, INT4, and binary formats. This variety lets developers choose the best precision for training or inference depending on accuracy and speed requirements.

Efficient use of tensor cores

Tensor core performance depends on matrix dimensions. The NVIDIA matrix multiplication guide recommends that the key dimensions M, N, and K be aligned to multiples of 16 bytes for best efficiency. For FP16 data, that means dimensions should be multiples of 8 elements, and on A100 GPUs, multiples of 64 are even better.

Arithmetic intensity, or the ratio of operations to memory accesses, also matters. The same guide provides an example: a GEMM of size 8192 x 128 x 8192 has an intensity of 124.1 FLOPS per byte, which is below the V100's ops-to-byte ratio, so it is memory limited. A larger GEMM of 8192 x 8192 x 8192 has an intensity of 2730, making it math limited. Choosing block sizes that keep arithmetic intensity high helps saturate tensor cores.

Relative math throughput of A100 GPU vs FP32FP32: 1 x (relative to FP32); TF32: 8 x (relative to FP32); FP16 / BF16: 16 x (relative to FP32)Relative math throughput of A100 GPU vs FP32FP321 x (relative to FP32)TF328 x (relative to FP32)FP16 / BF1616 x (relative to FP32)
Figures from NVIDIA
Recommended alignment for tensor core efficiency on recent cuBLAS and cuDNN versions
PrecisionAlways usableMost efficient multipleMost efficient on A100
INT8Yes16128
FP16Yes864
TF32Yes432
FP64Yes216

Frequently asked questions

What is a tensor core?

A tensor core is a specialized hardware unit in some GPUs that performs a matrix multiply-accumulate operation in one step, speeding up the matrix multiplications common in neural networks.

How does TF32 work?

TF32 is a precision format used on Ampere GPUs that has 8 exponent bits and 10 mantissa bits. It matches the range of FP32 but offers lower precision, which allows tensor cores to run faster while keeping accuracy sufficient for AI training.

Why use mixed precision training?

Mixed precision training uses FP16 for most operations and FP32 for critical accumulation and gradient scaling. This reduces memory usage and bandwidth while speeding up arithmetic on tensor cores, and loss scaling prevents small gradients from underflowing.

What alignment is needed for tensor cores?

Matrix dimensions should be multiples of 16 bytes for best efficiency on recent cuBLAS and cuDNN versions. For FP16, that is multiples of 8 elements, and on A100 GPUs, multiples of 64 elements yield even better performance.

How much faster are tensor cores than normal cores?

On the A100 GPU, tensor core throughput for TF32 is 8 times that of FP32, and for FP16 it is 16 times, according to NVIDIA. Actual speedups depend on the workload and how much of it consists of matrix multiplications.