Matrix Multiplication: Composition and the GEMM That Runs AI

In the previous topic you learned the dot product: multiply two equal-length lists entry by entry and add up the result — one number that measures how aligned two vectors are.

A whole neural layer is the same question asked over and over. Imagine you have a batch of 32 inputs, and each of 512 neurons wants to score every input. That's 32 × 512 = 16,384 dot products, all at once.

So the question becomes

The answer is matrix multiplication. A matrix (rank-2 tensor: a grid of numbers, from Topic 1) is just a convenient pile of rows and columns, and matmul is "do the dot product of every row with every column."

The rule in one bold line

Formally, given A ∈ R^{m×k} and B ∈ R^{k×n}, the product C = AB ∈ R^{m×n} has entries:

Cᵢⱼ = Σₗ Aᵢₗ · Bₗⱼ

Read Σₗ as "add over the shared position." Two constraints to burn in:

• Inner shapes must match — A has k columns, B must have k rows. Otherwise: matmul shapes error.
• Output takes the outer shapes — (m×k)(k×n) → (m×n).

Algebra facts worth memorizing:

• Matrix multiplication is associative (AB)C = A(BC) and distributive over addition.
• It is not commutative: AB ≠ BA in general (and BA is often not even defined).
• Order encodes the order you apply transformations — that's the whole story in section 4.

import torch

A = torch.randn(32, 128)   # batch of 32 vectors, dim 128
W = torch.randn(128, 512)  # layer: maps 128 -> 512
C = A @ W                  # (32,512): inner 128 matches
print(C.shape)             # torch.Size([32, 512])

try:
    W @ A                  # (128,512) @ (32,128): 512 != 32
except RuntimeError as e:
    print("shape mismatch")  # the most common ML error there is

# Batched: attention scores for 8 heads at once
Q = torch.randn(8, 64, 256)   # (heads, seq, head_dim)
K = torch.randn(8, 32, 256)
scores = Q @ K.transpose(1, 2)  # (8, 64, 32) batched matmul

Smallest interesting case: one row times one column.

A = [1, 2] (shape 1×2) and B = [3, 4]ᵀ (shape 2×1, i.e. the column with 3 on top of 4).

• one dot product: 1·3 + 2·4 = 3 + 8 = 11
• result: the 1×1 matrix [[11]]

Now add a second row to A: A = [[1, 2], [0, 5]] (2×2), with the same column B = [[3], [4]] (2×1).

• row 1 · column: 1·3 + 2·4 = 11
• row 2 · column: 0·3 + 5·4 = 20
• result: [[11], [20]] — shape (2,1). Outer sizes 2 and 1 survived; the inner 2/2 vanished.

That vanishing is the rule in action. A had 2 columns, B had 2 rows → the shared 2 cancels, output takes the remaining outsides. Try to reverse it and B @ A gives shape (2,2) but completely different numbers — because AB ≠ BA.

Here is the deep "why". A matrix is a linear transformation — a rule that stretches, rotates, and shears space, fixed entirely by where it sends the basis vectors. Multiplying matrices composes those transformations: applying A after B is the single matrix AB.

code
Copy
   x ──W₁──►  (W₁x) ──W₂──►  W₂(W₁x) ──W₃──►  W₃(W₂(W₁x))
  input        layer 1 out      layer 2 out       logits

  read right-to-left: the RIGHTMOST matrix acts FIRST

Consequences that pay off later:

• Stacking layers means composing maps; the deeper the stack, the more the product matrix can amplify or shrink certain directions (bridge to eigenvalues, Topic 7, and exploding gradients, Topic 12).
• Because AB applies B first, transformations read right to left: y = W₃(W₂(W₁x)).
• If a product equals the identity matrix, the transformations undo each other (Topic 6, inverses).

The Mermaid diagram at the top draws the production version: input → matmul W₁ → ReLU → matmul W₂ → logits.

Carry one image through everything below: a factory where each machine reshapes the parts coming off the previous machine.

• A raw part = your input vector x.
• A machine = a matrix (a transformation: stretch, tilt, rotate).
• Bolt machine A onto machine B's output → you get one combined machine BA. That's matrix multiplication: composing the factory line, not running it twice.
• The part enters the rightmost machine first — the line reads right to left, like y = W₃(W₂(W₁x)).
• Two machines in the wrong order → broken part. AB and BA are different factory lines.
• A machine and its inverse, placed side by side, un-do each other: the product is the identity — the part comes out unchanged.

Now the punchline of the whole topic: a neural network is this factory. Each layer is "one matrix machine, then one small non-linear tweak" (ReLU), bolted into a long line. And the line only fits together if each joint's widths match — the inner-shape rule.

Back to the factory: once every layer is a matmul, the machine that runs matmuls is the machine that runs AI.

A dense matmul of (m×k)(k×n) costs about 2·m·k·n FLOPs. Large language model training is dominated by a handful of gigantic matmuls per step — so much so that accelerator benchmarks are basically matmul benchmarks.

• The BLAS kernel GEMM (GEneral Matrix Multiply) and its batched/strided variants (cuBLAS, cuBLASLt, oneDNN) execute the bulk of FLOPs in every training and inference run.
• A 70B-parameter forward pass performs roughly 2×70×10⁹ FLOPs per token — a parade of (T×d)(d×4d) style products.
• Tensor cores exist because dot-product accumulation is the workload: they fuse multiply-add tiles for 16-bit and FP8 inputs (2024–2026: H100/B200-era mixed precision).
• Inference at batch size 1 is memory-bound: the matmuls are skinny (m = 1), so reading weights from HBM, not FLOPs, sets the speed — the reason for batching, and for techniques like continuous batching in vLLM and TensorRT-LLM.

One corollary from the factory picture, useful in interviews: stacking two linear layers with no activation between them is just one combined machine. Linear∘linear = linear — nonlinearity between the matmuls is what gives depth any value.

Real tensors are rank-N, so you multiply along chosen axes. NumPy/PyTorch's @ broadcasts over leading batch dimensions, and einsum makes axis intent explicit:

• np.einsum("ij,jk->ik", A, B) — plain matmul.
• np.einsum("bhd,ehd->bhe", Q, K) — attention scores: batch × query-heads dot embedding-heads.

When code review asks "what are the shapes?", they are checking that inner dimensions agree and that batch axes broadcast. Practice reading products as "sum over the shared index" — the Σₗ in the formula from section 2, spoken aloud.