Matrix multiplication is all you need: shapes, batching, and why GPUs love it

LEARN · MATHEMATICS FOR MACHINE LEARNING

Why matrix multiplication matters

If you learn one operation in modern machine learning, make it matrix multiplication (often shortened to matmul).

Almost every neural network eventually boils down to a handful of highly optimized matrix multiplications. Dense layers, attention, embeddings, projections, convolutions (after lowering or specialized kernels), and many scientific computing workloads spend most of their execution time multiplying matrices.

Understanding matrix multiplication is less about memorizing formulas and more about understanding shapes.

In this guide you’ll learn:

How matrix multiplication works from a shape perspective
How a linear layer is just one matrix multiplication
How batching works without changing the math
Why GPUs are exceptionally good at matmul
A simple CPU vs GPU benchmark you can run
A surprising fact about how much of an LLM’s computation is just matrix multiplication

Given:

A matrix A with shape (m, k)
A matrix B with shape (k, n)

Their product has shape:

(m, k) × (k, n) → (m, n)

The inner dimensions must match.

The output keeps the outer dimensions.

That’s the entire rule.

For example:

A = (3, 4)
B = (4, 2)

Output = (3, 2)

Visualized:

(3 × 4)
      ×
(4 × 2)

↓

(3 × 2)

Each output value is produced by taking:

one row from the first matrix
one column from the second matrix

Then computing their dot product.

Example:

A =
[[1, 2],
 [3, 4]]

B =
[[5, 6],
 [7, 8]]

The top-left element becomes:

1×5 + 2×7 = 19

The top-right element becomes:

1×6 + 2×8 = 22

And so on.

import numpy as np

A = np.array([
    [1, 2],
    [3, 4]
])

B = np.array([
    [5, 6],
    [7, 8]
])

C = A @ B

print(C)

Output:

[[19 22]
 [43 50]]

The @ operator is Python’s matrix multiplication operator.

Suppose we have:

Input:
(5, 3)

Weights:
(3, 4)

Then:

(5, 3)
    @
(3, 4)

↓

(5, 4)

Notice that we never needed to know the numbers.

Only the shapes matter.

This “shape-first” way of thinking is how deep learning practitioners debug models.

A fully connected layer computes:

output = input @ weights + bias

Nothing more.

Imagine:

Batch:
(32, 128)

Weights:
(128, 256)

Bias:
(256,)

Then:

(32,128)
      @
(128,256)

↓

(32,256)

+

(256,)

The bias is broadcast across every row.

import numpy as np

batch = np.random.randn(4, 3)

weights = np.random.randn(3, 5)

bias = np.random.randn(5)

output = batch @ weights + bias

print(batch.shape)
print(weights.shape)
print(bias.shape)
print(output.shape)

Output:

(4, 3)
(3, 5)
(5,)
(4, 5)

That is exactly what happens inside a dense neural network layer.

Instead of processing one example:

(128,)

we stack many together:

(batch_size, features)

Example:

64 images

↓

(64, 784)

Weights:

(784, 512)

Result:

(64, 512)

Every example is multiplied by the same weight matrix simultaneously.

The math stays identical.

Only the first dimension grows.

Deep learning libraries also support batched matrix multiplication.

Suppose:

(8, 32, 64)

This means:

8 batches
32 vectors
64 features

Multiply by:

(8, 64, 128)

Result:

(8, 32, 128)

Each batch performs its own matrix multiplication independently.

This pattern appears constantly in transformer models.

Inside self-attention, projections are implemented as matrix multiplications.

For example:

X
↓

Q = X @ Wq

K = X @ Wk

V = X @ Wv

Attention itself computes:

Attention(Q, K, V) = softmax(Q·Kᵀ / √dₖ) · V

Notice that the heavy operations are:

Q × Kᵀ
attention × V

Both are matrix multiplications.

The feed-forward network that follows attention is also two large matrix multiplications with a nonlinear activation in between.

Modern LLM inference therefore spends the overwhelming majority of its arithmetic inside highly optimized matmul kernels.

