You type "Write a story" into ChatGPT. After a pause, the response starts appearing: "Once upon a time..."

Longer prompts often mean a longer wait before the first piece of text appears. After that, the response tends to arrive at a fairly steady rate, and a longer response takes more time to finish.

Training and Inference

Training adjusts the numbers that control a model’s behavior, called its parameters or weights. During inference, those weights stay fixed while the model uses them to generate a response.

Part 1: How Inference Works

A Black-Box View

The interface exposes two things: the text you send and the text the model returns.

Detail 1: Tokens

The model computes with numbers. A tokenizer first divides the text into pieces and assigns each piece an integer ID.

One possible unit is a character. For a limited English alphabet, letters, digits, and punctuation give a vocabulary of around 100 entries. But "Write a story" already takes 13 characters, and a paragraph takes hundreds. Generating one character per step makes sequences long, while a character such as "t" says little on its own.

Whole words give shorter sequences and more meaning per entry, but the vocabulary has to cover technical terms, names, URLs, code, and multiple languages. A fixed list of words also needs a way to handle words it has never seen.

Subword units offer a compromise. Common words can have their own entries, while unfamiliar words can be split into smaller pieces already in the vocabulary. For example, "responding" can split into ["respond", "ing"]. The pieces are often larger than characters, keeping sequences shorter without requiring an entry for every possible word.

These pieces are called tokens.

The tokenizer returns integer IDs from a fixed vocabulary. TinyLlama has 32,000 entries. The splits and IDs in these examples are illustrative rather than verified TinyLlama tokenizations, and omit chat-template and special-token details. For "Write a story":

["Write", " a", " story"] → [8144, 264, 3446]

The vocabulary and split rules are prepared before language-model training. Algorithms such as BPE learn common pieces from a text corpus. With fixed tokenizer settings, inference-time tokenization is deterministic and doesn’t use the language model’s weights.

The model receives the token IDs rather than the original text.

Try it yourself: You can experiment with tokenization at tiktokenizer.vercel.app. Paste in some text and see how different models split it into tokens.

At the output end, a detokenizer converts token IDs back into text chunks.

Detail 2a: Embeddings

The IDs [8144, 264, 3446] identify tokens, but their numerical values don’t describe those tokens.

Average the IDs for "Write" (8144) and "story" (3446), and you get 5795, which may identify an unrelated token. Adjacent IDs can also have unrelated meanings. Arithmetic on the IDs tells us nothing useful about the text.

The model needs numbers it can combine and compare. Ideally, tokens with related meanings would have similar numerical representations: "lion" and "tiger" would be close, while "lion" and "teacup" would be farther apart.

An embedding maps each token ID to a vector, a list of floating-point numbers. You can think of that list as the coordinates of a point in a space with many dimensions.

To build intuition, imagine just 2 dimensions. One axis captures something like "danger level," the other "size." Plot a few animal words:

Lion and scorpion are high on the danger axis; elephant and lion are high on size. In this toy space, proximity tells us which animals share those properties.

TinyLlama uses 2,048 dimensions. The embedding for "Write" might be [0.21, -0.83, 0.45, ..., 0.12], with one number per dimension. Every vocabulary entry has its own vector, learned during training.

In the toy diagram, the direction from "mouse" to "elephant" means getting larger while staying on the harmless side of the danger axis. Add that direction to "scorpion" and you move toward "crocodile": large and dangerous. Vector addition lets us combine relationships expressed by the coordinates.

System implications: With a 32,000-token vocabulary and 2,048-dimensional embeddings stored in half-precision (FP16 or BF16, both 2 bytes per value):

32,000 × 2,048 × 2 bytes ≈ 131.1 MB

(MB and GB below use decimal units: 1 MB = 10⁶ bytes and 1 GB = 10⁹ bytes. KiB and MiB use powers of 1,024. Hardware capacity names and reported software usage don’t always use the same convention.)

This embedding table is a matrix of shape [32,000 × 2,048]. When the model sees token ID 8144, it looks up row 8,144 and retrieves that token's vector. That vector, not the original ID, is what flows through the rest of the model.

Detail 2b: Unembeddings

After the model processes the prompt, it returns a vector for each token position. To choose the next token, we use the vector at the last position. For TinyLlama, that leaves us with 2,048 numbers and a vocabulary of 32,000 possible outputs.

Scoring Every Token in the Vocabulary

We can give each vocabulary entry a score for how well it fits as the next token, then use those scores to choose an output.

Assign each output token a learned vector, as we did for input embeddings. Its dot product with the model’s output vector supplies the score, so both direction and magnitude affect the result. These output vectors can be learned separately from the input embedding table.

Place the output-token vectors in the columns of a matrix. Multiplying the model’s output vector by that matrix computes all 32,000 scores at once. Each score is a dot product: multiply corresponding coordinates, then add the products.

output_vector × W          = scores
[1 × 2,048]   [2,048 × 32,000]   [1 × 32,000]

This scoring matrix is called the LM head (language modeling head). It's learned during training alongside the rest of the model.

The scores can be any real number, positive or negative. A score of 5.2 for "Once" doesn’t mean a 5.2% probability; it ranks "Once" above a token with a score of 3.1. These raw scores are called logits.

Picking a Token: The Naive Approach

The simplest selection rule is to choose the highest-scoring token. If "Once" has a score of 5.2 and every other score is lower, choose "Once." This is greedy decoding.

Greedy decoding can produce repetitive text. A model may repeatedly favor a familiar continuation such as "The cat sat on the mat," even after it has already generated that sentence.

There are usually several plausible ways to continue a sentence. Always choosing the highest-scoring token gives the other continuations no chance, even when one would make a good response.

Instead, we can sample a token: make a random choice weighted by the model’s probabilities. A more likely token is chosen more often, but other candidates can still appear.

That requires probabilities between 0 and 1 that sum to 1. The logits don’t yet have those properties.

From Scores to Probabilities

Dividing each score by the sum of the scores would make them sum to 1, but negative scores could produce negative probabilities. We need to make the scores positive before normalizing them.

Exponentiation makes the scores positive: $e^{5.2}$ is a large positive number, while $e^{-1.4}$ is a small positive number. The function $e^x$ maps every real number to a positive value.

Now divide each exponentiated score by the sum of all exponentiated scores. Every value lands between 0 and 1, and they sum to 1. That's a valid probability distribution.

The difference between logits determines their probability ratio: with logits 5.2 and 3.1, the ratio is $e^{5.2} / e^{3.1}$, or about 8.2. Adding the same constant to every logit doesn’t change the probabilities, but changing the gaps between logits does.

This operation (exponentiate, then normalize) is called softmax.

The result is a probability distribution over the vocabulary. For example, "Once" might receive 0.12, "The" 0.08, and "In" 0.05, with the remaining probability spread across other tokens.

Sampling Strategies

Unrestricted sampling can select a very unlikely token and send a response off track. These controls let us adjust how much of the distribution we sample from:

• Temperature: Divide the logits by a temperature value before softmax. Temperature > 1 makes the distribution flatter (more random, more creative). Temperature < 1 makes it peakier (more deterministic, more focused). As temperature approaches 0, sampling converges to greedy. Most APIs treat temperature=0 as a special case meaning "always pick the top token" (since the literal math would be division by zero).
• Top-k sampling: Only consider the k highest-probability tokens, then sample among them. This filters out the long tail of unlikely tokens while preserving variety among the good options.
• Top-p (nucleus) sampling: Sort tokens by probability from highest to lowest, then keep only the smallest set whose cumulative probability reaches p (e.g., 0.9). Unlike top-k, this adapts to the distribution's shape. If the model is very confident, only a few tokens are considered. If it's uncertain, more tokens make the cut.

For one output token, the steps are:

1. The model's layers output a 2,048-dim vector for the last position
2. The LM head projects it to 32,000 logits
3. Softmax converts logits to probabilities
4. Sampling picks one token
5. The detokenizer converts that token ID back to text

Detail 3: Autoregressive Generation

Choosing "Once" gives us the beginning of a story. To continue with "upon," the model needs to use "Once" as part of its input.

Without caching, we append "Once" to the prompt, giving [8144, 264, 3446, 12483] ("Write a story Once"), and run the longer sequence through all 22 layers. The model produces new logits, from which we choose "upon." Appending each chosen token and repeating the computation extends the response one token at a time.

In this basic generation loop, the system hasn’t prepared a complete response to reveal word by word. The choice at step 6 depends on the token chosen at step 5, so it can’t finalize step 6 in advance.

This is autoregressive generation: each output becomes part of the input used to predict the next output.

The computation through the model is a forward pass. The LM head produces logits, then the selection rule chooses a token. On generation step N, the available context has P + N − 1 tokens: P prompt tokens and N − 1 output tokens.

How does generation stop? The vocabulary includes special tokens like <EOS> (end of sequence) or </s>. When the model samples one of these, generation terminates. The model learned during training that these tokens signal "I'm done responding." You can also set a maximum output length, and generation stops when that limit is hit, even if no EOS appeared.

What’s the context window? The context window is the supported number of token positions for a model and serving configuration. TinyLlama’s original configuration supports 2,048 positions. A server may reject an oversized request, truncate earlier content, or stop generation at its limit. Sliding-window attention is a separate architectural choice, not a general drop-in fix for exceeding a model’s trained context length.

With a 100-token prompt and a 200-token response, this uncached implementation makes 200 forward passes. The first processes 100 tokens, the second 101, and the last 299. Each pass runs through all 22 layers.

Each row in the diagram is a forward pass. The blue prompt tokens appear unchanged on every row, but we process them again each time. After the first pass, the green token at the end is the only new input.

At step 150 in the 100-token example, 248 of the 249 input tokens were already present at step 149. Reusing earlier results could remove much of this repeated computation.

Detail 4: Self-Attention

"A cave bat flew out into the night" and "She picked up the baseball bat" use the same word for different things. Yet the embedding table returns the same 2,048 numbers for "bat" in both sentences, because it looks up a token ID without seeing the surrounding words.

I like to think of the embedding table as a dictionary. Its entry for "bat" covers several uses, but it doesn’t choose the one that fits a sentence. The surrounding words do that: "cave" suggests the animal, while "baseball" suggests sports equipment.

We need a way to update each token’s vector using information from the surrounding tokens. Then the representation of "bat" can depend on the sentence it appears in.

Inside the box is a stack of transformer blocks. TinyLlama has 22 of them, each structurally identical but with different learned weights. The embedded vectors enter block 1, flow through all 22 blocks in sequence, and exit into the LM head for next-token prediction.

Each block has two main components:

1. Self-attention combines information from different token positions. It lets the representation of "bat" incorporate information from "cave".

2. Feed-forward network (FFN) transforms each token’s vector independently after attention has added context.

Building Attention from First Principles

For the phrase "A cave bat," suppose the embedding lookup returns these three vectors:

embed("A")    = [0.04, -0.11, 0.52, ..., 0.03]   (2,048 numbers)
embed("cave") = [0.58,  0.39,  0.21, ..., -0.14]
embed("bat")  = [0.61,  0.45, -0.33, ..., 0.27]

These are illustrative values, not measurements from TinyLlama. At this point, each vector depends only on its own token ID.

A first attempt is to compare "bat" with each token, then combine their vectors in proportion to their similarity. The geometry of embeddings gives us a way to make that comparison.

Measuring Similarity

Geometrically, vectors that point in similar directions have a small angle between them. Perpendicular vectors have a right angle. The angle gives us one way to compare the vectors without inspecting each coordinate separately.

The cosine of the angle is 1 for the same direction, 0 for perpendicular vectors, and −1 for opposite directions. This is cosine similarity, computed by dividing the dot product by the product of the vector lengths. These geometric relationships can capture learned similarities; opposite directions don’t automatically mean opposite dictionary meanings.

Standard attention uses a scaled dot product without normalizing each vector by its length. Vector magnitudes can therefore affect the score. Dot products are also efficient to batch: if $X$ holds token vectors as rows, $X \times X^T$ computes all pairwise scores. The $\sqrt{d_k}$ factor controls score scale; it isn’t a replacement for cosine normalization.

Focus on "bat." Compute its dot product with every token in the sequence:

# How much is the word "bat" related to each of the words in "A cave bat"?

# How related is "bat" with "A"?
score("bat", "A")    = embed("bat") · embed("A")    = 0.3

# How related is "bat" with "cave"?
score("bat", "cave") = embed("bat") · embed("cave") = 5.8

# How related is "bat" with itself?
score("bat", "bat")  = embed("bat") · embed("bat")  = 6.1

In this example, "bat" scores low against "A" (0.3), high against "cave" (5.8), and highest against itself (6.1). Its self-score is the squared length of its vector.

To turn these raw scores into proper weights, we normalize with softmax (same function from the LM head section) so they sum to 1:

weight("bat", "A")    = 0.002
weight("bat", "cave") = 0.425
weight("bat", "bat")  = 0.573

Now compute a new vector for "bat" as a weighted sum of all the embeddings:

new_bat = 0.002 × embed("A") + 0.425 × embed("cave") + 0.573 × embed("bat")

The weighted sum keeps 57.3% of the original "bat" vector and adds a 42.5% contribution from "cave." In the diagram’s simplified geometry, that contribution moves "bat" toward the animal meaning. "A" contributes very little.

(In reality, these embeddings live in 2,048 dimensions, not the 2D picture you might be imagining. The shift happens across many dimensions simultaneously. But the intuition of "movement toward the right meaning" holds.)

Repeating the calculation gives a context-dependent vector at each position. In a decoder, causal masking restricts each position to itself and earlier positions; the examples above focus on "bat," which is last in its phrase.

This weighted mixing is a simple form of attention. It works for the example because the relationship we need is captured by the raw embedding similarity. Other relationships need a different comparison.

Why Raw Dot Products Aren't Enough

In "The animal didn’t cross the street because it was too tired," the pronoun "it" refers to "animal." The model needs a way to connect the pronoun with a possible referent even though their embeddings needn’t be similar. At the position of "it," a causal model can use the preceding words; it can’t yet use "tired," which makes the reference clearer to us when we read the complete sentence.

A raw dot product can only compare the features already present in the embeddings. It has no separate way to ask whether one token could be the referent of another.

The relevant relationship here is grammatical: "it" needs a referent, and "animal" can supply one. We want the comparison to give that relationship a high score even if general embedding similarity doesn’t.

Other examples have the same problem. "Quickly" modifies "ran," and "not" changes the meaning of the verb it negates. Attention needs to recognize those relationships, not rely on the tokens being close in the original embedding space.

We can let the model learn what to compare by transforming the vectors before taking their dot product. Training can then shape a comparison suited to the relationships needed for next-token prediction.

Queries, Keys, and Values

Matching tokens and combining their information are different jobs. A transformation that helps identify "animal" as a possible referent needn’t preserve all the information we want to pass from "animal" to "it."

