Inference Representations

This repository is my attempt to learn inference engineering by building and implementing things myself. I am reading Philip Kiely's inference engineering book and, whenever I encounter something unfamiliar or do not feel confident about a concept, I dig deeper by reading more and implementing a small representation of it here.

The end goal is to build my own inference engine.

Attention Is All You Need

A compact encoder-decoder Transformer based on the Attention Is All You Need paper.

• config.py defines the model dimensions, including n_layers.
• layers.py implements positional encoding, multi-head attention, embeddings, feed-forward layers, and the language-model head.
• encoder.py and decoder.py each build n_layers independent blocks with nn.ModuleList.
• transformer.py connects encoder, decoder, and output head.

Run the example:

uv run python -m attention_paper.transformer

Resources

• Attention Is All You Need (PDF)

KV Cache

A focused representation of autoregressive KV-cache flow. For each decoder layer, cache.py stores projected keys and values during prefill, then appends only the new token during decoding. The query attends over the whole accumulated cache with a causal mask.

It intentionally omits surrounding Transformer components such as positional encodings, residual paths, norms, FFNs, and the output head. See the Hugging Face cache explanation for the production concept.

Run its prefill and token-by-token decode demonstration:

uv run python kv_cache/cache.py

Resources

• KV Cache
• KV Cache video