GPUs were built to execute many arithmetic operations simultaneously.

Matrix multiplication is nearly ideal because:

every output element is independent
thousands of threads can compute different outputs at once
memory access patterns are predictable
arithmetic intensity is high, meaning many calculations happen for each byte loaded from memory

Modern GPU architectures also include specialized hardware:

NVIDIA Tensor Cores
AMD Matrix Cores
Intel XMX units
Apple AMX acceleration on Apple Silicon CPUs and Metal-accelerated GPU kernels

These units are designed specifically to accelerate matrix multiplication and related tensor operations.

Deep learning frameworks automatically use them whenever possible.

The following example uses PyTorch.

If CUDA is available, it measures GPU performance.

import time
import torch

size = 4096

A = torch.randn(size, size)
B = torch.randn(size, size)

start = time.perf_counter()
C = A @ B
cpu_time = time.perf_counter() - start

print(f"CPU: {cpu_time:.3f} seconds")

if torch.cuda.is_available():
    device = torch.device("cuda")

    A = A.to(device)
    B = B.to(device)

    torch.cuda.synchronize()

    start = time.perf_counter()

    C = A @ B

    torch.cuda.synchronize()

    gpu_time = time.perf_counter() - start

    print(f"GPU: {gpu_time:.3f} seconds")

Your exact numbers will depend on hardware, but on a modern desktop GPU the matrix multiplication is often several times to well over an order of magnitude faster than on a typical CPU for large matrices.

Small matrices may not benefit because launching GPU work has overhead.

Inner dimensions don’t match

Incorrect:

(32, 128)
@
(64, 256)

The middle numbers differ.

128 ≠ 64

This raises a runtime error.

Mixing batch and feature dimensions

Many beginners accidentally write:

(features, batch)

instead of:

(batch, features)

Most machine learning frameworks expect:

(batch, features)

for dense layers.

Forgetting the transpose

Suppose:

Input:
(64, 512)

Weights:
(128, 512)

This cannot multiply directly.

Often the intended operation is:

input @ weights.T

giving:

(64,512)

@

(512,128)

↓

(64,128)

Always verify the expected orientation of your weight matrix.

You almost never implement matrix multiplication yourself.

Frameworks dispatch to highly optimized libraries such as:

cuBLAS on NVIDIA GPUs
oneDNN on Intel CPUs
BLIS or OpenBLAS on many CPU platforms
Metal Performance Shaders on Apple platforms

These libraries use techniques including:

cache blocking
vectorized instructions
fused operations
kernel autotuning
asynchronous execution
mixed-precision acceleration

As a result, calling a single @ operator may launch one of the most optimized numerical kernels on your system.

Here’s a surprising fact.

For today’s transformer-based large language models, the overwhelming majority of floating-point operations during both training and inference come from matrix multiplications.

The largest contributors are:

query, key, and value projections
attention matrix multiplications
output projections
feed-forward (MLP) layers

Everything else—activations, normalization, masking, residual additions, and element-wise operations—typically accounts for only a relatively small fraction of total FLOPs.

That is why GPU vendors invest so heavily in specialized matrix engines. Making matrix multiplication even a little faster has an outsized impact on end-to-end model performance.

Remember these rules:

Matrix multiplication follows (m, k) × (k, n) → (m, n)
Only the inner dimensions must match.
The output keeps the outer dimensions.
A dense neural network layer is simply input @ weights + bias.
Batching adds dimensions without changing the underlying mathematics.
GPUs excel because matrix multiplication exposes enormous parallelism.
Most of the compute in modern transformer models is spent performing highly optimized matrix multiplications.

Once you become comfortable reading tensor shapes, neural network architectures become dramatically easier to understand and debug.

What’s next?

Open a Python notebook, recreate every example in this guide, and experiment with the shapes. Change batch sizes, feature counts, and weight dimensions, then predict the output shape before running the code. Developing that intuition is one of the fastest ways to build confidence with linear algebra and modern deep learning.