A hiring analogy helps separate these jobs. You need an ML infrastructure engineer and have three candidates: an infrastructure engineer with five years at a big tech company, a generalist software engineer, and an ML researcher fresh out of a PhD.

For the analogy, treat the outcome as a weighted blend of the candidates’ contributions:

hiring_outcome = 0.7 × infra_engineer + 0.2 × generalist + 0.1 × researcher

The infra engineer is the strongest match for your need, so they contribute the most. The generalist brings something useful but less targeted. The researcher contributes least to this particular role.

Focus on one candidate. We need two things: how relevant the infra engineer is to your need, and what they actually contribute if hired.

hiring_outcome = (how_relevant_infra_engineer_is) × what_infra_engineer_actually_brings

Simplify:

hiring_outcome = (weight) × value_of_infra_engineer

The weight describes how well the candidate matches the role. The value describes what the candidate contributes. Both concern the same person, but they answer different questions.

The comparison needs three representations:

(a) Your job spec: "someone who combines ML knowledge with infrastructure experience." Structured for matching against candidates.

(b) Each candidate's resume: the infra engineer highlights cloud deployments, the generalist highlights system design, the researcher highlights ML publications. Optimized for comparison against job specs, not a description of what they'd actually do on the job.

(c) Each candidate's actual contribution if hired: the infra engineer writes the modules, debugs the 3 AM outages, and knows which Kubernetes configs silently break under load. The generalist spots architectural bottlenecks early and writes code that the rest of the team can actually maintain. The researcher catches a flawed training setup before it burns a week of GPU time.

The job spec is compared with a resume to determine the weight; the candidate’s work supplies the value:

hiring_outcome = match(your_job_spec, candidate_resume) × candidate_actual_contribution
               = weight                                 × value_of_infra_engineer

Here, weight comes from comparing (a) with (b), while value_of_infra_engineer corresponds to (c).

Same thing with tokens. Each token starts as one embedding vector, but attention needs three forms of it:

1. What it's looking for from context (like a job spec). What kind of surrounding information would help disambiguate this token? This is the Query.
2. How it describes itself for matching (like a resume). When other tokens are looking around for context, what should they see when they look at this token? This is the Key.
3. What it actually contributes when selected (like the real work a hired person does). The content it shares when pulled into the weighted sum. This is the Value.

Separate Query and Key transformations also make matching directional. Your job spec matched against a candidate’s resume is a different comparison from that candidate’s job spec matched against your resume. Attention can likewise assign different scores to token A attending to token B and token B attending to token A.

Each form is produced by multiplying the layer’s input vector by a learned weight matrix. In the first block that input comes from the embeddings; later blocks receive contextual representations:

Query = input × W_Q
Key   = input × W_K
Value = input × W_V

Three separate matrices ($W_Q$, $W_K$, $W_V$), all learned during training. The model discovers what aspects of each embedding to expose for each role.

Putting it all together for a given token:

1. Compute this token's Query. Compute every token's Key and Value.
2. Dot product of this token's Query against each Key to get relevance scores.
3. Softmax to normalize scores into weights summing to 1.
4. Weighted sum of Values to produce a new, context-aware embedding.

The Attention Formula

For one attention head in "A cave bat," each token has a Query, Key, and Value. The numbers below are a small toy example, with scaling temporarily omitted; TinyLlama’s actual heads have 64 dimensions.

Matching first. Compute the dot product of bat's Query against every token's Key:

score("bat" → "A")    = Q_bat · K_A    =  1.0
score("bat" → "cave") = Q_bat · K_cave =  5.0
score("bat" → "bat")  = Q_bat · K_bat  =  4.0

These scores differ from the raw embedding scores (0.3, 5.8, 6.1) because the learned Query and Key projections expose different features for comparison.

Stack all the Queries as rows of a matrix Q and all the Keys as rows of K, and every pairwise dot product falls out of a single matrix multiply:

$QK^T$

For 3 tokens, the result is a 3×3 grid. A 2,048-token prompt produces about 4.2 million scores per head.

Now these scores need to become proportions that sum to 1. Softmax (same operation from the unembeddings section) handles this:

softmax([1.0, 5.0, 4.0]) → [0.01, 0.72, 0.27]

In this unscaled example, "cave" receives 72% of the weight, "bat" 27%, and "A" 1%.

Use those weights to blend the Value vectors:

output_bat = 0.01 × V_A  +  0.72 × V_cave  +  0.27 × V_bat

The attention weights control how much each Value contributes. The resulting vector incorporates context from "cave" into the representation of "bat."

In matrix notation, the whole process for all tokens at once:

$\text{softmax}(QK^T) \cdot V$

Before using this formula, we need to account for how the size of the dot products changes with vector dimension.

The Scaling Problem

The illustrative scores were small: 1.0, 5.0, and 4.0. A dot product adds one product per dimension, so increasing the dimension can increase the spread of the scores. TinyLlama uses 64 dimensions per head.

Watch what happens:

With a few dimensions:
  scores = [1.0, 5.0, 4.0]
  softmax → [0.01, 0.72, 0.27]  ← blends three tokens

With 64 dimensions:
  scores = [8, 42, 35]
  softmax → [0.00, 1.00, 0.00]  ← collapsed to one token

Larger score gaps make softmax more concentrated. Here, $e^{42} / e^{35} \approx 1{,}097$, so the top token receives about 99.9% of the weight. Adding a constant to all scores wouldn’t have that effect; it’s their separation that matters.

We need to shrink the scores before they hit softmax:

$\text{softmax}\left(\frac{QK^T}{\text{???}}\right) V$

A fixed divisor such as 10 doesn’t account for the head dimension. We want the divisor to follow the expected spread of the dot products.

Call the number of dimensions per head $d_k$. A dot product adds $d_k$ terms: 3 when $d_k = 3$, or 64 when $d_k = 64$. Their typical combined spread depends on how those terms vary.

Think of the dot product as a random walk that adds one term per dimension. Under the simplifying assumption that Q and K components are independent, zero-mean, and have variance 1, each product also has variance 1. Adding $d_k$ independent products gives variance $d_k$. Variance measures spread in squared units, so taking its square root gives the standard deviation: $\sqrt{d_k}$. That is why the typical spread grows with the square root of the dimension.

For $d_k = 64$: $\sqrt{64} = 8$. Divide the inflated scores by 8:

Before scaling:  [8, 42, 35]       → softmax → [0.00, 1.00, 0.00]  ← collapsed
After  ÷ 8:     [1.0, 5.25, 4.38] → softmax → [0.01, 0.70, 0.29]  ← blends again

Scaling reduces the tendency for dimension alone to concentrate softmax. "Cave" still receives the most weight, and the model can still learn sharply peaked attention when that is useful.

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

In the formula, $QK^T$ computes the matching scores, softmax turns them into weights, and multiplying by $V$ combines the contributions. Dividing by $\sqrt{d_k}$ controls the score scale as the head dimension grows.

Each projection is a learned weight matrix. In standard attention with a 2,048-dimensional hidden state:

Q = input × W_Q    (2,048 → 2,048)
K = input × W_K    (2,048 → 2,048)
V = input × W_V    (2,048 → 2,048)

(TinyLlama actually uses smaller K and V projections than this, sharing them across groups of query heads. Detail 6 covers the exact shapes. The attention formula stays the same.)

Detail 5: Multi-Head Attention

Extend the "bat" sentence: "A cave bat flew out and landed on a nearby branch." Focus on "landed."

To interpret "landed," several earlier tokens matter: "bat" identifies who landed, "flew" supplies the preceding action, and "cave" gives the starting location.

A single attention computation produces one set of weights over the sequence. If "landed" assigns 0.4 to "bat," 0.3 to "flew," and 0.2 to "cave," those contributions enter the same weighted sum.

That sum mixes information about the subject, the motion, and the location. It may be useful, but the model can’t independently select a different mixture for each relationship within this one attention computation.

With several independent attention computations, the model could give "bat" a high weight in one mixture and "flew" a high weight in another, keeping those results separate until a later step.

The Fix: Ask Multiple Questions

Run attention several times in parallel, with a separate $W_Q$, $W_K$, $W_V$ for each computation.

One computation could focus on the subject and another on the preceding action. These are possible learned behaviors, not roles we assign in advance.

Each of these independent attention computations is called a head. Running multiple heads in parallel is multi-head attention.

The Cost, and the Trick

An immediate concern: TinyLlama uses 2,048-dimensional vectors. Running 32 separate full-size attention computations means 32 copies of Q, K, V projections, 32 score matrices, 32 weighted sums. That's 32× the parameters and 32× the compute.

We can keep the dimensions smaller for each head.

Use 32 heads with 64-dimensional vectors each. Their combined width is 32 × 64 = 2,048, so splitting into heads doesn’t multiply the projection parameter count by 32.

d_model = 2,048
n_heads = 32
d_head  = d_model / n_heads = 64

In standard multi-head attention, you compute Q, K, V at full width:

Q = input × W_Q    (2,048 → 2,048)
K = input × W_K    (2,048 → 2,048)
V = input × W_V    (2,048 → 2,048)

