LLM Inference under the hood - Part 1: From Prompt to KV Cache

Chat and coding agents simplify the general usage of AI models. When diving deeper, it is important to understand how these large language models (LLMs) work under the hood to truly optimize them and get the most out of them. This requires understanding the inference pipeline—how the model reads your text and generates a response.

To see how hardware limits apply in practice when hosting models locally, check out: Run AI Locally, It Works!.

💡 By understanding how words flow through your model’s layers and how memory works, you can unlock modern techniques to drastically speed up generation and save on costs. This single optimization can transform sluggish chatbots into near-instantaneous interfaces.

To understand why LLMs can sometimes be slow, we have to look under the hood at the core phases of processing.

1. Tokenization: Chopping up the puzzle

LLMs do not read words or characters directly. Instead, raw text is chopped up into small puzzle pieces called tokens. A token can be a whole word, part of a word, or just a single punctuation mark.

Every token in the model’s vocabulary is assigned a unique number. When you send a prompt, a tokenizer converts your text into a list of these numbers.

2. Embeddings: Turning pieces into meaning

Once the text is broken into numbered tokens, the model needs to understand what they mean. In the very first step, the model looks up each token number and converts it into a “meaning coordinate”.

Think of it as looking up a word in a dictionary, but instead of a definition, you get a location on a massive map of meaning. Tokens with similar vibes or meanings are placed close together on this map.

3. Contextual Layers: Understanding the sentence

Initially, each piece of meaning doesn’t know about the other words around it. To soak up surrounding context, these pieces travel through many layers:

• Attention Blocks: A “crowded room” where tokens look at each other to figure out context. For example, the word “model” will adjust its meaning depending on whether “machine learning” or “fashion” is sitting nearby.
• Processing Layers: Individual booths where the pieces are updated based on a series of questions.

After passing through these layers, the final piece contains the accumulated context of the entire sentence.

4. The KV Cache: The Short-Term Memory

As the model reads through your prompt, recalculating the context for the entire history every time would be very slow and wasteful.

To optimize this, the model saves the computed context for every token in memory as it goes. This is called the KV Cache (Key-Value Cache). Think of it as the model’s short-term memory (see this visual explanation of the KV Cache). When generating the next word, it simply looks at this cache instead of re-reading the entire prompt from scratch!

Now that you understand how a prompt is read and cached, continue to Part 2: Prefill, Decode, and Prompt Caching to see how the model generates text and how we can supercharge the speed!

Visual & Technical Learning Resources

• Matt Pocock / AI Hero: Most devs don’t understand how LLM tokens work
• ByteByteGo: Understanding KV Cache in LLMs