puffergpt

How this works

The playground runs GPT-2 small, the 124-million-parameter model OpenAI released in 2019. Its weight matrices are not on this server. Every row of every matrix is a document in turbopuffer, and every matrix multiply is a vector search.

What GPT-2 does for one token

1. It looks up a vector of 768 numbers for the token and another for its position, and adds them.
2. It passes that vector through 12 layers. Each layer has two parts. In attention, the token asks which earlier tokens matter to it and mixes in information from them. Then a small feed-forward network multiplies the vector by a matrix, squashes the result, and multiplies by a second matrix.
3. It compares the final vector with the vector of each of the 50,257 tokens. The closest matches are the likely next tokens.

It picks one, feeds it back in as the next input, and repeats. Text you type goes through the same steps, so reading a token costs the same as writing one.

A matrix multiply is a search

Multiplying a vector by a weight matrix means taking its dot product with every row. A vector search measures the distance from a query to every stored vector, which is nearly the same thing. Each row is stored as one document, with its bias, and padded so that every document of a matrix has the same length M:

document  d = (w, b, pad)
pad       √(M² − |w|² − b²)
query     q = (x, 1, 0)
distance  |d − q|²
          = M² + |q|² − 2·(w·x + b)

Here w is one row of the matrix, b is its bias and x is the layer's input. Because every document has length M, the distance depends only on w·x + b, and the layer's output comes straight out of the distances: y = (M² + |q|² − distance) / 2. A filter picks the matrix and the search asks for every row, so one request returns one whole multiply.

This relies on turbopuffer's exact search, which returns the true distance to every document that matches the filter.

One token is 76 requests

The requests are sent in the order shown, one row at a time. Each takes 20 to 30 ms, which adds up to a little under two seconds a token. The first token of a text has no earlier tokens to attend to, so it takes 52.

Attention

Attention has the shape of a search to begin with: a query is scored against stored keys. In each layer, 12 heads each turn the current token into three vectors of 64 numbers: a query, a key and a value.

The 144 keys are written to turbopuffer, in a namespace that belongs to the text, each tagged with its layer, head and position. Each query is searched against the keys of the earlier positions, filtered to its own layer and head. The distances become scores, and the scores set how much of each earlier token's value is mixed in. The values stay in this server's memory.

The next token

The last step is an ordinary nearest-neighbour search. The final vector is the query, the 50,257 token vectors are the documents, and the 40 nearest come back with their scores. The next token is picked from those 40, or from however many the Top-k box asks for. Click a generated token in the playground to see the top of that list.

What this server still does

The arithmetic that involves no weight matrix: normalising the vector before each step, the squashing function in each feed-forward network, turning attention scores into weights, the additions, and picking the token. The normalising uses 38,400 numbers that are also kept in turbopuffer and fetched when a text is opened.

Stepping back

A text's keys stay in its namespace, so trying again, branching from another candidate and cutting the text only pay for tokens the model has not read. Stepping back deletes nothing. Attention only searches positions before the current one, and new tokens overwrite the old keys.

What is stored

134,299 documents, 501 MB. The 770 and 3,074 are vector lengths: the input size, plus one for the bias and one for the padding. Each text adds a namespace for its keys, deleted when the text is closed.

In numbers

The costs are worked out from turbopuffer's published rates. The comparison ran a ten-token prompt through turbopuffer and through GPT-2 in numpy, and compared the 40 scores that came back. At temperature 0 the first starter prompt produces the same eight tokens either way.

Running GPT-2 the ordinary way is far faster. This is the same computation, done as searches.

Inspiration

• GPT in 500 lines of SQL, by Alex Bolenok. GPT-2 as a single PostgreSQL query, built on pgvector. It showed that a database can run the whole model.
• Accelerating LLM Inference via Vector Index Based Output Embeddings, by Martin Loretz and Sepp Hochreiter. The paper replaces a model's output layer with a nearest-neighbour search over its token vectors, which is the last step here.