How huge AI models fit, remember, and respond
A model can contain a trillion learned numbers without using all of them for every token. It can also reuse earlier calculations and keep parts of a long conversation in a compact memory.
To see how, follow three jobs: remember the conversation, choose which parts of the model to use, and divide the work across GPUs. Each saves a different resource.
Start with a small request: The key is blue. What color is the key? As the conversation grows, the
model still needs to find the relevant information and produce its answer one token at a time.
First, separate knowledge from conversation memory
Weights are the numbers learned during training (the model's parameters). They stay fixed during ordinary generation and are shared across requests.
Conversation state holds information computed from this particular input. It changes as the model reads and generates tokens. A token is a word, word piece, or punctuation mark.
A large model needs room for both, plus temporary working space. Using fewer weights per token helps computation, but does not automatically make the stored model smaller.
Attention looks back at useful information
For our color question, attention helps bring information from blue into the representation used
to answer. It scores how strongly positions in the text should contribute, then mixes their
information.
Each token is represented by a list of numbers (a vector). Attention operates on those vectors, and repeated layers transform them into the representation used to predict the next token.
Each attention layer runs several calculations in parallel (attention heads). A head makes three projections of each token's current representation:
Query (Q): what information this position is looking for.
Key (K): what another position can be matched against.
Value (V): what information that position contributes when it receives attention.
The query scores the keys. A normalization called softmax turns those scores into weights that sum to one. The weighted values are mixed together. These are learned numerical operations; “looking for” is an intuition for their role. Original Transformer paper.
Each layer receives the previous layer's output, so it computes its own keys and values. In a conventional KV cache, one processed token has one K vector and one V vector per KV head, per attention layer. A model with 32 attention layers and 8 KV heads stores 256 such pairs for that one token. Different requests have separate logical caches.
The same token has different K/V vectors at different layers. The animation below zooms into one head in one layer. How representations pass between layers.
This head passes on a mixture of value vectors. The rest of this layer and later layers still run before the model predicts a token.
Inspect
Reading the prompt (prefill) processes the known input and produces the first next-token prediction. Earlier positions cannot read later ones; a causal mask enforces that rule.
Continuing the answer (decode) processes each selected token to predict the next one. The model reuses previously computed keys and values from a KV cache, then appends the new K/V vectors at each full-attention layer. Selecting a token does not cache it; processing it on the following pass does.
That saves repeated computation. But the cache keeps growing, and full attention still reads the available history for each new token. Moving that data from GPU memory to its compute units can become a bottleneck. How KV caching works.
What is the actual cache shape?
Each full-attention layer has two tensors: one for keys and one for values. A common logical shape for each is:
[sequences in the batch, KV heads, processed token positions, numbers per vector]
For one request, 8 KV heads, 1,000 processed tokens, and 128 numbers per vector, each tensor has
shape [1, 8, 1000, 128]. A 32-layer model has 32 such K/V tensor pairs. At 2 bytes per number,
that is 125 MiB of cached data. Processing one more token extends the position dimension in every
layer from 1,000 to 1,001.
Grouped-query attention (GQA) lets several query heads share a KV head. Count KV heads when sizing the cache. Queries are computed for the current positions and are not retained in this cache. Compressed or sliding-window caches use different layouts or retention rules. Cache shape and ownership.
Does attention memory grow linearly or quadratically?
For one full-attention head with N tokens and key/value width d, the cached keys and values
contain 2 × N × d numbers. Doubling the conversation doubles this cache.
Processing the whole prompt involves roughly N² query-key comparisons. Doubling the prompt roughly
quadruples that attention work. A causal mask removes future-facing comparisons, but the growth is
still quadratic.
These describe different things: stored KV data versus attention computation. Efficient kernels such as FlashAttention avoid storing the entire attention score matrix, while still computing full attention. Multiple heads, grouped-query attention, and compressed KV representations change the memory calculation.
Keep a compact memory instead of every entry
Linear attention combines incoming information into a fixed-size table of numbers (a recurrent state). Each token updates that table. Its values change; its dimensions stay the same.
Full attention keeps individual keys and values available for later queries.
Linear attention combines information into a fixed-size state that later queries read.
After 16 tokens: 128 stored numbers in K + V, versus 16 in the recurrent state. The fixed state combines information; it does not retain a separate entry for every token.
Inspect
The fixed state saves memory as a sequence grows. But different associations can overlap inside it and interfere with each other. There is no promise that every detail remains separately recoverable.
Better fit: keep the current plan. Imagine a 100,000-token planning conversation where you need the latest destination, date, and budget. A linear-attention model can keep updating those few associations without storing another K/V entry for every token. This can make long-running generation cheaper in memory, if the model maintains the facts accurately enough for the task.
Worse fit: retrieve an exact old entry. Imagine a 100,000-token invoice log and the question "What was the total on invoice 8317?" A fixed state must distinguish many separate invoice-to-amount associations. Interference can lose the requested detail. Full attention retains token-level K/V entries that a query can revisit, though retaining them does not guarantee a correct answer. Retrieval limits and delta-rule improvements.
Sequence length changes the storage tradeoff too. For one head in one layer, with 128 numbers per key/value vector and 2 bytes per number:
| Processed tokens | Full-attention K + V | Linear-attention state matrix |
|---|---|---|
| 64 | 32 KiB | 32 KiB |
| 2,048 | 1 MiB | 32 KiB |
| 131,072 | 64 MiB | 32 KiB |
The K/V cache stores 2 × tokens × 128 numbers. This example's recurrent matrix stores 128 × 128
numbers. These are storage calculations, not a benchmark or a claim of equal quality; extra state
and runtime buffers are omitted. At short lengths, a fixed matrix may save little or no memory.
Actual speed also depends on the kernels, hardware, and batch size.
“Linear” describes how attention work scales with sequence length for fixed model dimensions. It does not mean the entire neural network is a linear function. Linear attention paper.
The matrix trick, and why softmax cannot simply be moved
Without softmax, matrix multiplication can be regrouped:
(Q × Kᵀ) × V = Q × (Kᵀ × V)
The left form builds token-to-token interactions. The right form first combines keys and values into a matrix whose size depends on their vector widths. During causal generation, a running state includes only the prefix processed so far.
But standard attention has a softmax between those multiplications. You cannot move it through and obtain the same result. Linear-attention architectures change the attention rule. Some use feature maps and a separate normalization state; later variants use learned memory updates. The model is trained with this different attention rule. Applying it to an arbitrary existing model would change that model's computation.
Update the memory when information changes
Suppose the conversation later says: Actually, the key is red.
An additive update adds the new key/value association to memory. The old blue association can remain, so a later read can mix blue and red.
A delta-rule update reads the value currently associated with the key, compares it with red, and writes the difference back. This helps revise an existing association, such as a changed color, date, or destination.
The key now retrieves red. The door still retrieves red. The memory has the same size.
Inspect
A learned write strength controls how much of that correction to apply. A full update replaces the retrieved value in our simplified example; a partial update blends old and new. DeltaNet's update and parallel algorithm.
A write gate controls how strongly to update the association addressed by the current key.
A forget gate controls how much existing state to retain before that update. This lets the model clear information beyond the currently addressed association.
Gated DeltaNet combines these write and forget gates with the delta-rule update.
The delta rule in one equation
With column vectors, let S map a key to a value. Read old = S × k, then update:
S_new = S + β × (v − old) × kᵀ
Here v is the desired value and β is the write strength. For a unit-length key and β = 1,
querying that same key retrieves v after the update. Orthogonal keys are unaffected; overlapping
keys can change too.
The diagram uses two orthogonal keys and two value coordinates, represented by blue and red. Real
systems learn their vectors and gates. Updating S during inference does not retrain the model's
weights.
Make the updates suit a GPU
GPUs process large blocks of matrix arithmetic efficiently. Updating memory in a literal loop, one token at a time, can leave much of that hardware underused.
During prefill, the prompt tokens are already known. Chunkwise algorithms group them and rewrite the dependent updates into matrix operations. Work inside a chunk becomes parallel; a state carries information between chunks. Future tokens remain masked.
Parallel work inside each chunk; memory carried between chunks. The model never needs to peek ahead.
Inspect
Grouping helps during prompt processing because the inputs are available together. Ordinary generation still needs the next token to be selected before continuing. Grouping computations does not make unknown future tokens available. Parallelizing DeltaNet.
Models can also combine the two memory approaches (hybrid attention). Some layers use a compact state; others keep token-level history. This preserves a way to revisit individual positions while reducing total memory use. The layers keeping that history still grow with the conversation.
Where Kimi Delta Attention fits
Kimi Delta Attention (KDA) refines the forget gate: different feature channels can retain different amounts of state. Channels are learned numerical features, not named topics such as “colors” or “people.”
Kimi Linear uses three KDA layers for each full-attention layer. The latter uses a compressed key-value representation called multi-head latent attention (MLA). The model therefore keeps both fixed-size recurrent states and a cache that grows with the conversation. Kimi Linear architecture.
How can dependent updates become parallel math?
Take a small example: new state = ½ × previous state + input. Starting from zero, the inputs
2, 4, 1 produce states 2, 5, 3.5.
Expanding the steps gives:
- First state:
x1 - Second state:
½ × x1 + x2 - Third state:
¼ × x1 + ½ × x2 + x3
All three expressions can be evaluated from the known inputs. Written as rows of a matrix, those coefficients let hardware perform the work together. Each row still excludes future inputs. Delta-rule kernels use a more involved version of this idea for matrix-valued states.
A huge model can use only a few experts per token
A model also transforms each token through a neural network inside a layer (a feed-forward network). Offering several of these networks lets the model choose which ones to use.
A dense layer sends every token through the same feed-forward network.
A mixture-of-experts layer (MoE) provides several feed-forward networks. A learned router chooses a few for each token, then combines their outputs.
Combine the two outputs using routing weights. A different token may select different experts.
Inspect
This separates total parameters, which determine how much weight data exists, from active parameters, which participate in processing one token. Sparse routing allows a model to have many more total parameters than it uses each time.
Experts are numerical networks, not separate chatbots. Their names do not guarantee tidy human specialties. Mixtral's architecture.
In fast GPU serving, the system generally keeps all experts available in GPU memory. Different tokens can choose different experts, and many requests run together (a batch). Constantly loading missing experts from slower storage would add delays.
Spread the remaining work across GPUs
If the model is too large for one GPU, several GPUs can share the job. What gets split changes how they need to communicate.
Tensor parallelism splits a layer's matrices across GPUs. They compute pieces of the same operation and communicate to assemble the result.
Expert parallelism places experts on different GPUs. Token representations travel to the selected experts, and their results travel back, as in the diagram above.
Pipeline parallelism places groups of layers on different GPUs. A request moves through those groups in order; multiple requests can occupy different stages.
A serving system can combine these forms of parallelism. Each changes how much data GPUs must exchange, what they store, and when they wait. Sending everything to one popular expert can leave other GPUs waiting. Model parallelism.
Fitting the weights is only the first requirement. The serving system also needs enough conversation memory for simultaneous requests and enough bandwidth to keep computation supplied.
Why serving adds checkpoints, separate workers, and speculation
State checkpoints save earlier recurrent states. A KV cache retains entries along a prefix; a recurrent state has been overwritten by later updates. Reusing an earlier prefix therefore needs a saved state at that boundary, or recomputation from an earlier checkpoint. More checkpoints consume more memory.
Separate prefill and decode workers let a service size those pools for its traffic. Long prompts and long answers stress different parts of the system. Transferring request state between workers adds a cost that the benefit must outweigh. Disaggregated serving.
Speculative decoding lets a cheaper draft model propose several tokens. The main model verifies them together, accepts a valid prefix, and corrects a rejection. Exact speculative sampling can preserve the main model's output distribution. The payoff depends on acceptance rate and the cost of drafting and verification. Speculative decoding paper.
When you see a model described as “huge,” ask three questions: How much runs per token? What memory grows with the conversation? What must move between GPUs? Parameter count alone cannot answer them. Together, these explain much more about how a model can fit, remember, and respond.
This explainer was prompted by Baseten's conversation, How trillion-parameter AI models work and how to serve them. The examples and diagrams here are simplified teaching examples; the linked papers supply the technical details.