Custom Transformer Decoder

This is a decoder-only transformer written in NumPy, with no machine learning framework. That covers the query, key, and value projections, scaled dot-product attention with a causal mask, layer normalization, residual connections, and the last projection onto vocabulary logits. I wrote the forward pass out so I could see the shapes.

It lives in a Jupyter notebook, with Matplotlib for the plots. The tokenizer is character-level, so the vocabulary stays small and the generated text is easy to read. A small demo lets you prompt the model and watch it emit one token at a time.

KV cache

During generation, each new token would otherwise run attention over the whole prefix again. The cost of a token grows with the context, and a full generation grows with the square of the length. The cache keeps the key and value tensors from earlier steps. The first token runs a full forward pass and fills it. Each later token computes only its own key and value, appends them, and attends over the tensors already stored.

The demo prints how long each generated token took. That time stays about flat as the context gets longer.

Shapes

Most of the time went into tensor shapes and the order of the multiplies. The notebook can draw attention heatmaps, and it plots how the cache grows at each step. I added both while debugging. Looking at the intermediate tensors showed mistakes faster than reading the loop.