Matmul tiling.

Matmul tiling

Matmul tiling.

Most of the arithmetic in a language model is multiplying tables of numbers. Tonight Pylon shows a classic trick for doing that fast on a graphics chip: carry a tray, not one can at a time.

One table times another

Multiplying two square tables, each number of the answer is a row of the first table times a column of the second: multiply pairs, add them up.

For tables 1,024 wide, one answer number needs 2,048 reads from the warehouse, 1,024 from each table.

One can at a time

The simple kernel gives each thread one answer number and lets it fetch everything itself. Every number of each table is then fetched 1,024 times, once for every thread that needs it.

In all, 2,147,483,648 reads from the warehouse, for one multiplication.

A tray to the bench

Now give a crew of threads a square of the answer. Together they carry a square tile of each table to their bench, the chip's small fast shared memory, and wait until every thread has put its part down.

Then each thread works through the tile on the bench. When every thread is done, the crew fetches the next pair of tiles. With tiles 16 wide, each number is fetched 64 times instead of 1,024.

Tiles 16 wide. How many times fewer reads from the warehouse? (the film pauses here)

Counting the reads

Reads from the warehouse: 2,147,483,648 for the simple kernel; 268,435,456 with tiles of 8; 134,217,728 with tiles of 16. The wider the tile, the fewer the trips.

Another way to say it: work done per byte fetched. The simple kernel does 0.25 operations per byte, counting each multiply and each add as one. Tiles of 16 do 3.97. Vasily Volkov and James Demmel used this counting in 2008 when they tuned matrix code for graphics chips.

Wait for the whole tray

The waiting matters. If a thread starts working before the crew has put the whole tile down, it can read a can from the last tile.

We broke the kernel that way on purpose. In our emulator, under the worst timing the missing wait allows, its answers came out wrong at every size we tried, and the checker, which compares each number with a careful answer computed on the main processor, caught it.

What the count leaves out

These are counts of reads asked for, not times. Real chips have caches that can serve repeated reads, so the simple kernel may not be as slow as its count says. How much it gains depends on the chip.

The lab runs both kernels on your own chip, if your browser allows it, and times them beside the count.

Goodnight

Carry a tray once, use it many times. Fast matrix code on graphics chips is built on the same idea.

Next in the cannery: attention itself, done in one trip. Goodnight.