(If you check TinyLlama's config, you'll notice its K and V projections are actually smaller, 2,048 → 256. That's because it uses grouped-query attention, which we'll cover in Detail 6. For now, we're building the standard picture.)

Each result is a 2,048-dimensional vector per token. Slice each into 32 pieces of 64 dimensions. Piece 1 gets dimensions 1 through 64. Piece 2 gets dimensions 65 through 128. And so on. Each piece is one head's workspace. Run the full attention computation ($QK^T / \sqrt{d_k}$, softmax, multiply by $V$) independently on each 64-dimensional slice. Head 1 uses only its slice of Q, K, V. Head 2 does the same with its slice. All 32 heads operate in parallel, each in its own 64-dimensional subspace.

This is where the $\sqrt{d_k}$ scaling becomes concrete: each head operates in 64 dimensions, so $d_k = 64$ and the scale factor is $\sqrt{64} = 8$.

The matrix-multiplication work is roughly comparable to one full-width attention computation, though each head has its own attention scores and softmax.

Each head produces a 64-dimensional output per token. To get back to full width, concatenate all 32 outputs end-to-end: 32 × 64 = 2,048. Then one final learned projection ($W_O$) mixes information across heads, letting the model combine what different heads discovered into a single output vector.

What Heads Learn

Because each head has its own projections, heads can specialize without anyone telling them to. The training signal (predict the next token well) is the only pressure, and it shapes each head into whatever role helps.

Trained heads can exhibit different patterns. Some attend to the previous one or two tokens; others connect distant subjects and verbs or attend to delimiters such as quotes and newlines. Some put substantial weight on the first token regardless of its content, a behavior called an "attention sink." Its interpretation is still an active research topic.

Some heads also learn redundant behavior. In the "landed" example, we can imagine head 7 assigning 90% of its weight to "bat," head 15 assigning 85% to "flew," and head 22 focusing on "cave." The point is that the architecture permits separate attention patterns; it doesn’t guarantee these particular ones.

Detail 6: Attention Variants (GQA, SWA, MLA)

Storage for 32 Heads

During generation, each new token uses the Keys and Values of earlier tokens. Inference engines keep those vectors in a KV cache so they can reuse them. Detail 10 explains why that reuse is valid; here, the storage cost matters.

Standard multi-head attention stores a separate Key and Value for each head, at every token position and every layer. With 32 heads, that means 32 sets of cached vectors.

For a hypothetical model with TinyLlama’s dimensions but standard multi-head attention:

Per position, per layer:
  32 heads × 64 dims × 2 (one K, one V) = 4,096 values

Across all 22 layers, at FP16 (2 bytes each):
  4,096 × 22 × 2 bytes = 180,224 bytes ≈ 180.2 KB per token position

At max context (2,048 tokens):
  2,048 × 180,224 bytes ≈ 369.1 MB per sequence

That’s about 369 MB for one full-length conversation. Larger models can require gigabytes per sequence, so serving many users can make the cache a substantial part of GPU memory.

Reducing the number of distinct Key and Value heads reduces the cache in the same proportion.

Not All Projections Are Created Equal

The heads attending from "landed" use different Queries to find the subject, preceding action, and starting location. They can keep those different Queries while sharing some of the Keys and Values they search.

Keys and Values describe the context tokens. Sharing them means the heads compare against the same representation of that context, while their separate Queries can still produce different attention weights.

Imagine 32 analysts researching different questions about a company. They can consult the same financial filings even though one is studying revenue and another is studying supply-chain risk.

The Extreme: One Shared K/V Set

Take this to its logical endpoint. What if all 32 query heads shared a single Key and a single Value? Each head keeps its own $W_Q$, so each can still attend to different tokens. But they all look into the same K/V.

MQA for TinyLlama:
  (1 K + 1 V) × 64 dims = 128 values per position per layer
  128 × 22 layers × 2 bytes = 5.6 KB per token position
  At 2,048 tokens: ~11.5 MB per sequence

Storing one K/V set instead of 32 gives a 32× reduction, from about 369.1 MB to 11.5 MB per full-length sequence.

This is called Multi-Query Attention (MQA).

Sharing one K/V set also constrains the heads. In the analogy, analysts lose the ability to organize separate copies of the filings to suit their own questions. In attention, all heads must match against the same Key representation and draw from the same Value representation.

That constraint can reduce quality. The effect depends on the model and task, so a useful compromise is to retain more than one K/V set.

Sharing Within Groups

Instead of going all the way from 32 K/V sets down to 1, keep a small number in between and divide the query heads into groups. Each group shares K/V.

TinyLlama uses GQA with 4 KV groups:

• 32 query heads ÷ 4 KV groups = 8 query heads per group
• Within each group, 8 heads share the same K and V
• Across groups, the K/V representations are different

This changes the projection shapes. The query projection stays full-width, but K and V shrink by 8×:

W_Q: 2,048 → 2,048  (32 query heads × 64 dims)
W_K: 2,048 → 256    (4 KV heads × 64 dims)
W_V: 2,048 → 256    (4 KV heads × 64 dims)

Each of the 4 KV heads produces a 64-dim Key and Value. Those get broadcast to the 8 query heads in their group. The attention math per head is identical to standard MHA. The only difference is where the Keys and Values come from: shared within the group, not unique per head.

GQA-4 for TinyLlama:
  4 groups × 64 dims × 2 (K+V) = 512 values per position per layer
  512 × 22 layers × 2 bytes = 22,528 bytes ≈ 22.5 KB per token position
  At 2,048 tokens: ~46.1 MB per sequence

Four groups use 8× less cache than 32 independent K/V heads. The groups can still learn different representations of context, though there are fewer distinct representations than in standard MHA.

This is Grouped-Query Attention (GQA). TinyLlama, Llama 2 70B, Llama 3, and Gemma 2 models use it. Other Gemma generations and sizes use different head layouts.

All three variants (standard MHA, GQA, MQA) are really points on the same spectrum, controlled by a single number: n_kv_heads, the count of distinct K/V sets.

The model’s configuration fixes this choice before serving. MHA gives each Query head its own K/V set, MQA shares one set across all Query heads, and GQA uses an intermediate number.

A Different Lever: How Far Back to Look

Standard causal attention lets a token attend to every earlier position. A smaller cache is possible if we restrict attention to a recent window.

Recent context often supplies what a token needs, while an occasional reference depends on something much earlier. Restricting the window saves memory but makes those long-range references harder to recover.

Sliding Window Attention (SWA) formalizes this observation. Instead of attending to all previous tokens, each token only attends to the nearest W tokens (the "window"), including itself. Everything older is masked out, as if it doesn't exist.

If W = 4,096 and the sequence has reached 10,000 tokens, token 5,000 can attend to tokens 905 through 5,000. Token 1 is invisible to it.

The trade-off is direct. A smaller window bounds memory because the KV cache only needs to hold the last W tokens' worth of K/V, regardless of how long the sequence gets. But it also caps reach: if important context lives 5,000 tokens back and your window is 4,096, you can't directly reach it.

But there's a subtlety. Information can still propagate beyond the window size by flowing through intermediate layers. Token 5,000 can't attend to token 1 directly, but if token 3,000 attended to token 1 in an earlier layer, and token 5,000 attends to token 3,000, the information has traveled indirectly. With window size $W$ and $L$ layers, the effective receptive field is bounded by roughly $O(W \times L)$, and only if intermediate tokens learn to relay the signal forward. In practice this propagation is lossy and unreliable: it is not a substitute for direct global attention.

In practice, many models use SWA in some layers and full attention in others. Gemma 3, for example, uses a 5:1 ratio: five sliding-window layers for every one full-attention layer. The sliding-window layers handle local relationships cheaply. The occasional full-attention layer provides a direct path for long-range information. This hybrid approach gets most of the memory savings while keeping long-range capability.

GQA reduces the data stored per position. SWA reduces the number of stored positions, so a model can use both.

Yet Another Lever: Compress What You Store

We can also reduce storage by compressing the cached representation at each position.

Multi-Head Latent Attention (MLA), used in DeepSeek V2 and V3, takes this approach. Instead of caching the full K and V vectors for each token, it compresses them into a smaller "latent" vector and avoids reconstructing a full per-head K/V cache on the hot path.

The cached latent is a shorter vector from which the model’s learned projections can obtain the information needed for attention. It avoids storing a full Key and Value for every head.

The savings depend on the latent width and head layout. DeepSeek V2’s large-model comparison gives 860.2K cached elements per token for an MHA baseline and 34.6K for MLA: about 24.9× fewer, a 96.0% reduction. The paper’s separate 93.3% claim compares with DeepSeek 67B, a different baseline.

GQA requires heads in a group to share identical K/V vectors. MLA instead stores a shared latent from which different heads can derive their own K/V representations using learned matrices.

DeepSeek V2’s MLA-versus-MHA ablations reported comparable or stronger scores on most evaluated benchmarks while using much less cache memory. Those results don’t establish a universal quality advantage, nor were they an equal-cache-size comparison of MLA with GQA.

You might expect a compute trade-off: reconstructing K and V from the latent for every past token sounds like it would add an expensive matrix multiply per layer per token. But MLA avoids an explicit full decompression pass through a linear algebra trick. Since the decompression is just a matrix multiply, and matrix multiplication is associative, you can absorb the decompression matrices into the Query and Output projection weights offline. Instead of reconstructing full K from the latent and then computing Q·K^T, you fold the key-side decompression into Q's projection. Same idea on the output side: fold the value-side decompression into the output projection. On the hot path, there is no separate "expand the whole cache back to per-head K/V" step.

Position-dependent operations complicate that absorption. They vary by token position and can’t be folded into a single static weight matrix.

DeepSeek V2 therefore caches a small positional component separately from the compressed content component. The RoPE section explains how those two paths work.

GQA saves cache space by sharing K/V across heads. MLA saves it by storing a compressed latent and arranging the projections so attention can use that latent without expanding the entire cache.

The Three Levers

For uncompressed K/V with equal Key and Value dimensions, the cache size follows:

KV cache memory ∝ n_layers × seq_len × n_kv_heads × d_head × 2 (K,V)

A model can combine head sharing with local attention. MLA uses a different storage formula: a latent vector plus a separate positional Key component per position. It can’t be modeled exactly by just reducing d_head in the uncompressed formula.

For a Llama 2 70B-like configuration at 2,048 tokens:

• Standard MHA: ~5.4 GB per sequence
• With GQA-8 (what it actually uses): ~0.67 GB per sequence

The 8× reduction leaves room for more concurrent sequences when KV storage is the capacity limit. Weights and other working memory still occupy part of the GPU.

Detail 7: Positional Encodings with RoPE

The same words can express different events:

"The dog bit the man." "The man bit the dog."

The ordering determines who bit whom. Token embeddings alone don’t encode that ordering.

Treating the words as tokens for this example, both sentences contain the same token identities. The embedding lookup returns the same vector for each identity regardless of its position.

Applying the same Q, K, and V projections to the same embeddings gives the same vectors, wherever the tokens occur.

Ignoring the causal mask for this comparison, the score between "dog" and "man" is $Q_{\text{dog}} \cdot K_{\text{man}}$. Swapping their positions doesn’t change these first-layer content vectors or $Q_{\text{dog}} \cdot K_{\text{man}}$, so unmasked attention weights and weighted sums simply follow the permutation.

A causal mask does distinguish earlier positions from later ones. We still want to encode position explicitly so the model can use the distance between tokens in its attention scores.

The First Attempt: Stamping Each Position

The most natural fix: give each position a unique identity. Token 0 gets a vector $p_0$, token 1 gets $p_1$, and so on. Add these to the token embeddings before block 1. Now "dog" at position 1 carries a slightly different vector than "dog" at position 4. The Q and K projections produce different results. The dot products change. Order becomes visible.

This is exactly what the original Transformer paper (Vaswani et al., 2017) did. They used fixed sinusoidal patterns: position-dependent sine and cosine values added to embeddings at the input. GPT-2 replaced the fixed patterns with learned vectors, one per position, but the idea is the same. Encode position once, at the input, by adding something to the embedding.

Adding position at the input works, but it leaves the model to carry that information through the later layers and use it to infer relative distances.

The encoding is attached to an absolute position. Token 10 gets the same position vector regardless of which other token it attends to. Sinusoidal encodings do provide structure from which relative offsets can be recovered; the model doesn’t have to discover subtraction from nothing. But after adding position to content and projecting to Q and K, the score isn’t guaranteed to depend on position only through relative distance.

The signal enters only once. The position vectors are added before block 1, so later attention and FFN operations transform them along with the content. There’s no fresh position encoding supplied to each layer.

For fixed content vectors, we’d like the positional part of the score to depend on their relative distance. Moving a pair from positions (2, 5) to (100, 103) should leave that part unchanged. Applying the encoding inside each attention layer would also give every layer direct access to position.

What Kind of Operation Gives Us This?

Think about the constraint precisely. We want to apply some function to Q and K based on their positions, such that when we take their dot product, the result depends on the position difference and nothing else.

Call the function $f$. We apply $f(q, \text{pos}_i)$ to the Query at position $i$, and $f(k, \text{pos}_j)$ to the Key at position $j$. Their dot product should contain $(i - j)$ but not $i$ or $j$ independently.

Adding position vectors $p_i$ and $p_j$ gives $(q + p_i)^T(k + p_j) = q^Tk + q^Tp_j + p_i^Tk + p_i^Tp_j$. The cross terms mix content with $i$ and $j$ separately. Without further constraints, the result doesn’t depend only on their difference.

Picture two clock hands on the same face. Each hand points in some direction: that's its "content" (what the token means). Now spin the entire clock by some angle. Both hands rotate. The angle between the two hands stays the same. The absolute orientation of each hand changed, but their relative angle is invariant.

Rotating one vector by $i\theta$ and the other by $j\theta$ adds $(i - j)\theta$ to their relative angle. Shifting both positions by the same amount leaves that angle unchanged.

Tracing Rotation in 2D

Let's make this concrete. Take a Query $q$ at position $i$ and a Key $k$ at position $j$, both 2D vectors. To encode position, we rotate each vector by an angle proportional to its position index. The rotation uses the standard 2D rotation matrix:

$R(\alpha) = \begin{pmatrix} \cos\alpha & -\sin\alpha \\ \sin\alpha & \cos\alpha \end{pmatrix}$

Apply $R(i\theta)$ to $q$ and $R(j\theta)$ to $k$, where $\theta$ is a fixed frequency. Then compute the dot product of the rotated vectors.

Using the rotation identities, the dot product becomes:

$q^T_{\text{rot}} \cdot k_{\text{rot}} = \|q\|\|k\|\cos(\phi_q - \phi_k + (i - j)\theta)$

Here $\phi_q$ and $\phi_k$ are the original vector angles. The position-dependent term is $(i - j)\theta$, which uses the relative offset.

Let's verify this with real numbers. Set $\theta = 30°$. Take $q = [1, 0]$ at position 3, and $k = [0.7, 0.7]$ at position 5.

Without rotation, the dot product is $1 \times 0.7 + 0 \times 0.7 = 0.7$, with no position term.

With rotation: $q$ gets rotated by $3 \times 30° = 90°$, and $k$ by $5 \times 30° = 150°$. Both vectors change. Their dot product comes out to roughly $-0.26$. The value now encodes both content similarity and the fact that these two tokens sit 2 positions apart.

Move the same vectors to positions (100, 102). The absolute rotations change to $100 \times 30°$ and $102 \times 30°$, but their relative rotation is still $2 \times 30° = 60°$. The dot product remains $-0.26$.

For these fixed content vectors, shifting both positions equally leaves the attention score unchanged. Rotation gives us the relative-position property directly.

This is Rotary Position Embedding (RoPE), the position encoding used by Llama-family models including TinyLlama. The name captures the mechanism: position is encoded through rotation, not addition.

Scaling to 64 Dimensions

A 2D rotation acts on two coordinates. To apply the same idea to TinyLlama’s 64-dimensional heads, RoPE pairs coordinates into 2D planes.

For illustration, pair dimensions (1, 2), (3, 4), and so on through (63, 64). Each pair rotates at its own frequency. Implementations can use a different coordinate ordering, such as pairing the two halves of the vector, with the same underlying operation.

Think of a clock with hands turning at different speeds. A fast hand changes noticeably between nearby positions but wraps around over longer distances. A slow hand changes little between adjacent positions and helps distinguish more distant offsets. Combining frequencies gives the model information at several distance scales.

Each coordinate pair has the same relative-angle property, so the full dot product combines content-dependent position signals at 32 frequencies. This supplies a useful pattern of relative offsets, without guaranteeing reliable behavior at distances beyond those supported by training.

The rotation angles are deterministic functions of position and dimension index, with no learned parameters. Unlike the original additive sinusoidal encoding, RoPE rotates Q and K inside every attention layer, supplying position information at each layer.

What Gets Rotated

RoPE rotates Queries and Keys, which determine the attention scores. It doesn’t rotate Values.

Values contain the information combined after the attention weights have been calculated. Leaving them unrotated lets position affect the weights without applying the same rotation to the information being passed between tokens.

Why This Matters for Inference

Because RoPE is applied inside each attention layer, the Keys stored in the KV cache (which we'll formalize in Detail 10) already have their positional rotations baked in. When a new token at position t arrives during decode, we only need to rotate the new Q and K vectors for position t. There is no extra cost to replay or recompute position information for cached tokens.

This also explains why positional information doesn't appear in the pipeline diagrams before the transformer stack: in a RoPE-based model, there is no separate "add position" step at the input. Position enters the computation inside each attention layer, 22 times rather than once.

Detail 8: Feed-Forward Network (FFN)

The attention example produced this weighted sum for "bat":

output_bat = 0.01 × V_A  +  0.72 × V_cave  +  0.27 × V_bat

The result combines the Value vectors using the attention weights. We also need a way to transform the resulting features at each token position.

A weighted sum stays within the space spanned by the vectors being combined. Adjusting the weights changes the mixture, but doesn’t by itself define an arbitrary transformation of the resulting vector.

The feed-forward network, or FFN, applies that transformation independently to each position. I wrote a separate explanation of the underlying network: Multi-Layer Perceptrons: How Neural Networks Bend Space to See. In a transformer, the same FFN runs on each token vector after attention.

The Structure

Step 1: Expand. TinyLlama’s token vectors have 2,048 dimensions. The FFN projects each one into a wider space with 5,632 dimensions, where it can compute more intermediate features before returning to the standard width.

Think of the wider space as a larger workbench. It gives the network more intermediate values to work with, though the model learns what those values encode rather than assigning one dimension to each named concept.

Step 2: Apply a nonlinearity. Expansion alone doesn’t give us a more expressive function.

Multiplying a column vector by matrix A and then matrix B is equivalent to multiplying it once by BA. Without a nonlinear operation between the projections, the 2,048 → 5,632 → 2,048 computation could be replaced by one 2,048 → 2,048 matrix.

TinyLlama applies an elementwise function called SiLU within its gated FFN. The nonlinearity prevents the two projections from collapsing into one matrix multiplication, allowing the network to compute feature interactions that a single linear transformation can’t express.

Step 3: Compress. The rest of the model expects 2,048-dimensional vectors. The expansion was working memory, not a permanent change in representation size. The final matrix multiplication projects back down from 5,632 to 2,048, packing the results into the standard width.

Same Function, Every Position

The FFN applies the same weights independently to every row of the [seq_len × d_model] matrix. It doesn’t exchange information between token positions.

The outputs still differ because the inputs differ. In "A cave bat," the vector for "bat" can already contain information from "cave," while "cave" has a different vector and earlier context. The same FFN transforms each according to the information it contains.

Where the Parameters Live

The FFN is where the majority of the model's parameters live. For each of TinyLlama's 22 layers:

FFN parameters per layer:
  Three matrices of ~2,048 × 5,632 each ≈ 34.6M parameters

Across all 22 layers: 34.6M × 22 ≈ 761M parameters

That’s roughly two-thirds of TinyLlama’s 1.1B parameters. Attention gets much of the discussion, but the FFNs account for most of the learned weights and a large share of the matrix-multiplication work.

Detail 9: The Complete Transformer Block

A transformer block takes one vector per token and returns vectors of the same width:

Inside, attention is followed by the FFN:

TinyLlama chains 22 blocks together. Passing information through that many transformations is easier if each sublayer can adjust an existing representation instead of producing its entire replacement.

Preserving Information Through the Stack

If every sublayer replaces its input, block 1’s result becomes the input to block 2, which transforms it again. Information needed much later must survive every intermediate transformation.

Suppose block 15 only needs to add a small detail about "bat." It would still have to reproduce all the useful information from earlier blocks in its output. An explicit path for carrying that information forward would let the block focus on its change.

The Fix: Add, Don't Replace

Add the sublayer’s output to its input:

output = input + sublayer(input)

The input has a direct path around the sublayer, whose output is added as an update. Attention and the FFN each contribute an update to this running vector.

Block 15 can contribute a small update or even zero without having to reconstruct its input. Later updates can still change or cancel earlier information, but preserving it no longer requires every sublayer to copy it explicitly.

This pattern is called a residual connection (or skip connection). The name comes from what the sublayer learns: the residual, the difference between what came in and what should go out, rather than the full output from scratch.

During training, residual connections also give gradients a direct path backward through the stack. This helps with the vanishing-gradient problem, where repeated transformations can make the training signal too small for early layers to learn effectively.

The Scale Problem

Repeated additions can change the overall scale of the running vector.

There are two updates per block, giving 44 across TinyLlama’s 22 blocks. Their accumulated values may grow even if each update is modest.

Suppose the original embedding for "bat" has values mostly in the range [-1, 1]. Block 1's attention adds a correction in a similar range. The FFN adds more. After a few layers, some dimensions have accumulated large positive values while others barely moved. By block 22, dimension 47 might sit at 42.3 while dimension 100 sits at -0.8.

A later sublayer then receives a vector whose scale may differ from the scale seen earlier in the stack. Controlling that variation helps keep training stable.

TinyLlama rescales the vector before each sublayer using RMSNorm, or root-mean-square normalization. It squares the elements, averages those squares, and takes the square root to measure the overall scale. It divides by that value, with a small numerical safeguard, then applies a learned scale factor to each dimension. This controls the vector’s overall scale; it doesn’t make all coordinates equal in magnitude.

TinyLlama normalizes before attention and before the FFN. This arrangement is called pre-norm. The original 2017 Transformer normalized after each sublayer; pre-norm is widely used because it can make deeper models easier to train.

The Complete Block

The data flow through one block is:

1. Input arrives (2,048-dim vector per token)
2. RMSNorm → Self-attention → Add to input (residual connection)
3. RMSNorm → FFN → Add to previous result (residual connection)
4. Output passes to the next block

Each sublayer receives a normalized copy of the running vector and adds its output back through a residual connection.

The Full Stack

TinyLlama repeats this block 22 times. Block 1 feeds block 2, block 2 feeds block 3. The wiring is identical, but the weights differ per layer, so each layer can specialize. After the final block, one last RMSNorm, and the vectors enter the LM head (from Detail 2b) for next-token prediction.

That's the complete pipeline: tokenizer → embedding → 22 transformer blocks (each with RoPE-augmented attention + FFN, wrapped in residuals and norms) → LM head → sampling → detokenizer.

Detail 10: The KV Cache

After generating "time" in "Write a story Once upon a time," the sequence contains seven tokens. An uncached implementation processes all seven again to predict the eighth.

At each layer, each of the seven tokens gets projected into Q, K, and V. Attention runs: token 7's Query dots against all seven Keys, softmax normalizes, the weighted sum collects from all seven Values. The FFN processes each result. Repeat across all 22 layers.

Only "time" is new. Causal masking prevents earlier positions from attending to it, so their inputs to each layer haven’t changed. With fixed weights, their Keys and Values are the same mathematical results as on the previous step.

The repeated work grows with the prefix. With a three-token prompt, generation step 100 processes 102 input tokens, though only one was added since the preceding step.

What Attention Actually Needs

At one attention layer, suppose we’re processing the newly appended token at position t to predict the following token. The attention calculation needs:

1. The new token’s Query, used to calculate attention scores.
2. Keys for positions 0 through t, including the new token.
3. Values for those same positions, used to form the weighted sum.

Token t's Query has to be computed fresh. It depends on the representation this token has built up through the preceding layers, which is new at every step.

But the Keys and Values for tokens 0 through t-1? Token 3's Key at layer 7 is determined by token 3's representation entering layer 7 and the weight matrix $W_K$ at that layer. Neither of these changes between step 4 and step 200. The result is identical every time we compute it.

Q is temporary. The Query for position t is used to compute that position’s output. Later positions have their own Queries and don’t need the old one.

K and V are reused. Every later position may attend to token t, so its Key and Value remain useful. Because the causal mask prevents later tokens from changing earlier representations, we can store those vectors once.

The Fix

Store each token’s K and V when they’re first computed, then read them for later attention calculations.

Here is what a single decode step looks like with this change:

1. Compute Q, K, V for only the new token (token t) at each layer.
2. Append token t's K and V to the stored collection.
3. For attention: dot token t's Q against all stored Keys (tokens 0 through t), softmax, weighted sum of all stored Values.

The mathematical attention operation is unchanged. Old K/V vectors now come from memory instead of being recomputed. Different kernels or batch shapes can produce small floating-point differences, so this doesn’t guarantee bit-for-bit identity.

This persistent store of past Keys and Values is the KV cache.

Attention still compares the new Query with the stored Keys and combines the stored Values, so that work grows with prefix length. But old tokens no longer need their projections, FFNs, or other layer computations rerun. In particular, multiplying by $W_K$ and $W_V$ happens once per token per layer.

Building and Extending the Cache

The prompt pass builds the initial cache. Later passes process one new input token per sequence and extend that cache. These are called prefill and decode, respectively.

The cache trades storage for avoiding repeated computation. It doesn’t remove the need to process each new token through all layers.

How Much Memory Does the Cache Take?

The cache stores K and V for every token position at every layer. It grows by one entry per decode step. What does each entry cost?

TinyLlama uses GQA with 4 KV heads, each 64 dimensions:

Per position, per layer:
  K: 4 heads × 64 dims = 256 values
  V: 4 heads × 64 dims = 256 values
  Total: 512 values

At FP16 (2 bytes per value):
  Per layer: 512 × 2 bytes = 1,024 bytes per position
  All 22 layers: 1,024 × 22 = 22,528 bytes per position

For max context (2,048 tokens):
  2,048 × 22,528 bytes ≈ 46.1 MB per sequence

For one sequence, that cache sits alongside roughly 2.2 GB of model weights.

Each additional active sequence needs its own cached state. Ten full-length sequences need roughly 460 MB; a hundred need about 4.6 GB. Cache memory grows with both context length and concurrency, and can eventually exceed weight memory.

The 8× cache reduction from GQA applies to every active sequence. Under standard MHA, a hundred full-length sequences would need roughly 37 GB of cache instead.

Detail 11: Prefill vs Decode

Two Different Jobs

For "Write a story," the prompt IDs [8144, 264, 3446] are all available before the model runs. Prefill processes those positions together through each layer, while causal masking controls which positions can attend to which. The pass fills the prompt’s KV cache and uses the last position’s logits to choose "Once."

Decode then processes "Once," using the prompt’s cached Keys and Values, and chooses "upon." The next step processes "upon" to choose the following token. Unlike prompt tokens, these inputs aren’t all known in advance: each step must wait for the preceding choice.

Why the GPU Cares

A longer prompt makes the hardware difference easier to see. Suppose a request has 1,000 prompt tokens and will produce 200 output tokens.

During prefill, each weight matrix is used for 1,000 token vectors. One full-width TinyLlama projection has shape [2,048 × 2,048], about 8.4 MB at FP16. Matrix-multiplication kernels load tiles of that matrix and reuse them across many prompt positions. This gives the GPU substantial arithmetic to perform for each byte read.

With enough prompt tokens, the GPU can spend much more time doing arithmetic than waiting for weights. Its compute throughput often becomes the main limit for these matrix multiplications.

For a single decode stream, each matrix is applied to one new token vector. The weights still have to be read, but they produce much less arithmetic per byte. The memory system can become the limit even though the GPU has spare compute capacity.

When arithmetic throughput limits an operation, we call it compute-bound. Large prefill matrix multiplications often operate in this regime.

When memory transfers limit an operation, we call it memory-bandwidth-bound. Single-stream decode commonly operates in this regime because it reads the weights for so little arithmetic.

These are workload-dependent descriptions. Short prefills, long-context attention, and large decode batches can behave differently.

Back to the User Experience

The wait until the first output token is time to first token (TTFT). Prefill contributes to that wait, along with any queueing and request-processing overhead. Longer prompts generally require more prefill work.

The time between successive output tokens is inter-token latency (ITL). During ordinary decode, each step reads a similar amount of weight data, which helps explain a steady streaming rate. The growing KV cache and other active requests can still change the time per step.

Decode Dominates

Using average ITL for the intervals after the first token, total response time is:

Total latency = TTFT + ITL × (output_tokens - 1)

For our 1,000-token prompt with a 200-token response: 1 prefill pass, then 199 decode steps. (The first output token comes from prefill; every subsequent token is a decode step.)

Even if prefill takes 10× longer than a single decode step, 199 decode steps overwhelm it. Prefill contributes less than 5% of the total latency. Decode contributes over 95%. Extend the response to 500 tokens and prefill drops below 2%.

Under these assumptions, decode accounts for most of the response time. A long prompt, a short output, or a long queue can change that balance.

Reducing decode cost can mean reading fewer weight bytes, sharing a read across more requests, or accepting several proposed tokens from one target-model pass.

Detail 12: Memory Bandwidth is the Bottleneck

One Decode Step, Traced

For one decode stream, each step applies the model to one new token.

The main weight reads come from the attention projections, the FFN matrices in all 22 layers, and the LM head. Using the rounded model size, 1.1 billion parameters at 2 bytes each gives roughly 2.2 GB. The actual read is somewhat smaller because the embedding layer looks up a row rather than scanning the whole table. Attention also reads the KV cache, which becomes more significant as context grows.

For a matrix multiplication, each weight contributes roughly one multiply and one add, or two floating-point operations (FLOPs). Using parameter count as an approximation gives:

~2 FLOPs × 1.1 billion parameters ≈ 2.2 billion FLOPs

Dividing roughly 2.2 billion FLOPs by 2.2 GB gives about one floating-point operation per byte loaded.

An A100 SXM 40GB has a peak FP16 tensor-core throughput of 312 trillion FLOPs/s and peak memory bandwidth of 1,555 GB/s. At 1 FLOP per byte, even saturating memory bandwidth would supply only about 1.6 trillion FLOPs/s, around 0.5% of that compute peak. This ratio describes the workload’s arithmetic intensity, not a measurement of how long the GPU is idle.

At those peak rates, 2.2 billion FLOPs would take about 7 microseconds, while transferring 2.2 GB takes about 1.4 milliseconds. Real matrix-vector kernels don’t reach the tensor-core peak, but the comparison shows why adding arithmetic capacity alone won’t remove this bottleneck.

The weight traffic repeats at every decode step. An A100’s roughly 40 MB L2 cache can’t hold a 2.2 GB model, so a 200-step calculation can involve roughly 440 GB of weight reads under this simplified model.

Caches can retain some data, and actual traffic depends on the kernels. Treat one model-sized read per step as a useful approximation for this workload, rather than exact traffic accounting or an unconditional latency bound.

The Bandwidth Floor

If a step transfers 2.2 GB through a memory interface with peak bandwidth B GB/s, that transfer needs at least 2.2 / B seconds. The bound applies to that assumed traffic volume; reducing the traffic changes the bound.

These are transfer-time bounds for the assumed 2.2 GB. Kernel overhead and additional traffic can increase latency, while a lower actual transfer volume changes the bound. Form factor matters too: H100 PCIe is specified at about 2,000 GB/s, compared with 3,350 GB/s for H100 SXM.

Under this model, data movement alone takes fractions of a millisecond to a few milliseconds. Quantization reduces the weight bytes, while tensor parallelism can split their transfer across several GPUs.

Arithmetic Intensity

The ratio of computation to data movement is called arithmetic intensity:

For batch-size-1 decode, the weight-only estimate is about 1 FLOP per byte. Real arithmetic intensity also includes cache, activation, and intermediate traffic.

For a 1,000-token prefill, the weight-only estimate becomes:

• FLOPs: roughly 2.2 trillion from the parameter-count approximation
• Weight bytes: roughly 2.2 GB under ideal reuse
• Weight-only intensity: roughly 1,000 FLOPs per byte

The roofline model compares peak compute throughput with the throughput allowed by memory bandwidth at a given arithmetic intensity. Other limits, such as kernel-launch latency, can put measured performance below both.

The ratio of peak compute throughput to peak memory bandwidth gives the point where these limits meet. For the A100 figures above, it’s about 201 FLOPs per byte. Below that point, bandwidth limits the roofline estimate; above it, compute does. The exact threshold depends on the GPU and numerical format.

Single-stream FP16 decode at about 1 FLOP per byte is far below that threshold. A sufficiently large prefill can be above it. This explains the different optimization priorities without implying that a bandwidth-bound GPU consumes little power or has every arithmetic unit idle.

Reusing Weights Across Requests

The weights are shared by all requests using the same model. Applying one loaded matrix to 32 token vectors instead of one raises the weight-only arithmetic intensity:

With batch size 1:

• Read ~2.2 GB weights
• Produce 1 token
• Arithmetic intensity: ~1 FLOP/byte

With batch size 32:

• Read ~2.2 GB weights (same bytes)
• Produce 32 tokens (one per sequence)
• Arithmetic intensity: ~32 FLOPs/byte

Each request still contributes its own KV-cache traffic, and the additional arithmetic eventually matters. But reusing weight reads gives batching a substantial advantage at small batch sizes.

The Inference Pipeline and Memory Budget

For TinyLlama serving one full-length sequence, weights and KV storage give this partial memory budget. Allocator reservations, kernel workspaces, and runtime state add to it:

Model weights (FP16):     2.2 GB
KV cache (2,048 context): ~46.1 MB (GQA keeps this small)
Working buffers/runtime: implementation-dependent
Weights + KV cache:      ~2.25 GB, before other allocations

The Three Bottlenecks

A serving system has to manage three limits:

1. Memory bandwidth: single-stream decode repeatedly reads model weights and cached attention state.
2. Memory capacity: weights, active KV caches, and working memory must fit on the available devices.
3. Sequential generation: each output choice conditions the next, limiting how much work one request can do in parallel.

Part 2: Production Optimizations

Detail 13: Batching Requests

A single decode stream leaves compute capacity available while the GPU reads weights. Another request using the same model can reuse those weight reads.

The two requests have different prefixes and separate KV caches, but they use the same projection and FFN matrices. We can apply each loaded matrix to both requests’ token vectors.

This is batching: processing several sequences in one forward pass. Reusing the weights across their token vectors raises arithmetic intensity.

The batch-size-32 example in Detail 12 raises weight-only intensity to about 32 FLOPs/byte. That leaves compute headroom on the A100, though KV traffic still adds to the transfer time.

Throughput can initially grow close to linearly with batch size. KV-cache traffic, compute throughput, and memory capacity eventually limit that growth.

Static Batching

A simple implementation collects N requests, prefills them, then decodes them in lockstep. Each step generates one token per unfinished sequence, and the group stays fixed until its longest request finishes.

The scheme above works if all requests arrive at the same time and generate the same number of tokens. In production, neither is true.

Three requests arrive together:

Request A: "Write a poem"      → generates 50 tokens
Request B: "Explain gravity"   → generates 120 tokens
Request C: "Hello"             → generates 20 tokens

For the first 20 steps, all three requests are active.

At step 20, C finishes and returns its result, but its slot remains in the fixed batch as padding while B continues.

A finishes at step 50, leaving two padded slots for the final 70 steps. During that interval, each weight read produces only one useful token.

A request arriving at step 30 can’t join the group, even though C’s slot is empty. It waits until step 120, when the fixed batch completes: 90 steps of avoidable queueing.

There is also a cost before decoding even begins. You need N requests to form a batch of N. Under low traffic, you wait for them to accumulate (or use a timeout, which produces smaller, less efficient batches). Either way, users experience added latency from the queueing alone.

This is static batching: requests can finish at different times, but new requests can’t replace them until the batch completes.

Static batching reuses weights, but fixed membership leaves padded slots and delays admission of new requests. Waiting to form a batch can add latency before execution even starts.

Detail 14: Continuous (In-Flight) Batching

C’s empty slot doesn’t contribute anything to A or B’s computation. Each request needs its own prefix and KV cache; none depends on which other requests share the weight read.

That independence lets the scheduler choose a new batch at every iteration. Requests don’t have to start or finish together.

When C finishes, remove it and make its memory available. The scheduler can admit another request after preparing its prompt cache.

With requests D and E already waiting and their prefill accounted for, a schematic schedule looks like this:

Step  0: [A, B, C]     all three decoding
Step 20: [A, B, _]     C finishes, result returned, slot freed
Step 21: [A, B, D]     Request D admitted, joins decode loop
Step 50: [_, B, D]     A finishes, slot freed
Step 51: [E, B, D]     Request E admitted
         ...

C’s slot can do useful work on the next scheduled iteration. D doesn’t need to wait for A and B to finish, though it must complete prefill before decoding.

Under a steady queue of requests, replacing finished sequences keeps more batch slots useful. Reported gains over static batching can be several-fold, but depend on request lengths, arrivals, and scheduling overhead.

This is continuous batching, also called in-flight batching because the batch composition changes while generation is in flight. Most production inference servers implement some version of it. vLLM, TGI, and TensorRT-LLM all use continuous batching as their default scheduling strategy.

What Happens When a New Request Joins

A new request can't simply appear in the decode loop. It has no KV cache yet. The decode step expects to read cached Keys and Values for every past token position (Detail 10), and for a brand-new request there is nothing to read.

So when the scheduler decides to admit Request D, it first runs prefill for D's prompt. All of D's prompt tokens get processed through the 22 transformer layers in a single parallel pass, just like the prompt processing from Detail 11. This populates D's KV cache. Once prefill finishes, D joins the decode loop alongside everyone else, advancing one token per step.

D still goes through prefill and decode. The scheduler now has to fit its prefill work around the sequences already decoding, which can affect their latency.

The Scheduler

The engine that makes this work is a scheduling loop that runs once per decode step:

1. Evict finished sequences. Any sequence that sampled an end-of-sequence token or hit its maximum length gets removed. Its output is returned to the user. Its KV cache memory is freed back to the pool.

2. Admit waiting requests. If there are queued requests and enough KV cache memory to hold a new sequence, run prefill and insert the new sequence into an active slot. This is the step that keeps the batch full.

3. Run one decode step for active sequences. Apply the model to each sequence’s latest input token, append that processed position’s K/V to its cache, and sample the following token. Batched matrix multiplications reuse weight reads across the sequences.

4. Repeat.

Detail 15: The KV Cache Memory Problem

Admitting a request allocates a KV cache that will grow as generation continues. Even when the GPU has compute capacity available, the scheduler may have to leave a request queued because there isn’t enough cache space.

The Allocation Problem

Requests arrive at different times with different prompt lengths. Each cache starts at a different size, grows during decode, and is released when that request finishes.

This is a memory-allocation problem involving variable-sized objects with unpredictable lifetimes. Two straightforward allocation policies show where space gets wasted.

Approach 1: Reserve for the Worst Case

A simple policy reserves contiguous space for the maximum context length as soon as a request arrives. For TinyLlama at 2,048 positions, that’s about 46.1 MB, whether the request eventually needs that many positions or not.

Reserving the full block avoids reallocations and copies as the cache grows, but prevents other requests from using its spare capacity.

Suppose a short request uses 50 processed positions before it finishes. It needs about 1.13 MB of KV data, but the maximum-length allocation reserves about 46.14 MB.

Reserved:  46.14 MB (2,048 positions)
Used:       1.13 MB (50 processed positions)
Wasted:    45.01 MB (97.6% of the allocation)

The unused capacity stays reserved for this request, so other requests can’t use it.

If ten requests reserve full-length caches but use only about 20 MB in total, hundreds of megabytes remain allocated without holding useful state.

This unused space inside an allocation is internal fragmentation.

Approach 2: Grow As You Go

An alternative is to allocate only enough contiguous space for the prompt, then extend it as the sequence grows. This reduces unused space inside allocations, but requires room next to each growing cache or a copy to a larger region.

Three sequences arrive. A needs 100 MB, B needs 150 MB, C needs 80 MB. They get placed contiguously:

[  A: 100 MB  ][  B: 150 MB  ][  C: 80 MB  ][    free: 670 MB    ]

The allocations fit in a 1,000 MB pool. When A finishes, its 100 MB becomes free:

[  free: 100  ][  B: 150 MB  ][  C: 80 MB  ][    free: 670 MB    ]

Sequences D and E arrive. D needs 60 MB, fits in the gap left by A. E needs 200 MB, takes from the end:

[ D: 60 ][ free: 40 ][  B: 150 MB  ][  C: 80 MB  ][ E: 200 MB ][ free: 470 ]

B finishes:

[ D: 60 ][ free: 40 ][  free: 150  ][  C: 80 MB  ][ E: 200 MB ][ free: 470 ]

A new 200 MB request can use the 470 MB region at the end. The adjacent 40 MB and 150 MB regions can be merged into 190 MB, but still can’t hold that request. Total free space is 660 MB; only some of it is usable for this allocation.

A separate snapshot of a fragmented 1,000 MB pool could look like this:

[D: 60][free: 40][C: 80][free: 150][E: 200][free: 30][G: 90][free: 350]

This layout has 570 MB free, but its largest gap is 350 MB. A request needing 400 MB would fail despite enough free bytes in total.

This is external fragmentation: free memory is divided into regions that are individually too small for an allocation.

Two Failures, One Root Cause

Reserving for maximum length wastes space inside allocations. Growing contiguous allocations can leave unusable gaps or require expensive copies. Both policies insist that a sequence’s cache occupy one contiguous region, even though attention only needs a way to find each token’s K/V.

A cache doesn’t have to be physically contiguous if the kernel can map logical token positions to their stored locations.

Detail 16: PagedAttention (vLLM's Innovation)

Operating-system virtual memory provides a useful model for removing that contiguity requirement.

How Virtual Memory Works

With a contiguous physical allocation, programs that grow and terminate leave gaps like the KV-cache allocations above.

Virtual memory separates the addresses a program uses from physical memory locations. The OS divides the virtual address space into fixed-size pages, often 4 KiB, and maps them to available physical frames through a page table.

For example, address 5,000 is in virtual page 1 when pages hold 4,096 bytes. A page-table mapping can place that page in physical frame 37. Hardware performs address translation, usually aided by a translation cache, so the program doesn’t need to know where the page is stored.

Within a pool of uniform frames, any free frame can back a page, avoiding the need for a large contiguous physical allocation. Internal fragmentation remains in partially filled pages, but the unused tail is bounded by one page for an allocation.

From RAM Pages to KV Blocks

The mapping is direct:

• A sequence is the program
• Its logical KV positions (token 0, token 1, ...) are the virtual addresses
• Physical frames become fixed-size KV blocks in VRAM (say, 16 tokens each)
• A per-sequence block table plays the role of the page table

Walking Through a Request

The earlier 50-position example reserved 2,048 positions under the maximum-length policy. Paging allocates storage as positions are needed.

With blocks of 16 positions, the first three blocks store positions 0–47. Positions 48–49 occupy two slots in a fourth block, leaving 14 unused. The four blocks can be anywhere in the KV pool:

Block table for Sequence A:
  Logical 0 → Physical 7     (tokens 0-15)
  Logical 1 → Physical 2     (tokens 16-31)
  Logical 2 → Physical 14    (tokens 32-47)
  Logical 3 → Physical 5     (tokens 48-49, 14 slots empty)

From the sequence's perspective, its KV cache is a contiguous stretch of 50 positions. Physically, it's four scattered chunks.

Bounded Waste

Block 3 has 14 empty slots. That's real waste, the same internal fragmentation from Detail 15. But paging caps it at 15 slots per sequence (one block minus one token). Compare that to 1,998 under static pre-allocation.

Any free block in the KV pool can satisfy the next block allocation. The cache doesn’t require adjacent physical blocks, so gaps between them no longer prevent growth.

The vLLM paper measured substantial reservation waste in the systems it compared. With paging, unused tail space is bounded by one partial block per sequence; its percentage depends on block size and sequence lengths.

Growth and Cleanup

If the sequence continues, new processed positions fill the last block. After its 14 free slots are used, the next position requires a new block, such as physical block 11. Existing blocks don’t move; the block table gains an entry mapping logical block 4 to physical block 11.

When the sequence finishes, its private blocks return to the pool. Shared prefix blocks may remain referenced by other requests or a prefix cache. Neither case requires compacting the other live allocations.

What Changes in the Kernel

The attention kernel still reads the Keys and Values of the positions available to the current token. It follows a block-table mapping to find them. For token 35:

block_size    = 16
logical_block = floor(35 / 16)           # = 2
offset        = 35 % 16                   # = 3  (4th slot in that block)
physical      = block_table[2]            # = 14  (from our table above)

Token 35’s Key is in the fourth slot of physical block 14. Once loaded, it participates in the same attention calculation as a contiguously stored Key.

The kernel must handle that indirection efficiently. It can arrange contiguous accesses within each block, but block-table lookups and the chosen physical layout affect memory-access efficiency.

In the vLLM paper’s comparisons, memory utilization jumps from 20-38% to over 96%. Better allocation lets the scheduler fit more requests, contributing to the reported throughput improvements. The gain depends on the baseline allocator and workload.

Back-of-the-Envelope (TinyLlama)

Suppose 50 active sequences each have 500 processed positions. Each position stores 22,528 bytes across all layers.

With a full 2,048-position reservation per sequence:

• Reserved: 50 × 2,048 × 22,528 bytes ≈ 2.307 GB
• Used: 50 × 500 × 22,528 bytes = 563.2 MB

With 16-position blocks, each sequence uses 32 blocks, or 512 slots. Across all 50 sequences, that is 50 × 512 × 22,528 bytes = 576.7 MB. The final-block slack is 13.5 MB in this example, bounded by fewer than 16 positions per sequence.

Paging frees about 1.73 GB relative to the maximum-length reservation, which the scheduler can allocate to other sequences.

This combination of paged KV storage and attention kernels that read it is PagedAttention, introduced by vLLM.

Detail 17: FlashAttention

During prefill, attention can spend substantial bandwidth on temporary score and probability matrices, even when the model’s linear layers have enough work to be compute-bound.

An unfused implementation writes those matrices to GPU memory and reads them back for the next operation. For a 2,048-token prompt, that traffic is much larger than the input and output vectors.

Two Tiers of Memory

A GPU has several memory levels. Two matter for this explanation: device memory and the smaller working memory near the compute units.

HBM (High Bandwidth Memory) is the device memory used by the A100. The SXM 40GB model has about 1.6 TB/s of peak bandwidth. HBM sits on separate chips in the same package as the GPU. Other GPUs, such as the RTX 4090, use GDDR memory instead.

On-chip SRAM includes the storage used for registers and shared memory near the compute units. The FlashAttention hardware model uses roughly 20 MB of on-chip working memory across an A100’s streaming multiprocessors, with much higher aggregate bandwidth than HBM. A single computation can use only a portion of that capacity.

This smaller working memory is where kernels keep tiles and intermediate results while computing. It can’t hold a full model or arbitrarily large attention matrices.

Separate GPU kernels generally exchange intermediate arrays through global device memory. Those accesses may hit on-chip caches, but a large array can still incur substantial HBM traffic. Fusing operations lets one kernel consume intermediate results while they remain in registers or shared memory.

(Not all inference hardware follows this two-tier pattern, but GPUs dominate inference today. The rest of this section addresses their specific bottleneck.)

Standard Attention: Where the Bytes Go

For one head and a 2,048-token prompt, the Q, K, and V matrices have these shapes:

Q: [2,048 × 64]   (one 64-dim query per token)
K: [2,048 × 64]   (one 64-dim key per token)
V: [2,048 × 64]   (one 64-dim value per token)

Each input matrix has 131,072 FP16 values, or 262,144 bytes: 256 KiB, about 0.262 MB. Q, K, and V together occupy about 0.786 MB.

Now run the attention formula from Detail 4. Step one: compute the score matrix $S = QK^T / \sqrt{d_k}$.

Q is [2,048 × 64]. K transposed is [64 × 2,048]. The result: [2,048 × 2,048]. One attention score for every pair of tokens. $2{,}048 \times 2{,}048 \approx$ 4.2 million entries.

S: [2,048 × 2,048] = 4,194,304 numbers
At FP16: 4,194,304 × 2 bytes ≈ 8.389 MB (8 MiB)

The three inputs total about 0.786 MB, while S alone is about 8.389 MB. Doubling the sequence length quadruples S because it contains one score per token pair.

Softmax produces P with the same [2,048 × 2,048] shape, about 8.389 MB if stored in FP16. Some implementations use higher precision for intermediate values.

Multiplying P by V gives [2,048 × 64], about 0.262 MB. The output is much smaller than the intermediate score and probability matrices.

In this simplified unfused implementation, the three kernels exchange materialized S and P arrays through global memory. Ignoring cache hits and repeated tile loads, the traffic is:

1. Scores: read Q and K, then write S (8.389 MB).
2. Softmax: read S (8.389 MB), then write P (8.389 MB).
3. Output: read P (8.389 MB) and V, then write the output (0.262 MB).

Writing and reading S and P moves 4 × 8.389 MB ≈ 33.554 MB. Reading Q, K, V and writing the output adds about 1.049 MB, for roughly 34.603 MB per head.

Most of this traffic comes from writing and reading S and P. The final output doesn’t require retaining either matrix once its contribution has been computed.

Across 32 heads and 22 layers, this per-head accounting gives about 24.36 GB of traffic, around 15.7 ms at 1,555 GB/s. It models full dense matrices; causal kernels can skip masked regions, and GQA can reuse K/V reads. The corresponding dense matrix products take about 0.756 trillion FLOPs across the stack, an ideal 2.4 ms at the A100’s peak FP16 tensor-core rate.

(The estimate counts K/V reads separately for each Query head. GQA can reduce that traffic through sharing, but S and P still dominate this materialized-matrix example.)

Under these assumptions, unfused prefill attention is bandwidth-bound because of intermediate arrays. Single-stream decode has a different source of traffic: it repeatedly reads weights and the KV cache.

The two dense matrix products require about 1.074 billion FLOPs per head. Dividing by 34.603 MB gives about 31 FLOPs/byte, below the A100 threshold. This comparison explains the incentive to avoid the intermediate transfers; actual timings depend on the kernels and cache behavior.

The Fix: Never Materialize the Full Matrix

Fusing score calculation, softmax, and the Value multiplication could avoid writing S and P to HBM. The constraint is fitting the working data into on-chip memory.

The full score matrix for one head is about 8.4 MB, far larger than the shared memory available to a thread block. Fusion therefore needs a way to compute the output without retaining all the scores.

Once we have normalized the scores in row i and used them to combine the Values, we have that Query’s output. The score row is no longer needed. We can process small tiles and accumulate the output as we go.

FlashAttention organizes this as follows:

1. Divide Q into blocks of rows (say, 128 at a time). Divide K and V into blocks of the same size.
2. Load one Q block into SRAM: [128 × 64] = 16 KiB at FP16. Fits easily.
3. Stream K/V blocks through one at a time. For each K block, compute a small score tile: [128 × 128] = 32 KiB at FP16. Update the running softmax state, multiply by the corresponding V block, and accumulate into a running output. All of this stays in SRAM.
4. After all K/V blocks have streamed through, the output for these 128 query rows is complete. Write it to HBM.
5. Move to the next Q block. Repeat.

K and V are read for each Query block in this loop ordering. With $\lceil N / B_r \rceil$ Query blocks, this repeats input reads, but avoids writing and reading the much larger S/P arrays. In the IO model, the resulting transfer bound is $O(N^2 d^2 / M)$ elements, where $M$ is on-chip capacity measured in elements; standard materialization incurs $O(N^2)$ intermediate elements.

With this tiled implementation, only a [128 × 128] score tile needs storage at a time. That’s about 32.8 KB in FP16, 256× smaller than the complete score matrix.

The Softmax Obstacle

Softmax normalizes each row by the sum of exponentials across all keys:

$\text{softmax}(s_i) = \frac{e^{s_i}}{\sum_{j=1}^{N} e^{s_j}}$

A tile contains only part of the row, so its local sum can’t yet normalize scores over all keys.

After processing keys 0–127, we have a partial sum of exponentials. Later keys add to the denominator and change the normalized weights for the earlier keys. We need a way to update their contribution without keeping every earlier score.

One option is to scan the scores once for the denominator, then scan them again to calculate the output. That repeats score computation and Key reads. An incremental normalization avoids that second pass.

FlashAttention avoids the second pass entirely. The idea: maintain three running quantities per query row as key blocks stream through:

1. Running maximum $m$: the largest score seen so far. (Subtracting the max before exponentiating prevents overflow. Standard practice in any stable softmax implementation, not specific to FlashAttention.)
2. Running sum of exponentials $\ell$: $\sum e^{s_j - m}$, corrected whenever the maximum changes.
3. Partial output vector $O$: the weighted sum of value vectors accumulated so far.

With an unnormalized running output, changing the maximum from $m$ to $m_{\text{new}} > m$ requires multiplying the old sum $\ell$ and output $O$ by $e^{m - m_{\text{new}}}$ before adding the new contributions. If the maximum is unchanged, the factor is 1. Dividing by the final sum produces the normalized output.

After all blocks, the result is mathematically equivalent to applying softmax to the full row. Floating-point rounding can differ because the operations occur in a different order. This is online softmax, which makes single-pass tiled attention possible.

Memory Traffic After Tiling

FlashAttention keeps score tiles and running softmax state on-chip rather than materializing the full [N × N] intermediates in HBM. The running state consists of a maximum, an exponential sum, and an output vector for each query row.

FlashAttention computes the same attention function with the same main matrix products, plus bookkeeping for the incremental softmax. Its benefit comes from reducing HBM traffic. That raises arithmetic intensity and can move attention closer to the compute limit, especially for longer prompts.

PyTorch’s scaled_dot_product_attention can select a fused attention backend when the hardware, data type, shapes, and other constraints permit it. Serving engines also use specialized attention kernels for their cache layouts.

The N × N intermediate arises during prefill. During single-token decode, the score array is one row long, so the opportunity is different. Decode kernels still benefit from fusion and careful KV reads, but weight and cache traffic remain major costs.

Detail 18: Quantization

TinyLlama’s FP16 weights occupy about 2.2 GB. A dense 70B model needs about 140 GB, exceeding the 80 GB capacity of an A100 80GB before any cache or working memory is allocated.

Even when a model fits, each decode step has to move its weights. Reducing the bytes per weight can therefore help both capacity and bandwidth.

Do We Need All 16 Bits?

Suppose a weight is 0.0217. FP16 stores it in 16 bits, with roughly three to four decimal digits of precision. A matrix multiplication multiplies that weight by an activation and adds the product to many others.

Rounding it to 0.024 changes the weight by 0.0023, about 10.6% of its original value. The effect on the output depends on the activation it multiplies and how that result propagates. A small absolute weight error may have little effect, but we can’t infer that from the weight alone.

Many models tolerate lower-precision weights well enough to make the storage savings useful. The engineering question is how to choose the representation and scales while keeping the output quality acceptable.

Fewer Bits, Coarser Grid

Think of representable numbers as points on a number line. FP16 has 65,536 bit patterns, including encodings for infinities and NaNs. Its finite values are unevenly spaced: the spacing grows with magnitude.

A scaled INT8 representation has at most 256 levels, while INT4 has 16. Their coarser spacing introduces more rounding error, but storing each weight takes half or a quarter of the FP16 bits.

For a symmetric INT8 quantizer, divide the largest weight magnitude by 127, the largest positive code used here. Divide each weight by that scale and round to an integer. This uses 255 codes from −127 through 127. Multiplying the code by the scale reconstructs an approximation.

scale = max(abs(W)) / 127
q     = round(W / scale)   # INT8 values in [-127, 127]
W_hat = q * scale           # approximate reconstruction

If the largest magnitude is 1.0, scale = 1/127 ≈ 0.00787. The weight 0.0217 maps to round(0.0217 × 127) = round(2.7559) = 3. Reconstructing gives 3/127 ≈ 0.02362, an absolute error of about 0.00192, or 8.9% of the original weight. The output effect still depends on how that weight is used.

INT4 has 16 integer levels, [-8, 7] in signed two’s-complement, using 4 bits per weight. Different quantizers may use a symmetric subset or an offset to map those levels to the weight range.

One Scale Factor Isn't Enough

An outlier can make a shared scale too coarse for the other weights. If a group contains a weight of 10.0 while the rest lie below 0.5 in magnitude, covering the outlier spreads the integer levels across a much wider range. The smaller weights then have fewer distinct reconstructed values available.

Quantizers can choose scales at different granularities:

• Per-tensor: one scale for the entire matrix. Simplest, but one outlier stretches the range for everyone.
• Per-channel: one scale per output dimension. Better, because magnitudes often vary across channels.
• Per-group: one scale per small block of weights, typically 64 to 128 values. Each block gets its own scale, tuned to its local range.

Smaller groups can give their weights a better-fitting scale, at the cost of more scale metadata. GPTQ and AWQ implementations commonly use groupwise quantization, which helps preserve quality at 4 bits.

This whole process, mapping floating-point weights to low-bit integers plus a scale factor, is quantization.

What Gets Quantized

An "INT4 model" often means weight-only quantization. The weight matrices are stored at low precision while activations, the values computed from each input, remain in FP16 or BF16.

Weights are fixed, so their scales can be chosen offline. Activations change with the input and can contain outliers much larger than the other values. Low-precision activation formats need scales and methods that account for that variation.

In a common weight-only implementation, the kernel reconstructs weight tiles as it loads them, multiplies them by FP16 or BF16 activations, and accumulates the results. It doesn’t first expand the entire model into a second full-precision copy.

FP8, an 8-bit floating-point format, is also used for weights and activations. Its exponent provides more dynamic range than a fixed integer grid, but it still needs suitable scales and can still clip outliers. GPUs including NVIDIA Hopper and Ada support FP8 operations; Blackwell also supports low-precision FP4 formats. The best choice depends on hardware, kernels, and measured quality.

The Bandwidth Payoff

Under the weight-transfer model from Detail 12, reducing bytes per weight lowers the time needed to transfer them. Using the A100 SXM 40GB’s 1,555 GB/s bandwidth:

These numbers are approximate. Real quantized formats store scale (and sometimes zero-point) metadata per group, so the file is slightly larger than the naive calculation. And the speedups are bandwidth ceilings, not guarantees. Kernel overhead, dequantization cost, KV-cache traffic, and batch size all affect the actual latency.

For bandwidth-bound kernels, reading fewer weight bytes creates an opportunity to reduce ITL. The kernel implementation determines how much of that opportunity becomes a measured speedup.

Scale the math to a 70B model. FP16: ~140 GB of weights. One A100 (80 GB) can't hold it. You'd shard across at least two GPUs, paying inter-GPU communication overhead on every forward pass.

INT4: ~35 GB of weights. Add scale-factor metadata, call it 38 to 40 GB. Fits on a single 80 GB GPU with room for KV cache and activations. A multi-GPU problem becomes a single-GPU problem.

What Weight-Only Quantization Does Not Change

First: the KV cache. If you only quantize weights, the KV cache stays in FP16 or BF16. For short contexts and small batches, weights dominate memory and quantization helps a lot. But as context grows or batch size increases, KV cache becomes a larger share of total memory. Weight-only quantization shrinks the fixed cost (weights). It doesn't touch the per-sequence cost (KV cache).

The KV cache can also be quantized, for example from FP16 to FP8 or INT8. This roughly halves the stored data before metadata; INT4 can reduce it further. Quality depends on the model, scales, and workload, especially for long contexts, so the cache format needs its own evaluation.

Quantizing the KV cache reduces memory pressure (more sequences fit in VRAM) and bandwidth demand (each decode step reads less KV data per sequence). The latter matters directly in Detail 21 when we look at how KV traffic scales with batch size.

Second: the number of forward passes. You still read the (now smaller) weights once per token. The cost per read drops, but the count stays the same. Each decode step still produces exactly one token.

The Cost: Quality

In the INT8 example, 0.0217 reconstructs as about 0.0236. Many such errors enter a matrix multiplication. Some may cancel, while others can reinforce each other or affect sensitive features.

With the example scale of 1/127, the midpoint between codes 1 and 2 is about 0.011811. A weight near that boundary can round to a different code after a small change. Most such differences have little effect, but one can shift nearly tied logits enough to change a greedy choice and therefore the following context.

INT8 generally permits smaller errors than INT4, but bit width alone doesn’t determine quality. Calibration, grouping, and the model’s sensitivity matter too. Conversational output may tolerate a change that causes failures in exact code generation or multi-step arithmetic, so evaluate the tasks the model will actually serve.

At 3 or 2 bits, there are fewer levels available and preserving quality becomes harder. Results depend heavily on the quantization method and model.

Calibration runs representative inputs through the model to measure activation ranges or the effect of quantization errors. It can guide scales and the treatment of sensitive channels without updating the model through a full training run. Its cost depends on the method, model size, and calibration data.

Detail 19: Speculative Decoding

Quantization reduces the cost of a weight read, and batching shares it across requests. For one request, the basic decode loop still chooses one token per forward pass.

The Dependency That Keeps Us Slow

The obvious wish: run the model once and produce 4 tokens instead of 1.

Autoregressive generation won't allow it. Say the prefix is "Once upon a" and we want the next 4 tokens. Token 1 might be "time." But token 2 depends on token 1. If token 1 is "time," token 2 is probably ",". If token 1 were "hill," token 2 might be "there" instead. You can't compute the right token 2 without first knowing token 1. Token 3 depends on token 2. Token 4 on token 3. Each token is conditioned on every token before it. That's the constraint from Detail 3.

If we already have candidate tokens, however, the target model can evaluate their conditional probabilities together in one pass.

The candidates supply the inputs that ordinary generation would have to wait for. A causal forward pass can score each candidate using the prefix and the earlier candidates, just as prefill processes known tokens together. Acceptance must still be checked from left to right: after a rejection, later scores use a prefix we won’t retain.

The Cost of Verification

For the single-stream estimate in Detail 12, one TinyLlama decode step:

• Reads ~2.2 GB of weights from VRAM
• Performs ~2.2 billion FLOPs
• Arithmetic intensity: ~1 FLOP/byte

Processing five candidate positions uses more arithmetic than processing one, but can reuse the same weight reads. When weight bandwidth is the main limit and compute capacity is available, that extra arithmetic can add relatively little latency.

Verification can therefore cost much less than five separate decode steps. It isn’t free, and the benefit shrinks if the batch is already using the available compute throughput.

Generating Candidates With a Smaller Model

Familiar phrasing gives a small model a chance to agree with a larger one. After "Once upon a," both may assign a high probability to "time." Structured or repetitive text can offer many such continuations.

Suppose we have a compatible 125M-parameter draft model using the target’s tokenizer. Its FP16 weights occupy about 250 MB, roughly one-ninth of TinyLlama’s. This is a hypothetical size comparison, not a claim that a particular 125M model will draft well for TinyLlama.

For greedy decoding, the procedure is:

1. The draft model generates K candidate tokens autoregressively.
2. The target evaluates the candidates together, reusing its cached prefix and processing the new positions.
3. Check candidates from left to right. Accept each one that matches the target’s highest-scoring token. At the first mismatch, emit the target’s choice and discard the remaining candidates.

Tracing Through an Example

Prefix: "Write a story. Once upon a"

Draft phase. The 125M model generates ["time", ",", "there", "was"] in four passes, each involving roughly 250 MB of weights in this estimate.

Verification phase. TinyLlama reuses its cached prefix. It processes the needed new positions together, including the preceding token if that token hasn’t yet been cached, and obtains the target scores for the draft:

Position 1: draft = "time"  → target agrees   ✓ accept
Position 2: draft = ","     → target agrees   ✓ accept
Position 3: draft = "there" → target disagrees ✗ → use target's token: "in"
Position 4: draft = "was"   → discarded (conditioned on wrong token 3)

Result: 3 tokens from one target forward pass. Two accepted plus one from the target at the rejection point. Without this approach, those 3 tokens would have cost 3 separate reads of 2.2 GB each.

If the first candidate is rejected, we emit one replacement token. If all four are accepted, we can emit five: the four drafts plus a bonus token predicted from the last draft position. At the first rejection, later predictions must be discarded because they depend on the rejected prefix.

Ignoring end-of-sequence and output limits, each target pass emits the accepted prefix of drafts plus one target-selected token.

This is speculative decoding. The name borrows from CPU design, where processors speculatively execute instructions before knowing if a branch condition is true, discarding the work if the guess was wrong.

The Output Doesn't Change

A smaller model is proposing tokens. We're sometimes accepting them. The natural worry: does this pull quality toward the draft model?

Under the greedy acceptance rule, each accepted token agrees with the target’s choice at that position. At a mismatch, we emit the target’s choice instead. The draft determines which continuations can be checked together.

With identical target logits and tie-breaking, the emitted tokens match greedy decoding from the target alone. Numerical differences between execution shapes can still affect near-tied logits.

Sampling requires a different acceptance rule to preserve the target distribution. A draft token isn’t accepted merely because it matches the target’s highest score.

The rejection-sampling rule preserves the target’s distribution in exact arithmetic. Relaxing it to accept tokens that are merely plausible would change that distribution. Different random draws can produce different text even when two procedures sample from the same distribution.

What It Costs

Two models means memory for two weight sets and two KV caches:

Target weights (TinyLlama): 2.2 GB
Draft weights (125M):       ~250 MB
Total:                      ~2.45 GB (plus KV caches for both)

After a rejection, both caches must discard any entries for the rejected draft and later positions. The replacement token is processed before generation continues. Implementations can write tentative entries during verification, but those entries must not be reused after their prefix has changed.

Total time includes all draft passes and the target verification. If few candidates are accepted, the additional work can make speculative decoding slower than ordinary decoding.

When It Helps (and When It Doesn't)

The acceptance rate is the probability that a candidate passes the acceptance rule. With K = 4 and a constant conditional acceptance probability $a$, the expected number of output tokens per target pass is $1 + a + a^2 + a^3 + a^4$: one guaranteed output, plus each additional output weighted by the probability that all required earlier drafts were accepted.

Works well when:

• The draft model has high measured agreement with the target
• Output is locally predictable: boilerplate, common phrasing, structured text, code
• Batch sizes are small, so decode is bandwidth-bound and the spare compute that makes verification cheap is actually available

Works poorly when:

• The draft model is weak relative to the target (low acceptance means most draft work is wasted)
• Large batches already saturate the GPU's compute (verification stops being "free")
• Memory is tight (two models plus two KV caches may not fit)

Published experiments report multi-fold speedups for suitable low-batch workloads. Acceptance rate alone doesn’t predict wall-clock speedup; draft cost and verification latency also matter.

Variants

The draft-model approach works, but loading a second model has costs. Several alternatives eliminate that overhead.

Self-speculative decoding uses a cheaper path through the target model to draft, then verifies with the full model. Some methods skip selected layers or use trained early exits. Simply stopping after an arbitrary number of layers may not produce useful draft logits; the method must make that shorter path effective.

Medusa. Attach lightweight prediction heads to the target model's final hidden layer. Head 1 predicts t+1, Head 2 predicts t+2, and so on. Each head proposes several candidates, forming a tree of possible continuations verified in one forward pass. The limitation: the heads are independent. Head 3 doesn't know what Heads 1 and 2 predicted. It's guessing 3 steps ahead without knowing the intermediate tokens, which limits accuracy for later positions.

EAGLE uses a small autoregressive module to predict future hidden features, conditioned on both features and token information. Conditioning on the selected token helps distinguish possible continuations. EAGLE variants report multi-fold speedups under their evaluation settings; those results depend on the target model and workload.

Lookahead decoding uses iterative refinement, called Jacobi iteration, to form candidate continuations for target-model verification. It avoids a separate draft model or prediction heads, but still uses additional computation and candidate storage. Repeated patterns can make those candidates more useful.

Detail 20: Prefix Caching

A KV cache avoids recomputing earlier positions within a request. Separate requests can also share work when they begin with exactly the same tokens.

Suppose a chatbot prepends a 500-token system prompt containing instructions and tool descriptions. For illustration, count the first user message and its formatting as 12 more tokens. Without prefix caching, the server prefills all 512 positions.

A second request uses the same system prompt with an 8-token user suffix. Its 508-position prefill repeats the shared prefix.

With 100 requests per minute and an average 20-token unique suffix, the server processes 50,000 shared-prefix token positions and about 2,000 unique positions. Roughly 96% of the positions repeat, although that percentage isn’t a direct estimate of FLOPs or latency.

We could retain the shared prefix’s KV entries and attach them to later requests. That requires the entries to remain valid when a different suffix follows.

Can We Actually Reuse Those KV Entries?

The causal mask allows position $i$ to attend only to positions up to and including itself. The Key and Value at position $j$ therefore depend on tokens at positions $0$ through $j$, together with the model and positional configuration. Later tokens can’t change those states.

If the first 500 token IDs and all relevant model settings match, the corresponding cached states are mathematically reusable. Floating-point results may differ slightly from a fresh computation with a different batch shape, but the suffix introduces no new dependency into the prefix.

Reusing that state across requests is prefix caching.

Walking Through Two Requests

Request 1 has a 500-token system prompt and a 12-token user suffix, for 512 tokens total. With an empty cache, it needs full prefill. The shared prefix occupies about 11.3 MB of TinyLlama KV data. A paged engine can retain references to those blocks for later requests.

Request 2 has the same 500-token prefix and an 8-token suffix. With 16-token blocks and a policy that caches only complete blocks, it can reuse 31 blocks covering 496 positions. Prefill processes the remaining 4 prefix tokens and 8 suffix tokens: 12 positions instead of 508. Those positions still attend to the cached prefix, so the reduction in newly processed positions doesn’t translate directly into the same reduction in TTFT.

How Much Does It Help?

Long shared prefixes offer more reuse. Ignoring block boundaries, a 500-token prefix with a 10-token suffix leaves about 2% of positions to process anew; a 100,000-token prefix with a 200-token suffix leaves about 0.2%. Suffix attention still reads the cached prefix in both cases.

This can reduce TTFT substantially, depending on the suffix, cache lookup, and scheduling costs. It doesn’t remove the prefix from subsequent decode attention: each new token still reads the relevant cached Keys and Values.

Exact Tokens Only

Cache hits require matching token IDs in the same positions under the same model configuration. Similar meaning isn’t sufficient.

An extra space, a different quote style, or a timestamp can change token IDs and break reuse from that point onward. Reusing the exact prompt template is a practical way to keep IDs stable. Put dynamic fields after the shared prefix when possible. Byte-identical text is sufficient for a deterministic tokenizer, but identical token IDs are the actual requirement.

Production Reality

The cache consumes the same scarce GPU memory discussed in Detail 15, so servers use LRU or TTL eviction to keep only the most frequently reused prefixes resident. SGLang's RadixAttention organizes this cache as a radix tree (trie) over token sequences, enabling efficient longest-prefix matching, insertion, and eviction in a single data structure. Think of it as the prefix-caching analog of PagedAttention for KV memory management.

Block granularity explains the 496-position cache hit in the example. Cache keys must also distinguish settings that change computed states, such as adapters. Deployments may isolate cached user content by tenant.

Detail 21: Latency vs Throughput Trade-offs

Batching improves aggregate throughput by sharing weight reads. Each added sequence also brings its own KV-cache reads, and newly admitted requests need prefill. Those costs affect the users already receiving tokens.

The Two Clocks

TTFT includes queueing and prefill; ITL measures the spacing between output tokens. Suppose, purely for a timing example, prefill takes 5 ms and each decode step takes 1.4 ms.

A 200-token response: 5 + (199 × 1.4) ≈ 284 ms end-to-end.

These are illustrative timings, not a TinyLlama benchmark. Real timings include kernel overhead, cache traffic, and other work omitted from the bandwidth estimate.

A scheduler can improve one metric while worsening the other. Prioritizing a new prompt may lower that request’s TTFT while interrupting active decode streams.

Throughput counts output tokens per second across all active requests. A user’s ITL and the server’s throughput describe different things: increasing the number of active requests can raise aggregate throughput even when each individual stream slows down.

Batch Size and KV Traffic

In the simplified traffic model, a decode step shares roughly 2.2 GB of weight reads across the batch. Each sequence also contributes its own KV-cache reads, so adding sequences increases bytes transferred per step.

TinyLlama stores 22,528 bytes of K/V per position. At 300 processed positions, that’s 6.7584 MB per sequence. Using 2.2 GB for weights and 1,555 GB/s peak bandwidth:

In this weight-and-cache traffic estimate, batch 128 transfers about 39% more data per step than batch 1 while producing 128 tokens. Aggregate output rises by roughly 92×, but each stream’s estimated ITL increases. Actual timings can be higher because the estimate omits compute and scheduling limits.

Weight traffic is shared, while KV traffic grows with the number of sequences. At about 326 sequences with this context length, total KV bytes per step would equal the weight bytes. Beyond that, adding sequences produces diminishing throughput gains in the bandwidth-only model.

That weight/KV crossover is different from the roofline’s compute/bandwidth crossover. Decode can remain bandwidth-bound even after KV reads exceed weight reads. A sufficiently large batch can also reach a compute limit before or after that point.

The weight/KV crossover is the weight bytes divided by the cache bytes per sequence. It doesn’t depend on bandwidth; bandwidth sets the transfer time. Longer contexts increase the cache bytes and bring the crossover to a smaller batch size.

Fewer KV heads and lower cache precision reduce bytes per sequence, allowing a larger batch before KV traffic dominates. Neither removes the compute or latency constraints of serving that batch.

When Prefill Interrupts Decode

The table assumes every active sequence is decoding. Continuous batching also has to schedule prefill for incoming requests.

If a scheduler runs a long prefill without interleaving decode work, active streams wait for it to finish. The prefill and decode kernels share GPU resources even when their main bottlenecks differ.

Suppose 32 sequences decode at 1.55 ms per step, and an incoming prefill occupies the GPU for 50 ms. Those are illustrative timings. If decode waits throughout that prefill, the next token interval exceeds 50 ms instead of 1.55 ms.

Limiting the prefill work scheduled at once can protect active streams, but may make the new request wait longer for its first token.

Chunked prefill divides the prompt into smaller pieces, such as 128 tokens, and schedules those pieces alongside or between decode iterations.

This bounds the amount of prefill work that can delay one decode iteration. Some schedulers also combine prefill tokens and decode tokens in the same batch, reusing weights across both. The useful chunk size depends on kernel efficiency, available memory, and the latency targets.

Smaller chunks can reduce ITL spikes while increasing the incoming request’s TTFT. Choosing that trade-off requires measuring both, rather than assuming the two workloads overlap without cost.

The Queue

Everything above assumed the request was already running on the GPU. In a real service, requests arrive over time and compete for slots in the decode batch.

Admission depends on the scheduler’s token budget, memory capacity, and concurrency limits. When capacity is available, the request can begin prefill. Otherwise it waits for work to finish or resources to be reclaimed.

Without queueing, TTFT is mainly prefill and request overhead. Under load, queue time can exceed prefill time.

How fast does queue time grow? Think about it from the arriving request's perspective. At low utilization, most slots are open. A new request almost always finds one immediately. As utilization climbs, open slots become rare. Near full capacity, almost every request waits, and each waiting request makes the line longer for the next one.

A simple single-server queue illustrates how sharply waiting can grow: mean waiting time relative to mean service time scales as $\frac{\rho}{1 - \rho}$ under that model’s assumptions, where $\rho$ is utilization. The ratio is 1 at 50%, 9 at 90%, and 19 at 95%. A batched inference server needs a richer model, but it can show the same rapid growth near saturation.

• At moderate load, a request often starts promptly.
• Near capacity, bursts and long requests cause queues to persist.
• Tail latency can worsen while average throughput still looks healthy.

The utilization percentage alone doesn’t predict p99 latency, the delay exceeded by the slowest 1% of requests. Arrival patterns, output lengths, and scheduling policy determine how much headroom a service needs.

Goodput

Consider two server configurations:

• Config A: 50,000 tok/s throughput, p99 TTFT = 4 seconds
• Config B: 20,000 tok/s throughput, p99 TTFT = 200 ms

If the service requires p99 TTFT below 500 ms, Config B meets that target and Config A doesn’t. These two p99 values alone don’t tell us how many requests meet the target in either configuration.

Goodput measures useful throughput under latency requirements. One common serving definition is the maximum request rate that meets specified TTFT and ITL targets for a required fraction of requests. A comparison needs that fraction and both targets, not just aggregate token throughput.

A server can increase raw token throughput while failing more latency targets. Measuring goodput makes that trade-off explicit; whether it rises or falls depends on the observed latency distribution.

Picking an Operating Point

For an interactive chatbot, the scheduler needs headroom for incoming requests and must limit interruptions to active streams. A concurrency cap and chunked prefill can help, though the settings should follow the application’s TTFT and ITL targets.

An overnight job summarizing a million documents may accept much longer waits in exchange for lower cost per token. It can run larger batches as long as memory and completion deadlines allow.

Most real systems serve a mix. The usual solution: separate the pools. Give interactive traffic a strict concurrency cap with headroom for fast admission. Route batch traffic to a throughput-maximizing pool where high utilization is acceptable.

Choose batch and admission limits using measured throughput and tail latency at representative context lengths. If the model can’t fit or meet the latency target on one GPU, distributing it introduces another set of trade-offs.

Detail 22: Multi-GPU Inference

A dense 70B model at FP16 needs about 140 GB for weights alone. It can’t fit on one 80 GB GPU.

INT4 reduces the raw weight storage to about 35 GB, before metadata. That may fit, but moving those bytes can still exceed a single-request latency budget.

Using a 35 GB weight-transfer estimate and an A100 SXM 80GB’s 2,039 GB/s peak bandwidth:

Under the 35 GB traffic estimate and the bandwidth specified above, a single device can’t reach a 15 ms transfer time. Reducing traffic further or using several memory interfaces would change that limit.

Distributing the model can address capacity, latency, or both. The result depends on how the work is split.

The Assembly Line

A transformer is a stack of identical layers. The most obvious split: give each GPU a contiguous chunk and process tokens in sequence, like stations on a factory line.

4 GPUs, 32-layer model:

• GPU 0: layers 1 to 8
• GPU 1: layers 9 to 16
• GPU 2: layers 17 to 24
• GPU 3: layers 25 to 32

A token enters GPU 0, passes through 8 layers, and the resulting activation tensor ships to GPU 1. GPU 1 runs layers 9 to 16, ships to GPU 2. And so on. Communication is light: one small tensor at each stage boundary.

For one microbatch, these stages execute in sequence.

While GPU 0 runs layers 1–8, the other three GPUs wait. The activation then moves to GPU 1, and GPU 0 waits. If the four stages take equal time, only one GPU is active at a time for this microbatch.

This idle time is the pipeline bubble.

The fix comes from the same intuition as any assembly line: don't wait for one item to finish before starting the next. If multiple sequences are in flight (and continuous batching from Detail 14 provides them naturally), GPU 0 can start processing sequence 2 while GPU 1 handles sequence 1. This is microbatching.

For a balanced pipeline processing a finite wave of $m$ microbatches across $p$ stages, the ideal fill-and-drain bubble fraction is $(p - 1) / (m + p - 1)$. This ignores communication and stage imbalance; a continuously running decode schedule can behave differently.

More microbatches reduce the idle fraction in this balanced-stage model.

Each individual microbatch still passes through all four stages in sequence. Overlapping different microbatches improves utilization and throughput, but doesn’t divide a single token’s computation across simultaneous stages.

This is pipeline parallelism (PP). It solves capacity: models that won't fit on one GPU can spread across many. It improves throughput: more total tokens per second. But it usually does not improve single-request latency, and can add a bit of stage-boundary communication on top.

The Latency Wall

PP can distribute weights that don’t fit on one device. But dividing 35 GB of weights across four sequential stages doesn’t divide the total transfer time for one token by four: the stage times still add. Microbatching primarily improves throughput here.

Reducing that token’s latency requires several GPUs to work on it simultaneously.

Splitting Weight Matrices

Instead of giving each GPU different layers, give every GPU a slice of every layer. All of them work on the same token, at the same time.

Use the row-vector convention $y = xW$. Split the columns of $W$ across two GPUs and replicate $x$. Each GPU computes a different portion of $y$.

Those local multiplications need no communication once each GPU has its input. Gathering the output portions into a complete vector would require communication, but the next operation can sometimes consume the distributed portions directly.

The original Megatron-LM example uses a two-projection MLP: $\text{output} = \text{GeLU}(x \cdot A) \cdot B$. Split $A$ by columns, apply GeLU independently to each GPU’s output slice, then split $B$ by rows so each GPU consumes its local hidden slice. TinyLlama’s SwiGLU has an additional gate projection, which can be partitioned alongside the up-projection before the same row-parallel down-projection.

The row-parallel multiplication by B produces partial sums for the final output. An all-reduce adds those partial results and distributes the sum to the GPUs that need it.

The classic Megatron-style forward pass uses two all-reduces per transformer layer: one after attention’s output projection and one after the FFN. For 22 layers, that’s 44 collectives per decode step. Other implementations may use different collective arrangements.

This is tensor parallelism (TP).

The Bandwidth Payoff

Using the earlier A100 SXM 40GB estimate, transferring 2.2 GB takes at least about 1.41 ms.

With ideal 2-way tensor parallelism, each GPU reads half the weight bytes at the same time. The weight-transfer component falls to about 0.71 ms before communication and other overhead.

TP can reduce a single token’s weight-transfer time because the transfers happen in parallel. PP transfers smaller sets of weights at each stage, but those stages execute sequentially for that token.

For 35 GB of weights on A100 SXM 80GB GPUs at 2,039 GB/s each, the ideal weight-transfer floors are about 17.2 ms on one device, 8.6 ms with 2-way TP, 4.3 ms with 4-way TP, and 2.1 ms with 8-way TP. These exclude quantization metadata and communication; they aren’t predictions of end-to-end ITL.

The Communication Tax

Frequent collectives can offset the transfer savings, especially for a small model or batch. Their cost depends on message size, collective algorithms, and interconnect latency as well as bandwidth.

Within a node, H100 NVLink provides up to about 900 GB/s of bidirectional bandwidth per GPU. Inference collective messages can be small, so startup and synchronization latency matter even when the peak link bandwidth is high. That peak alone doesn’t establish how much of a decode step communication will take.

A 400 Gb/s network link carries at most 50 GB/s before protocol overhead; a 100 Gb/s link carries at most 12.5 GB/s. Multi-rail systems combine links, but effective collective performance also depends on topology and contention. Cross-node TP can lose much of its benefit if communication costs exceed the weight-transfer savings.

Pipeline parallelism, by contrast, communicates only at stage boundaries: one point-to-point activation transfer between adjacent GPUs per step. Far less frequent than TP's per-layer all-reduces. PP tolerates slow interconnects.

A common deployment uses tensor parallelism within a node and pipeline parallelism across nodes. Fast local links support frequent TP collectives, while PP needs fewer transfers between nodes.

A deployment can combine 8-way TP within a node, PP across nodes, and several replicas of that whole group. Replication serves independent requests; model parallelism makes each replica fit or meet its latency target.

The Third Axis: Splitting the Sequence

Long contexts create another capacity problem: the KV cache grows with the number of token positions, even after weight storage has been distributed.

TP can split the KV cache across head partitions, but the reduction is limited by the number of KV heads and the implementation. If the TP degree exceeds the KV-head count, some systems replicate KV heads. Each partition’s cache also continues to grow with sequence length.

Consider a 70B model with GQA (8 KV heads, $d_{\text{head}}$ = 128, 80 layers) serving a 1M-token context in FP16. Each layer stores keys and values: $2 \times 8 \times 128 \times 1{,}000{,}000 \times 2$ bytes. That is roughly 4.1 GB per layer. Across 80 layers: ~327 GB of KV cache for a single sequence. Quantized to INT8: still ~164 GB. An 80 GB GPU can’t hold this cache. Even with 8-way TP splitting across heads, each GPU's share of the KV cache is ~41 GB, for one sequence alone.

Context parallelism (CP) splits the sequence dimension itself across GPUs. With 8-way CP on a 1M-token context, each GPU holds a 125K-token slice and stores only its portion of the KV cache.

A Query must attend to all allowed Key positions, including positions stored on other GPUs. In Ring Attention, each GPU computes against its local K/V chunk and then passes that chunk around a logical ring. With P GPUs, P − 1 transfers expose all chunks; causal masking excludes future positions. Online softmax combines the partial results without storing the full attention matrix.

CP transfers K/V chunks between devices during prefill. Implementations overlap those transfers with attention on another chunk. The $O(N^2)$ arithmetic can help hide communication, though this depends on chunk size, hardware, and load balance.

Causal masking gives earlier query positions fewer Keys to attend to, which can unbalance contiguous partitions. Interleaving positions or using paired early/late chunks can distribute the work more evenly.

During prefill, quadratic attention work can help hide communication. Decode has one new Query per sequence, so circulating full KV chunks would be expensive relative to its arithmetic. Distributed decode can instead send Queries to cache partitions and combine partial attention results. The latency benefit depends on that implementation and its interconnect.

Context parallelism is orthogonal to TP and PP. A large long-context deployment might use all three: TP within a node (split weights for low ITL), CP across nodes (split the KV cache so million-token contexts fit), and PP if the model exceeds even one node's aggregate weight capacity. Each axis addresses a different dimension: TP handles weights and heads, CP handles sequence length, PP handles layers.

Matching Optimizations to Bottlenecks

The techniques affect different costs. Reported speedups depend on the baseline, model, and workload; they aren’t factors that can simply be multiplied together.

Some techniques combine well because they change different costs. Quantization reduces weight bytes, while GQA reduces cache bytes. Others interact: larger batches can use the compute capacity that speculative verification would otherwise exploit. After an optimization, measure which resource limits the new workload.

Reading an Inference Benchmark

A claim of "2× faster" needs a metric. It could describe TTFT, ITL, aggregate throughput, or tail latency under load. Those improvements affect different users and can require different trade-offs.

For a claim that a model runs on a single 4090, parameters × bytes per weight gives a first check against its 24 GB capacity, before cache and working memory. Estimated bytes transferred ÷ memory bandwidth gives a transfer-time bound under the assumed traffic. For a throughput comparison, check batch size, context length, precision, and the latency targets used in the benchmark.

I find these estimates useful because they give a concrete starting point: what data must be stored, what data must move, and how much work can share each transfer.

Other Compression and Serving Techniques

References & Further Reading

Primary sources for the concepts discussed above, plus a few good overviews.

Blog posts & videos

• Aleksa Gordić, Inside vLLM: Anatomy of a High-Throughput LLM Inference System. https://www.aleksagordic.com/blog/vllm
• Vizuara, How the VLLM Inference Engine Works. https://www.youtube.com/watch?v=QyHHbeXqgrQ
• Lilian Weng, Large Transformer Model Inference Optimization. https://lilianweng.github.io/posts/2023-01-10-inference-optimization/
• NVIDIA Developer Blog, Mastering LLM Techniques: Inference Optimization. https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/
• Jay Alammar, The Illustrated Transformer. https://jalammar.github.io/illustrated-transformer/
• Andrej Karpathy, Let's build GPT: from scratch, in code, spelled out. https://www.youtube.com/watch?v=kCc8FmEb1nY

Papers (mostly arXiv)

Core architecture

• Vaswani et al. (2017). Attention Is All You Need. https://arxiv.org/abs/1706.03762
• Bahdanau et al. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. https://arxiv.org/abs/1409.0473
• Su et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. https://arxiv.org/abs/2104.09864
• He et al. (2016). Deep Residual Learning for Image Recognition. https://arxiv.org/abs/1512.03385
• Ba et al. (2016). Layer Normalization. https://arxiv.org/abs/1607.06450
• Shazeer (2020). GLU Variants Improve Transformer (SwiGLU). https://arxiv.org/abs/2002.05202

Tokenization

• Sennrich et al. (2016). Neural Machine Translation of Rare Words with Subword Units (BPE). https://arxiv.org/abs/1508.07909
• Kudo & Richardson (2018). SentencePiece: A Simple and Language Independent Subword Tokenizer and Detokenizer for Neural Text Processing. https://arxiv.org/abs/1808.06226

Models referenced

• Zhang et al. (2024). TinyLlama: An Open-Source Small Language Model. https://arxiv.org/abs/2401.02385
• Touvron et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. https://arxiv.org/abs/2307.09288
• Llama Team (2024). The Llama 3 Herd of Models. https://arxiv.org/abs/2407.21783
• Jiang et al. (2023). Mistral 7B. https://arxiv.org/abs/2310.06825
• Gemma Team (2024). Gemma: Open Models Based on Gemini Research and Technology. https://arxiv.org/abs/2403.08295
• Gemma Team (2025). Gemma 3 Technical Report. https://arxiv.org/abs/2503.19786

Inference + serving

• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM). https://arxiv.org/abs/2309.06180
• Yu et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI). https://www.usenix.org/conference/osdi22/presentation/yu

Attention kernels

• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. https://arxiv.org/abs/2205.14135
• Dao (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. https://arxiv.org/abs/2307.08691
• Milakov & Gimelshein (2018). Online Normalizer Calculation for Softmax. https://arxiv.org/abs/1805.02867

KV-cache shape tricks

• Shazeer (2019). Fast Transformer Decoding: One Write-Head Is All You Need (MQA). https://arxiv.org/abs/1911.02150
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. https://arxiv.org/abs/2305.13245
• DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (MLA). https://arxiv.org/abs/2405.04434

Speculative decoding

• Leviathan et al. (2022). Fast Inference from Transformers via Speculative Decoding. https://arxiv.org/abs/2211.17192
• Chen et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. https://arxiv.org/abs/2302.01318
• Li et al. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. https://arxiv.org/abs/2401.15077
• Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. https://arxiv.org/abs/2401.10774
• Fu et al. (2024). Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. https://arxiv.org/abs/2402.02057

Quantization

• Frantar et al. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. https://arxiv.org/abs/2210.17323
• Xiao et al. (2022). SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models. https://arxiv.org/abs/2211.10438
• Lin et al. (2023). AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration. https://arxiv.org/abs/2306.00978
• Dettmers et al. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. https://arxiv.org/abs/2208.07339
• Dettmers et al. (2023). bitsandbytes: Accessible large language models via k-bit quantization for PyTorch. https://github.com/bitsandbytes-foundation/bitsandbytes

Parallelism

• Shoeybi et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. https://arxiv.org/abs/1909.08053
• Huang et al. (2018). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. https://arxiv.org/abs/1811.06965

Performance modeling

• Williams et al. (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM, 52(4), 65-76. https://doi.org/10.1145/1498765.1498785

Alternate attention (context for long sequences)

• Kitaev et al. (2020). Reformer: The Efficient Transformer. https://arxiv.org/abs/2001.04451
• Wang et al. (2020). Linformer: Self-Attention with Linear Complexity. https://arxiv.org/abs/2006.04768
• Beltagy et al. (2020). Longformer: The Long-Document Transformer. https://arxiv.org/abs/2004.05150

Alternative hardware & memory economics

• Groq. Inside the LPU: Deconstructing Groq's Speed (blog). https://groq.com/blog/inside-the-lpu-deconstructing-groq-speed
• Groq. Groq LPU Inference Engine Crushes First Public LLM Benchmark (blog, February 2024). https://groq.com/blog/groq-lpu-inference-engine-crushes-first-public-llm-benchmark
• Groq. Groq Sets New LLM Performance Record of 300 Tokens per Second per User on Llama-2 70B (press release, November 2023). https://www.prnewswire.com/news-releases/groq-sets-new-large-language-model-performance-record-of-300-tokens-per-second-per-user-on-meta-ai-foundational-llm-llama-2-70b-301980280.html
• Patel & Nishball. Groq Inference Tokenomics: Speed, But At What Cost? SemiAnalysis (February 2024). https://newsletter.semianalysis.com/p/groq-inference-tokenomics-speed-but
• Patel et al. The Memory Wall: Past, Present, and Future of DRAM. SemiAnalysis (September 2024). https://semianalysis.com/2024/09/03/the-memory-wall/
• Cerebras. Cerebras Systems Unveils World's Fastest AI Chip (press release, March 2024). https://www.cerebras.ai/press-release/cerebras-announces-third-generation-wafer-scale-engine
• Cerebras. Cerebras Inference now 3x faster: Llama3.1-70B breaks 2,100 tokens/s (blog, October 2024). https://www.cerebras.ai/blog/cerebras-inference-3x-faster
• Prabhakar et al. (2024). SambaNova SN40L: Scaling the AI Memory Wall with Dataflow and Composition of Experts. https://arxiv.org/abs/2405.07518
• d-Matrix. Corsair Product Overview. https://www.d-matrix.ai/product/

Compression beyond quantization (optional)

• Sanh et al. (2019). DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. https://arxiv.org/abs/1910.01108
• Frantar & Alistarh (2023). SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot. https://arxiv.org/abs/2301.00774
• Shazeer et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. https://arxiv.org/abs/1701.06538
• Fedus et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. https://arxiv.org/abs/2101.03961

Implementations & docs

• vLLM: https://github.com/vllm-project/vllm
• FlashAttention: https://github.com/Dao-AILab/flash-attention
• Hugging Face TGI: https://github.com/huggingface/text-generation-inference
• NVIDIA TensorRT-LLM: https://github.com/NVIDIA/TensorRT-LLM
• LMCache (prefix caching): https://github.com/LMCache/LMCache
• XGrammar (structured generation): https://arxiv.org/abs/2411.15100
• TinyLlama model card: https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0
• Tiktokenizer (interactive tokenizer): https://tiktokenizer.vercel.app/