Matrix Multiplication on a GPU

The idea in one minute#

Almost all the time in a neural network is spent multiplying matrices. The textbook three-loop algorithm is correct and hopelessly slow, because it re-reads the same data from far-away memory over and over. The fast version cuts the matrices into tiles small enough to sit in fast memory, and reuses each tile many times before fetching the next.

Every fast matmul — on CPU or GPU, in any library — is that one idea applied carefully.

An analogy#

You are cooking from 100 recipes that share the same 20 ingredients. The slow way: for each step of each recipe, walk to the pantry, fetch one ingredient, walk back. The fast way: bring the 20 ingredients to the counter once and make all the recipes from there.

Tiling brings the ingredients to the counter.

A picture#

flowchart LR
  subgraph G["Global memory (slow, huge)"]
    A[("Matrix A")]
    B[("Matrix B")]
    C[("Matrix C")]
  end
  subgraph S["One block on one SM"]
    TA["Tile of A<br/>in shared memory"]
    TB["Tile of B<br/>in shared memory"]
    MUL["Threads multiply tiles<br/>tensor cores, registers"]
    TA --> MUL
    TB --> MUL
  end
  A -->|"load tile once, coalesced"| TA
  B -->|"load tile once, coalesced"| TB
  MUL -->|"write result tile once"| C
  class A,B,C memory
  class TA,TB queue
  class MUL compute

How it really works#

The operation#

C = A × B, where A is m×k and B is k×n. Each output element is a dot product:

C[i][j] = Σ over p of A[i][p] × B[p][j]

Work: 2·m·n·k FLOPs. Data: m·k + k·n + m·n numbers. For square matrices that is 2n³ FLOPs on 3n² numbers — intensity grows with n, so large matmuls are compute-bound in principle.

Libraries call this operation GEMM (general matrix-matrix multiply). Matrix-vector is GEMV.

Why the naive version is slow#

In the naive loop, computing one row of C walks across all of B. If B is larger than the cache, every element is fetched from far memory again for the next row. The algorithm should be compute-bound, but it has turned itself into a memory-bound one through bad access order.

Tiling#

Split A, B and C into tiles of, say, 32×32. For each output tile:

Load one tile of A and one of B into shared memory — two clean, coalesced reads.
Every thread in the block computes with them; each loaded number is used 32 times.
Move to the next pair of tiles along the row/column, accumulating in registers.
Write the finished output tile once.

Global memory traffic falls by roughly the tile size. Now the operation is as compute-bound as the mathematics says it should be, and tensor cores (which consume small tiles natively) can run flat out.

What a real library adds#

cuBLAS, and the kernels inside frameworks, add: multiple tile levels matched to registers, shared memory and L2; tensor-core instructions; different kernels for different shapes; and heavy per-GPU tuning. You should never write production GEMM yourself. You should understand it well enough to know when you are not getting it — for example when odd dimensions or FP32 inputs silently push you onto a slow path.

Why batch-1 LLM decoding is different#

Decode for one sequence multiplies each weight matrix by one vector (GEMV). There is nothing to reuse: every weight is needed exactly once. Tiling cannot help, intensity is ~1, and the operation is memory-bound by nature (IV.01). Batching turns GEMV back into GEMM — each weight loaded is now used once per sequence — which is why it recovers the GPU’s arithmetic.

Code#

Naive versus tiled, in Go. Same FLOPs, same answer.

// matmul.go — the same multiplication with bad and good memory behaviour.
package main

import (
	"fmt"
	"time"
)

const n = 1024

// naive: for each output element, walk a row of A and a COLUMN of B (a cache miss every step).
func naive(a, b, c []float32) {
	for i := 0; i < n; i++ {
		for j := 0; j < n; j++ {
			var sum float32
			for p := 0; p < n; p++ {
				sum += a[i*n+p] * b[p*n+j]
			}
			c[i*n+j] = sum
		}
	}
}

// tiled: work on T×T tiles so the data in use stays in fast memory and is reused.
func tiled(a, b, c []float32) {
	const T = 64
	for i0 := 0; i0 < n; i0 += T {
		for p0 := 0; p0 < n; p0 += T {
			for j0 := 0; j0 < n; j0 += T {
				for i := i0; i < i0+T; i++ {
					for p := p0; p < p0+T; p++ {
						aip := a[i*n+p]
						row, out := b[p*n+j0:p*n+j0+T], c[i*n+j0:i*n+j0+T]
						for j := range row {
							out[j] += aip * row[j]
						}
					}
				}
			}
		}
	}
}

func main() {
	a, b := make([]float32, n*n), make([]float32, n*n)
	for i := range a {
		a[i], b[i] = float32(i%7), float32(i%5)
	}
	for _, v := range []struct {
		name string
		fn   func(a, b, c []float32)
	}{{"naive", naive}, {"tiled", tiled}} {
		c := make([]float32, n*n)
		t0 := time.Now()
		v.fn(a, b, c)
		dt := time.Since(t0).Seconds()
		fmt.Printf("%s: %6.0f ms  %5.2f GFLOP/s  (c[12345]=%v)\n", v.name, dt*1e3, 2*n*n*n/dt/1e9, c[12345])
	}
}

Expect several-fold between the two on one CPU core. A BLAS library adds SIMD and threads for another 50–100x; a GPU with tensor cores another 100x beyond that. The idea is the same at every level.

Remember this#

Neural networks are mostly GEMM; fast GEMM is tiling plus reuse.
Tiles live in shared memory / registers; global memory is touched once per tile.
Tensor cores consume tiles directly. Friendly shapes and low-precision formats keep you on the fast path.
Batch-1 decode is GEMV: no reuse possible, memory-bound. Batching makes it GEMM again.

Try it#

Run matmul.go. Try tile sizes 8, 32, 64, 256. Where is the sweet spot, and what does it say about your CPU’s cache sizes?
Parallelize tiled across goroutines by i0. Why is that safe without locks?
For d = 4096, compute FLOPs and bytes for one GEMV and for a GEMM with 64 vectors. What are the two intensities?

Check yourself#

Why is naive matmul memory-bound in practice although it is compute-bound in theory?
What does tiling change?
Why can tiling not help a single-sequence decode step?