In 2021 a small idea slipped into the transformer with almost no fanfare. A paper called RoFormer proposed encoding a token’s position not by adding something to it, but by rotating it. The idea, Rotary Positional Encoding (RoPE), spread quickly. Within two years it had become the default in nearly every serious open model: GPT-NeoX, PaLM, LLaMA and its descendants, Mistral, Qwen, DeepSeek, Gemma. If you have used a modern LLM, you have used RoPE.

It is also the quiet reason your model can sometimes read a whole book. Almost any conversation about long context, whether that means 128K windows, million-token prompts, or the needle-in-a-haystack test, eventually runs into RoPE. RoPE is the thing you have to stretch to make long context work, and the thing that breaks when you stretch it wrong.

This is where to start. Before we can talk about how models reach for longer and longer context, the subject of the posts that follow this one, we need to actually understand the small, elegant trick at the heart of it. By the end of this post, you should be able to look at this formula.

and find it obvious.

We will keep one example in hand the whole way: the sentence “the dog chased the cat,” and in particular how the word chased should relate to the word dog.

The problem: attention sees an unordered bag

The attention mechanism has an uncomfortable property. By itself, it has no idea what order the words came in.

Attention works by comparing every token with every other token and taking weighted sums, and a sum does not care about order. Shuffle the inputs and you get the same answer back. To raw attention, “the dog chased the cat” and “the cat chased the dog” are the same bag of vectors. One sentence has the dog doing the chasing and the other has it being chased, and attention cannot tell them apart.

That is a problem, because meaning lives in order. We have to inject position somehow, giving the model a way to know that in our sentence, chased sits one step after dog, and that this adjacency is part of what the sentence means.

The obvious fix, and its hidden flaw

The original 2017 Transformer solved this the obvious way. It built a position vector out of sines and cosines and added it onto each word’s embedding, like stamping a timestamp onto a letter before you mail it. The token at position 0 gets stamp #0, position 1 gets stamp #1, and so on. The network is then left to untangle “content plus stamp” back into “content” and “where.”

It helps to see those sines and cosines directly. Each row below is one dimension oscillating at its own frequency, fast at the top and slow at the bottom. Move the slider, and the column the cursor marks out is the stamp that gets added to the word sitting at that position:

Two things stand out. The stamp depends only on the position, never on the word underneath it. And neighbouring positions get very similar stamps, since the waves move smoothly. Hold that picture, because it is what rotation is about to improve on.

It works, but it carries two flaws that, once you see them, motivate everything RoPE does.

Flaw one: it smears content and position together. Adding a vector moves the point. A word’s embedding encodes its meaning; the stamp we add shifts that vector somewhere new, so the same word at two different positions becomes two genuinely different vectors, with different length and direction. Meaning and position now sit tangled in the same numbers, and the model has to spend capacity pulling them back apart.

Flaw two: it is absolute, but attention wants relative. The stamp records where a token sits counting from the start of the sequence. That is almost never what matters. What matters for chased is that dog is one token back, not that dog happens to be the second word in this particular sentence. Prepend the word “Yesterday,” to our sentence and every absolute position shifts by one, yet the relationship between chased and dog has not changed at all. Absolute encodings force the model to learn how to turn “position 2 versus position 3” into “one step apart,” and to relearn that for every pair of positions. It is work we should not have to do.

Keep both flaws in mind. RoPE fixes them at once, with a single geometric move.

The insight: position is not a number you add, it’s a rotation you apply

There is another way to think about it. Instead of adding a position vector, what if we rotated the token’s vector by an angle that grows with its position?

Take the token’s vector and chop it into pairs of coordinates. Each pair is just a point on a plane, an arrow from the origin. To encode position $m$, spin that arrow by an angle $m\theta$: position 0 gets no turn, position 1 turns by $\theta$, and position $m$ turns by $m\theta$. In two dimensions this is exactly the rotation matrix from high-school geometry:

Drag the slider below and watch what happens to a single pair as its position climbs:

Notice the one thing that never changes: the arrow’s length. A rotation changes direction, never magnitude, and that is what makes it a rotation. That single property already takes care of the first flaw. A word’s meaning lives in the length and shape of its vector, and spinning it leaves all of that alone. Position ends up written purely into the angle, kept separate from content.

Length preservation is only the warm-up, though. The real payoff is what rotation does to the dot product, and to see it we need one fact about how attention scores tokens.

Why a dot product only feels the angle between

Attention decides how much chased should attend to dog by taking the dot product of chased’s query vector with dog’s key vector. A big dot product means strong attention.

The dot product has a clean geometric meaning. For any two vectors $q$ and $k$,

where $\phi$ is the angle between them. The lengths $\|q\|$ and $\|k\|$ are fixed properties of the two words, a kind of loudness. Everything about how the two words relate is carried by that single $\cos\phi$ term. Vectors pointing the same way score high ($\cos 0 = 1$), perpendicular vectors score zero, and opposed vectors score negative.

So an attention score is really a question about the angle between two arrows. That is the sentence to hold onto. If position is an angle, and attention only responds to angles, then position and attention are speaking the same language.

The magic: relative position, for free

Now put the two halves together. Rotate chased’s query by its position $m\theta$ and dog’s key by its position $n\theta$. What is the angle between them afterwards?

Rotating one vector by $m\theta$ and comparing it against another rotated by $n\theta$ composes into a single rotation by the difference. Writing it out with the matrices, and using the identity $R(a)^{\top}R(b) = R(b-a)$, the rotated dot product collapses:

Look at the right-hand side. The two absolute positions $m$ and $n$ are gone. Only their difference $n-m$ remains. The attention score between chased and dog depends only on how far apart they are, not on where the pair happens to sit in the sentence.

That takes care of the second flaw. We never asked the model to learn how to convert absolute positions into relative ones; the geometry does it on its own. Relative position falls out of the structure for nothing, with zero parameters spent.

Try it. Move the query and key positions on their own, then press “shift both +1,” which is the same as prepending “Yesterday,” to the sentence. Both arrows spin, but the offset between them, and the score, stay put:

This is the payoff that made RoPE win. Shift the whole sentence and chased’s relationship to dog survives intact, because that relationship was stored as the angle between the two vectors, and shifting both simply rotates them together. One thing the demo quietly does, though, is hold the two words’ content fixed so the positional effect stands on its own. In a real attention head the content match is the dominant, learned signal, and rotation only modulates it. The next section puts content back in.

How rotation fits inside attention

It is easy to leave that demo thinking the rotation is the whole story. It is not. To see why, it helps to look at where RoPE actually sits inside an attention head, and at how much it leaves alone.

Attention scores a query against a key with a dot product, and before any position is involved that dot product is pure content matching. Think of it as a tiny search engine running in every layer. Each token sends out a query, a search for what it wants (“I am a verb, I am looking for my subject”). Every other token offers a key, an advertisement for what it is (“I am a noun, I could be a subject”). The dot product scores how well the advertisement answers the search. All of this is learned, by the projection matrices $W^Q$ and $W^K$, and it is the main event. It is what lets chased know it wants nouns and not commas.

RoPE does not touch any of that. It rotates the query and key after they are built, and because a rotation changes direction but not length, the learned content survives untouched. For a single pair of dimensions the score works out to

where $\alpha$ is the angle between the two words’ content directions. Read the cosine as taking in two things at once. The content angle $\alpha$ is small when the words genuinely match and large when they do not. The positional term $(m-n)\theta$ is a fixed turn that depends only on how far apart they are. The score is highest when the content matches and the relative distance is one the head cares about. A perfect content match at an unwanted distance gets pulled down, and a favoured distance with mismatched content still scores low. Both have to agree.

Our running example needs both halves. Content alone tells chased to look at nouns, but dog and cat are equally nouns, so content cannot say which is the subject. Position alone would prefer whatever sits one step back, but it cannot tell a noun from a comma. Put them together and chased attends to dog because dog is both a content match and at the relative offset the head has learned to read as a subject. RoPE did not pick dog. Content narrowed the field to nouns, and rotation broke the tie by distance.

And everything downstream is left alone. The value vectors are never rotated, the softmax is the same softmax, and the weighted sum that produces the head’s output is unchanged. RoPE is a small, surgical edit to one quantity, the query-key score, not a new attention mechanism.

One frequency isn’t enough: a clock with many hands

A single rotation speed has a catch, because a circle wraps around. If every pair spun at the same rate $\theta$, two positions a full turn apart would land in the same place, and you could not tell them apart. One spinning hand cannot tell you the time on its own.

A clock fixes this with several hands at different speeds. The second hand resolves fine detail, the hour hand tracks the long sweep, and together they pin down a single moment. RoPE does the same thing, giving each coordinate pair $i$ its own rotation speed:

The first pairs spin fast, so they resolve fine, local distances like “one token apart.” The last pairs spin slowly, tracking coarse, long-range position across thousands of tokens. Stack them all and every position gets a unique multi-frequency fingerprint, with no wraparound ambiguity across the range that matters.

Watch the bank of dials below. The leftmost races while the rightmost barely moves.

The full RoPE rotation is just this whole set of 2-D rotations stacked into one large block-diagonal matrix, with each pair of dimensions spun at its own frequency. Conceptually it is the picture above, repeated $d/2$ times.

Two things are worth noticing now, because they come back when we talk about long context. The base $b = 10000$ is a knob you can turn. And the fast pairs are the ones that wrap around soonest.

A free locality prior: nearby leans in, distant fades

The many frequencies do one more thing, almost by accident. When you add up the cosines across all the pairs, they reinforce at zero distance and start to interfere as the distance grows. The result is that the raw attention score between two identical vectors is high when they sit close together and decays, with a gentle ripple, as they move apart.

Slide the dimension count and watch the curve: more frequencies, smoother decay, sharper peak at $\Delta = 0$.

So RoPE quietly hands the model a sensible default, namely “pay more attention to what is near,” without spending a single parameter on it. The model can override that prior when it needs to reach far, but it starts from a reasonable place.

Adding moves the point; rotating keeps it honest

We can now see, side by side, why rotation beats addition. Sinusoidal encoding adds a position vector, so the point drifts off its circle: its length changes, and content gets tangled up with position. RoPE rotates instead, so the point glides along its circle, its length perfectly preserved and position kept separate from meaning.

Sweep the position. On the left the length wobbles, which is the word’s meaning being disturbed as it moves through the sentence. On the right it stays rock-steady. Same goal of encoding position, very different treatment of the content.

Why it won

Pull back, and the list of advantages is long for something so simple:

• Relative position, for free. The attention dot product depends only on $m-n$. The model never has to learn to subtract positions, because it is guaranteed by construction.
• Meaning stays intact. Rotation preserves length, so a token’s content is not corrupted by where it sits, unlike additive encodings, which blur the two together.
• Applied where it matters. RoPE rotates the queries and keys inside every attention layer, right where the comparison happens, instead of being bolted once onto the input embedding and left to fade.
• Zero extra parameters. It is a fixed geometric operation. There is nothing to train, almost nothing to compute, and it composes cleanly with efficient-attention kernels like FlashAttention.
• A built-in locality prior. Scores naturally taper with distance, a free and sensible default.
• It stretches. Because it encodes relative distance through smooth, tunable frequencies, RoPE can be rescaled to longer sequences far more gracefully than anything before it, which is the whole reason it underpins modern long-context models.

That last point is where this post ends and the next one begins.

The bridge: from rotation to long context

This last point is also where the trouble starts. The same frequency structure that makes RoPE so elegant puts a hard ceiling on context length.

A model trained with a context window of, say, 4K tokens has only ever seen rotation angles up to $4096 \cdot \theta_i$ for each pair. The fast pairs will have swept through their whole range many times inside those 4K tokens, while the slow pairs have turned only a fraction of a circle. The network has learned to read positions inside that envelope of angles, and nowhere else.

Now feed it a 100K-token prompt at inference. The fast pairs are suddenly spinning to phase angles the model has never seen in training. As far as the network is concerned, the positions have gone out of distribution. Attention destabilizes, and quality falls off a cliff long before the prompt ends.

This is why context extension has become its own discipline. Every major technique is, underneath, a way of manipulating the angles and frequencies we just built up. Position Interpolation squeezes the positions back into the trained range. NTK-aware scaling turns that base-frequency knob $b$ up so the fast pairs slow down. YaRN interpolates each frequency band differently. None of them make much sense until you can see position as a rotation, which you now can.

That is the subject of the next post in this series.

RoPE also shows up in surprising places elsewhere on this blog. The evolution of attention post shows how DeepSeek’s Multi-head Latent Attention has to do some delicate surgery, a decoupled form of RoPE, to stay compatible with rotary embeddings while compressing the KV cache. The state-space models post shows Mamba-3 reusing the same rotary machinery with data-dependent angles. The idea travels a long way.

Lessons for builders

A few takeaways that generalize beyond RoPE:

1. Positional information wants to be relative. When you catch yourself making a model re-derive the same relationship at every absolute offset, look for a representation where that relationship is built in rather than learned.
2. Magnitude-preserving operations keep signals clean. Rotation works partly because it refuses to touch the content’s length. When you have to inject one kind of information into a vector that already carries another, prefer transforms that leave the existing signal undisturbed.
3. The base frequency is a real knob, not a constant. That $10000$ is not sacred; long-context models routinely raise it to $500{,}000$ or $1{,}000{,}000$ to slow the fast pairs down. When a hyperparameter is set “because the paper said so,” it is worth knowing what it actually controls.
4. The elegant default and its failure mode are two sides of one coin. The frequencies that give RoPE its clean relative encoding are the same ones that go out of distribution past the training length. The mechanism and its breaking point cannot be pulled apart, so understanding one means understanding the other.

Conclusion

RoPE comes down to one choice: to encode position, rotate the query and key instead of adding to them. Rotation keeps each vector’s length, so a word’s meaning stays intact, and since an attention score depends only on the angle between two vectors, the score ends up tracking how far apart two tokens are rather than where they sit.

If you want the original source, it is the RoFormer paper by Su et al. It is short and readable, and after this it should be easy to follow. It is also the groundwork for the next question in this series: how a model trained on a few thousand tokens manages to read much longer inputs.

This is Part 1 of a two-part series on how the transformer’s attention mechanism has evolved. Every attention variant shipped in production since 2019 is fighting one number: the bytes of KV cache you have to carry per token. That number controls how many concurrent users fit on a GPU, how long a context you can serve, and ultimately whether your model is economically viable to deploy.

Part 1 walks the first wave of answers: the variants that attack the cache by changing what gets stored per token. We start with the bottleneck, recap multi-head attention, look at the stepping stones (MQA and GQA), then spend most of the post on Multi-head Latent Attention as introduced in DeepSeek-V2. By the end you will see how a single low-rank bottleneck plus a clever bit of algebra collapses the cache by nearly two orders of magnitude without giving up the expressivity of standard softmax attention.

Part 2 picks up at the question MLA cannot answer: once each cached token is about as small as it gets, can we cache fewer tokens? That is sparse attention (DSA, NSA, MoBA), linear-attention hybrids, and the V4-Pro synthesis where compression and sparsity stack. (Coming soon.)

The audience is engineers who deploy models. FlashAttention has its own dedicated post on this blog and we will not re-cover it here.

Part 1: The KV Cache Wall

Modern transformer inference is, in practice, a memory bandwidth problem. During autoregressive decoding, each new token must attend to every prior token, which means every prior token’s key and value vectors must already sit in fast memory. That stash is the KV cache, and it grows linearly with sequence length, linearly with batch size, and linearly with the number of layers.

For a model with $L$ layers, $n_h$ heads, per-head dimension $d_h$, sequence length $T$, and batch size $B$ in float16, the standard MHA cache is:

That factor of 2 is for keys and values. For a Llama-3-style 70B model at BF16 with 80 layers, 64 heads, and $d_h = 128$, each token consumes about 2.5 MB of cache per layer. At 128K context that is 320 GB of KV cache for a single sequence. The H100 has 80 GB of HBM. You cannot serve a single 128K-context request on one H100 without doing something to shrink that number. The cache, not the parameters, is what limits how many concurrent users fit on a GPU and how long a context you can serve.

Several approaches have chipped away at this. Multi-Query Attention (MQA) shares a single K/V across all heads, a brutal compression that visibly degrades quality. Grouped-Query Attention (GQA) is the negotiated middle ground that ships in Llama, Mistral, and Qwen. Multi-Head Latent Attention, introduced in DeepSeek-V2 in May 2024, takes a different tack: it caches a single low-rank latent vector per token and reconstructs full-rank K and V on the fly.

The trick, and the whole point of the rest of this post, is that you never actually have to reconstruct them.

Part 2: Notation

A few symbols recur throughout. None of them are unusual; the table is a quick reference so the equations below read fast.

Convention: all vectors are row vectors, so $W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$ and $h W$ produces an output row vector. This matches how PyTorch nn.Linear behaves and how most modern transformer code reads.

Part 3: Recap: Standard Multi-Head Attention

Before getting to MLA, it is worth pinning down exactly what we are trying to replace. Standard MHA at position $t$ takes the hidden state $h_t \in \mathbb{R}^{d}$ and produces three projections:

where $W^Q, W^K, W^V \in \mathbb{R}^{d \times n_h d_h}$. The result is split into $n_h$ heads of dimension $d_h$. Attention is computed head-wise, with $k_t$ and $v_t$ cached for every past position $t$:

The cost is paid not in the projection matrices (those live in HBM regardless) but in the running cache of all past $k_t, v_t$. Per token per layer, that is $2 \cdot n_h \cdot d_h$ floats. For DeepSeek-V2 scale (128 heads, 128 per-head dim) it is 32,768 floats per token per layer. Everything MLA does is in service of shrinking that cache while keeping the per-head expressivity.

Part 4: Stepping Stones: MQA and GQA

MLA is most legible when contrasted with what came before. Both MQA and GQA attack the cache by reducing the number of distinct K and V projections.

Multi-Query Attention (Shazeer 2019) keeps $n_h$ query heads but uses a single shared key projection and a single shared value projection across all of them. The cache shrinks by a factor of $n_h$. On a 64-head model that is a 64x reduction. MQA worked in production for PaLM and Falcon-40B, but most successors retreated. With one shared K and one shared V, every query head looks at the same key/value subspace, the model loses head-level specialization on the recall side, and quality regressions at scale were measurable.

Grouped-Query Attention (Ainslie et al. 2023) is the practical middle ground. Instead of all-or-nothing sharing, you partition the $n_h$ query heads into $n_g$ groups, and each group shares one K/V head. MHA is $n_g = n_h$. MQA is $n_g = 1$. GQA-8 (the Llama-3 default) is $n_g = 8$. For our 70B reference at GQA-8 the cache is 320 KB per token, an 8x reduction, and quality regression vs MHA is in the noise.

The K/V width is the lever each method pulls. MHA keeps full per-head K/V. MQA collapses to a single shared pair. GQA picks a comfortable middle. All three keep K and V as the cached objects, though.

MLA changes the question. Instead of asking “how many K/V projections do we keep?”, it asks: what if the cache is not K or V at all, but a compressed representation we expand on demand?

Part 5: MLA’s Core Insight

The hidden state $h_t \in \mathbb{R}^{d}$ already contains everything needed to compute that token’s keys and values. The standard pipeline burns most of its width on a high-dimensional intermediate ($n_h d_h \approx 16{,}000$ for DeepSeek-V2) that we store, when really we could store a compact summary and reproject when needed.

Concretely: introduce a latent bottleneck $\mathbf{c}_t^{KV} \in \mathbb{R}^{d_c}$ with $d_c \ll n_h d_h$. Cache only this latent. Recover K and V via dedicated up-projections at attention time:

If you stop reading right here, you would be tempted to ask: doesn’t reconstructing K and V at every attention step add a huge amount of compute? The honest answer is yes, naively. The whole punchline of MLA, which we get to in §8, is that during inference you do not have to materialize K and V at all. The up-projection matrices can be absorbed into Q and the output projection. The cache stays small and the FLOPs stay manageable.

A low-rank cache only saves memory. A low-rank cache that you never have to up-project saves memory and compute. That is the actual MLA result.

But before we get to the absorption, there is a wrinkle: RoPE. Rotary positional embeddings break the clean factorization above, and dealing with that gracefully is what gives MLA its slightly baroque final form. We will first walk through the no-RoPE version step by step, then patch it.

Part 6: Matrix Walkthrough, Step by Step

We will work through one token’s forward pass. Position $t$, hidden state $h_t \in \mathbb{R}^{d}$. Numbers in parentheses are the DeepSeek-V2 values, so the shapes feel concrete.

Step 1: KV down-projection

A single linear layer projects the residual stream down into the latent space:

The resulting $\mathbf{c}_t^{KV} \in \mathbb{R}^{512}$ is the only KV-related thing we cache for this token. It is shared across all heads and contains the information from which both K and V will eventually be reconstructed.

Step 2: K and V up-projection

The latent fans out into the full multi-head K and V via two more linear layers:

The superscript $C$ marks the content portion of K (we add a separate RoPE portion in §7). After this step, $k_t^C$ and $v_t^C$ live in $\mathbb{R}^{n_h d_h}$ and split into $n_h$ heads of width $d_h$. Since $d_c \le n_h d_h$, these are low-rank reconstructions: every head’s K and V is constrained to lie in a $d_c$-dimensional subspace of the full $d_h$-space. This rank constraint is the price of compression, and empirically it is the right trade.

Step 3: The query path

Symmetrically, and primarily to save training-time activations rather than KV cache, the query is also routed through a bottleneck:

The query bottleneck $d_c' = 1536$ is wider than the KV bottleneck because queries are not cached. There is no inference benefit from making them narrower. The reason to compress them at all is parameter and activation memory during training.

$d_c = 512$ is sized for cache miniaturization. $d_c' = 1536$ is sized for representational room. They are independent design knobs and DeepSeek-V2 chose them to be quite different.

Part 7: The RoPE Complication and the Decoupled Fix

So far the story is clean: cache a latent, up-project on demand, profit. But all modern transformers, DeepSeek-V2 included, use rotary position embeddings, and RoPE is exactly the kind of thing that ruins clean factorizations.

Why naive RoPE breaks MLA

RoPE applies a position-dependent rotation matrix $\mathcal{R}_t$ to the query and key after they are projected. The attention dot product becomes:

That last simplification, the rotation depending only on the relative position $s-t$, is the whole reason RoPE works. But now imagine we tried to apply RoPE to our reconstructed key $k_s = \mathbf{c}_s^{KV} W^{UK}$. The pre-rotated key cached as $\mathbf{c}_s^{KV}$ would need to also be rotated by $\mathcal{R}_s$. And $\mathcal{R}_s$ depends on $s$, the actual position. Different tokens use different rotations. So we would need to store the rotated reconstruction per token, defeating the cache.

The deeper algebraic problem: in §8 we will want to absorb $W^{UK}$ into $W^Q$. But if RoPE applies between them, in $q^\top W^{UQ\top} \mathcal{R}_t^\top \mathcal{R}_s W^{UK} \mathbf{c}^{KV}$, the position-dependent $\mathcal{R}$ blocks any precomputed absorption.

RoPE and low-rank absorption are fundamentally incompatible when applied to the same vector. MLA’s fix is to give RoPE its own, separate vector.

The decoupled RoPE construction

Each token gets two key vectors per head:

1. A content key $k_t^C$ (no RoPE), reconstructed from the cached latent as in §6.
2. A small RoPE key $k_t^R$, a separate, narrow tensor produced directly from $h_t$, to which RoPE is applied. Crucially, it is shared across all heads.

And on the query side, the rotation also lives on its own piece, but here it is per-head, since queries are not cached:

The final per-head query and key are the concatenation of the content part and the RoPE part:

Note: $k_t^R$ has no head subscript. Every head’s RoPE-key part is the same vector. This is exactly an MQA-style sharing, isolated to the RoPE 64-wide tail.

Attention is then computed on the concatenated vectors:

The dot product splits cleanly into a content term and a RoPE term, exactly because we concatenated rather than added. The content term will get the absorption treatment in §8. The RoPE term stands alone and is computed directly from the (small) cached $k_s^R$.

The output is then the usual per-head weighted sum of $v_{s,i}^C$ (the value has no RoPE, it never did), and the heads are concatenated and pushed through $W^O$ as usual:

The KV cache now holds $\mathbf{c}_t^{KV}$ (512 floats) plus $k_t^R$ (64 floats) per token per layer. Total: 576 floats, or 1,152 bytes in BF16. That is the headline 57x reduction versus MHA’s 32,768 floats.

Part 8: The Absorption Trick

We have reduced the cache from 32,768 to 576 floats per token. But the up-projections $W^{UK}$ and $W^{UV}$ are huge ($512 \times 16{,}384$ each), and computing them per attention step looks alarming. The trick is that we never compute them at inference time. We absorb them into the surrounding matrices.

Absorbing $W^{UK}$ into the query

Look at one head’s content-attention term:

The product $\widetilde{W}_i^Q = W_i^{UQ\top} W_i^{UK}$ is two parameter matrices multiplied together. It depends on no input, so we precompute it once. The score becomes a single bilinear form between the cached query latent and the cached key latent, both of width at most 1536.

After absorption, “computing the content key” disappears. The query latent talks to the key latent directly through a precomputed bridge matrix. The cache really is the key.

Absorbing $W^{UV}$ into the output projection

The value side gets the same treatment, but from the other end. The per-head output is a weighted sum of $v_{s,i}^C = \mathbf{c}_s^{KV} W_i^{UV}$, which then gets multiplied by $W_i^O$:

Again, the bracketed matrix is parameter-only: precompute and store. At runtime, the attention weights multiply directly against the cached $\mathbf{c}_s^{KV}$, and the result projects directly to the output dimension. $v$ is never materialized.

Subtle point: the absorption changes the effective compute layout. The per-head $\widetilde{W}^Q$ matrices are dense and head-specific, so you do not get a literal FLOPs reduction in the obvious places, but you do collapse what was a stream of large reconstructions into a single small inner product. The combined effect (small cache plus small inner product) is what makes long-context decoding cheap.

The RoPE part is computed the old way: the cached $k_s^R$ is dotted with the freshly computed $q_{t,i}^R$. It is small (64 wide) and outside the absorbed factorization, which is exactly why we paid the cost of decoupling it.

Part 9: DeepSeek-V2 by the Numbers

Plug the actual hyperparameters in. Drag the slider to see how the per-layer KV cache footprint scales with context length, per request, in fp16.

At 128K context, the kind of regime DeepSeek-V2 was designed for, each MLA layer holds about 144 MB of cache, versus around 8 GB for naive MHA. Multiply across DeepSeek-V2’s 60 layers and the difference is the difference between a context that fits and one that does not.

Part 10: Comparison and Takeaways

What is actually new

MLA is not the first low-rank attention proposal. Linformer (2020), Performer, and various Nyström approximations long predate it. What makes MLA practical and distinct is:

1. The low-rank object is the cache, not the attention pattern. Earlier work approximated the attention matrix itself. MLA keeps softmax attention exact and only compresses what gets stored across timesteps.
2. The absorption trick keeps inference compute bounded. Without §8, MLA would just be a memory-vs-compute tradeoff. With it, both move in the right direction at once for the autoregressive decoding regime.
3. The decoupled-RoPE construction shows how to make a low-rank cache compatible with rotary embeddings. This is the part most easily glossed over but is what makes the technique deployable in modern transformer stacks. Future low-rank schemes that ignore positional embeddings are nice in a paper and broken in practice. MLA paid the engineering tax.

The cost

MLA is not free. The decoupled RoPE adds parameters and a slightly more elaborate forward path. The bilinear absorbed matrices $\widetilde{W}_i^Q$ are larger than the raw $W^Q$ pieces would be. For very short context inference, plain GQA may be faster wall-clock, because the absorbed matmul is the dominant cost rather than memory bandwidth. MLA’s wins compound with context length and batch size, exactly the regimes that matter when serving a frontier model.

In the long-context, large-batch regime that production inference cares about, MLA pulls the KV-cache lever further than any other published trick, while keeping standard softmax attention semantics intact.

If you remember three things

1. The cache is a low-rank latent $\mathbf{c}^{KV}$, not K and V. K and V are reconstructed by up-projections that you never actually run at inference.
2. RoPE rides on a tiny separate vector that bypasses the latent compression, so the rotation can stay position-aware without contaminating the absorbable content path.
3. The whole construction is engineered around the fact that during decoding, you do not need K and V. You need scores. Storing scores requires only enough information to compute them, and the latent is sized to exactly that.

What comes next

Part 1 ended at the point where each cached token is about as small as it gets. The next pressure point is not the size of each token but the number of tokens you carry. Once your model can hold 1M tokens of context, do you really need to read every one of them at every decode step?

Part 2 will pick up there. Three different bets on how to answer that question shipped in 2025: sparse top-k selection driven by a cheap relevance scorer (DSA, deployed in DeepSeek V3.2), three-branch sparsification that is trainable from scratch (NSA, ACL 2025 best paper), and mixture-of-block routing that retrofits existing dense checkpoints (MoBA, powering Kimi K2’s 1M context). A different bet entirely is the linear-attention hybrids that change the math so no $O(L^2)$ matrix ever materializes (Lightning Attention, in MiniMax M1). The synthesis at the end is DeepSeek V4-Pro, released last month, where every thread we have followed so far gets composed at once into a 61-layer stack that runs at 2% of GQA’s cache footprint.

Part 2 is in progress. check back soon.

References

1. Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS 2017. The original Transformer paper. The baseline MHA that every variant in this post is reacting to.
2. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. The MQA paper. Concise and worth reading in full; the whole argument is six pages.
3. Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. The GQA paper. Includes the mean-pooling conversion recipe that made GQA easy to adopt.
4. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. The MLA paper. Section 2.1.3 has the decoupled RoPE derivation; almost everyone missed it the first time around.
5. Ji, F. et al. (2025). TransMLA: Multi-Head Latent Attention Is All You Need. Proves that MLA has strictly greater expressive power than GQA at the same cache budget, and gives a conversion recipe from GQA to MLA.

References for sparse attention (DSA, NSA, MoBA), linear hybrids (Lightning Attention), and the V4-Pro report will appear in Part 2.

By April 2026, the adoption numbers are staggering: 90% of developers use AI at work and over 80% say it’s made them more productive. Gartner projects 40% of enterprise applications will feature task-specific AI agents by the end of 2026, up from less than 5% in 2025.

Then you read the other column. CIO reports that 95% of enterprises see zero return on their AI investments. McKinsey’s maturity model puts only around 11% of enterprises in the “AI-native” tier. The 2025 DORA report is more uncomfortable still: AI raises throughput and raises change-failure rate. PR size is up 154%. 30% of engineers don’t trust the code their own agents produce.

The gap isn’t the model. Frontier models are a commodity. They get swapped every six months and the next one is better. The gap is everything around the model: the control plane that routes requests, attributes cost, enforces policy, retrieves context, evaluates quality, and measures outcomes. It’s the platform that turns “we rolled out Copilot” into “we shipped a 30.8% reduction in PR cycle time across 1,900 repos,” which is what Atlassian did with Rovo Dev and published at ICSE 2026.

This post is for the architect who has been asked to lead that platform. Not to choose between Copilot and Cursor; that’s a week of spreadsheets. To design what sits around whatever agent you pick, so that a year from now your CFO knows what AI is costing and your CTO knows what it’s earning.

Chapter 1: What “Platform” Actually Means Here

When a VP of Engineering says “we have an AI platform,” they might mean one of three things:

1. We bought Copilot Enterprise. Everyone has access.
2. We stood up a chat UI in front of a couple of models.
3. We run an internal control plane that mediates every AI request our engineers make, attributes cost per team, enforces policy per repo, evaluates quality continuously, exposes a curated surface of tools and skills, and lets any engineer publish a repeatable workflow that triggers on a schedule, a webhook, or a repository event.

Only the third one is a platform. The first two are procurements.

Shopify made this distinction concrete. Per Bessemer’s write-up of their AI-first engineering playbook, Shopify runs an LLM proxy. Every AI request from every tool, every engineer, every script, goes through one internal gateway. Engineers can pick their harness (Claude Code, Cursor, Copilot), but the proxy is non-negotiable. That single architectural choice is what gives them centralised cost control, usage analytics, model flexibility, and the ability to swap a model provider in a day instead of a quarter.

Block took a different turn at the same fork. Rather than wrapping third-party agents, they built Goose internally and open-sourced it. The stated reason, per CTO Dhanji Prasanna on the Sequoia Training Data podcast, was that “data leaving our infrastructure” was unacceptable. The outcome: engineers save 8-10 hours a week, Goose is on track to reclaim 25% of manual hours company-wide, and 100% of Goose’s own PRs are now written by Goose.

Two tech companies. Two legitimate answers to the same architectural prompt. Both built a platform. Neither bought one.

The platform you build, whether Shopify-shaped (gateway + BYO-harness) or Block-shaped (build-the-harness), owns five responsibilities:

• Capability: what tools and skills agents can reach, and how they’re discovered.
• Identity & Policy: who’s acting, with what scope, under what guardrails.
• Context: how org knowledge gets ingested, permissioned, and served.
• Evaluation: how you know the agent is actually getting better, not just shipping faster.
• FinOps & Observability: what it costs, who paid, what it produced, where it broke.

Those five wire up into a system shape worth drawing. The harness runs on each engineer’s laptop (in a Docker container) or in a CI runner. Everything the harness reaches into over the network is the platform. Everything the platform calls out to is a model provider, a curated MCP server, or a source system the platform has already indexed.

When multiple agents cooperate within this system, the topology that ships is orchestrator-worker, not swarm. Airbnb’s 3,500-file Enzyme-to-RTL migration proved this: per-file parallel workers, central orchestration, brute-force retries with dynamic prompts. 97% automated, 6 weeks, 6 engineers. Swarms, by contrast, are the dominant source of silent failure because a single hallucination in shared memory propagates to every peer that reads it. Chapter 5 shows how Goose sub-recipes implement the orchestrator-worker pattern concretely.

If any of these responsibilities isn’t owned by a named team with a roadmap, you don’t have a platform. You have a shadow-IT problem that will compound. These five exist to enable one thing: the workflow lifecycle. Chapter 5 names it; everything else makes it safe, cheap, and measurable.

Chapter 2: The Capability Surface

The first architectural argument inside every platform team is about how agents get things done. It usually gets framed as a choice between the Model Context Protocol (MCP), now stewarded by the Linux Foundation’s Agentic AI Foundation as of December 2025, and Agent Skills, the behavioural-instruction packages that started shipping with Claude.

Framing them as competitors is a category error. MCP is an execution fabric: a standardised RPC for tools, resources, and prompt templates, with bidirectional comms and dynamic tool discovery. Skills are a knowledge layer: portable instructions that encode the how-to of a specific job. MCP tells an agent what it can do. Skills tell an agent what it should do in a specific situation.

The context-bloat failure is the sharper risk. A single enterprise-grade GitHub MCP server exposes 90+ tools. Loaded naively, that’s 50,000+ tokens of schema entering the context window before the model has read a single line of user intent. Add Jira, the cloud provider, a feature-flag platform, your incident system, and your agent is spending six figures a year on tokens that are just tool catalogue. The overhead scales linearly with the number of services you connect.

The three-tier model that scales:

1. Built-ins: the small primitive set the harness ships (file ops, shell, code execution). Always loaded.
2. Curated MCP registry: a governed list of approved MCP servers with progressive disclosure. The agent sees metadata first, loads full tool schemas on semantic match.
3. Skills library: org-specific playbooks in a searchable registry, discovered by description, expanded on demand.

Block’s Goose is the cleanest public expression of this. The operational equation is Goose = LLM + MCP + Agent, but the load-bearing piece isn’t MCP. It’s Goose’s Recipes and Sub-recipes. Recipes are declarative YAML workflows that encode a repeatable piece of work; sub-recipes run in isolated sub-sessions with their own context windows. That isolation keeps token cost linear in the work done rather than quadratic in conversation depth. The result is the 30-40% of code Block’s top engineers now get from Goose in legacy codebases, per the Sequoia interview.

Chapter 3: Identity, Policy, and the Execution Boundary

Every agentic action has two identities the audit team cares about: the human on whose behalf the agent is acting, and the agent’s own service identity. Conflate them and compliance review kills your rollout.

The human identity provides authorisation scope. The agent identity provides attribution and accountability (which agent, which version, which session). Every tool invocation carries both. Tokens are short-lived (minutes, not days) and scoped to the specific task, not the session. Policy is enforced at the MCP gateway, as code, so auditors can diff it and engineers can review it.

How do you keep humans in the loop without drowning them in approval prompts? The pattern that works is the asynchronous wait-state. When an agent hits a high-risk decision (production deploy, financial transaction, irreversible write), the workflow suspends, persists its state externally, and emits an approval event. Reviewers act on their own clock, often hours later. On approval, the signal routes back and the workflow resumes exactly where it left off.

The anti-pattern is approval fatigue. The fix is graduated trust: scope approvals by blast radius.

• Read-only on scoped data: no approval.
• Mutations inside a sandbox or personal branch: no approval, full audit.
• PR against the main branch: standard code review.
• Production-shaped actions (deploys, config changes, prod data reads): explicit, async approval with a named owner.
• Irreversible (delete, drop, disable safety): two-person review.

GitHub’s Copilot Enterprise surface has become the most concrete public implementation. Per the December 2025 Enterprise roundup and Microsoft’s DevBlogs on agentic platform engineering, admins get fine-grained permissions, explicit MCP control, audit-log review, and policy-based gating of model upgrades.

Policy tells the agent what it’s allowed to try. The sandbox decides what happens when it tries the wrong thing. Three isolation tiers map to the graduated-trust model:

Match the isolation tier to the trust tier. A read-only retrieval agent does not need a microVM. A production migration worker rewriting other teams’ code absolutely does.

Chapter 4: Context Engineering at Platform Scale

The frontier model isn’t your moat. The agent harness isn’t your moat. Your context is your moat: the graph of your repos, the runbooks nobody wrote down, the incident history, the ADRs, the style guides, the org’s service catalog.

MCP is connectivity, not context

A common mistake is assuming that connecting MCP servers to your agent solves the context problem. MCP gives your agent a standardised way to query any single system. Four things it does not do:

1. Cross-source retrieval. “Find everything relevant to this migration across code, tickets, docs, and incidents” requires a unified index. No single MCP server spans all your sources.
2. Pre-indexing. MCP queries are live. For a 500k-file monorepo, live search on every agent call is slow and expensive.
3. Governance. PII redaction, access-aware filtering, staleness SLOs. Each MCP server returns raw data under its own auth model.
4. Token-budget management. Fitting retrieved context to the model’s window is orchestration the pipeline owns, not the protocol.

The clean architecture: build the context pipeline, expose it as an MCP server. The agent queries one context endpoint. The pipeline behind it handles ingest, index, govern, serve.

The pipeline

Ingest. Connectors to the authoritative sources: repos, docs wiki, ticket system, incident tracker, service catalog. Each with an idempotent, versioned schema and an owner.

Index. Hybrid retrieval is the production default: BM25 for lexical recall, dense embeddings for semantic similarity, graph for structural relationships. No single index is sufficient.

Govern. Staleness SLOs per source. PII and secret redaction before indexing, not after retrieval. Access-aware retrieval: the retriever filters by the caller’s permissions before ranking. If your agent can see secrets its invoking user can’t, you have a data-exfiltration vulnerability wearing a productivity tool’s clothes.

Serve. A token-budget manager (compression, summarisation, eviction) that fits retrieved context to the model’s window and the task’s importance.

Augment Code’s Context Engine is the clearest public reference for this in 2026. It indexes up to 500,000 files across multiple repositories with roughly 100ms retrieval latency, building semantic dependency graphs. The telling move: Augment recently shipped the Context Engine as an MCP server, the exact pipeline-behind-protocol pattern. Sourcegraph’s Cody takes a three-layer approach (local file, local repo, remote repos), handling 300k+ repositories for enterprise customers. Stripe’s agent harness takes the curation angle: each “minion” gets scoped context per task, not the whole repo. Context curated, not copied.

The metric to watch: context hit rate per task type. If your hit rate is under 30%, your pipeline is ornamental.

Chapter 5: Workflows, the Unit That Ships

Four chapters described infrastructure. This chapter is about what the infrastructure produces. The deliverable is the workflow: a versioned, parameterised unit of work any engineer can build once, evaluate, and hand to other engineers (or to CI runners) who invoke it on a trigger they didn’t author.

Authoring

Four patterns; the choice follows who the author is:

• Recipe / YAML: Goose Recipes, GitHub Agentic Workflows (Feb 2026 preview). Structured, diff-reviewable, CI-friendly. The enterprise default.
• Prompt-as-code: Claude Skills. Flexible, closer to prose, weaker composition.
• DSL / real code: Temporal, LangGraph, Kestra. Maximum control; needs engineer authors.
• Low-code: Atlassian Rovo Studio. Natural-language authoring for non-engineers.

A Goose Recipe is the concrete shape most architects will end up writing:

name: pr_security_review
recipe:
  version: 1.0.0
  title: PR Security Review
  description: OWASP-informed review of a pull-request diff.
  settings:
    goose_provider: anthropic
    goose_model: claude-sonnet-4-5
  parameters:
    - key: pr_url
      input_type: string
      requirement: required
      description: "Pull request URL to review"
  extensions:
    - type: builtin
      name: developer
    - type: streamable_http
      name: github
      uri: https://api.githubcopilot.com/mcp/x/pull_requests/readonly
  instructions: |
    You are a security reviewer. Check the diff for OWASP Top-10
    issues, secrets, and unsafe patterns. Be specific and sparing.
  prompt: |
    Review PR {{ pr_url }}. For each finding, cite the file,
    line, severity, and suggested fix. Post findings as a single
    PR comment. If nothing is found, say so.

Every primitive the last four chapters described is visible here. settings routes through the LLM gateway. extensions declares which approved MCP servers the capability surface exposes. parameters is how a non-author reuses the workflow. instructions vs prompt separates policy from task, which is what makes a Recipe testable.

Parameterisation and sub-workflows

A Recipe without parameters is a one-off. With parameters, it’s a product. The sharper Goose primitive is the sub_recipes array: each sub-recipe runs in its own isolated subagent session with its own context window, and sequential_when_repeated: true/false picks parallel vs sequential execution. This is the orchestrator-worker pattern from Chapter 1, made concrete. It’s what makes the Airbnb migration topology possible: 3,500 files fan out across parallel sub-recipe invocations, each with fresh context, orchestrated by one parent.

Triggers

Cron. goose schedule add recipe.yaml --cron '0 9 * * 1-5'. Nightly lint, weekly security audit, daily stale-PR report. The built-in scheduler is single-machine; for distributed schedules, wrap with a Kubernetes CronJob or a Temporal worker pool.

Event-driven. PR opened, issue labelled, incident created, build failed. Atlassian’s Rovo Dev fires on every PR. The Goose GitHub Action wraps the same pattern: label an issue with goose and a PR opens. Event-driven is where agents stop being assistants and start being automation.

Manual / API. goose run -i recipe.yaml --param pr_url=https://... from a CI step, or goose serve running as a webhook receiver inside the cluster.

Runtime, observability, and governance

Triggered workflows run on ephemeral CI runners (GitHub Actions, Buildkite) for sub-five-minute PR-shaped work, or on dedicated agent pools for long-running stateful work. Match runtime to the trust tier from Chapter 3.

Every triggered run is a first-class object: trigger source, parameters, spans with retry counts, final status, cost, trace ID. Kestra recorded over two billion workflow executions in 2025, up from one hundred million in 2024. That twenty-fold increase signals the direction of travel. If your platform cannot answer “what ran when, triggered by what, with what outcome?” in two clicks, it is opaque.

Shared workflows need product discipline. The GitHub Actions governance model (internal org, SHA-pinned versions, PR-reviewed contributions) is the pattern most enterprises borrow.

Chapter 6: Evaluation and Economics

Most platform teams skip evaluation and then wonder why their rollout plateaus. Evaluation is not a phase of delivery; it is the product that determines whether the other five chapters compound.

Silent failure

An agent completes its run without any software error (no exception, no crash, no red log line) and produces output that looks plausible and is wrong. The PR passes review because the diff looks reasonable. The test the agent wrote passes because it tests the buggy behaviour it introduced. Every DORA-2025 data point on increased change-failure rate is a silent-failure story that got written to disk.

The evaluation stack that catches silent failure has three layers.

Unit-level. Tool schemas, prompt templates, and system prompts each get their own regression suite. Every change runs a deterministic test set before it can ship.

Task-level. A curated golden set of real tasks, graded by LLM-as-judge with a rubric that includes business-outcome correctness, not just style. This is eval-as-CI.

Production. Shadow traffic and online signals: thumbs-up/down, PR accept rate on agent-authored code, downstream defect escape rate. The production signals feed back into the golden set. Every thumb-down becomes a candidate regression test.

Atlassian’s Rovo Dev Code Reviewer ran a year-long evaluation across more than 1,900 internal repos before general availability. The result, published at ICSE 2026, was a 30.8% reduction in PR cycle time and a 35.6% reduction in human-written review comments. The same three eval layers apply at the Recipe level: shadow-run the candidate against live triggers before promoting; canary to a subset before broad ship.

Token economics

By the time you have 5,000 engineers on your platform, token cost is non-linear in three dimensions: context depth, fan-out, and retry depth.

Tiered routing. Simple classification and extraction routes to a cheap model (Haiku-class). Standard code generation routes to mid-tier (Sonnet-class). Hard planning and architectural synthesis reserves to the frontier (Opus-class). Defaulting every call to the most expensive model is the single largest source of cost inflation.

Prompt caching as an SLI. Structured prompts should cache at 90%+ hit rate. A 90% cache hit translates to roughly 10x cost reduction on the cached portion. Cache hit rate deserves a dashboard, an owner, and an alert when it drops.

Attribution at every level. Per-team, per-repo, per-task, per-session. Without attribution there’s no chargeback; without chargeback there’s no incentive for teams to care about efficiency.

Shopify’s LLM proxy, mentioned in Chapter 1, is the artefact that makes all of this possible. You cannot attribute cost you don’t see. You cannot route by complexity if requests bypass your router. Per First Round’s write-up, the proxy is what let Shopify’s engineering dashboard correlate AI usage with shipping impact, which in turn gave VP Eng Farhan Thawar the evidence to support the ~20% productivity gain the org now claims.

What to measure

The most common failure mode in “AI productivity” reporting is Goodhart’s Law in a lab coat. A measurement stack that survives scrutiny operates in four families: proxy (acceptance rate, session count), activity (DORA: PR count, lead time, CFR), outcome (defect escape, rework, dev-reported friction), and economic (hours saved, cost per merged PR). An architect reporting to leadership needs at least one number from each.

Consider the published record: Uber reports ~10% PR-velocity lift (Pragmatic Engineer), an activity metric. Shopify claims ~20% productivity accompanied by a public refusal to measure it in LOC, an outcome claim. Block’s 8-10 hours saved per engineer per week is a clean economic metric. Airbnb’s 18 months to 6 weeks is a sharp outcome metric with a legible counterfactual. Same reality. Different slices.

Chapter 7: The Build Sequence

The platform described above is not a weekend project. It also does not require a three-year transformation program. The sequence that has worked in the public record collapses into three horizons.

Days 0-90. Stand up the minimum viable control plane.

• Pick one harness. Don’t debate it for a quarter. Any of them is fine; the harness is replaceable.
• Stand up the LLM gateway. Every agent request flows through it. Day-one cost attribution.
• Ship one Recipe. Not twelve. Pick one repeatable task (PR security review, migration shard, on-call triage). Versioned, parameterised, triggered by one event, observable end-to-end. Everything else is scaffolding for the next Recipe.
• Stand up one golden eval set with an LLM-as-judge rubric. Wire it into CI. Refuse to promote prompts or Recipes that regress.
• Turn on OpenTelemetry tracing end-to-end.

Months 3-6. Build the moat.

• Context pipeline for your top-five repos: ingest, index, govern, serve. Measure hit rate.
• Policy-as-code at the gateway. Scoped tokens. Async approvals for production actions.
• Expand the eval harness to workflow-level: golden sets of Recipe invocations, shadow-mode promotion.
• First KPI dashboard: one proxy, one activity, one outcome, one economic metric.

Months 6-12. Compound.

• Orchestrator-worker topology for the hard workloads: migrations, cross-repo refactors, bulk compliance work.
• Recipe registry self-service with SHA-pinned versions. Teams contribute; the platform team curates.
• Progressive autonomy tiers. Graduate teams through read-only, sandboxed, PR, and production as their eval and incident track record earns it.
• Per-team chargeback. The budget conversation changes the usage conversation.

Fund internal DevRel from day one. Uber’s coursework moved Claude Code adoption from 32% to 63% of engineers in three months. Block’s engineers found Goose through Slack channels, not mandates. Shopify paired a top-down AI-first memo with bottom-up tool freedom through the LLM proxy. The technical platform and the organisational motion need to ship together.

In twelve months, when your CFO asks what AI is costing and what it’s earning, you have an answer, because you built a platform rather than bought a license. That’s the answer the 11% have. It’s not because they picked a better model.

References

1. Google Cloud / DORA. 2025 State of AI-Assisted Software Development Report. Source for 90% adoption, 30% distrust, PR size +154%, and the stability/throughput tension.
2. Faros AI. Key Takeaways from the DORA Report 2025. Practitioner analysis of the DORA findings.
3. McKinsey / KPMG. AI at Scale: Q4 2025 AI Pulse. Source for the four-stage maturity model and the ~11% AI-native figure.
4. OneReach / CIO. What Shapes Enterprise AI Agents in the Future. Source for the 95% zero-ROI and 14% change-management figures.
5. Block. Block Open Source Introduces “codename goose” and Goose on GitHub.
6. Sequoia. Training Data podcast with Dhanji Prasanna. Source for Block’s 8-10 hours/week, 25% target, and 30-40% legacy-code figures.
7. All Things Open. Meet Goose: The open source AI agent built for developers.
8. Bessemer Venture Partners. Inside Shopify’s AI-First Engineering Playbook.
9. First Round Review. From Memo to Movement: Shopify’s Cultural Adoption of AI.
10. Augment Code. Context Engine and Context Engine MCP now live. Source for the 500k-file indexing, ~100ms retrieval, and pipeline-behind-MCP pattern.
11. Pragmatic Engineer. How Uber Uses AI for Development. Source for the 84% agentic-coding adoption, Claude Code 32% to 63%, and DevRel investment.
12. Sourcegraph. How Cody understands your codebase and How Cody provides remote repository awareness. Source for the three-layer context architecture and 300k+ repo scale.
13. Atlassian. 30.8% Faster PRs: How AI-Driven Rovo Dev Code Reviewer Improved Developer Productivity. Source for the ICSE 2026 publication figures.
14. GitHub. December 2025 Enterprise Roundup. Source for Copilot Enterprise governance features.
15. Microsoft DevBlogs. Agentic Platform Engineering with GitHub Copilot.
16. Airbnb Engineering. Accelerating Large-Scale Test Migration with LLMs.
17. Anthropic. Model Context Protocol.
18. Gartner. 40% of Enterprise Apps Will Feature Task-Specific AI Agents by 2026.
19. Block. Goose Recipes reference and Goose Recipes cookbook.
20. Pulse MCP. Configure your agent with Goose Recipes.
21. Block. Goose AI Developer Agent GitHub Action.
22. GitHub. Automate repository tasks with GitHub Agentic Workflows.
23. Kestra. Kestra 1.0 launch. Source for the 2B+ workflow executions in 2025.
24. Temporal. Orchestrating Ambient Agents with Temporal.
25. MindStudio. Stripe Minions vs Shopify Roast. Source for Stripe’s scoped-context agent pattern.
26. GitHub. Building organization-wide governance for CI/CD with GitHub Actions.

NVIDIA’s Nemotron-3-Super is not a Transformer. Not entirely. It is a hybrid architecture that interleaves Mamba-2 SSM layers with select attention layers, using the majority of its compute on state space operations rather than self-attention. It ships in production on NVIDIA’s inference stack and competes with pure Transformer models at the same scale. NVIDIA is not alone. IBM’s Granite 4.0 uses a 9:1 SSM-to-Transformer ratio. AI21’s Jamba uses 1:7. Zyphra’s Zamba, Google’s Griffin, Microsoft’s Phi-4-mini-flash-reasoning: all hybrid architectures, all in production.

Something shifted. For years, the Transformer was the only architecture that mattered for language. Now every major AI lab is replacing most of their Transformer layers with SSM layers and getting better results at lower inference cost. If you deploy models, manage GPU clusters, or care about inference latency, this is worth understanding deeply.

This post builds State Space Models from zero. I start with the simplest possible differential equation: one variable, one parameter. From there, I build up to the full SSM formulation, explain the key breakthroughs (HiPPO, S4), and walk through the three generations of Mamba. By the end, you will understand the math well enough to explain to your team why these hybrid architectures are winning, what trade-offs they make, and what it means for your inference stack.

Part 1: Why SSMs? The Transformer’s Inference Problem

You already know the Transformer’s self-attention mechanism scales quadratically with sequence length: $O(L^2)$ in both time and memory. But the pain runs deeper than asymptotic notation.

During autoregressive decoding, the Transformer generates one token at a time. For each new token, it must load the entire KV cache from GPU HBM into SRAM, compute a single attention score against every previous token, and write the new KV entry back. The GPU spends the vast majority of its time moving data, not computing. On an H100 generating tokens from a 70B model, the Tensor Cores that deliver 989 TFLOPS of BF16 matmul sit almost entirely idle during decoding. The bottleneck is memory bandwidth, not compute.

This is why you need PagedAttention to manage fragmented KV cache memory. This is why vLLM exists: to batch requests efficiently despite variable KV cache sizes. This is why context windows beyond 128K tokens start requiring multi-GPU setups just to hold the KV cache.

State Space Models offer a fundamentally different deal. Instead of caching every token’s key-value pair (lossless but expensive), they compress the entire sequence history into a fixed-size hidden state (lossy but cheap). Processing each new token during inference takes $O(1)$ time and memory. No growing KV cache. No PagedAttention. Constant memory per sequence regardless of whether you have processed 100 tokens or 100,000.

The question has always been whether a compressed, lossy state can match the quality of the Transformer’s lossless KV cache. For years, the answer was no. SSMs excelled on audio, time series, and synthetic long-range benchmarks, but they lagged on language. The Mamba line of work changed that. To understand how, we need to start from scratch.

Part 2: State Space Models from Scratch

A Single Differential Equation

Forget matrices, vectors, and neural networks for a moment. Start with a single number $h(t)$ that changes over time:

$h'(t)$ is the rate of change of $h$ at time $t$. If you know calculus, this is the derivative. If not, think of it as: “how fast is $h$ changing right now?” When $h'(t)$ is positive, $h$ is increasing. When negative, $h$ is decreasing. When zero, $h$ is holding steady.

The constant $a$ controls everything:

• $a > 0$: $h$ grows exponentially. Think compound interest. A bank account with interest rate $a = 0.05$ earns interest on its interest, accelerating upward forever. Unstable.
• $a < 0$: $h$ decays exponentially. Think radioactive decay. A substance with decay rate $a = -0.5$ loses half its remaining mass roughly every 1.4 time units. The more you have, the faster it drains, but it never quite reaches zero. Stable.
• $a = 0$: nothing happens. $h$ is constant forever.

For building sequence models, we want $a < 0$. Our hidden state should be a fading memory, not an explosion.

Adding an Input

A decaying state by itself is useless. We need to feed information in:

Now $x(t)$ is an input signal (think: a stream of token embeddings arriving over time), and $b$ controls how strongly the input drives the state.

Picture a leaky bucket with a tap. The water level $h(t)$ is the state. The hole in the bottom drains water at rate $a \cdot h(t)$: the more water in the bucket, the faster it leaks (more pressure = faster drain). The tap $b \cdot x(t)$ pours water in at a rate proportional to the input signal. The water level at any moment reflects a fading, weighted average of all the water that has ever been poured in, with recent additions contributing more because older ones have partially leaked away.

This is the core intuition for the entire SSM line of work. The hidden state $h(t)$ is a running, compressed summary of the input history, where old inputs fade at a rate controlled by $a$.

Adding an Output

We read out the state with simple scaling:

The output $y(t)$ is just a weighted view of the state. Together, these two equations form the complete scalar SSM:

Three parameters. One input, one hidden state, one output. This is the entire architecture, in its simplest form.

Why One Bucket Is Not Enough

A single leaky bucket has one leak rate, which means one timescale of memory. If $a = -0.5$, the state “forgets” with a half-life of about 1.4 time units. It cannot simultaneously maintain a short-term memory (last few tokens) and a long-term memory (paragraph-level context).

The fix: use $N$ buckets, each with a different leak rate.

This is where scalars become vectors. The scalar state $h(t)$ becomes an $N$-dimensional vector $\mathbf{h}(t) \in \mathbb{R}^N$. The scalar parameters become matrices:

• $A \in \mathbb{R}^{N \times N}$ (state to state): governs how each of the $N$ state dimensions evolves and potentially interacts with the others. It is $N \times N$ because each state dimension can influence every other state dimension.
• $B \in \mathbb{R}^{N \times 1}$ (input to state): fans a scalar input out into $N$ state dimensions. It is $N \times 1$ because it needs to distribute one input value across $N$ state slots. Think of it as an adapter between a narrow input pipe and a wide state vector.
• $C \in \mathbb{R}^{1 \times N}$ (state to output): narrows the wide state back down to a scalar output. It is $1 \times N$ because it takes a weighted combination of all $N$ state dimensions to produce one output value.

In practice, $A$ is almost always diagonal. A diagonal $A$ means each state dimension evolves independently. No cross-talk between buckets. Dimension 1 decays at its own rate, dimension 2 at its own rate, and so on. This simplification works just as well empirically (the S4D paper proved this) and is much cheaper to compute.

Eigenvalues: The Retention Rates

For a diagonal $A$, the diagonal entries ARE the eigenvalues. No linear algebra required to understand this. Each eigenvalue $\lambda_i$ is simply the decay rate of one state dimension. Think of them as $N$ different bank account interest rates running simultaneously:

• $\lambda_i = -0.01$: very slow decay. This dimension remembers inputs from thousands of timesteps ago. It is the long-term savings account.
• $\lambda_i = -0.5$: moderate decay. This dimension tracks information over dozens of timesteps.
• $\lambda_i = -2.0$: fast decay. This dimension mostly tracks the last few inputs. It is the checking account that turns over quickly.
• $\lambda_i > 0$: growth. Unstable. The state explodes. We never want this.

By having $N$ state dimensions with different eigenvalues, the model simultaneously maintains memory at multiple timescales. Some dimensions track recent tokens (large negative eigenvalues, fast decay), others preserve long-range context (small negative eigenvalues, slow decay).

This is why the initialization of $A$ matters enormously. If you set eigenvalues randomly, you get random memory timescales, and the model struggles to learn useful representations of sequences. Random eigenvalues might cluster all your memory in one timescale, leaving gaps in others. This is exactly the problem that HiPPO solved. But we need to cover discretization first.

Part 3: Discretization, Making It Computable

Why Discretize

The continuous ODE $\mathbf{h}'(t) = A\mathbf{h}(t) + Bx(t)$ processes smooth, continuous signals. But an LLM does not receive continuous signals. It receives a discrete sequence of tokens, one after another. We need to convert the continuous dynamics into a step-by-step recurrence: given the previous state and the current token, compute the next state.

The Step Size $\Delta$

Discretization introduces a learnable parameter $\Delta$ that controls “how much continuous time passes between tokens.” Small $\Delta$ means the model takes fine-grained steps, preserving detailed temporal structure. Large $\Delta$ means coarse steps, compressing more time into each token. Each channel in the model can learn its own $\Delta$, so different parts of the network can operate at different temporal resolutions.

The Euler Discretization

The simplest approach: approximate the derivative as constant over each timestep, following the Euler method from numerical analysis. This gives us discrete parameters:

This is first-order accurate, with local truncation error $O(\Delta^2)$. There are more accurate methods (zero-order hold, bilinear transform), but Euler is the one that matters for the Mamba story because Mamba-3 directly improves on it.

The Discrete Recurrence

The discretized system gives us a step-by-step formula. At each timestep $k$:

This is now a simple step-by-step formula. Given the previous state and the current input, compute the next state and output. No calculus needed at runtime.

Numerical Walkthrough

Let me make this concrete. Take a scalar system with $a = -0.5$, $b = 1.0$, $c = 1.0$, and step size $\Delta = 0.1$.

Now run 5 timesteps with the input sequence $[1, 1, 0, 0, 1]$, starting from $h_0 = 0$:

# a_bar = 0.95, b_bar = 0.1

# k=0: input=1  h = 0.95 * 0      + 0.1 * 1 = 0.1000   y = 0.1000
# k=1: input=1  h = 0.95 * 0.1    + 0.1 * 1 = 0.1950   y = 0.1950
# k=2: input=0  h = 0.95 * 0.195  + 0.1 * 0 = 0.1853   y = 0.1853
# k=3: input=0  h = 0.95 * 0.1853 + 0.1 * 0 = 0.1760   y = 0.1760
# k=4: input=1  h = 0.95 * 0.176  + 0.1 * 1 = 0.2672   y = 0.2672

The state accumulates when inputs arrive (steps 0-1, step 4) and decays when they stop (steps 2-3). You can verify every number with a calculator. There is nothing hidden in the SSM recurrence: it is a multiply-and-add, repeated.

The Dual Computation Modes

This is the SSM’s defining superpower. In the formulation above, $A$, $B$, and $C$ are constants that do not change over time. This property is called Linear Time-Invariance (LTI), and it unlocks something powerful. Because the recurrence $\mathbf{h}_k = \bar{A}\mathbf{h}_{k-1} + \bar{B}x_k$ is linear with fixed parameters, we can unroll it algebraically:

The output $y_k = C\mathbf{h}_k$ is then a weighted sum of all past inputs, with weights $K = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \ldots)$. This sequence of weights is a convolution kernel.

This means we can compute the output of the SSM in two completely different ways:

Training mode (convolution): Compute the kernel $K$ once, then convolve it with the entire input sequence via FFT in $O(L \log L)$. Fully parallel, like a CNN. The GPU processes all $L$ tokens simultaneously.

Inference mode (recurrence): Step through $\mathbf{h}_k = \bar{A}\mathbf{h}_{k-1} + \bar{B}x_k$ one token at a time in $O(1)$ per step. The state $\mathbf{h}$ is a fixed-size vector regardless of how many tokens have been processed. No KV cache.

“Train like a CNN, infer like an RNN.” This is the fundamental efficiency proposition. During training, you get the parallelism of convolutions. During inference, you get the constant-time, constant-memory decoding of RNNs, without the KV cache that makes Transformer inference expensive.

This duality is only possible because the system is LTI: the parameters $A$, $B$, $C$ are fixed, so the same convolution kernel $K$ applies to every input. When parameters become input-dependent (which is what Mamba does), there is no single kernel for the whole sequence. The duality breaks, and new algorithms are needed.

Part 4: HiPPO, The Initialization That Made SSMs Work

Before HiPPO, SSMs initialized the state matrix $A$ randomly. Random eigenvalues produce random memory timescales. On Sequential MNIST (classifying a handwritten digit fed one pixel at a time, 784 steps), random initialization achieved about 60% accuracy. Barely above chance for some digit classes.

Albert Gu’s HiPPO framework (2020) solved this by deriving $A$ matrices from a mathematical objective: at every timestep, the state should store the best polynomial approximation of the entire input history. Each state dimension corresponds to one polynomial coefficient, with low-order coefficients capturing broad trends (long-range memory) and high-order coefficients capturing fine details (short-range memory). The resulting $A$ matrix has eigenvalues arranged to cover multiple timescales without redundancy.

The concrete impact: switching from random $A$ to HiPPO improved Sequential MNIST from 60% to 98%. Same architecture, same training. Only the initialization of $A$ changed.

Part 5: S4 and S4D, Making SSMs Practical

S4 (Gu, Goel, and Re, 2022) was the first architecture to make deep SSMs work at scale by finding an efficient algorithm to compute the convolution kernel from a HiPPO-initialized $A$ matrix. It was the first model to solve long-range tasks at sequence lengths of 16,000+, a result no Transformer or RNN had achieved. S4 also fully exploited the recurrent-convolutional duality: convolution mode for training, recurrence mode for inference.

A key simplification followed quickly. S4D (2022) showed that restricting $A$ to a fully diagonal matrix matched S4’s performance while dramatically simplifying the implementation. Independent state dimensions with well-chosen eigenvalues were sufficient. This diagonal restriction became the standard for all subsequent work, including Mamba.

Part 6: Mamba-1, Selectivity Changes Everything

The LTI Problem

S4 and its variants excelled on continuous signals and synthetic long-range benchmarks. On language modeling, they consistently lagged behind Transformers of the same size.

The reason is exactly the LTI limitation I described earlier. In an LTI system, the matrices $A$, $B$, $C$ are fixed constants. Every token receives identical treatment. The Mamba paper demonstrated this failure precisely with two diagnostic tasks:

Selective Copying: Given “A B _ _ _ C _ _ A _ _”, copy only A, B, C while ignoring the padding underscores. An LTI system cannot distinguish content tokens from padding because it applies the same transformation to everything.

Induction Heads: Given “A B … A ?”, recall that B followed A earlier and predict B. This requires content-based lookup: comparing the current token (A) against stored tokens to find what came after it. An LTI system has no mechanism for content comparison.

Language is full of these patterns. The word “not” should be remembered differently from the word “the.” A name mentioned once in a document needs to be retrievable later. All of this requires the model to make content-dependent decisions about what to store and what to forget.

The Fix: Input-Dependent Parameters

The December 2023 paper “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” by Albert Gu and Tri Dao introduced a single, elegant idea: make $B$, $C$, and $\Delta$ functions of the input token.

Now the model dynamically modulates its behavior on a per-token basis. When the network encounters an important token, it can predict a large $\Delta_t$ to reset the state and absorb the new information. When it encounters filler, it can predict a tiny $\Delta_t$ to preserve existing memory and let the filler leak away.

The roles of each parameter are clear. $\Delta$ controls the gate: large $\Delta$ resets the state and focuses on the current input; small $\Delta$ persists the state and ignores the current input. $B$ controls what enters the state (content-based filtering of what to remember). $C$ controls what exits (content-based modulation of what to read out).

Note that the state matrix $A$ itself remains fixed. This is intentional. $A$ affects the discrete recurrence only through its interaction with $\Delta$ via $\bar{A} = \exp(\Delta A)$, so making $\Delta$ input-dependent is sufficient to make the entire system input-dependent.

The Cost: Convolution Mode Breaks

Input-dependent parameters mean the system is no longer LTI. $B_t$ and $C_t$ change at every timestep, so there is no single convolution kernel $K$ that describes the entire sequence. The FFT-based parallel training mode is gone.

Naively, this forces sequential computation: process token 1 to get $h_1$, then token 2 to get $h_2$, and so on. This would be catastrophically slow on GPUs, which need parallel workloads to achieve decent utilization.

Mamba-1 solved this with a hardware-aware selective scan algorithm, directly inspired by FlashAttention’s approach to the memory hierarchy. The key idea: fuse all SSM operations (discretization, recurrence, output) into a single GPU kernel that runs entirely in SRAM, avoiding expensive HBM round-trips. The recurrence is parallelized using a parallel scan that exploits the associativity of the multiply-add operation, and intermediate states are recomputed in the backward pass rather than stored. The result: 40x faster than a naive implementation, with the same memory footprint as an optimized Transformer with FlashAttention.

The Mamba Block

A common misconception is that Mamba replaces only the attention layer. It replaces both attention AND the MLP. A standard Transformer decoder block has two sub-layers: multi-head self-attention (which mixes information across sequence positions) and a feed-forward network (which mixes information across feature channels). The Mamba block handles both in a single, unified structure.

Here is how the Mamba block works. The input ($B \times L \times D$) passes through a LayerNorm and is linearly projected to expand the feature dimension by a factor of $E = 2$. This expanded representation is then split into two parallel branches:

The SSM branch (left): Processes through a short 1D causal convolution (width 4) to capture immediate local patterns between neighboring tokens. Then a SiLU activation. Then three parallel linear projections produce the token-specific $\Delta_t$, $B_t$, and $C_t$. The selective SSM recurrence runs using these dynamic parameters. This branch handles sequence mixing: how information flows across token positions.

The gate branch (right): Takes the other half of the expanded input and passes it through a SiLU activation. This branch serves as a dynamic gate that controls which channels of the SSM output are passed through and which are suppressed.

The two branches merge via element-wise multiplication. If elements in the gating vector are near zero, the corresponding SSM information is suppressed. The result passes through a linear projection back to dimension $D$ and is added to the input via a residual connection.

The entire block is one homogeneous module. No separate attention layer. No separate MLP.

Inference: The Fundamental Trade-off

The trade-off is fundamental. The Transformer’s KV cache stores every token (lossless, but $O(L)$ memory). Mamba’s hidden state compresses all history into a fixed vector (lossy, but $O(1)$ memory). The question is always whether the compressed representation is good enough for the task.

Mamba-1 demonstrated that it was. Mamba-2.8B matched or exceeded Pythia-6.9B (a model more than twice its size) on zero-shot downstream evaluations. On the Pile dataset, Mamba-1.4B achieved 59.7% average across common-sense reasoning benchmarks, matching Pythia-2.8B (59.1%). At batch size 16, Mamba-2.8B completed generation in 18.6 seconds versus GPT-Neo-2.7B’s 65.9 seconds (3.5x faster). GPT-Neo ran out of memory at batch size 32 on a 64GB GPU; Mamba continued scaling to batch 128+.

Part 7: Mamba-2, Maximizing GPU Utilization

Mamba-1 had an embarrassing practical problem: it was 2-3x slower than equivalently sized Transformers during training. The root cause is that modern GPUs deliver roughly 16x more throughput for matrix multiplication (via Tensor Cores) than for general arithmetic. Transformers are pure matmul. Mamba-1’s selective scan was not.

Tri Dao and Albert Gu’s May 2024 paper “Transformers are SSMs” solved this by proving that unrolling the SSM recurrence produces a structured matrix that can be computed via matrix multiplications. The resulting algorithm (SSD) splits the sequence into chunks: within each chunk, the computation runs as dense matmuls on Tensor Cores; between chunks, a short scan passes state forward. Training speed improved 2-8x over Mamba-1.

The trade-off: to fit this matrix framework, $A$ is restricted from a diagonal matrix (Mamba-1) to a scalar-times-identity (Mamba-2), meaning all state dimensions within a head share one decay rate. Mamba-2 compensates with a multi-head structure and increases the state dimension from $N = 16$ to $N = 64\text{-}256$.

Part 8: Mamba-3, Three Innovations from Classical SSM Theory

Published at ICLR 2026, Mamba-3 asks a different question than its predecessors. Mamba-2 optimized for training speed by simplifying the SSM to leverage Tensor Cores. But with the rise of RL post-training, agentic workflows, and test-time compute scaling, inference efficiency has become the primary bottleneck.

Here is the problem Mamba-3 targets. During autoregressive decoding, Mamba-2’s simplified recurrence is memory-bound. The GPU loads the state from HBM to SRAM, performs a trivially small computation (the scalar-times-identity update is cheap), and writes the result back. The arithmetic intensity is roughly 2.5 ops/byte. The H100 needs $\sim$295 ops/byte to be compute-bound. More than 99% of GPU compute sits idle during token generation.

Mamba-3’s overarching philosophy is to increase arithmetic intensity during decoding by making the state update mathematically richer, spending more compute per byte of memory traffic, filling idle GPU cycles rather than adding new ones. Three innovations accomplish this.

Innovation 1: Exponential-Trapezoidal Discretization

The problem. Mamba-1 and Mamba-2 used what the Mamba-3 authors retroactively classify as “Exponential-Euler” discretization: the exact formula $\bar{A} = \exp(\Delta A)$ paired with the first-order Euler approximation $\bar{B} = \Delta B$. This is a hybrid: exact for the state decay, but approximate for how the input enters the state. The local truncation error is $O(\Delta^2)$.

In numerical analysis terms, the Euler method approximates the area under a curve using a rectangle aligned to one endpoint. It captures the value at the start of the interval but ignores how the curve changes across the interval. This crude approximation struggles with fast-moving temporal dependencies, producing “jerky” transitions in the state.

In practice, prior Mamba models compensated by adding an explicit short 1D causal convolution (Conv1d, width 4) before the SSM. This Conv1d smoothed out immediate local token interactions that the imprecise discretization missed. It worked, but it was an architectural bandage for a mathematical shortcoming. And it added latency at inference: one more sequential operation per token.

The intuition. The trapezoidal rule approximates the area under a curve using a trapezoid instead of a rectangle. A rectangle uses only one endpoint’s value. A trapezoid uses both endpoints and draws a straight line between them, capturing the slope of the curve across the interval. This gives second-order accuracy: the local error drops from $O(\Delta^2)$ to $O(\Delta^3)$.

The math. Applying the generalized trapezoidal rule to the SSM’s state equation produces a three-term recurrence, where the old Euler formula had only two:

Term 1 is identical to the old formula: decay the previous state. Term 3 is similar to the old Euler input term, but weighted by $\lambda_t$ instead of 1. The new addition is Term 2: the previous timestep’s input $x_{t-1}$, projected through $B_{t-1}$, scaled by $(1 - \lambda_t)$, and decayed by the same exponential factor as the state.

The parameter $\lambda_t$ is a data-dependent convex combination weight. The $(1-\lambda_t)$ and $\lambda_t$ coefficients on consecutive inputs are the weights of the trapezoid: $(1-\lambda)$ on the left endpoint (previous input), $\lambda$ on the right endpoint (current input).

Let me verify the connection to the old formula. When $\lambda_t = 1$: Term 2 vanishes entirely because its coefficient $(1-\lambda_t) = 0$. Term 3 becomes $1 \cdot \Delta_t \cdot B_t \cdot x_t = \Delta_t B_t x_t$, which is exactly the Euler formula $\bar{B}x_t$. So the old Mamba-1/2 discretization is the special case $\lambda = 1$ of this more general formula.

A numerical walkthrough. Take $a = -0.5$, $\Delta = 0.1$, $b = 1$, and process the input sequence $[1, 1, 0]$ with $\lambda = 0.5$ (balanced trapezoidal blending):

# Parameters: a = -0.5, delta = 0.1, b = 1.0, lambda = 0.5
# exp(delta * a) = exp(-0.05) = 0.9512

# k=0: input=1 (no previous input, Term 2 uses x_{-1}=0)
#   Term 1: 0.9512 * 0 = 0                              (no state yet)
#   Term 2: (1-0.5) * 0.1 * 0.9512 * 1.0 * 0 = 0        (no prev input)
#   Term 3: 0.5 * 0.1 * 1.0 * 1 = 0.05
#   h_0 = 0 + 0 + 0.05 = 0.0500

# k=1: input=1, prev_input=1
#   Term 1: 0.9512 * 0.05 = 0.04756                     (decayed state)
#   Term 2: 0.5 * 0.1 * 0.9512 * 1.0 * 1 = 0.04756     (prev input, decayed)
#   Term 3: 0.5 * 0.1 * 1.0 * 1 = 0.05                  (current input)
#   h_1 = 0.04756 + 0.04756 + 0.05 = 0.14512

# k=2: input=0, prev_input=1
#   Term 1: 0.9512 * 0.14512 = 0.13804                  (decayed state)
#   Term 2: 0.5 * 0.1 * 0.9512 * 1.0 * 1 = 0.04756     (prev input=1, decayed)
#   Term 3: 0.5 * 0.1 * 1.0 * 0 = 0.0                   (current input=0)
#   h_2 = 0.13804 + 0.04756 + 0.0 = 0.18560

# Compare to Euler (lambda=1):
# k=2 with Euler: h = 0.9512 * 0.1904 + 0.0976 * 0 = 0.18107

At step 2, the trapezoidal version ($h = 0.1856$) is higher than the Euler version ($h = 0.1811$). The difference comes from Term 2: even though the current input is 0, the trapezoidal rule still accounts for the previous input ($x_1 = 1$) via the left endpoint of the trapezoid. The Euler method ignores this entirely. For fast-changing input sequences, this difference matters.

The implicit convolution and the death of Conv1d. Here is the subtle and important consequence. Because the state update at time $t$ depends on both $x_t$ and $x_{t-1}$, the trapezoidal recurrence contains an implicit convolution of width 2. The $(1-\lambda_t)$ and $\lambda_t$ weights on the consecutive inputs play the role of a learned, data-dependent convolution filter operating on pairs of adjacent tokens.

For years, SSM architectures (H3, RWKV-4, Mamba-1, Mamba-2) required an explicit external Conv1d (width 4) before the SSM to handle immediate local token interactions. The Conv1d was considered essential for capturing “induction head” copying behaviors and local patterns. Mamba-3 found that the implicit width-2 convolution from the trapezoidal discretization, combined with learnable bias terms on $B$ and $C$ (constant vectors added after normalization), is expressive enough to replace the external Conv1d entirely. Mamba-3 is the first Mamba variant to drop the Conv1d without performance loss.

Why it matters for your H100. Fewer sequential operations per token at inference. No Conv1d kernel launch, no Conv1d memory traffic, no Conv1d compute. The architecture is simpler and the discretization is now theoretically justified ($O(\Delta^3)$ error) instead of a heuristic patched by an external convolution.

Innovation 2: Complex-Valued SSMs via RoPE

The problem. Real-valued SSMs with non-negative eigenvalues can only decay monotonically. The state gets smaller over time, or stays the same, but it cannot oscillate. Mathematically: if $\bar{a} \in [0, 1]$, then $\bar{a}^k$ is a monotonically decreasing sequence. The state can only move in one direction (toward zero).

This means real-valued SSMs cannot solve simple state-tracking tasks that require flipping between states. Consider parity: given a stream of bits, track whether the running count of 1s is even or odd. Every time a 1 arrives, the parity flips. This requires the state to toggle between two values indefinitely. A monotonically decaying state cannot do this. On the bit sequence parity task, Mamba-2 scored no better than random guessing.

The intuition. Real eigenvalues restrict the state to movement along a line: it can only grow or shrink. Complex eigenvalues enable rotation: the state can cycle through values, oscillate, and flip. For even/odd tracking, you need the state to flip sign every time a 1 appears. A 180-degree rotation achieves exactly this. Real, non-negative arithmetic cannot.

The clever trick. Implementing complex arithmetic on GPUs is painful. Complex numbers double memory requirements, break existing CUDA kernel optimizations, and introduce alignment issues. Mamba-3 avoids all of this through a mathematical equivalence.

The key theoretical result (Proposition 3 in the paper): a discretized complex-valued diagonal SSM is mathematically equivalent to a real-valued SSM with data-dependent Rotary Positional Embeddings (RoPE) applied to $B$ and $C$. The decomposition works as follows:

• The real part of the complex eigenvalue controls decay. This is handled by the existing SSD machinery, exactly as in Mamba-2. No changes needed.
• The imaginary part controls rotation. This is factored out and implemented as rotary embeddings applied to the $B$ and $C$ projection vectors.

The rotation angles are produced dynamically via projections from the current input token $x_t$, rather than using static positional indices as in standard Transformer RoPE. This is why it is called “data-dependent” RoPE. The rotation applied to $B$ and $C$ changes based on what token is being processed.

No complex number ever appears in the GPU kernels. The real-valued SSD computation runs at the same speed as before. The rotational dynamics are absorbed into $B$ and $C$ via the same RoPE infrastructure that Transformers already use for positional encoding. Existing Transformer tooling (rotary embedding kernels, fused attention implementations) can be reused directly.

The result. On synthetic state-tracking benchmarks:

Mamba-3 solves parity perfectly and near-perfectly executes complex modular arithmetic. Standard (non-data-dependent) RoPE does not help. Static positional rotation angles cannot implement content-dependent state flipping. The data-dependent version can, because the rotation is a function of the input.

These are tasks that were mathematically impossible for real-valued SSMs, regardless of scale or training budget. The real-complex boundary is a hard expressivity ceiling, not a soft scaling issue.

Innovation 3: MIMO (Multi-Input Multi-Output)

The problem. As discussed above, standard SSMs waste more than 99% of GPU compute during decoding because each state update is a trivial rank-1 operation: the GPU loads the entire state from memory, performs a single multiply-add, and writes it back. The computation is too cheap relative to the memory transfer.

The fix. Instead of processing one input and producing one output per SSM (SISO), process $R$ inputs and $R$ outputs simultaneously (MIMO). The scalar input $x_t$ is linearly projected into a matrix $X_t$ with rank $R$. The projection vectors $B_t$ and $C_t$ are correspondingly expanded to rank-$R$ structures. The state update becomes a matrix multiplication instead of an outer product:

With $R = 4$, the model performs 4x the floating-point operations for the same amount of memory traffic. The arithmetic intensity jumps from $\sim$2.5 to $\sim$10 ops/byte. Still not enough to fully saturate the H100, but a 4x improvement in GPU utilization during the memory-bound decode phase.

Crucially, only the SSM-specific parameters ($B_t$, $C_t$, and the state $H_t$) grow with $R$. The main input projections, the output projections, and the residual gate all remain at their original sizes. This contains the parameter increase to the SSM core.

Why it does not hurt latency. The extra compute fills idle GPU cycles. During decoding, the bottleneck is the time it takes to load the state from HBM to SRAM. While that data transfer is in flight, the Tensor Cores have nothing to do. MIMO gives them work. The wall-clock time per decode step is dominated by memory transfer time, not compute time, so adding compute within the transfer window is effectively free. MIMO with $R = 4$ matches Mamba-2’s decode speed while delivering substantially better accuracy.

The result. At 1.5B scale with Chinchilla-optimal training: the base Mamba-3 (SISO) outpaces Gated DeltaNet (the previous state-of-the-art sub-quadratic model) by 0.6 percentage points on average downstream accuracy. Adding MIMO with $R = 4$ adds another 1.2 points, for a total gain of 1.8 points over Gated DeltaNet, 1.9 over Mamba-2, and 2.2 over equivalently-sized pure Transformers. The 1.5B MIMO variant achieves 57.6% average accuracy across benchmarks.

A Mamba-3 MIMO model with state dimension $N = 64$ matches the perplexity and downstream accuracy of a Mamba-2 model with $N = 128$. Halving the state size while maintaining quality doubles inference throughput within the same hardware footprint.

Architectural Changes

Two more structural changes round out Mamba-3:

Normalization. Mamba-3 replaces post-gate RMSNorm (Mamba-2) with QKNorm, also called BCNorm: RMS normalization applied directly to the $B$ and $C$ projections before mixing. This stabilizes variance and activation spikes during large-scale pretraining, which is especially important with the added mathematical complexity of trapezoidal recurrence and MIMO updates.

Block structure. Mamba-1 and Mamba-2 fused the sequence mixer (SSM) and channel mixer (MLP) into a single homogeneous block. Mamba-3 reverses this decision. It adopts an interleaved architecture that matches the Llama family: alternating Mamba-3 SSM blocks with standard SwiGLU MLP blocks. Each Mamba-3 block handles sequence mixing; each SwiGLU MLP handles channel mixing. This Llama-compatible topology makes it straightforward to create hybrid models by swapping some SSM blocks for attention blocks.

Evolution Comparison

Part 9: The Bigger Picture

Hybrid Architectures Are the Production Standard

The field has converged on an empirical finding: hybrid architectures combining SSM layers with a small fraction of attention layers outperform both pure approaches. Albert Gu has articulated the fundamental reason clearly. Transformers are like databases: they cache every token for future reference (perfect recall, but linear memory growth). SSMs are like brains: they compress all history into a fixed-size state (infinite context, but lossy).

Pure SSMs struggle with two specific capabilities. Exact retrieval: finding a specific fact buried in a long context degrades as context grows, because the fixed-size state cannot perfectly memorize arbitrary content. In-context learning: few-shot pattern matching from prompt examples requires comparing the current token against specific stored tokens, which is fundamentally an attention operation.

The solution adopted by every major lab: use SSM layers for the vast majority of the network and sprinkle in a few attention layers for precise retrieval. The exact ratio varies (5:1 in Mamba-3’s recommended config, 9:1 in Granite), but the pattern is universal.

What This Means for Your Inference Stack

For inference infrastructure teams, the implications are concrete. With only 10-15% of layers using attention, you manage KV cache for those few layers, not the entire network. SSM layers need no KV cache management at all: no PagedAttention, no eviction policies, no memory fragmentation. And throughput advantages grow with context length. Pure Transformers are faster at short sequences ($<$2K tokens), but SSM-based models cross over quickly and the gap widens: at 57K tokens, Mamba-2 outperforms Transformers by 4x. SSM decode cost is constant per token; Transformer decode cost grows linearly.

The Fundamental Trade-off Persists

SSMs compress. Attention caches. Mamba-3 makes the compressed memory more expressive through complex dynamics, higher-order discretization, and MIMO. But it cannot eliminate the compression. If your workload requires perfect verbatim retrieval of a specific sentence from a 100K-token document, you need attention layers for that.

The Transformer monopoly has ended. But Transformers are not dead. They are becoming a specialized, strategically-placed component within a larger hybrid architecture, used precisely where their lossless memory is needed and nowhere else.

References

1. Gu, A., Dao, T., Ermon, S., Rudra, A., & Re, C. (2020). HiPPO: Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020.

2. Gu, A., Goel, K., & Re, C. (2022). Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022. (S4)

3. Gu, A., Gupta, A., Goel, K., & Re, C. (2022). On the Parameterization and Initialization of Diagonal State Space Models. NeurIPS 2022. (S4D)

4. Gu, A. & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. COLM 2024.

5. Dao, T. & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024. (Mamba-2 / SSD)

6. Lahoti, A., Li, A., Chen, Y., Wang, Z., Bick, T., Kolter, J. Z., Dao, T., & Gu, A. (2026). Mamba-3: Improved Sequence Modeling using State Space Principles. ICLR 2026.

7. Together AI. Mamba-3 Blog Post. Technical overview and benchmark results.

8. Goomba Lab. Blog series on Structured State Space Duality and Mamba-3 mathematical foundations.

9. Tri Dao. Mamba-3 Part 2: Methodological Deep Dive. Detailed derivation of the exponential-trapezoidal discretization and RoPE equivalence.

10. Princeton Language and Intelligence. Mamba-2: Algorithms and Systems. Technical walkthrough of the SSD algorithm.

11. NVIDIA. Nemotron-3-Super: hybrid Mamba-2 + MoE + Attention architecture for production deployment.

12. IBM. Granite 4.0: 9:1 Mamba-to-Transformer ratio with >70% memory reduction vs. conventional LLMs.

In our post on speculative decoding, we covered how a small draft model proposes tokens that a large target model verifies in parallel, achieving 2-3x speedups without changing the output distribution. That technique exploits idle GPU compute during memory-bound inference.

This post examines a follow-up question: can we make speculative decoding itself faster? The answer is yes. A recent paper by Kumar, Dao, and May (ICLR 2026) identifies a sequential bottleneck within standard speculative decoding and eliminates it through a technique called Speculative Speculative Decoding (SSD). Their algorithm, Saguaro, achieves up to 5x speedup over autoregressive decoding and roughly 2x over optimized speculative decoding.

We will walk through the bottleneck SSD targets, the core idea of speculating about verification outcomes, and the three engineering challenges that Saguaro solves: cache construction, cache-aware sampling, and batch-size scaling. Each section explains why the naive approach fails before presenting the solution.

The Hidden Bottleneck in Speculative Decoding

Standard speculative decoding runs in a loop: the draft model generates K tokens, the target model verifies them in a single forward pass, and the process repeats. This is faster than autoregressive decoding because verification amortizes the expensive memory read of the target model’s weights across multiple tokens.

But there is a sequential dependency hiding in plain sight. Let’s trace what happens on the draft model’s GPU during one round:

1. The draft model generates K tokens (draft phase)
2. The draft model sends tokens to the target model
3. The target model runs verification (the draft model sits completely idle)
4. The target model returns the verification outcome
5. The draft model generates K new tokens for the next round
6. Go to step 2

The draft model does nothing during step 3. If verification takes $T_v$ time units, the draft model wastes $T_v$ time units every round. Since the target model is much larger than the draft, $T_v$ dominates the round duration. The draft model’s GPU is idle for the majority of each round.

This is the bottleneck SSD eliminates. Instead of waiting for verification to finish, the draft model spends that idle time doing something useful: predicting what the verification outcome will be and pre-computing the next round’s speculation for each likely outcome.

The Core Idea: Speculate About the Speculation

The idea draws from CPU speculative execution. When a CPU encounters a conditional branch, it does not wait for the condition to resolve. Instead, it predicts the likely outcome and begins executing instructions along that predicted path. If the prediction was correct, the results are kept. If wrong, they are discarded and the correct path is executed.

SSD applies this same principle to speculative decoding. While the target model verifies round $T$’s draft tokens, the draft model:

1. Predicts what the verification outcome will be
2. Pre-computes speculations for each likely outcome
3. Stores these in a speculation cache

When verification finishes, the actual outcome is compared against the cache. If it matches a cached prediction (a cache hit), the next round’s speculation is returned instantly with zero drafting latency. If it doesn’t match (a cache miss), the system falls back to standard synchronous drafting.

What Is a Verification Outcome?

To understand what we need to predict, let’s define the verification outcome precisely. When the target model verifies K draft tokens, two things are determined:

• $k$: the number of accepted draft tokens (ranging from 0 to K)
• $t^*$: the bonus token, sampled from the target distribution at the first disagreement point (or at position K+1 if all tokens are accepted)

The verification outcome is the pair $v^T = (k, t^*)$. This fully determines the context from which the next round’s speculation must begin. If we can predict $v^T$ before verification completes, we can pre-compute the next K draft tokens starting from that context.

The Speculation Cache

The speculation cache $S^T$ is a dictionary that maps predicted verification outcomes to pre-computed speculations:

Each entry contains K draft tokens generated autoregressively by the draft model starting from the context implied by $(k, t^*)$.

When verification returns the actual outcome $v^T$:

• Cache hit ($v^T \in S^T$): Return the cached speculation immediately. Zero drafting latency for this round.
• Cache miss ($v^T \notin S^T$): Fall back to generating a fresh speculation synchronously.

The key question is obvious: how do we choose which outcomes to cache? The space of possible outcomes is $(K+1) \times |V|$ where $|V|$ is the vocabulary size (typically 32,000-128,000). We cannot pre-compute speculations for all of them. This is the first of three challenges Saguaro addresses.

The Speedup Formula

Before diving into the three challenges, let’s quantify the potential. The expected speedup of SSD over autoregressive decoding (Theorem 7 from the paper) is:

Walking through each term: $p_{\text{hit}}$ is the probability of a cache hit. $E_{\text{hit}}$ and $E_{\text{miss}}$ are the expected number of tokens generated per round on a hit and miss respectively. $T_p$ is the latency of the primary speculator (the neural draft model) relative to the verifier, and $T_b$ is the latency of the backup speculator used on cache misses.

The numerator is the expected tokens per round, weighted by hit/miss probabilities. The denominator is the expected wall-clock time per round.

Two corollaries follow directly:

Corollary 8: SSD strictly outperforms standard speculative decoding whenever $p_{\text{hit}} > 0$. Any nonzero cache hit rate improves performance because cache hits eliminate drafting latency entirely, and cache misses simply revert to standard SD behavior.

Corollary 9: The maximum speedup is bounded by $(1 + T_{\text{SD}}) \cdot (E_{\text{hit}} / E_{\text{SD}})$, where $T_{\text{SD}}$ and $E_{\text{SD}}$ are the drafting time and expected tokens for standard SD. When the cache hit rate approaches 1, all drafting latency vanishes and we gain a factor of $(1 + T_{\text{SD}})$ in the denominator.

Challenge 1: Building the Cache (Saguaro Cache Construction)

The Problem

A verification outcome is a pair $(k, t^*)$. The acceptance length $k$ ranges from 0 to K (typically K=7 in SSD), giving K+1 possibilities. The bonus token $t^*$ comes from a vocabulary of size $|V|$. The total outcome space is $(K+1) \times |V|$, which is roughly 250,000 for K=7 and $|V|$=32,000.

We have a budget of $B$ speculations we can pre-compute during the verification window. Each speculation requires the draft model to run K autoregressive steps. With the verification latency of a 70B target model on 4 H100s, we can fit roughly $B = 20\text{-}30$ speculations. We need to choose wisely.

Decomposing the Problem

Saguaro decomposes this into two subproblems:

1. For each acceptance length $k$, how many bonus tokens should we cache? Call this the fan-out $F_k$.
2. For a given fan-out $F_k$, which bonus tokens should we cache?

The second question has a straightforward answer: use the top-$F_k$ tokens from the draft model’s own probability distribution at that position. The draft model has already computed logits during the current round’s speculation, so the most likely tokens are immediately available. Empirically, this predicts the actual bonus token with up to 90% accuracy at moderate fan-out.

The first question, how to allocate fan-out across positions, is where the interesting optimization happens.

Power-Law Cache Hits

The authors make a key empirical observation: cache miss probability follows a power law in the fan-out:

for some exponent $r > 0$ that depends on the draft-target alignment. This means that doubling the fan-out does not halve the miss rate. Instead, miss rate decreases polynomially, with diminishing returns as fan-out grows. This finding (confirmed across multiple model pairs and datasets) drives the allocation strategy.

Geometric Fan-Out

Given the power-law structure and a total budget $\sum_{k=0}^{K} F_k \leq B$, Saguaro solves a constrained optimization problem using Lagrange multipliers. The result (Theorem 12) is a geometric allocation:

where $\alpha_p$ is the per-token acceptance rate and $F_0$ is determined by the budget constraint. The formula allocates more fan-out to earlier positions (small $k$) and less to later positions.

The reasoning: position $k=0$ (first token rejected) is more probable than $k=5$ (five tokens accepted before rejection) because each acceptance is an independent event with probability $\alpha_p < 1$. The probability of reaching acceptance length $k$ is roughly $\alpha_p^k \cdot (1 - \alpha_p)$, a geometric distribution. Allocating fan-out proportionally to the probability of each outcome maximizes the expected cache hit rate.

There is one exception: position $k=K$ (all tokens accepted) receives a boost. When all K draft tokens are accepted, the bonus token comes directly from the target model’s distribution $p_{\text{target}}$ rather than the residual distribution. The target distribution is sharper and more concentrated, making the top-$F$ prediction more accurate. Saguaro accounts for this with a multiplicative correction:

Challenge 2: Trading Acceptance for Cache Hits (Saguaro Sampling)

The Tension

There is a fundamental tension between two objectives in SSD. On one hand, we want the draft model to closely match the target model so that acceptance rates stay high. On the other hand, we want to predict the bonus token $t^*$ accurately so that cache hit rates stay high.

These objectives conflict. Here is why.

Recall from our speculative decoding post that the bonus token is sampled from the residual distribution:

When the draft model closely matches the target ($p_{\text{draft}} \approx p_{\text{target}}$), the acceptance rate is high but the residual $\max(p_{\text{target}} - p_{\text{draft}}, 0)$ is spread thinly across many tokens. A thin residual means the bonus token could be almost anything, making it hard to predict and reducing cache hit rates.

When the draft model diverges from the target, the residual concentrates on tokens where $p_{\text{target}} \gg p_{\text{draft}}$, making the bonus token more predictable. But acceptance rates drop, meaning fewer tokens per round.

The Solution: Intentional Misalignment

Saguaro sampling resolves this tension by deliberately modifying the draft model’s sampling distribution. For a set of cached tokens (the top-$F$ tokens at each position), Saguaro suppresses the draft model’s probability on those specific tokens:

Here $z$ is the vector of draft model logits, $F$ is the fan-out, and $C \in [0,1]$ is a downweighting constant. When $C=1$, this is the standard softmax with no modification. When $C < 1$, the cached tokens receive reduced probability in the draft’s distribution.

Why does this help? Let’s trace the effect on the residual.

When the draft model assigns less probability to a cached token, the gap $p_{\text{target}}(\cdot) - p_{\text{draft}}(\cdot)$ becomes larger for that token. A larger gap means the residual distribution assigns more probability to that token. Saguaro is steering the entire residual distribution to concentrate on the exact tokens it has cached.

Theorem 15 from the paper confirms this formally: the cache hit rate $p_{\text{hit}}$ increases monotonically as $C \to 0$. Push $C$ all the way to zero and the residual distribution is forced entirely onto cached tokens (guaranteeing a hit), but the acceptance rate collapses because the draft distribution diverges maximally from the target.

The optimal $C^*$ balances these competing effects and depends on the sampling temperature:

• Temperature 0 (greedy decoding): $C = 1$ is optimal. The bonus token is deterministic (the argmax of the target distribution), so the top-1 draft prediction already has high accuracy. No need to sacrifice acceptance rate.
• High temperature: $C \ll 1$ becomes advantageous. The bonus token is sampled from a flatter distribution, making it harder to predict without help. Saguaro sampling concentrates the residual, recovering cache hit rates that would otherwise be low.

In practice, Saguaro sampling provides up to 50% additional end-to-end speedup at high temperatures compared to using $C=1$.

Challenge 3: Scaling to Large Batches (Saguaro Fallback)

The Batch Size Problem

Everything so far has assumed batch size 1 (a single sequence). At larger batch sizes, a new problem emerges.

With batch size $b$, the system can only proceed to the next round when all $b$ sequences have speculations ready. The probability that every sequence in the batch gets a cache hit is:

Even with $p_{\text{hit}} = 0.9$ per sequence, a batch of 16 sequences gives $P(\text{all hit}) = 0.9^{16} \approx 0.19$. At batch size 32, it drops to $0.9^{32} \approx 0.03$. The probability of at least one cache miss grows exponentially, and a single miss stalls the entire batch.

The Naive Fallback Fails

The intuitive solution is simple: when a cache miss occurs, have the primary draft model generate a fresh speculation on the spot. But this is catastrophically bad at scale.

The primary draft model is still a neural network that generates tokens autoregressively. Generating K tokens takes non-trivial time. While this one stalled sequence catches up, every other sequence in the batch (including those with cache hits) waits. The batch is only as fast as its slowest member.

Corollary 16 in the paper formalizes this: at large batch sizes, the SSD speedup becomes overwhelmingly bounded by the fallback latency. If the fallback speculator is slow, the theoretical gains from caching vanish.

Dual-Tier Fallback

Saguaro solves this with a dual-tier fallback strategy controlled by a critical batch size $b^*$:

Below $b^*$ (low batch regime): Cache misses are infrequent. The primary draft model serves as its own fallback, generating a fresh speculation synchronously. The latency penalty is acceptable because misses are rare and each one affects only one sequence.

Above $b^*$ (high batch regime): The system switches to an ultra-fast backup speculator with minimal latency. This could be an n-gram model, random tokens, or a token frequency model. The backup’s speculations will likely be rejected during verification (random tokens have near-zero acceptance rate). But the strategic insight is that the latency cost of feeding bad speculations to one sequence is vastly less than the latency cost of making the entire batch wait for a neural draft model.

Theorem 17 proves this formally: accepting a single sequence’s quality penalty (poor speculation) is strictly better than inflicting the primary drafter’s latency across the entire batch.

The critical batch size $b^*$ is derived analytically from the speedup equation and depends on $p_{\text{hit}}$, $T_p$ (primary drafter latency), and $T_b$ (backup latency).

The Full Algorithm

SSD runs three concurrent processes: a main coordinator, a verifier, and a speculator. Here is the algorithm:

# Main: launches speculator asynchronously, runs verifier
def main(prompt, target, primary_draft, backup_draft):
    launch_async(speculator, prompt, primary_draft, backup_draft)
    return verifier(prompt, target)

# Verifier: runs on target model's GPUs (e.g., 4x H100)
def verifier(prompt, target):
    target.prefill(prompt)
    spec_tokens = RECEIVE()              # wait for first speculation
    generated = []

    while True:
        outcome = target.verify(spec_tokens)   # standard SD verification
        generated.extend(outcome.tokens)

        if EOS in outcome.tokens:
            return generated

        SEND(outcome)                          # send (k, t*) to speculator
        spec_tokens = RECEIVE()                # wait for next speculation

# Speculator: runs on draft model's GPU (e.g., 1x H100)
def speculator(prompt, primary_draft, backup_draft):
    primary_draft.prefill(prompt)
    spec_tokens = primary_draft.speculate(prompt)  # initial speculation

    while True:
        SEND(spec_tokens)                      # send to verifier

        # While verification runs, build the cache
        cache = build_speculation_cache(
            spec_tokens, primary_draft           # §4.1: geometric fan-out
        )

        outcome = RECEIVE()                     # get actual (k, t*)

        if EOS in outcome:
            return

        if outcome in cache:                    # CACHE HIT
            spec_tokens = cache[outcome]        # instant return
        else:                                   # CACHE MISS
            spec_tokens = fallback(             # §4.3: dual-tier
                outcome, primary_draft, backup_draft
            )

Step-by-Step: One Round

Let’s trace through a single round to see how the pieces fit together.

Step 1: The speculator sends K draft tokens to the verifier and immediately begins predicting outcomes. It examines the draft model’s logits at each position to determine the most likely bonus tokens.

Step 2: Using geometric fan-out (Challenge 1), the speculator allocates its budget across positions. For each predicted outcome $(k, t^*)$, it generates K new draft tokens autoregressively from that context.

Step 3: These speculations are stored in the cache. With K=7 and fan-out F=3, the cache contains roughly 24 entries (8 acceptance lengths, with 3 bonus token predictions each, weighted by geometric allocation).

Step 4: The verifier finishes and sends back the actual outcome $(k, t^*)$.

Step 5: Cache lookup. If the outcome matches, the cached speculation is returned instantly. If not, the fallback mechanism kicks in.

The critical property is that on a cache hit, the round latency equals only the verification time (because the speculation was pre-computed in parallel). In standard SD, the round latency equals verification time plus drafting time. This is where the speedup comes from.

Correctness: SSD Is Lossless

An important guarantee: SSD produces the same output distribution as standard speculative decoding (which itself matches autoregressive decoding exactly).

On a cache hit, the cached speculation is verified using the same rejection sampling mechanism as standard SD. The fact that the speculation was pre-computed rather than computed just-in-time changes nothing about verification correctness.

On a cache miss, the system falls back to standard synchronous SD, which is known to be lossless.

The speculation cache is a performance optimization that never affects what tokens the target model accepts or rejects. Pre-speculation only changes when the draft tokens are computed, not what they are or how they are verified.

Hardware Setup

SSD requires the draft and target models to run on separate GPUs so they can operate concurrently. The typical configuration:

• Target model (verifier): 4x H100 80GB GPUs with tensor parallelism (for Llama-3.1-70B)
• Draft model (speculator): 1x H100 80GB GPU on a separate device
• Total: 5 GPUs for SSD vs. 4 GPUs for standard SD/AR

This is a 25% increase in hardware. The question is whether the speedup justifies the extra cost. At batch size 1, SSD achieves roughly 2x higher throughput than SD with the same target model, meaning the throughput per GPU still improves substantially.

Results

The authors benchmark SSD (with Saguaro optimizations) against autoregressive decoding and standard speculative decoding across four datasets: HumanEval (code), Alpaca (chat), GSM8K (math), and UltraFeedback (general).

Setup: Llama-3.1-70B-Instruct as the target model on 4x H100 GPUs, Llama-3.2-1B-Instruct as the draft model on 1x H100 GPU. K=6 for SD, K=7 for SSD, fan-out F=3 for SSD.

Key findings:

• SSD vs. autoregressive: Up to ~5x faster (e.g., 255.8 tok/s vs. 54.7 tok/s on some benchmarks)
• SSD vs. standard SD: Up to ~2x faster (255.8 vs. 161.8 tok/s)
• At larger batch sizes: SSD still provides ~20% improvement over SD, even as cache hit rates drop. The Saguaro optimizations push the throughput-latency Pareto frontier across all batch sizes.
• Temperature sensitivity: Cache hit rates decrease with sampling temperature, but Saguaro sampling (low $C$) compensates effectively, maintaining gains at high temperatures.

The results also confirm on Qwen-3-32B (target) with Qwen-3-0.6B (draft): 203.8 tok/s for SSD vs. 136.8 for SD vs. 88.8 for AR.

Where SSD Fits in the Landscape

SSD is not the first attempt to overlap drafting and verification. Several concurrent methods target the same bottleneck:

SSD is also orthogonal to several other inference optimizations, meaning they can be combined:

• EAGLE/EAGLE-2: Feature-level draft prediction. SSD could use an EAGLE-style drafter as its speculator.
• Tree-based verification (Sequoia, SpecInfer): Verify multiple candidates in one pass. SSD parallelizes the draft-verify loop itself, a different axis.
• Non-neural speculators (SuffixDecoding, Infini-gram): Could serve as SSD’s fast backup speculator for cache misses.

When to Use SSD (and When Not To)

SSD is a strong fit when:

• Latency matters more than throughput (real-time chat, interactive coding)
• Batch sizes are small to moderate ($b \leq b^*$)
• You have spare GPU capacity for a separate draft model
• You are already using speculative decoding and want to push further

SSD is not ideal when:

• You are throughput-bound (large-scale RL, offline batch generation). SSD optimizes per-request latency, not aggregate throughput at high concurrency.
• Hardware is constrained. The extra GPU for the draft model is not available.
• Batch sizes are consistently very large. Cache hit rates decay exponentially with batch size, and the gains narrow.

Looking Forward

SSD makes a compelling case that the sequential dependencies in LLM inference are not fixed constraints but engineering surfaces that can be optimized. The draft-verify loop in speculative decoding seemed inherently sequential. It turned out that the sequential part (waiting for verification before starting the next draft) could be hidden by speculating about the verification outcome.

This pattern of applying a technique recursively to itself is worth paying attention to. The authors frame SSD as “nested speculation,” and the natural question is whether another level of nesting could help. The answer is likely no for now (the overhead of a third speculation level would exceed the marginal benefit), but the thinking is instructive: whenever two stages of a pipeline are sequential, ask whether one stage can predict the other’s output and pre-compute accordingly.

The practical significance is clear. For latency-sensitive applications running large models, SSD with Saguaro optimizations roughly doubles the speedup of speculative decoding at modest hardware cost. As inference frameworks like NVIDIA Dynamo adopt disaggregated architectures (separate prefill and decode stages on different hardware), SSD’s separate-GPU design fits naturally into that direction.

References

1. Kumar, T., Dao, T., & May, A. (2026). Speculative Speculative Decoding. ICLR 2026.

   • The SSD paper introducing Saguaro and the three core optimizations.

2. Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023.

   • The original speculative decoding paper with distribution preservation proofs.

3. Chen, C., Borgeaud, S., Irving, G., Lespiau, J. B., Sifre, L., & Jumper, J. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv preprint.

   • Independent discovery of speculative decoding at DeepMind.

4. Li, Y., Cai, T., Zhang, Y., Chen, D., & Dai, D. (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. ICML 2024.

   • Feature-level speculation achieving superior speedups through hidden state prediction.

5. Chen, Z., Yang, X., Lin, J., Sun, C., Huang, J., & Chang, K. W. (2024). MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding. arXiv preprint.

   • Analysis of speculative decoding performance at high batch sizes.

6. Spector, B. & Ré, C. (2023). Accelerating LLM Inference with Staged Speculative Decoding. arXiv preprint.

   • Multi-stage speculation with cascaded draft models.

7. Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., & Jia, Z. (2024). SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification. ASPLOS 2024.

   • Tree-based speculative inference with parallel verification.

8. GitHub: tanishqkumar/ssd. Saguaro Implementation.

   • Open-source SSD implementation with custom inference engine, supporting Llama-3 and Qwen-3 families.

In The Anatomy of Agentic Code Assist, we looked at how agents like OpenHands work: event streams, sandboxed execution, tool use, the CodeAct framework. That post covered the agent itself, what it does and how it’s built. This post covers a different layer: the infrastructure that keeps agents running reliably in production.

When an agent runs for hours, makes hundreds of tool calls, and interacts with flaky LLM APIs, a whole class of infrastructure problems emerge that application-level code cannot solve:

1. State loss on process crashes: a worker dies mid-workflow and hours of accumulated context disappear. The agent restarts from scratch, re-executing every LLM call and tool invocation.
2. LLM API rate limits and timeouts: 429s, 500s, socket timeouts, multi-minute latencies. A reflexion loop running 10 cycles can consume 50x the tokens of a linear pass if any step fails and forces a restart.
3. Debugging non-deterministic behavior: the same prompt produces different outputs, different tool call sequences, different results. Without a complete execution trace, reproducing production bugs is close to impossible.
4. Tasks exceeding server timeouts: agent sessions lasting minutes to hours die on deployments, fail during scaling events, and exceed web server timeout limits.
5. Ambiguous recovery after parallel fan-out crashes: the agent launches ten parallel tool calls. The process crashes after seven complete. Which results were already obtained? Which need re-execution?
6. Losing context during human-in-the-loop waits: the agent pauses for human approval, potentially for hours or days. The server holding that state needs to remain available, or all accumulated context is lost.
7. Error cascades across multi-agent systems: a single failure in one agent propagates downstream without corrective mechanisms. Simple retry logic at the tail end is inadequate because the agent may have already deviated significantly from the intended path.

Temporal is an orchestration platform built around durable execution. We’ll walk through its architecture, understand why each design decision exists, and look at how OpenAI’s Codex team uses it in production.

The core idea can be expressed as a state transition: $S_{t+1} = f(S_t, M(S_t, T_t))$. Agent state evolves through deterministic orchestration ($f$) of non-deterministic operations ($M$ = LLM response, $T$ = tool results). Temporal separates these two concerns at the infrastructure level. The deterministic part goes in workflows. The non-deterministic part goes in activities.

Workflows and Activities

The fundamental design decision in Temporal: split all code into two categories based on determinism.

Workflows

A Workflow is the agent’s control flow, the logic that decides which tools to call, in what order, what to do with results, and when to wait for human input. Workflows run as ordinary code in Python, TypeScript, Go, or Java, with one hard constraint: they must be deterministic. Given the same inputs and the same activity results, a workflow must produce the same sequence of commands every time.

A Workflow Execution can run for seconds, hours, or years. It persists through infrastructure failures. The workflow doesn’t know or care about crashes; from its perspective, execution is continuous.

Activities

Activities are where all side effects live: LLM API calls, tool executions, database writes, HTTP requests. Anything that can fail, timeout, or produce different results on re-execution. Temporal records every activity result in a persistent Event History, an append-only log that serves as the authoritative record for the entire workflow’s state.

Why This Split Matters

The determinism requirement is what enables replay-based recovery (which we’ll cover in the next section). Here’s the reasoning: if we know the workflow logic is deterministic, and we have a recorded log of all activity results, we can reconstruct the exact workflow state after a crash. We don’t need developer-written checkpoint code. We don’t need serialization logic. We just replay the deterministic code with the previously recorded results, and we arrive at the same state.

This raises an obvious question: LLMs are non-deterministic, so how does this work? The answer maps directly to how agents already operate. The LLM call goes in an activity – it’s non-deterministic, its result gets recorded. The logic deciding what to call and when goes in the workflow – it’s deterministic. The agent loop says “if the LLM returned a tool call, execute that tool; if it returned a final answer, return it.” That orchestration logic doesn’t change between runs.

A Complete Agent Loop

Here’s what a complete agent workflow looks like in Python:

from temporalio import workflow, activity
from temporalio.common import RetryPolicy
from datetime import timedelta
from dataclasses import dataclass

@dataclass
class LLMRequest:
    goal: str
    history: list
    available_tools: list

@activity.defn
async def call_llm(request: LLMRequest) -> dict:
    # Non-deterministic: LLM API call lives here
    response = await llm_client.chat(
        messages=request.history,
        tools=request.available_tools,
    )
    return {"action": response.action, "params": response.params}

@activity.defn
async def execute_tool(tool_name: str, params: dict) -> str:
    # Non-deterministic: tool execution lives here
    return await tool_registry.execute(tool_name, params)

@workflow.defn
class AIAgentWorkflow:
    @workflow.run
    async def run(self, user_goal: str) -> str:
        conversation_history = []
        llm_retry = RetryPolicy(
            initial_interval=timedelta(seconds=1),
            backoff_coefficient=2.0,
            maximum_interval=timedelta(seconds=60),
            maximum_attempts=10,
        )

        while not self.is_goal_achieved(conversation_history):
            # Deterministic: this decision logic is the workflow
            next_action = await workflow.execute_activity(
                call_llm,
                LLMRequest(
                    goal=user_goal,
                    history=conversation_history,
                    available_tools=self.get_available_tools(),
                ),
                start_to_close_timeout=timedelta(seconds=120),
                retry_policy=llm_retry,
            )

            if next_action["action"] == "tool_call":
                # Parallel tool execution when multiple tools requested
                results = await asyncio.gather(*[
                    workflow.execute_activity(
                        execute_tool,
                        tool["name"], tool["params"],
                        start_to_close_timeout=timedelta(seconds=30),
                    )
                    for tool in next_action.get("tool_calls", [])
                ])
                conversation_history.extend(results)
            else:
                conversation_history.append(next_action)

        return self.format_final_result(conversation_history)

Deterministic Replay

Replay is the mechanism that makes Temporal’s fault tolerance work. Let’s walk through it in detail, because understanding replay is the key to understanding why the rest of the architecture looks the way it does.

The Event History

Every workflow execution has an Event History: an append-only log stored in Temporal’s persistence layer. When an activity completes, Temporal records both the request and the result.

What Happens on a Crash

Here’s a concrete scenario. An agent workflow is at step 4 of 7. It has completed three LLM calls and tool executions, and is partway through the fourth:

1. The worker process crashes (OOM, deployment, hardware failure)
2. The Temporal server detects the failure (heartbeat timeout or task timeout)
3. Another worker picks up the workflow from the task queue
4. Temporal re-executes the workflow code from the beginning
5. When the code reaches activity calls that already completed (steps 1–3), Temporal returns the previously recorded results from the event history instead of re-executing them
6. The workflow code deterministically reaches the exact same state it was in before the crash: same local variables, same loop counter, same conversation history
7. Forward execution resumes from step 4. Only now does an actual activity get dispatched

Because the workflow code is deterministic, replaying it with the same activity results always produces the same sequence of commands. The entire call stack and state are reconstructed with no developer-written checkpoint code. This is different from simple checkpointing because the developer never has to decide what to checkpoint or when – the replay mechanism reconstructs everything automatically from the event history.

The Determinism Contract

The determinism requirement imposes hard constraints on workflow code. You cannot use:

• random() – use workflow.random() instead
• datetime.now() – use workflow.now() instead
• time.sleep() – use workflow.sleep() or timers instead
• Direct I/O (network calls, file reads) – these must go in activities
• Threading or subprocess creation – use activities or child workflows

For AI engineers, this constraint is less restrictive than it sounds. LLM calls and tool executions are inherently side effects, so they already belong in activities. The orchestration logic that decides what to call and when – “call the LLM, check if it returned a tool call, execute the tool, loop” – doesn’t use random numbers or system clocks.

Here’s what non-determinism violations look like in practice:

# WRONG: non-deterministic workflow code
@workflow.defn
class BadAgentWorkflow:
    @workflow.run
    async def run(self, goal: str) -> str:
        if random.random() > 0.5:        # different result on replay
            strategy = "aggressive"
        else:
            strategy = "conservative"

        timestamp = datetime.now()         # different on replay
        await asyncio.sleep(5)             # blocks the event loop

# CORRECT: deterministic workflow code
@workflow.defn
class GoodAgentWorkflow:
    @workflow.run
    async def run(self, goal: str) -> str:
        if workflow.random().random() > 0.5:   # deterministic across replays
            strategy = "aggressive"
        else:
            strategy = "conservative"

        timestamp = workflow.now()              # deterministic across replays
        await workflow.sleep(5)                 # durable timer, survives crashes

Contrast with OpenHands

Both Temporal and OpenHands use event sourcing, but for different purposes. OpenHands records events (CmdRunAction, FileWriteAction, observations) for debuggability and observability. You can replay the event sequence to understand what the agent did. Temporal records events so the workflow can be reconstructed after a crash as if nothing happened. Same architectural pattern, different goals.

Formalization

If History = $[(a_1, r_1), (a_2, r_2), \ldots, (a_k, r_k)]$ records completed activities, then replay returns $r_1 \ldots r_k$ from history and only executes $a_{k+1}$ forward. The workflow’s determinism guarantees that replaying with recorded results produces the same sequence of activity commands, so the state at step $k$ is identical to the state before the crash.

Server Architecture

Temporal runs as four server-side services plus a persistence layer, with user-managed workers running externally.

The Four Services

Frontend Service: a stateless gRPC gateway. All client and worker communication flows through it. Handles rate limiting, routing, and authorization. Horizontally scalable because it holds no state.

History Service: owns workflow state and persists event histories. This is the most important component. Manages state transitions across configurable History Shards, which are the unit of concurrent throughput scaling. Each shard handles a subset of workflows. More shards = more concurrent workflows.

Matching Service: hosts Task Queues and dispatches work to workers. When a workflow needs an activity executed, the Matching Service places it on the appropriate task queue. When a worker polls for work, the Matching Service assigns a task.

Workers: stateless external processes that you deploy and manage. Workers long-poll task queues via gRPC, execute workflow or activity code, and report results back. Because workers hold no state, they can be killed, restarted, or scaled horizontally without any coordination. The Temporal server is always the authoritative record.

Task Queues

Task Queues provide a routing layer that becomes important for agent workloads. Workflow tasks and activity tasks flow through separate queues. You can route activities to specialized worker pools (GPU workers for inference, lightweight workers for API calls) by assigning them to different task queues. This lets teams scale heterogeneous agent workloads independently.

Primitives for Agent Patterns

Beyond workflows and activities, Temporal provides several primitives that map to common agent coordination problems.

Signals

Signals are asynchronous messages sent to a running workflow. The workflow can react at any point in its execution. This is the mechanism for human-in-the-loop: the agent reaches a decision point, calls workflow.wait_condition(), and a signal carrying the human’s approval resumes it.

The workflow can wait hours or days. It consumes no compute while waiting because its state lives in the event history, not in a running process. No worker is tied up, no server is keeping a connection open. The state is persisted in the database and can be reconstructed on demand when the signal arrives.

Queries

Queries let external systems read workflow state without modifying it. This powers dashboards and monitoring: “What step is the agent on? What was the last LLM response? How many tokens has it consumed?” The query handler runs against the in-memory workflow state and returns immediately.

Updates

Updates combine a signal and a query: send a command to the workflow and get a response. This is useful for interactive agent control (“redo step 2 with different parameters”) where you need to both modify the workflow’s behavior and confirm the modification was accepted.

Replit, for example, uses Workflow Updates for human-in-the-loop consent. When their agent wants to perform a destructive action, it pauses and waits for the user to accept or reject via an Update.

ContinueAsNew

Each workflow execution is limited to 51,200 events or 50MB of event history. For agents making hundreds of tool calls, history grows fast; each activity generates roughly 3 events. If activities return large LLM payloads (500KB+), the 50MB limit becomes binding well before the event count limit.

ContinueAsNew addresses this by atomically starting a fresh execution with the same Workflow ID, carrying forward essential state while resetting the history. The old history is archived. For long-running agents, this is how you keep the workflow alive indefinitely.

Human-in-the-Loop Pattern

@workflow.defn
class AgentWithHumanApproval:
    def __init__(self):
        self.approved = False
        self.current_step = "initializing"
        self.pending_action = None

    @workflow.signal
    async def approve(self, decision: str):
        self.approved = decision == "yes"

    @workflow.query
    def get_status(self) -> dict:
        return {
            "step": self.current_step,
            "pending_action": self.pending_action,
            "approved": self.approved,
        }

    @workflow.run
    async def run(self, goal: str) -> str:
        while not self.is_complete():
            action = await workflow.execute_activity(
                call_llm, goal,
                start_to_close_timeout=timedelta(seconds=120),
            )

            if action.requires_approval:
                self.pending_action = action.description
                self.current_step = "awaiting_approval"
                # Workflow state persists in DB -- no compute cost while waiting
                await workflow.wait_condition(lambda: self.approved)
                self.approved = False  # reset for next approval

            self.current_step = "executing"
            result = await workflow.execute_activity(
                execute_tool, action.tool, action.params,
                start_to_close_timeout=timedelta(seconds=60),
            )
        return self.format_result()

The workflow.wait_condition(lambda: self.approved) line is where the agent pauses. It can sit there for minutes, hours, or days. If the server restarts, if workers are redeployed, the workflow’s state survives. When the signal arrives, any available worker picks it up and resumes execution.

Retry Policies and Error Handling

LLM APIs fail routinely. Rate limits (429), server errors (500), socket timeouts, multi-minute latencies. These are the norm for agents making hundreds of calls, and different activities need different retry strategies.

Declarative Retry Policies

Retry policies are configured per activity with several parameters: initial interval, backoff coefficient, maximum interval, maximum attempts, and non-retryable error types. The important part is that retries happen at the infrastructure level. If a worker crashes during a retry cycle, another worker picks up with the retry state intact. The developer writes no retry logic.

Why Different Activities Need Different Strategies

LLM calls need aggressive retry with exponential backoff. Rate limits are transient, and the cost of not retrying (losing all accumulated context and starting the agent run from scratch) far outweighs the cost of waiting 30 seconds for capacity. Configure high maximum attempts (10+) with a long maximum interval.

Tool executions need limited retries. Tools may not be idempotent – running git commit twice produces different results. Blindly retrying could cause duplicate side effects. Configure low maximum attempts (2–3) and mark certain error types as non-retryable.

Human notifications often need no retry at all. Fire-and-forget: if the Slack message fails, don’t block the workflow.

llm_retry = RetryPolicy(
    initial_interval=timedelta(seconds=1),
    backoff_coefficient=2.0,
    maximum_interval=timedelta(seconds=60),
    maximum_attempts=10,
    non_retryable_error_types=["InvalidPromptError"],
)

tool_retry = RetryPolicy(
    initial_interval=timedelta(seconds=2),
    maximum_attempts=3,
    non_retryable_error_types=["PermissionDenied", "NotIdempotent"],
)

# Heartbeating for long-running activities
@activity.defn
async def execute_long_tool(task: dict) -> str:
    result = ""
    for i, chunk in enumerate(process_chunks(task)):
        activity.heartbeat({"progress": i, "last_chunk": chunk.id})
        result = await process(chunk)
    return result

Heartbeats

For long-running activities, the worker periodically reports progress via heartbeats. If the heartbeat stops (worker crashed), Temporal reschedules the activity on another worker. The new worker can read the last heartbeat details to resume from the last checkpoint rather than starting over. This matters for activities processing large datasets or running multi-step tool executions.

Saga Patterns for Multi-Agent Systems

When multiple agents coordinate, failure handling gets complex. Temporal supports saga patterns where compensation logic runs when a step fails. If a planning agent fails, downstream execution agents’ pending activities can be cancelled rather than left hanging. If the response agent produces an unsatisfactory draft, compensation logic can route back to the research agent for additional context.

Production Case Study: OpenAI Codex

OpenAI’s Codex, their cloud-based coding agent that writes, tests, and iterates on code, uses Temporal as its core orchestration backbone. Will Wang, a software engineer on the Codex team, confirmed publicly that “Temporal is a critical part of the infrastructure powering Codex, responsible for executing our core control flows.” He described it as enabling the team to “easily reason about concurrency, correctness, and fault tolerance” while scaling a complicated distributed system.

Codex sessions run for 6+ hours on complex tasks. The entire agent loop (prompt construction, model inference, tool calls, result observation, loop back) runs as a Temporal Workflow. Each LLM call and tool execution is an Activity with its own retry policy and timeout. A single “turn” can involve hundreds of tool calls.

The Codex harness manages three conversation primitives: Items (atomic I/O units like messages or diffs), Turns (one unit of agent work from user input), and Threads (the durable container for an ongoing session, with persisted event history supporting resume, fork, and archive operations). Thread persistence – OpenAI describes threads as “durable containers” with “persisted event history” supporting reconnection – aligns directly with Temporal’s Event History.

Codex has a self-review pattern internally called the “Ralph Wiggum Loop”: the agent reviews its own changes, requests additional agent reviews, and iterates until all reviewers are satisfied. In Temporal terms, the review results arrive as signals, and the workflow decides whether to iterate or complete.

The relationship extends beyond Codex. In July 2025, OpenAI and Temporal launched a formal integration adding durable execution to the OpenAI Agents SDK. Every agent invocation runs as a Temporal Activity, orchestration runs as a Temporal Workflow. Temporal also processes millions of ChatGPT Images generation workflows. Venkat Venkataramani (OpenAI’s VP of App Infrastructure) reinforced this at Temporal’s Series D announcement: “Durable execution is a core requirement for modern AI systems.”

Framework Integrations

Temporal integrates with existing agent frameworks so teams don’t have to rewrite their agent logic from scratch. The pattern is the same across integrations: Temporal provides the durability layer, the framework provides the agent logic.

PydanticAI + Temporal

PydanticAI has first-class Temporal support via a TemporalAgent wrapper that preserves PydanticAI’s type-safety while offloading non-deterministic model requests and tool calls to Temporal activities. The orchestration logic lives in a deterministic workflow, and all I/O-bound tasks are automatically wrapped as activities.

One significant design decision: thread-based workflows. Each conversation thread gets its own Temporal workflow that persists for the lifetime of the conversation. This is more efficient than stateless approaches because the system only processes new messages, maintaining context within workflow state rather than re-sending the entire history for every inference.

from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIModel
from temporalio.client import Client

# Define the agent with PydanticAI's type-safe interface
support_agent = Agent(
    model=OpenAIModel("gpt-4o"),
    system_prompt="You are a customer support agent.",
    result_type=SupportResponse,  # Pydantic model for type-safe output
)

@support_agent.tool
async def lookup_order(ctx, order_id: str) -> OrderDetails:
    return await db.get_order(order_id)

# Wrap with Temporal for durability
from pydantic_ai_temporal import TemporalAgent

temporal_agent = TemporalAgent(
    agent=support_agent,
    client=await Client.connect("localhost:7233"),
    task_queue="support-agents",
)

# Each conversation gets a durable workflow
result = await temporal_agent.run(
    "What's the status of order #12345?",
    thread_id="customer-session-abc",
)

OpenAI Agents SDK + Temporal

The OpenAI Agents SDK integration centers on the activity_as_tool helper. This function automatically generates OpenAI-compatible tool schemas directly from Temporal activity signatures. The agent reasons about and invokes activities as tools, with every tool call backed by durable execution.

import { activityAsTool } from "@temporalio/openai-agents";
import { OpenAIAgentsPlugin } from "@temporalio/openai-agents";

// Temporal activities become tools the agent can call
const searchTool = activityAsTool(searchDocuments, {
  startToCloseTimeout: "30s",
  retryPolicy: { maximumAttempts: 3 },
});

const writeTool = activityAsTool(writeDocument, {
  startToCloseTimeout: "60s",
  retryPolicy: { maximumAttempts: 1 },
});

// Agent orchestration runs as a Temporal Workflow
// Each tool call is a durable Activity
const plugin = new OpenAIAgentsPlugin({
  client: temporalClient,
  taskQueue: "agent-workers",
  tools: [searchTool, writeTool],
});

Developers use the OpenAIAgentsPlugin to configure the Temporal client and worker, enabling integrated tracing that provides visibility through both the Temporal UI and OpenAI dashboards.

When Temporal Adds Unnecessary Complexity

Temporal is not always the right choice. Here’s where it adds more complexity than value:

• Simple agents: a single LLM call followed by one tool call doesn’t benefit from durable execution infrastructure. One comparison found that adding Temporal to a simple document indexing pipeline required “rearchitecting the app, splitting it into two services, adding a runtime dependency on a third service, and adding over 100 lines of code” where a lighter-weight approach achieved the same with 7 lines.
• Prototyping and experimentation: when you’re iterating on agent architecture, the determinism constraints and operational overhead slow you down.
• Sub-30-second agents: if the agent completes before infrastructure failures become likely, the cost of durable execution exceeds the benefit.
• Teams without infrastructure engineering capacity: self-hosted Temporal requires operating four services plus a database. If you don’t have the team to manage this, the operational burden may outweigh the reliability gains.

Trade-offs

Temporal’s guarantees come with trade-offs that shape day-to-day development experience.

Operational Complexity

Self-hosted Temporal requires deploying four independent services plus a persistence database (PostgreSQL, MySQL, or Cassandra) and optionally Elasticsearch for advanced visibility. This is not a single process with a single run command.

Learning Curve

Engineers must internalize: workflows vs activities, determinism rules, event history mechanics, signals, queries, updates, ContinueAsNew, versioning strategies, worker configuration.

The determinism constraint confuses newcomers, especially because LLMs are inherently non-deterministic. The resolution (LLM calls go in activities, not workflows) is simple once understood, but the documentation framing perpetuates the misconception.

Event History Limits

Each workflow execution is limited to 51,200 events or 50MB. An activity generates roughly 3 events. If activities return large LLM payloads (500KB+), the 50MB limit becomes binding well before the event count limit. The mitigation – ContinueAsNew, which atomically starts a fresh execution with carried-over state – works but adds architectural complexity. Teams building agents with many LLM calls must implement payload offloading (store large payloads in S3, pass references) and proactively manage history growth.

Latency

Temporal Cloud’s minimum end-to-end latency is roughly 100ms per workflow step, with a single activity round-trip taking approximately 220ms. Local Activities save ~50ms per call but sacrifice heartbeating and independent retry capabilities. For agents where sub-second interactivity matters (chatbot-like interactions), this overhead accumulates across many steps. Agents with 50+ steps per interaction may see 5–10 seconds of pure infrastructure overhead.

Versioning

Code changes to workflow logic can cause non-determinism errors during replay of running workflows. If a running workflow was started with version 1 of the code and a worker running version 2 picks it up, the replay may produce different activity commands, causing a non-determinism exception. Temporal provides patching APIs and worker versioning, but patches accumulate in code and “need to be removed with extreme care.” Airbyte documented struggles with non-determinism exceptions, ultimately deciding to fail affected workflows rather than attempting recovery. Safe deployment requires replay testing against production event histories in CI.

Closing Thoughts

We covered a lot of ground here: the workflow/activity split, deterministic replay, server architecture, coordination primitives, retry strategies, and how OpenAI’s Codex team puts it all together.

The core design insight is the separation of deterministic orchestration from non-deterministic execution. Once you accept that split, replay-based recovery falls out as a consequence – and with it, most of the infrastructure problems we listed at the top of this post.

OpenAI, Replit, Block, NVIDIA, and others have independently converged on durable execution for their agent workloads. Temporal’s recent $300M Series D at a $5B valuation, with 380%+ year-over-year revenue growth driven substantially by AI workloads, suggests this is a real pattern. The company joined the Agentic AI Foundation (under the Linux Foundation) alongside Anthropic, OpenAI, and Block.

For most teams, the practical path is: prototype with something lighter (LangGraph, CrewAI), validate the agent architecture, and migrate when the agents run long enough and matter enough that you can’t afford to lose state on a crash. The operational investment is real, but so is the cost of rebuilding reliability from scratch.

References

1. Temporal Documentation. Core Concepts – Workflows, Activities, Workers. Temporal Technologies.

2. Temporal. Temporal for AI. Overview of Temporal’s AI-specific capabilities and customer stories.

3. Wang, W. (2025). Codex and Temporal Integration. Will Wang’s public statements on Codex’s use of Temporal for core control flows.

4. OpenAI. Harness Engineering: Leveraging Codex in an Agent-First World. OpenAI engineering blog on the Codex harness architecture.

5. Temporal. Build Durable AI Agents with Pydantic AI and Temporal. PydanticAI integration guide.

6. Temporal. Of Course You Can Build Dynamic AI Agents with Temporal. Temporal’s architecture for dynamic AI agent loops.

7. Quo (formerly OpenPhone). How We Built a Real-Time AI Voice Agent with Temporal. Production case study on Temporal primitives for voice agents.

8. Temporal. Production-Ready Agents with the OpenAI Agents SDK + Temporal. OpenAI Agents SDK integration announcement.

9. Temporal. AI Cookbook – OpenAI Agents SDK. Code examples and patterns for the OpenAI integration.

10. PydanticAI Documentation. Temporal Durable Execution. Official PydanticAI guide for Temporal integration.

11. Vanlightly, J. Explanations of deterministic replay mechanics and the determinism contract in Temporal workflows. Referenced via Temporal community resources.

12. Wang, X., et al. (2025). The OpenHands Software Agent SDK. arXiv preprint arXiv:2511.03690. The predecessor post’s primary reference for event sourcing comparison.

If you have been following this blog, you know how LLMs generate text: attention mechanisms, KV caches, speculative decoding, quantized weights moving through GPU memory hierarchies. We have spent considerable time understanding what happens after the model exists. This post asks a different question: how did the model learn to be useful in the first place?

A pretrained language model is a remarkable thing. It can complete sentences, mimic writing styles, and recite facts absorbed from trillions of tokens of internet text. But it cannot follow instructions. Ask it to summarize an article and it might continue the article instead. Ask it to refuse a harmful request and it will cheerfully comply. The gap between “can predict the next token” and “can be a helpful assistant” is enormous, and closing it is the job of reinforcement learning from human feedback (RLHF) and its descendants.

This post covers three techniques that bridge this gap: PPO (Proximal Policy Optimization), the original workhorse that proved RL could align language models; DPO (Direct Preference Optimization), an elegant reformulation that eliminates the reward model entirely; and GRPO (Group Relative Policy Optimization), the technique behind DeepSeek-R1’s reasoning capabilities. Each optimizes the same underlying objective (maximize reward while staying close to a reference policy) but they make fundamentally different engineering trade-offs.

We will not cover pretraining, nor will we survey every RLHF variant (KTO, SimPO, ORPO, and others exist but are beyond our scope). Instead, we will go deep on these three methods: the math, the intuition, the practical trade-offs, and the reasons each one was invented.

Where RL Fits: The Model Training Lifecycle

Before we touch any equations, let’s establish context. Training a modern LLM that can follow instructions involves three distinct phases, each with a different objective.

Phase 1: Pretraining. The model learns language by predicting the next token on a massive corpus: books, articles, code, web pages. This produces a powerful text completion engine that knows facts, grammar, and reasoning patterns but has no concept of a “conversation” or “helpfulness.” This phase consumes the vast majority of compute (months on thousands of GPUs) and produces what we call the base model.

Phase 2: Supervised Fine-Tuning (SFT). Humans write demonstration data: pairs of (instruction, ideal response). The model is trained to reproduce these demonstrations using standard cross-entropy loss:

This is the same next-token prediction objective from pretraining, just applied to curated instruction-response pairs. The model learns the format of helpful responses: how to structure answers, when to use code blocks, how to handle multi-turn conversations. SFT typically requires only thousands of examples and a few hours of training.

But SFT has a fundamental limitation: it teaches the model to mimic demonstrations without understanding why one response is better than another. The model learns that a particular answer to “explain quantum computing” was in the training data, but it cannot distinguish between a clear explanation and a subtly misleading one. It learns format, not judgment.

Phase 3: RL-based Alignment. This is where the techniques in this post come in. Instead of showing the model what to produce, we teach it what better means. The model generates its own responses, receives feedback on quality, and updates its parameters to produce higher-quality outputs. This is reinforcement learning: the model (agent) generates text (actions), receives scores (rewards), and improves its generation strategy (policy).

The SFT model serves double duty here: it initializes the policy we will optimize, and it becomes the frozen reference policy $\pi_{\text{ref}}$ that prevents the RL-trained model from drifting too far from coherent language. This reference is crucial. Without it, the model can find degenerate ways to maximize reward that produce nonsensical text.

The three methods we will examine differ in how they implement Phase 3. PPO trains a separate reward model, then runs RL against it. DPO skips the reward model by extracting the reward signal directly from preference data. GRPO replaces learned value estimates with group statistics and pairs naturally with verifiable rewards. Let’s start with the foundation they all share: the reward signal.

The Reward Signal: Bradley-Terry and What Makes It Work

The fundamental challenge of aligning language models is this: we cannot write a reward function for “helpfulness.” Unlike game-playing AI where the score is clearly defined, the quality of a text response is subjective, contextual, and multidimensional. But humans can do something simpler: given two responses to the same prompt, they can usually say which one is better.

This observation is the foundation of RLHF. Collect pairwise comparisons, then train a model to predict which response humans prefer. The mathematical framework for this is the Bradley-Terry model, originally developed in 1952 for ranking chess players:

The elegant property: only the difference in rewards matters. Adding a constant to all reward scores leaves preferences unchanged. This means the reward model only needs to learn a relative ranking, not absolute quality scores.

We train this reward model by maximizing the log-likelihood of observed human preferences:

This is binary cross-entropy: we are training a classifier that says “response A is better than response B.” Architecturally, the reward model is typically the same transformer as the language model, with the language modeling head replaced by a single linear layer that maps the final hidden state to a scalar reward.

In practice, InstructGPT used a 6B parameter reward model to guide a 175B policy, trained on approximately 33,000 prompts with 4-9 ranked completions each. The reward model is trained for only a single epoch to avoid overfitting to the preference data, a detail that matters more than it might seem.

With a trained reward model in hand, we can now define what “better” means mathematically. The question becomes: how do we actually optimize the language model to produce higher-reward outputs?

PPO: The Four-Model Pipeline

The RLHF Objective

The goal of RLHF is captured in a single objective. Let’s walk through it symbol by symbol:

Reading left to right: $\max_{\pi_\theta}$ means “find the policy parameters that maximize the following expression.” $\mathbb{E}$ is the expected value, averaging over many prompts and responses. $x \sim \mathcal{D}$ means prompts are drawn from the training distribution. $y \sim \pi_\theta(\cdot \mid x)$ means responses are sampled from the current policy, not taken from a fixed dataset. $r_\phi(x, y)$ is the reward model’s score. $\beta$ is a coefficient controlling constraint strength. $D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})$ is the KL divergence measuring how far the policy has drifted from the reference.

This objective contains two opposing forces. The first term pushes toward human-preferred outputs: maximize the expected reward. The second term, the KL divergence penalty, is the guardrail that prevents reward hacking.

Reward hacking is not a theoretical concern. Without the KL constraint, models learn to game the reward model: they produce longer responses (reward models often prefer length), use confident language and bullet-point formatting (which correlates with higher human ratings), and can even produce convincing fabrications that fool the reward evaluator. Wen et al. (2024) showed that RLHF without proper regularization increases human approval ratings while simultaneously decreasing actual correctness. The KL penalty keeps the optimized policy close enough to the reference that these degenerate strategies remain unlikely.

The PPO Clipped Surrogate

The RLHF objective tells us what to optimize. PPO tells us how. The challenge is that policy gradient methods are notoriously unstable. A single large update can destroy the policy, and recovery is difficult. PPO solves this with a clipping mechanism that limits how much any single update can change the policy.

First, we define the probability ratio, how much the policy’s opinion of a particular token has changed:

When $r_t = 1.0$, the policy hasn’t changed its probability for this token. When $r_t = 1.5$, the token is 50% more likely under the new policy. When $r_t = 0.6$, it is 40% less likely. The ratio tells us the direction and magnitude of the policy shift.

The PPO clipped surrogate objective is:

Here $\hat{A}_t$ is the advantage estimate: how much better (positive) or worse (negative) this token was compared to the expected baseline. Concretely, a critic network $V_\psi(s)$ estimates the expected future reward from each state; the advantage is the difference between the actual return and this estimate. If the model produced a token that led to higher reward than expected, $\hat{A}_t > 0$ (“good token, do more of this”); if the reward was lower than expected, $\hat{A}_t < 0$ (“bad token, do less of this”). The advantage is what tells PPO which direction to push; the clipping mechanism controls how far.

The clip function is a simple three-case clamp: if $r_t < 1-\varepsilon$, return $1-\varepsilon$; if $r_t > 1+\varepsilon$, return $1+\varepsilon$; otherwise return $r_t$ unchanged. With the standard $\varepsilon = 0.2$, the ratio is constrained to the range $[0.8, 1.2]$.

The behavior of this objective follows a 2×2 matrix that is worth internalizing:

The pattern reveals something important: clipping only constrains the policy when it is already moving in the right direction too aggressively. When the policy increases probability for a good token (top-left), clipping says “that’s enough reinforcement for one update.” When it decreases probability for a bad token (bottom-right), clipping says “that’s enough suppression.” But when the policy is moving in the wrong direction — decreasing a good token or increasing a bad one — the full gradient signal flows through. Clipping never protects wrong moves.

Why $\min$ and not $\max$? The $\min$ operator takes the pessimistic bound. If the clipped version yields a lower objective than the unclipped version, we take the clipped (lower) one, preventing overconfident updates. If the unclipped version is already lower (meaning the policy moved in a harmful direction), we take that instead, allowing the full corrective gradient.

Let’s trace through concrete numbers. With $\hat{A}_t = +2.0$ (a good token) and $\varepsilon = 0.2$:

• At $r_t = 1.1$: Unclipped = $1.1 \times 2.0 = 2.2$. Clipped = $1.1 \times 2.0 = 2.2$ (within bounds). $\min = 2.2$. Full gradient.
• At $r_t = 1.3$: Unclipped = $1.3 \times 2.0 = 2.6$. Clipped = $1.2 \times 2.0 = 2.4$ (capped at $1+\varepsilon$). $\min = 2.4$. Gradient is reduced.
• At $r_t = 2.0$: Unclipped = $2.0 \times 2.0 = 4.0$. Clipped = $1.2 \times 2.0 = 2.4$. $\min = 2.4$. The objective plateaus. No matter how much more likely this token becomes, the gradient contribution is capped.

This plateau is the key mechanism. The objective becomes flat beyond the clip boundary, which means the gradient is zero, so the optimizer receives no signal to push the ratio further. The policy can only change by $\pm 20\%$ per update, ensuring training stability.

An important subtlety for LLM training: clipping operates per-token, not globally. In a 512-token response, some tokens might have $r_t$ well within bounds (contributing full gradients) while others hit the clip boundary (contributing zero gradient). The overall update is a blend of these per-token signals, which produces remarkably stable training even without careful learning rate tuning.

One notable exception: DeepSeek-R1 uses $\varepsilon = 10$, which effectively disables clipping. Their group-normalized advantages (which we will see in the GRPO section) are already well-scaled, reducing the need for a tight trust region.

The Four-Model Problem

Running PPO for LLM alignment requires four models simultaneously in GPU memory:

1. The policy $\pi_\theta$ — the model being trained. Requires gradients and optimizer states (2-3× the weight memory).
2. The reference policy $\pi_{\text{ref}}$ — a frozen copy of the SFT model. Only forward passes, but still occupies full weight memory.
3. The reward model $r_\phi$ — scores generated responses. Frozen during PPO, forward passes only.
4. The value/critic network $V_\psi$ — estimates expected future reward to compute advantages $\hat{A}_t$. Requires gradients and optimizer states.

For a 7B parameter model in fp16, weights alone consume approximately 14GB per model, roughly 56GB across all four, before accounting for optimizer states (Adam stores two additional copies of the policy’s and critic’s parameters). With a batch of generated sequences in memory, the total easily exceeds 100GB for a single 7B model. Running PPO on a 70B model requires multi-node setups that only frontier labs can afford.

Beyond memory, PPO faces two systemic challenges. Distribution shift: as the policy improves, the reward model’s training data (collected from an earlier, weaker policy) becomes stale. The proxy reward keeps climbing while true human preference plateaus or declines. Gao et al. (2022) formalized this as “reward model overoptimization.” Hyperparameter sensitivity: learning rates, KL coefficients, clipping parameters, and even Adam’s epsilon require careful tuning. Huang et al. (2023) found that reward scores and loss values are poor indicators of training health; practitioners should monitor KL divergence, response length distributions, and perplexity instead.

Despite all this complexity, PPO produced the first convincing result: InstructGPT showed that a 1.3B parameter model trained with RLHF was preferred by human evaluators over the 175B parameter base GPT-3. A 130× smaller model, made more useful through alignment. The engineering was expensive, but the result was undeniable.

The Question That Sparked DPO

PPO demonstrated that RL could align language models with human preferences. But the engineering complexity was severe: four models, meticulous hyperparameter tuning, and infrastructure that only a handful of organizations could afford. Researchers began asking: could we achieve similar results without the reward model entirely?

The mathematical observation that makes this possible: the RLHF objective has a closed-form optimal policy. If we can express the reward in terms of the policy itself, we can optimize directly on preference data without a reward model, RL loop, or critic network. This insight leads to DPO.

DPO: Your Language Model Is Secretly a Reward Model

The Reparameterization That Changes Everything

Let’s start from the same KL-constrained RLHF objective we defined for PPO. Using variational calculus (or, more practically, by expanding the KL divergence and completing the algebra), we can derive the optimal policy in closed form:

where $Z(x) = \sum_y \pi_{\text{ref}}(y \mid x) \cdot \exp(r(x, y) / \beta)$ is the partition function that ensures the distribution sums to 1.

The intuition here is direct: the optimal policy is the reference distribution “warped” by an exponential reward function. Responses with high reward get boosted in probability; responses with low reward get suppressed. The parameter $\beta$ controls how aggressive this warping is. When $\beta \to 0$, the policy collapses toward pure reward maximization; only the highest-reward response gets any probability mass. When $\beta \to \infty$, the exponential flattens and the policy stays frozen at the reference.

We can rearrange this to express the reward in terms of the policy:

This says something remarkable: the reward is fully determined by the log-ratio of optimal policy to reference policy, plus a prompt-dependent constant $Z(x)$. The reward is hiding inside the policy all along.

But $Z(x)$ is intractable. It requires summing over all possible responses to prompt $x$, every possible sequence of tokens the model could produce. For a vocabulary of 50,000 tokens and responses of even modest length, this is an astronomically large set. PPO avoids computing $Z(x)$ by using iterative approximate optimization. DPO avoids it through algebraic cancellation.

From RL to Classification in One Substitution

Here is where DPO’s elegance emerges. We substitute the implicit reward expression into the Bradley-Terry preference model. For a preferred response $y_w$ and dispreferred response $y_l$ given the same prompt $x$:

The $\beta \log Z(x)$ terms cancel exactly. This is the critical step. $Z(x)$ depends only on the prompt, not the response, so it appears identically in both terms and drops out of the difference. The intractable partition function vanishes.

Substituting into the Bradley-Terry model and replacing the theoretical optimal policy $\pi^*$ with our trainable policy $\pi_\theta$, we get the DPO loss:

This is a binary cross-entropy loss. The “logit” is the difference in implicit rewards between the preferred and dispreferred responses. Each implicit reward $\beta \log(\pi_\theta / \pi_{\text{ref}})$ measures how much the current policy has shifted its probability relative to the reference, a direct proxy for how much the policy “values” that response.

During training, this loss simultaneously increases the relative probability of preferred completions and decreases the relative probability of dispreferred ones. The $\beta$ parameter controls how sharply: low $\beta$ (e.g., 0.1) allows aggressive optimization away from the reference, while high $\beta$ (e.g., 0.5) keeps updates conservative. No explicit KL penalty is needed because the reference policy appears directly in the loss function. Deviating too far automatically reduces the gradient signal through the sigmoid saturation.

A Worked Numerical Example

Let’s make this concrete. Consider a prompt $x$ = “What is the capital of France?” with two responses:

• $y_w$ (preferred): “The capital of France is Paris.”
• $y_l$ (dispreferred): “France’s capital is Berlin, a beautiful city.”

Suppose the reference policy assigns $\pi_{\text{ref}}(y_w \mid x) = 0.15$ and $\pi_{\text{ref}}(y_l \mid x) = 0.12$. The trainable policy $\pi_\theta$ starts as a copy of the reference, so initially $\pi_\theta = \pi_{\text{ref}}$. Let $\beta = 0.1$.

At initialization:

The implicit rewards are both zero:

The reward difference is $0 - 0 = 0$. The loss is $-\log \sigma(0) = -\log(0.5) = \log 2 \approx 0.693$. The model has no preference, exactly what we’d expect before any training.

After one gradient step:

The gradient pushes $\pi_\theta(y_w \mid x)$ up to $0.20$ and $\pi_\theta(y_l \mid x)$ down to $0.08$:

The reward difference is $0.029 - (-0.041) = 0.069$. The loss drops to $-\log \sigma(0.069) \approx 0.676$.

The model is learning without ever computing an explicit reward. The reward signal emerges from the probability shift relative to the reference. And notice the structural KL constraint at work: as $\pi_\theta$ pushes probabilities further from $\pi_{\text{ref}}$, the log-ratio grows, which eventually saturates the sigmoid and produces diminishing gradient signal. The policy naturally resists extreme deviations.

The insight, elegantly stated by the DPO authors: “The reward model was never eliminated — it was absorbed into the policy itself.”

Where DPO Shines and Where It Falls Short

DPO’s practical advantages are substantial. Training requires only two models (policy and reference), not four. The implementation is roughly 20 lines of PyTorch on top of a standard language modeling pipeline. HuggingFace’s TRL library provides a DPOTrainer that handles the details. Major models adopted it quickly: Llama 3, Zephyr-beta, and Tulu 2 all used DPO in their alignment pipelines. DPO democratized alignment research. Any lab with a GPU and preference data could train an aligned model.

But DPO has limitations that become apparent at scale. The most fundamental is its offline nature: DPO trains on a fixed dataset of preference pairs, with no mechanism for the model to explore and discover new behaviors. As training progresses, the policy drifts from the distribution that generated the training data, but the training data cannot adapt. This is particularly problematic for tasks where the model needs to discover novel reasoning strategies.

Xu et al. (ICML 2024) conducted a systematic comparison and found that PPO consistently surpasses DPO across all tested benchmarks when properly tuned, especially on challenging code generation tasks (on CodeContest, PPO-34B achieved 22.4% while DPO-34B scored significantly lower). The gap widens on tasks that require exploration and long-horizon reasoning.

There is also a subtler issue: DPO assumes the Bradley-Terry preference model perfectly fits the data. Real human preferences can be intransitive (A > B, B > C, but C > A), context-dependent, and noisy. When these assumptions break down, DPO’s loss function can produce misleading gradients.

DPO traded RL’s complexity for supervised learning’s simplicity. The next technique we’ll examine takes a different path: keep the online RL loop, but find a cheaper way to run it.

GRPO: Grading Responses on a Curve

The Insight: Eliminate the Critic, Keep the RL Loop

DPO eliminated RL entirely but lost online exploration. GRPO takes a different approach: retain the online RL loop (the model generates responses, gets feedback, and updates) but eliminate the critic network, which is the most expensive component of PPO after the reward model.

Recall that PPO needs a critic $V_\psi(s)$ to compute advantages $\hat{A}_t$, estimating how much better each token was compared to baseline expectations. This critic is a full-sized neural network with its own gradient computation and optimizer states. GRPO’s key observation: instead of learning this baseline, we can estimate it empirically by sampling multiple responses to the same prompt and comparing them to each other.

The Mechanism: Sample, Score, Normalize

For each prompt $q$, GRPO samples $G$ completions $\{o_1, o_2, \ldots, o_G\}$ from the current policy $\pi_\theta$. Each completion is scored by a reward function, producing rewards $\{r_1, r_2, \ldots, r_G\}$. The advantage for each completion is computed via z-score normalization:

This is “grading on a curve.” Instead of evaluating each response against an absolute rubric (the learned critic), we evaluate it relative to its peers. A response that scores 0.8 when all other responses also score around 0.8 gets a near-zero advantage because it was average for this prompt. The same score of 0.8 when peers score around 0.3 earns a strongly positive advantage because it was exceptional.

The group mean serves as an empirical Monte Carlo estimate of the expected reward for this prompt, playing the same role as the learned value function $V(s)$ in PPO. More samples mean a better estimate. In practice, $G = 8$ to $G = 64$ provides sufficient accuracy without excessive compute.

The standard deviation in the denominator does something subtle but important. It acts as a curvature-adaptive gradient mechanism. For easy prompts where the model consistently scores well (low reward variance), the std is small and the advantage magnitudes are amplified, but since the raw rewards are already clustered near the mean, the actual advantages remain small. For hard prompts where reward variance is high, the std normalizes away the scale differences, producing moderate advantages regardless of the raw reward range. This provides automatic per-prompt learning rate adaptation without any additional hyperparameters.

The full GRPO objective incorporates this advantage into a clipped surrogate structure that should look familiar:

This is structurally identical to PPO’s clipped surrogate. The probability ratio $\rho_t^{(i)} = \pi_\theta(o_{i,t} \mid q, o_{i,\lt t}) / \pi_{\theta_{\text{old}}}(o_{i,t} \mid q, o_{i,\lt t})$ is the same per-token ratio. The clipping mechanism works identically. The only difference is where the advantage $\hat{A}_i$ comes from: PPO estimates it with a learned critic, GRPO estimates it with group statistics. The double averaging (over group members and over tokens) combined with clipping and KL penalty gives GRPO the stability of PPO without the critic’s memory cost.

GRPO + RLVR: The Reasoning Revolution

GRPO’s natural partner is Reinforcement Learning from Verifiable Rewards (RLVR). These are tasks where correctness can be checked deterministically: math problems have right answers, code must pass test cases, logic puzzles have verifiable solutions. For these tasks, the reward function is a simple rule (correct or incorrect) requiring no learned reward model at all.

Rule-based rewards are immune to reward hacking. There is no neural network to exploit, no proxy to overoptimize. The reward is ground truth. This makes GRPO + RLVR an extraordinarily clean training setup: sample responses, check if they are correct, normalize advantages within the group, update the policy. Two models in memory (policy and reference), a deterministic reward function, and online exploration.

DeepSeek-R1 demonstrated how powerful this combination can be. Its reward function was remarkably simple:

where $R_{\text{accuracy}}$ is binary (1 if the final answer matches the ground truth, 0 otherwise, verified by regex matching) and $R_{\text{format}}$ enforces structured reasoning with <think>...</think> and <answer>...</answer> tags. That’s it. No neural reward model, no human preference data for the RL stage.

The results with DeepSeek-R1-Zero — trained with GRPO and no supervised fine-tuning at all — were striking: 71.0% on AIME 2024 (matching OpenAI’s o1-preview), 97.3% on MATH-500, and a 2,029 Elo rating on Codeforces. Perhaps most remarkable was the emergent behavior: the model spontaneously developed self-correction strategies (“Wait, let me reconsider this step…”) — without any explicit training signal for reflection. This self-verification behavior emerged purely from the pressure to produce correct final answers.

DeepSeek-R1’s practical configurations: $G = 16$ responses per prompt, batches of 32 unique questions, and notably $\varepsilon = 10$, which effectively disables clipping entirely. The group-normalized advantages are already well-scaled, reducing the need for a tight trust region constraint.

Connection to REINFORCE and Variance Reduction

GRPO is best understood as a variant of REINFORCE, the simplest policy gradient algorithm, with a group-based baseline for variance reduction. Vanilla REINFORCE computes policy gradients as:

This has notoriously high variance because raw returns fluctuate enormously between episodes. The standard fix is to subtract a baseline $b$ from the return: $\nabla_\theta J = \mathbb{E}[(R - b) \cdot \nabla_\theta \log \pi_\theta]$. Any baseline that does not depend on the action is unbiased. PPO learns this baseline with an expensive critic network. GRPO uses the group mean instead, a sample-based estimate that improves with larger group sizes, costs no additional parameters, and requires no additional training.

Choosing the Right Technique

The choice between PPO, DPO, and GRPO depends primarily on the nature of your reward signal and your computational constraints.

Use PPO when you are training on open-ended tasks (creative writing, general helpfulness, safety) where reward must come from a learned model, and you have the compute budget to support four models in memory. Despite its complexity, PPO with proper tuning remains the highest-ceiling approach. Xu et al. (ICML 2024) showed it consistently outperforms DPO on challenging benchmarks. It is the choice of frontier labs training flagship models.

Use DPO when you have high-quality paired preference data, want simple and stable training, and are working with limited compute. DPO matched or exceeded PPO on summarization benchmarks (61% GPT-4 win rate vs. PPO’s 57% on TL;DR) and is implementable with a standard supervised training pipeline. It is ideal for quick alignment passes and situations where the preference dataset covers the intended use distribution well.

Use GRPO when your task has verifiable rewards (math, code, logic, factual QA with ground-truth answers). It combines online exploration (like PPO) with low memory footprint (like DPO), and rule-based rewards eliminate reward hacking entirely. It is the standard for training reasoning models.

In practice, these techniques are often used in combination, not isolation. Llama 3 used a pipeline of SFT → Rejection Sampling → PPO → DPO, where each stage addressed different aspects of alignment. DeepSeek-R1 alternated between SFT stages, RLVR with GRPO for reasoning, and RLHF stages for general helpfulness. The techniques are complementary.

The landscape continues to expand. On the DPO side, IPO removes the Bradley-Terry assumption, KTO works with binary feedback (thumbs up/down) instead of pairwise preferences, and SimPO simplifies the reference model dependency. On the GRPO side, DAPO addresses training instabilities with dynamic sampling, and Dr. GRPO provides variance reduction to the gradient estimates. Each builds on the foundations covered here.

From Explicit Rewards to Emergent Reasoning

Let’s step back and trace the arc we have followed. PPO established that RL could align language models, using an explicit reward model to score responses and a learned critic to estimate advantages. DPO showed the reward model was unnecessary. The reward signal was implicit in the probability ratios, waiting to be extracted through a clever reparameterization. GRPO showed the critic was unnecessary too. Group statistics could replace learned value functions, especially when paired with verifiable rewards.

Each step eliminated a component that turned out to be inessential for the task at hand. What remained was the core objective: maximize expected reward while staying close to the reference. And progressively simpler ways of optimizing it.

But the most interesting result came from the simplest setup. DeepSeek-R1-Zero, trained with GRPO and binary correct/incorrect rewards, spontaneously developed multi-step reasoning, self-correction, and solution verification, capabilities that were not explicitly trained. The model learned how to think from the sole pressure to be correct. No demonstrations of reasoning. No reward for intermediate steps. Just final-answer accuracy and the group-relative advantage signal.

This suggests that the path to capable reasoning models may be less about sophisticated reward engineering and more about giving models the right optimization framework and letting them discover strategies on their own. The field is still learning which components are genuinely necessary and which are engineering artifacts of earlier approaches. These three techniques (PPO, DPO, and GRPO) represent the progression of that understanding.

References

1. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv preprint.

   • The original PPO paper introducing the clipped surrogate objective.

2. Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., & Finn, C. (2023). Direct Preference Optimization: Your Language Model is Secretly a Reward Model. NeurIPS 2023.

   • The DPO paper showing preference optimization can be reduced to classification.

3. Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Zhang, M., … & Guo, D. (2024). DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv preprint.

   • Introduces GRPO and demonstrates its effectiveness for mathematical reasoning.

4. DeepSeek-AI. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv preprint.

   • DeepSeek-R1 and R1-Zero results using GRPO with verifiable rewards.

5. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., … & Lowe, R. (2022). Training language models to follow instructions with human feedback. NeurIPS 2022.

   • The InstructGPT paper establishing the SFT → RM → PPO pipeline.

6. Christiano, P. F., Leike, J., Brown, T., Marber, M., Legg, S., & Amodei, D. (2017). Deep reinforcement learning from human preferences. NeurIPS 2017.

   • Foundational work on learning reward models from human preferences.

7. Gao, L., Schulman, J., & Hilton, J. (2022). Scaling Laws for Reward Model Overoptimization. ICML 2023.

   • Formalizes reward hacking and overoptimization in RLHF.

8. Xu, J., Xie, T., Zhao, A., Song, J., Wang, J., & Zhang, Y. (2024). Is DPO Superior to PPO for LLM Alignment? A Comprehensive Study. ICML 2024.

   • Systematic comparison showing PPO outperforms DPO when properly tuned.

9. Wen, Y., Zhang, Z., Jiao, H., Yang, M., Zhang, H., & Wang, G. (2024). From RLHF to RLHF: The Dilemma of Improving Human Alignment. arXiv preprint.

   • Analysis of reward hacking showing approval ratings can increase while correctness decreases.

10. Huang, H., Zhong, H., Li, S., Yang, K., & Zitnik, M. (2023). The N Implementation Details of RLHF with PPO. arXiv preprint.

    • Practical insights on PPO hyperparameter sensitivity and training diagnostics.

11. Bradley, R. A., & Terry, M. E. (1952). Rank Analysis of Incomplete Block Designs: I. The Method of Paired Comparisons. Biometrika, 39(3/4), 324-345.

    • The original paired comparison model adapted for preference learning.

12. Yu, Q., Zhang, H., Shao, Z., Guo, D., Zhu, Q., & Lu, H. (2025). DAPO: An Open-Source LLM Reinforcement Learning System. arXiv preprint.

    • Addresses GRPO training instabilities with dynamic sampling and clip-higher strategy.

OpenClaw is a 430K-line TypeScript project that turns any messaging platform into an interface for an autonomous AI agent. Instead of reading about it, explore the architecture below.

Three-Layer Architecture

The hub-and-spoke topology collapses into three clean layers, each with a single responsibility.

Message Flow

Seeing the layers in isolation is one thing. Here’s what happens when an actual message flows through them.

System Prompt Filesystem

The runtime’s “brain” isn’t code — it’s a filesystem of markdown documents that compose into the system prompt.

Memory System

The most distinctive primitive: long-term memory that lives on disk, not in the context window.

The entire system is an exercise in composition: message queues, schedulers, filesystems, and virtual memory, familiar abstractions from operating systems, recomposed into an AI agent.

Why this matters: new model architectures appear constantly, and vLLM needs to serve them. The interesting engineering question is how — because the optimizations that make inference fast (fused kernels, CUDA Graphs, tensor parallelism) require deep model-specific restructuring. You can’t just import a model and get peak performance. vLLM resolves this tension with a tiered system: immediate support through a compatibility layer, then a clear path to fully optimized native integration.

This post is structured into 4 parts:

1. The Gateway — how vLLM decides what to do with a model it receives
2. The Transformers Fallback — the zero-day mechanism and its trade-offs
3. Native Integration — what it takes to make a model truly fast in vLLM
4. The Execution Core — forward pass, weight loading, and distributed execution

We’ll build on concepts from previous posts. If you’re not familiar with PagedAttention and FlashAttention, or the hidden software stack beneath inference, those are worth reading first. We also won’t re-explain the Engine-Worker orchestration layer in full — just enough to ground the model integration story.

Here’s the starting point:

vllm serve some-brand-new/model-7B --dtype auto
# This works. Even for a model vLLM has never seen before.

That command succeeds for models that vLLM has no dedicated code for. Let’s understand why.

Part 1: The Gateway

The Engine-Worker-Model Hierarchy

vLLM enforces a strict separation of concerns in how it handles models. Before we get into the model-specific details, let’s establish the high-level architecture, since it determines where new model code actually lives.

There are four levels:

1. LLMEngine — the control plane. Handles scheduling, manages the BlockSpaceManager (which tracks physical GPU memory blocks), and decides which requests get processed in each iteration. The Engine is completely agnostic to model architecture.
2. Worker — one per GPU. Manages the GPU device, holds its slice of model weights, and coordinates with other Workers for distributed execution.
3. ModelRunner — sits inside each Worker. Responsible for converting logical request data (token IDs, sequence lengths) into the physical tensors the model needs. This is where input flattening happens.
4. Model — the neural network itself. Whether it’s a native LlamaForCausalLM or a wrapped TransformersModel, this is the only layer that changes when you add a new model.

The key property here: the Engine only needs the KV cache element size — derived from num_layers, hidden_size, and num_attention_heads in the model config to make scheduling decisions. It never touches the model’s forward pass. This means adding support for an entirely new architecture only changes the bottom layer of this stack. Everything above it stays the same.

The VllmConfig object is how information flows across these levels. It aggregates several sub-configs:

Registry Mechanics and the Architecture Lookup

When you run vllm serve <model>, the first thing that happens is a config.json resolution — either pulled from the Hugging Face Hub (if you pass a model ID like meta-llama/Llama-2-7b-hf) or read from disk (if you pass a local path, as you would in an air-gapped deployment). The architectures field — for example, ["LlamaForCausalLM"] — is the primary lookup key for the entire loading sequence.

This key gets checked against the _VLLM_MODELS dictionary, the core of vLLM’s ModelRegistry. It maps architecture strings to (module_name, class_name) tuples:

_VLLM_MODELS = {
    "LlamaForCausalLM":       ("llama", "LlamaForCausalLM"),
    "MistralForCausalLM":     ("mistral", "MistralForCausalLM"),
    "DeepseekV2ForCausalLM":  ("deepseek_v2", "DeepseekV2ForCausalLM"),
    "Qwen2ForCausalLM":       ("qwen2", "Qwen2ForCausalLM"),
    # ... hundreds of other architectures
}

The module_name is a relative path within vllm.model_executor.models — so "llama" resolves to vllm/model_executor/models/llama.py. The class_name is the specific nn.Module subclass to instantiate.

One important detail: vLLM does NOT import all model classes at startup. Instead, it uses _LazyRegisteredModel wrappers. When the ModelConfig requests a specific architecture, the registry:

1. Checks if the architecture string exists in _VLLM_MODELS
2. Retrieves the module path and class name
3. Dynamically imports the module using importlib
4. Returns the class constructor to the ModelLoader

This lazy loading matters for dependency isolation. A user running Llama shouldn’t need the specific kernels required for an audio-processing model. If those kernels aren’t installed and the audio model is loaded eagerly at startup, vLLM crashes for everyone.

Three things can happen when the registry receives an architecture string:

1. Found in registry → native path (optimized, model-specific code)
2. Registered by plugin → external native path (optimized, third-party code)
3. Not found → Transformers Modeling Backend fallback (compatibility shim)

The Plugin System

This is a significant evolution in vLLM’s architecture. External packages can register models without modifying vLLM core.

The mechanism uses Python’s vllm.general_plugins entry point. During vLLM’s initialization, it discovers and executes all registered plugins. A plugin can invoke ModelRegistry.register_model() to inject a new architecture mapping at runtime:

# In your package's plugin entry point
def register():
    from vllm import ModelRegistry

    if "MyNewModel" not in ModelRegistry.get_supported_archs():
        ModelRegistry.register_model(
            "MyNewModel",
            "my_package.models:MyNewModel"
        )

This decouples the vLLM release cycle from model release cycles. Model creators — Mistral, DeepSeek, Google — can ship a “vLLM adaptation package” alongside their weights. Users pip install that package, and vLLM recognizes the new architecture immediately. No PRs to vLLM core, no waiting for a new release.

Part 2: The Transformers Fallback (Zero-Day Support)

The Transformers Backend

When the registry lookup fails — or when you explicitly set model_impl="transformers" — vLLM resolves to its Transformers backend. This is a family of mixin-composed classes (TransformersForCausalLM, TransformersMoEForCausalLM, TransformersMultiModalForCausalLM, etc.) defined in vllm/model_executor/models/transformers/. They sit between vLLM’s scheduler and standard Hugging Face model code, and they’re the reason that vllm serve command works on day zero for new models.

There are two initialization steps worth understanding:

1. Config-Based Instantiation. The wrapper uses transformers.AutoModel.from_config(...) to build the model architecture on a meta device — meaning no GPU memory is allocated yet, just the module structure with placeholder parameters. Weights are loaded separately later through vLLM’s load_weights() pipeline. This two-phase approach (structure first, weights later) is critical for distributed loading: each GPU can load only its weight shard, rather than loading everything and then discarding what it doesn’t need.

2. Attention Backend Injection. Before instantiation, vLLM modifies the model’s text configuration:

# vLLM sets this before calling from_config()
text_config._attn_implementation = "vllm"

This is the critical mechanism. Modern Hugging Face models are written to be attention-backend-agnostic. They check _attn_implementation and query a registry of attention functions. vLLM populates ALL_ATTENTION_FUNCTIONS with its own PagedAttention-backed implementation:

# vLLM registers its attention backend into HF's registry
ALL_ATTENTION_FUNCTIONS["vllm"] = vllm_attention_forward

When the HF model reaches its attention layer and calls the registered function, it gets vLLM’s implementation instead of the default eager/SDPA/FlashAttention backend. The model doesn’t know the difference.

Let’s trace the data flow step by step:

1. Engine generates block_tables and slot_mapping → packs them into an AttentionMetadata object
2. TransformersForCausalLM.forward() receives flattened inputs + attn_metadata
3. The wrapper passes vLLM metadata as **kwargs into the HF model’s forward method
4. The HF model propagates **kwargs down through its layers (this is a convention in Transformers — unused kwargs flow through)
5. At each attention layer, the injected vLLM backend receives Q, K, V tensors + the attn_metadata
6. The PagedAttention CUDA kernel executes — storing K/V into paged blocks, computing attention scores using block tables

The result: even “unoptimized” models benefit from PagedAttention’s memory virtualization. No more OOM from naive KV cache pre-allocation. The KV cache is managed efficiently through fixed-size blocks, regardless of whether the model itself was designed for it.

The Trade-offs (What You Lose)

The Transformers backend enables immediate serving, but it sits in what we might call an “unoptimized valley.” Let’s be specific about the costs:

CUDA Graph Capture. In vLLM V1, the Transformers backend supports torch.compile with piecewise CUDA graph capture (via the @support_torch_compile decorator), closing what was historically the largest performance gap. However, models with dynamic RoPE scaling still fall back to eager mode. And native models can leverage more aggressive graph capture strategies that cover a larger fraction of the computation graph, since their code is explicitly written with static control flow in mind.

Kernel Fusion. Native vLLM models use fused kernels — LayerNorm + activation in one kernel, RoPE computation fused with the QKV projection, SiLU and gate multiplication combined. The fallback uses separate PyTorch operations for each step. Every separate operation means an extra round-trip to HBM: write intermediate result, read it back for the next op. On a memory-bandwidth-bound workload (which LLM decode always is), these extra reads and writes add up fast.

Parallelism Limitations. Basic Tensor Parallelism can sometimes be inferred automatically via the model’s base_model_tp_plan, but this doesn’t cover every case. Mixture-of-Experts routing, novel attention patterns, or architectures with unusual layer structures may not shard correctly — restricting you to single-GPU execution.

Note: The fallback is not meant to be the final state — it’s the starting point. It gives you a working, servable model while the community works on native integration. Think of it as a bridge: useful immediately, but you cross it to get somewhere better.

Part 3: Native Integration

The Model Interface and Prefix Protocol

To go from “supported” to “optimized,” a model must be implemented natively. This means creating a Python class that mirrors the original model structure but substitutes standard layers with vLLM’s distributed primitives.

Every module in a native vLLM model accepts a prefix="" argument during initialization. This string represents the module’s fully qualified name in the state dictionary — for example, model.layers.0.self_attn.q_proj.

class LlamaAttention(nn.Module):
    def __init__(self, config, prefix=""):
        super().__init__()
        self.qkv_proj = QKVParallelLinear(
            ...,
            prefix=f"{prefix}.qkv_proj"
        )
        self.o_proj = RowParallelLinear(
            ...,
            prefix=f"{prefix}.o_proj"
        )

The prefix serves two purposes:

1. Weight loading: maps checkpoint tensors to the correct layer instance. When the load_weights method receives a tensor named model.layers.0.self_attn.q_proj.weight, the prefix tells it exactly which module to route it to.
2. Non-uniform quantization: the QuantizationConfig can specify different quantization schemes per layer. Some layers might be FP16 while others are INT8. The prefix is how the config identifies which kernel to instantiate for each specific layer.

Parallel Layer Primitives

For models that won’t fit on a single GPU (70B+), vLLM provides distributed primitives that replace standard nn.Linear and nn.Embedding layers:

ColumnParallelLinear splits the weight matrix along the output dimension. Each GPU computes a fraction of the output features. This is used for QKV projections (each GPU computes a subset of attention heads) and MLP up-projections (each GPU computes a portion of the intermediate dimension). No inter-GPU communication is needed for this operation.

RowParallelLinear splits along the input dimension. Each GPU computes a partial result, then an AllReduce sums the partial results across all GPUs. This is used for the attention output projection and MLP down-projection — the operations where partial results need to be recombined.

VocabParallelEmbedding splits the embedding table (often 128k+ tokens for modern models) across GPUs. Each GPU holds a slice of the vocabulary and performs lookups only for tokens in its range.

The VllmConfig provides tensor_parallel_size during initialization, and each layer auto-configures its sharding based on the worker’s rank. A model developer doesn’t write explicit GPU assignment code — they use these primitives and the infrastructure handles partitioning.

Input Flattening and the 1D Computation Graph

This is one of the more interesting design decisions in vLLM. In standard PyTorch, inputs are 2D tensors of shape [batch_size, sequence_length]. This requires padding to align sequences of different lengths — if you’re processing three requests with lengths 5, 12, and 3, you pad everything to length 12. That means 8 wasted positions out of 20, nearly 40% of compute thrown away on padding tokens.

vLLM eliminates padding entirely. The ModelRunner concatenates all tokens from all concurrent requests into a single 1D tensor of shape [total_num_tokens]. A separate positions tensor (also 1D) provides the sequence position for each token:

# Three concurrent requests:
#   Request A: tokens [101, 204, 305]        (3 tokens, positions 0,1,2)
#   Request B: tokens [42, 55, 67, 89, 12]   (5 tokens, positions 0,1,2,3,4)
#   Request C: tokens [700, 801]             (2 tokens, positions 0,1)

# Flattened input — no padding, no wasted compute:
input_ids = [101, 204, 305, 42, 55, 67, 89, 12, 700, 801]  # shape: [10]
positions  = [0,   1,   2,   0,  1,  2,  3,  4,  0,   1]   # shape: [10]

Every layer in a native vLLM model is written to process this 1D stream. Embeddings do lookups on the 1D tensor. RoPE uses the positions tensor for correct positional encoding. The attention layer uses block_tables to reconstruct the logical sequence structure — knowing which tokens belong to which request and where their KV cache blocks live in physical memory.

Weight Loading — From Disk to Device

One of the trickier parts of native integration: implementing load_weights(self, weights). This method receives an iterator of (name, tensor) pairs from AutoWeightsLoader and must map checkpoint weights into the model’s parameters.

The parameter mismatch problem is why this isn’t trivial. vLLM often fuses layers that are separate in the Hugging Face checkpoint. For example, a standard Llama MLP has separate gate_proj and up_proj linear layers. In vLLM, these become a single gate_up_proj to reduce kernel launches. The load_weights logic must handle this:

def load_weights(self, weights):
    # Stacking mapping: which HF weights get concatenated into which vLLM param
    stacked_params = {
        "gate_proj": ("gate_up_proj", 0),  # goes into first half
        "up_proj":   ("gate_up_proj", 1),  # goes into second half
    }

    for name, loaded_weight in weights:
        if "gate_proj" in name or "up_proj" in name:
            # Buffer the tensor, wait for its partner, then concatenate
            param = self.state_dict()[name.replace("gate_proj", "gate_up_proj")
                                          .replace("up_proj", "gate_up_proj")]
            # Load into the correct slice of the fused parameter
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, name)
        else:
            param = self.state_dict()[name]
            param.copy_(loaded_weight)

Two utilities make this process more manageable:

AutoWeightsLoader abstracts away the routing of weights to child modules. It recursively discovers sub-modules that have their own load_weights methods and delegates the appropriate (name, tensor) pairs to each one, so the top-level model doesn’t need to manually dispatch weights. The upstream shard iteration — walking through model-00001-of-00005.safetensors through model-00005-of-00005.safetensors and presenting a unified (name, tensor) stream — happens in vLLM’s weight loading utilities (weight_utils.py), which feed into AutoWeightsLoader.

WeightsMapper provides declarative renaming rules. Instead of writing string manipulation inside load_weights, you define a mapping:

mapper = WeightsMapper(orig_to_new_prefix={
    "model.decoder.layers.": "model.layers.",
    "norm.weight": "model.norm.weight"
})

The loader applies these rules on the fly, letting the vLLM model structure diverge from the Hugging Face structure while maintaining compatibility with official checkpoints.

For quantized models, weight loading has an additional layer of complexity. In 4-bit quantization schemes like AWQ, eight 4-bit weights are packed into a single int32. The loader must recognize that the destination parameter is quantized and load the packed tensor directly, no casting to float16 first. If the config specifies quantization, the linear layer initializes a specialized “quantized parameter” object that overrides the default loading behavior.

Part 4: The Execution Core

The Attention Switchboard (ForwardContext)

In standard PyTorch, an Attention module is self-contained — it receives Q, K, V and computes the output. In vLLM, the Attention layer acts as a client to a global context. When the model executes a forward pass, a ForwardContext is established containing the AttentionMetadata generated by the scheduler.

The AttentionMetadata object is essentially a page table for the KV cache. Here’s what it contains, with concrete values for a batch of 3 requests:

# AttentionMetadata for a batch with 3 requests:
#   Request A: 128 tokens, KV spread across blocks [4, 17, 23, 8]
#   Request B: 64 tokens, KV in blocks [1, 12]
#   Request C: 256 tokens, KV in blocks [0, 5, 9, 14, 22, 31, 7, 19]

block_tables = [
    [4, 17, 23, 8, 0, 0, 0, 0],   # Request A (padded to max_blocks)
    [1, 12, 0,  0, 0, 0, 0, 0],   # Request B
    [0, 5,  9, 14, 22, 31, 7, 19], # Request C
]
# shape: [3, 8] — each entry is a physical block index in GPU memory

slot_mapping = [512, 65, 1024]
# For decode: maps each new token to its physical slot (block_idx * block_size + offset)

context_lens = [128, 64, 256]
# Sequence length per request, for correct attention masking

The attention layer’s forward() does not compute QK^T * V directly. Instead, it dispatches to different backends depending on the phase:

• Prefill (processing prompts) → FlashAttention variant, optimized for parallel computation over many tokens. All Q, K, V tokens are known upfront, so we can exploit parallelism across the sequence.
• Decode (generating tokens) → PagedAttention kernel. The query is a single new token. The kernel uses block_tables to gather K and V vectors from non-contiguous physical blocks, compute attention scores, and scatter the result. This is the operation that makes virtual memory for KV cache work, tokens don’t need to be stored contiguously.

The split between prefill and decode backends is important for performance. Prefill is compute-bound (large matrix multiplications), so FlashAttention’s tiling strategy works well. Decode is memory-bound (loading many cached K/V vectors for a single query), so PagedAttention’s gather-based approach is the right fit.

Distributed Execution Contracts

Supporting 405B-class models means multi-GPU execution across potentially many nodes. This introduces specific contracts that new model implementations must satisfy:

Tensor Parallelism requires precise synchronization. In a standard Transformer block, AllReduce happens exactly twice after the attention output projection (RowParallelLinear) and after the MLP down-projection (RowParallelLinear). These are the points where partial results from each GPU must be summed. Adding extra synchronizations (say, an unnecessary AllReduce after the QKV projection) doesn’t produce wrong results, but it degrades throughput. Each AllReduce is a blocking collective, so all GPUs wait.

Pipeline Parallelism splits the model vertically by layers. The native model’s forward method must handle an intermediate_tensors argument:

def forward(self, input_ids, positions, attn_metadata, intermediate_tensors=None):
    if intermediate_tensors is not None:
        # We're not the first pipeline stage — skip embedding
        hidden_states = intermediate_tensors["hidden_states"]
        start_layer = self.start_layer  # e.g., layer 16
    else:
        # First pipeline stage — process from the embedding
        hidden_states = self.embed_tokens(input_ids)
        start_layer = 0

    for layer in self.layers[start_layer:self.end_layer]:
        hidden_states = layer(hidden_states, positions, attn_metadata)

    return hidden_states

Rank 0 executes layers 0–N, outputs the hidden state as intermediate_tensors. Rank 1 receives it, skips the embedding layer, and resumes from layer N+1. If a developer forgets to implement this check, the model works fine in single-node TP mode but silently breaks in PP mode. it tries to re-embed already-processed hidden states.

CUDA Graph Compatibility requires static control flow. Dynamic Python branching based on tensor values like if tensor.sum() > 0: breaks CUDA Graph capture because the graph records a fixed execution path. This is particularly relevant for Mixture-of-Experts models where expert routing is inherently data-dependent. The routing logic must use masked tensor operations (scatter, gather with masks) rather than Python if/else, so the computation graph remains static even though different experts activate for different tokens.

Putting It All Together — The Optimization Ladder

To summarize the full lifecycle, model support in vLLM is a progression through increasingly optimized tiers:

1. Transformers Fallback — works immediately. PagedAttention memory management. No fusion, no graphs, limited parallelism. This is where every new model starts.
2. Plugin Registration — an external package provides a native model class. pip install and go. Model creators control their own release timeline.
3. Native Model Class — upstreamed into vLLM. Parallel primitives, 1D flattened computation, CUDA Graph compatible. This is where the performance lives.
4. Quantization Support — AWQ, GPTQ, FP8 weight loading tested and working. Packed tensor handling, per-layer quantization configs. Unlocks deployment on smaller hardware.
5. Full Production — Pipeline Parallelism support, custom attention patterns if needed, benchmarked against reference implementations. Ready for large-scale serving.

The plugin system represents the future direction — federated model support where model creators can ship “vLLM-ready” code independently of core releases. Instead of waiting for the vLLM team to implement every new architecture, the ecosystem moves toward model creators owning their integration path.

Closing

The process of introducing a new model into vLLM is a systems engineering exercise. It requires transforming a static model definition, essentially a recipe for matrix multiplications — into a dynamic, distributed execution graph that manages its own memory, shards its own weights, and coordinates across GPUs. The Transformers fallback bridges the gap for immediate access; native integration is where the performance lives.

There are four core contracts a model must satisfy for full integration: registry updates (mapping architecture strings to code), class restructuring (parallel primitives, 1D flattening), weight loading (handling mismatches between checkpoint and runtime structure), and PagedAttention integration (routing attention through the block-table-based memory system). Understanding these four contracts gives you a mental model for reasoning about model support in any inference engine, not just vLLM.

References

1. Kwon, W. et al. “Efficient Memory Management for Large Language Model Serving with PagedAttention.” SOSP 2023. arXiv:2309.06180
2. vLLM Documentation — Architecture Overview. docs.vllm.ai
3. vLLM Documentation — Adding a New Model. docs.vllm.ai
4. vLLM Documentation — Plugin System. docs.vllm.ai
5. vLLM Source — Model Registry (registry.py). GitHub
6. vLLM Source — Transformers Backend (transformers/). GitHub
7. vLLM Documentation — Class Hierarchy. docs.vllm.ai
8. Gordic, A. “Inside vLLM: Anatomy of a High-Throughput LLM Inference System.” aleksagordic.com
9. Zalt, M. “The Hidden Switchboard Behind vLLM Attention.” zalt.me
10. Prerepa, A. “ZvLLM: Zigzag forward pass with vLLM.” adiprerepa.github.io
11. El Shafie, H. “Paged Attention from First Principles: A View Inside vLLM.” hamzaelshafie.bearblog.dev
12. vLLM Documentation — Paged Attention Design. docs.vllm.ai

vllm serve meta-llama/Llama-3.1-70B-Instruct --tensor-parallel-size 8

One line. Eight GPUs. A 70-billion parameter model ready to serve requests. But this hides significant complexity.

Behind this command, three distinct software systems spring into action. Kubernetes allocates pods and manages node resources. Ray spawns actors, creates placement groups, and coordinates distributed execution. vLLM initializes workers, establishes NCCL communication rings, and begins orchestrating the token-by-token dance of autoregressive generation.

The interesting part is the choreography between systems. Each layer operates at a different granularity, speaks a different language, and solves a different class of problem. Kubernetes thinks in pods and nodes. Ray thinks in actors and tasks. vLLM thinks in requests and tokens. Yet when you hit that endpoint with a prompt, all three coordinate to produce a coherent response.

The question worth asking: How do these systems know when to hand off control to each other?

This post traces that coordination. We’ll follow the cascade from kubectl apply to the moment NCCL rings form and tensor data starts flowing. We’ll examine why placement groups matter more than you’d expect, why your network configuration can make or break performance, and how the industry is evolving toward disaggregated architectures that split inference across specialized pools.

If you’ve read the previous deep-dive on the hidden software stack behind inference, this builds on that foundation. We won’t revisit PagedAttention or continuous batching fundamentals. Instead, we’re zooming out to the orchestration layer, the software that transforms a rack of GPUs into something resembling a programmable supercomputer.

Prerequisites: This post assumes familiarity with PagedAttention, continuous batching, and basic Kubernetes concepts. If you’re new to the inference stack, start with The Hidden Software Stack Behind Fast LLM Inference.

Three Systems, Three Granularities

This stack works because of its division of labor. Each system operates at a different level of abstraction, handling the problems it’s best suited to solve.

Kubernetes sees the world in pods and nodes. It manages the lifecycle of containers, handles service discovery, and ensures workloads get scheduled onto machines with available resources. Its scheduling decisions happen at the coarse granularity of “does this node have 8 GPUs available?” Kubernetes has no concept of what happens inside those containers once they’re running.

Ray operates one level deeper. It sees actors (long-lived Python objects that can hold state and process messages) and tasks, which are stateless function invocations. Ray’s Global Control Store (GCS) maintains a distributed view of cluster resources, and its Raylets (one per node) handle local scheduling and object management. Ray also understands placement constraints: it can ensure that a group of actors lands on the same physical node, or spreads across nodes in a specific pattern.

vLLM cares about requests and tokens. It manages the KV cache, schedules which requests get processed in each iteration, and coordinates the actual tensor operations across GPU workers. vLLM’s scheduler operates at millisecond granularity, making decisions every inference step about which tokens to generate next.

Kubernetes has no understanding of GPU topology. It can count GPUs, but it cannot distinguish between eight GPUs connected via NVLink at 900 GB/s and eight GPUs scattered across nodes connected via Ethernet at 10 GB/s. Without additional tooling, Kubernetes might schedule your tensor-parallel workload across two nodes, a configuration that would perform 40-90x slower than necessary.

This is where KubeRay enters the picture. KubeRay is a Kubernetes operator that bridges the gap between Kubernetes’ pod-centric worldview and Ray’s actor-centric model. It introduces three Custom Resource Definitions (CRDs):

• RayCluster is the foundation. It defines head and worker node configurations, resource requirements, and cluster topology. Use this when you need a persistent Ray cluster for interactive development or long-running services.

• RayService builds on RayCluster to add Ray Serve deployments. It handles zero-downtime upgrades, health checking, and automatic recovery. This is the production choice for serving workloads.

• RayJob handles batch workloads. It spins up a cluster, runs a job, then tears everything down. Useful for fine-tuning runs or batch inference over large datasets.

The operator watches these CRDs and reconciles cluster state: creating pods, configuring networking, managing the Ray head node’s GCS, and ensuring workers connect properly. It’s the translation layer that lets Kubernetes manage Ray clusters without understanding Ray’s internal semantics.

The Reconciliation Dance

When you kubectl apply a RayService manifest, you trigger a cascade that touches every layer of the stack. Understanding this sequence reveals how control flows through the system.

Phase 1: KubeRay Operator Activation

The KubeRay operator runs as a deployment in your cluster, watching for changes to Ray CRDs. When it detects your new RayService, its reconciliation loop activates. The operator compares desired state (your manifest) against actual state (what’s running) and generates a plan to converge them.

Phase 2: Head Node Creation

First, the operator creates the Ray head pod. This pod runs the Global Control Store (GCS) on port 6379, a distributed metadata store that tracks cluster membership, resource availability, and actor locations. The head also exposes the Ray Dashboard on port 8265 for observability.

The head pod needs to be running and healthy before workers can join. KubeRay handles this sequencing automatically, using Kubernetes’ built-in readiness probes to gate worker creation.

Phase 3: Worker Pod Launch

Once the head is ready, the operator creates worker pods. Each worker’s entrypoint executes ray start --address=<head-ip>:6379, connecting to the head’s GCS. This is where the Kubernetes and Ray worlds first touch: Kubernetes schedules the pod, but Ray handles what happens inside.

Phase 4: Resource Discovery

Inside each worker pod, the Raylet process inspects its environment. It discovers available GPUs through CUDA, determines memory capacity, and inventories other resources. This information flows back to the GCS, which maintains a global resource table.

Phase 5: Cluster Ready

When all workers have connected and advertised their resources, the Ray cluster is ready. The GCS now has a complete picture: which nodes exist, what resources each has, and how to reach them. Ray Serve can start accepting deployment requests.

When vLLM initializes with --tensor-parallel-size 8, it needs to transform this general-purpose Ray cluster into a coordinated inference machine.

vLLM Initialization Sequence:

1. Cluster Connection: vLLM’s RayGPUExecutor calls initialize_ray_cluster(), connecting to the existing Ray cluster or starting a new one.

2. Placement Group Creation: vLLM creates a placement group with the specification [{"GPU": 1}] * 8, which means eight bundles, each requiring one GPU. The placement strategy is STRICT_PACK, meaning all bundles must land on a single node.

3. GCS Scheduling: The GCS consults its resource table. Can any single node satisfy eight GPU bundles? If yes, it reserves those resources atomically. If no, the placement group creation fails. Better to fail fast than scatter actors across nodes.

4. Actor Spawning: vLLM spawns RayWorkerWrapper actors inside the placement group. Each actor gets assigned to a specific bundle, guaranteeing GPU affinity. Ray sets CUDA_VISIBLE_DEVICES appropriately so each worker sees only its assigned GPU.

5. Process Group Initialization: Each worker calls torch.distributed.init_process_group(backend='nccl'). This creates the NCCL communicator that will handle all tensor data movement.

6. NCCL Ring Formation: NCCL establishes its communication topology (typically ring or tree patterns optimized for the underlying hardware). From this point forward, tensor data flows through NCCL, completely bypassing Ray’s object store.

Here’s how the handoff works: Ray’s job is setup and supervision. Once the NCCL rings form, Ray steps aside for the performance-critical path. Tensor data never touches Ray’s object store. It flows directly between GPUs over NVLink or the network fabric. Ray remains involved for health monitoring, actor lifecycle management, and metrics collection, but it’s out of the hot path.

Why STRICT_PACK Changes Everything

Placement groups are Ray’s mechanism for expressing scheduling constraints that go beyond “find me a node with resources.” For distributed inference, they determine whether your system performs at full speed or crawls.

Consider what happens without placement constraints. You request 8 GPU actors. Ray’s default scheduler might place 4 on Node A and 4 on Node B. Both nodes have available GPUs, the request is satisfied, everyone’s happy. Except they’re not.

The Disaster Scenario:

With tensor parallelism, every transformer layer requires an AllReduce operation to synchronize partial results across all GPUs. For a Llama-70B with 80 layers, that’s 160 AllReduce calls per forward pass. Each AllReduce must move data between every pair of GPUs.

When all 8 GPUs are on one node connected via NVLink:

• Bandwidth: ~900 GB/s bidirectional
• AllReduce latency: microseconds

When 4 GPUs are on Node A and 4 on Node B, connected via datacenter Ethernet:

• Bandwidth: ~10-25 GB/s (even with 100GbE)
• AllReduce latency: milliseconds

The performance difference is stark. You’re looking at a 40-90x bandwidth reduction for every AllReduce. For interactive inference where you need responses in hundreds of milliseconds, this makes the system unusable. A 50ms operation becomes a 2-second operation.

STRICT_PACK to the Rescue:

The STRICT_PACK placement strategy provides an atomic guarantee: “Reserve all N bundles on a single node. If no single node can satisfy the request, schedule none of them.”

# Conceptual placement group specification
placement_group = ray.util.placement_group(
    bundles=[{"GPU": 1}] * 8,
    strategy="STRICT_PACK"
)

This is all-or-nothing. Either you get all 8 GPUs on one node with NVLink connectivity, or you get an error telling you no suitable node exists. No silent degradation to a broken configuration.

SPREAD for Pipeline Parallelism:

Not all parallelism strategies want STRICT_PACK. Pipeline parallelism deliberately spans multiple nodes, with each node handling different layers of the model. Here, SPREAD makes sense: you want actors distributed across nodes to maximize aggregate memory capacity.

The communication pattern differs too. Pipeline parallelism uses point-to-point sends between adjacent stages, not AllReduce across all participants. This is less latency-sensitive because you’re overlapping computation with communication: while stage N processes micro-batch B, stage N-1 can send micro-batch C.

The placement group abstraction is what lets vLLM express “I need these actors to be co-located” without knowing anything about Kubernetes node topology. Ray’s GCS has that knowledge (from Raylet resource advertisements), and the placement group mechanism lets vLLM leverage it declaratively.

Two Interfaces, Two Purposes

Even with correct placement, there’s another way to tank your inference performance: letting NCCL traffic flow over the wrong network interface.

Production GPU nodes typically have multiple network interfaces:

• eth0: The standard Kubernetes pod network. Usually an overlay network (Calico, Cilium, Flannel) that provides cluster connectivity, DNS, and service discovery. Fine for control plane traffic: health probes, metrics scraping, Ray GCS heartbeats.

• net1/ib0/bond0: A high-performance interface connected to InfiniBand or RoCE fabric. This is your data plane, purpose-built for moving large tensors between nodes at 100-400 Gb/s with microsecond latencies.

The problem: NCCL doesn’t automatically know which interface to use. By default, it may discover eth0 first and decide that’s the interface for collective operations. Your carefully provisioned InfiniBand fabric sits idle while tensor data crawls through the overlay network.

The key environment variable:

NCCL_SOCKET_IFNAME=net1

This tells NCCL explicitly which interface to use for socket-based communication. For InfiniBand with RDMA, you’d also set:

NCCL_IB_HCA=mlx5_0

In Kubernetes, you expose multiple interfaces to pods using Multus CNI, a meta-plugin that lets you attach additional networks beyond the default pod network. Your pod spec includes annotations requesting attachment to the high-speed network:

annotations:
  k8s.v1.cni.cncf.io/networks: high-speed-net

The result is a pod with two interfaces: eth0 for Kubernetes integration, and net1 for NCCL traffic. Control plane and data plane are cleanly separated.

Why This Matters for Multi-Node:

For single-node tensor parallelism, NVLink handles everything and network configuration is less critical. But the moment you scale beyond one node, whether for pipeline parallelism, larger tensor-parallel groups, or disaggregated serving, network configuration becomes essential.

A properly configured InfiniBand fabric can deliver 400 Gb/s (50 GB/s) per port with single-digit microsecond latencies. The Kubernetes overlay network, even with modern CNIs, typically maxes out at 10-25 Gb/s with millisecond-scale latencies. For operations that happen 160 times per forward pass, this difference compounds dramatically.

Choosing Your Communication Pattern

We’ve seen how network configuration can make or break performance. The reason network matters so much depends on which parallelism strategy you’re using, and each strategy creates fundamentally different communication patterns.

Parallelism isn’t one-size-fits-all. Different strategies create different communication patterns, and understanding these patterns reveals why orchestration decisions matter.

Tensor Parallelism: The AllReduce Pattern

Tensor parallelism shards weight matrices across GPUs within a layer. Each GPU computes a partial result, then all GPUs synchronize via AllReduce to combine their contributions.

Ray’s responsibilities:

• Create STRICT_PACK placement group
• Spawn workers with correct GPU assignments
• Set CUDA_VISIBLE_DEVICES per worker
• Monitor actor health, restart on failure

What Ray doesn’t do:

• Manage AllReduce operations (that’s NCCL)
• Move tensor data (that flows through NVLink/NCCL)
• The object store is bypassed entirely for the hot path

The Communication Reality:

For an 80-layer model, tensor parallelism requires 160 AllReduce operations per forward pass (2 per layer—one after attention, one after FFN). Each AllReduce synchronizes tensors sized [batch_size, seq_len, hidden_dim]. With Llama-70B’s hidden dimension of 8192 and a batch of 32 sequences at 2048 tokens, you’re moving ~1 GB per AllReduce.

AllReduce has ring and tree implementations. Ring AllReduce on 8 GPUs requires each GPU to send and receive 7/8 of the data, essentially 7 full tensor transfers per operation. The only way this is fast is with NVLink’s 900 GB/s bandwidth.

Pipeline Parallelism: The Point-to-Point Pattern

Pipeline parallelism assigns different layers to different GPUs (or groups of GPUs). Data flows through stages sequentially: Stage 0 processes the input, sends activations to Stage 1, which processes and sends to Stage 2, and so on.

Orchestration Differences:

Ray creates a placement group that may span nodes (SPREAD rather than STRICT_PACK). Each stage gets its own bundle, and stages communicate via point-to-point sends rather than collective operations.

The Bubble Problem:

Pure pipeline parallelism has a fundamental inefficiency. While Stage 0 processes micro-batch 1, Stages 1-7 sit idle. While Stage 7 processes micro-batch 1, Stages 0-6 may be idle waiting for backward pass dependencies.

The bubble ratio quantifies this waste:

Where $p$ is the number of pipeline stages and $m$ is the number of micro-batches. With 8 stages and 8 micro-batches, you lose 7/15 ≈ 47% of potential throughput to bubbles.

Think of it this way: p is how many slices you cut the model into, m is how many requests you’re processing in parallel. With 8 pipeline stages but only 1 micro-batch, 7 out of 8 stages are always waiting, meaning 87.5% of compute wasted to bubbles. With 64 micro-batches, that drops to ~10%. The lesson: pipeline parallelism only pays off with large batches.

Continuous batching helps by keeping the pipeline fed with new requests, but the fundamental tradeoff remains: pipeline parallelism trades AllReduce bandwidth requirements for pipeline bubbles.

Expert Parallelism: The AllToAll Pattern

Mixture of Experts (MoE) models introduce a third communication pattern. Instead of every GPU needing data from every other GPU (AllReduce), or sequential point-to-point (pipeline), MoE requires AllToAll: each GPU sends different data to different destinations based on which expert each token routes to.

The orchestration complexity increases significantly. Expert assignments are dynamic (determined by a router network), so communication patterns vary per batch. Some experts may be hot (receiving many tokens) while others are cold.

Expert parallelism is its own orchestration challenge. Unlike tensor parallelism’s predictable AllReduce or pipeline’s sequential handoffs, MoE communication is dynamic. A router network decides which tokens go to which experts at runtime, so the communication pattern changes every batch. Some experts receive hundreds of tokens while others get none.

This dynamic routing breaks placement assumptions. You can’t pre-plan which GPUs need to talk to which. Solutions like expert replication (placing hot experts on multiple GPUs) and capacity factors (limiting tokens per expert) add orchestration complexity. MoE deserves its own treatment, but the key insight here is: AllToAll with dynamic routing is fundamentally harder to orchestrate than static patterns.

Prefill and Decode Don’t Have to Live Together

The inference phases we’ve discussed (prefill and decode) have fundamentally different computational profiles:

• Prefill: Process the entire prompt in parallel. Compute-bound. Benefits from high FLOPS.
• Decode: Generate tokens one at a time. Memory-bound. Benefits from high memory bandwidth.

For most of LLM inference history, both phases ran on the same hardware. But there’s no law of physics requiring this. Disaggregated serving splits them apart.

The Architecture:

1. Router receives incoming request, examines the prompt
2. Prefill Pool (optimized for compute: H100s with maximum FLOPS) processes the prompt, generates initial KV cache
3. KV Transfer moves the KV cache to the decode pool
4. Decode Pool (optimized for memory bandwidth: could be A100s or L40S) generates tokens autoregressively
5. Response streams back to client

Why Bother?

Different hardware, different economics. Prefill can run on fewer, more powerful GPUs because it’s compute-bound. You’re not paying for memory bandwidth you don’t use. Decode can run on more, cheaper GPUs optimized for memory bandwidth.

The pools also scale independently. A sudden spike in long prompts? Scale up prefill. Many concurrent users generating responses? Scale up decode. The tight coupling of traditional serving forces you to scale both together.

The KV Transfer Challenge:

The catch is moving the KV cache. For Llama-70B with 128K context, the KV cache can reach 40+ GB per request. Moving that between pools is non-trivial.

Two approaches are emerging:

• NIXL (NVIDIA Inference Transfer Library): GPU-to-GPU RDMA transfers over InfiniBand/RoCE. Keeps data on GPU memory throughout, avoiding PCIe bottlenecks.

• LMCache / Shared Storage: Write KV cache to a fast shared storage layer (think distributed NVMe or GPU memory pooling). This enables “context caching”: compute popular prompts once, reuse across millions of requests.

Context caching is particularly powerful for system prompts. If every request to your coding assistant starts with the same 8K token system prompt, why recompute that KV cache for every request? Compute it once, cache it, and let decode instances reuse it.

The Request Knows Where to Go

Traditional load balancing treats all requests as fungible. Round-robin, least-connections, random: they all assume any backend can handle any request equally well. For LLM inference with caching, this assumption is expensive.

If Request A and Request B share a common prefix (same system prompt, same few-shot examples), and Request A already warmed the KV cache on Pod 1, sending Request B to Pod 2 wastes the cache hit opportunity. You’ll recompute the shared prefix unnecessarily.

Prefix-Aware Routing:

Ray Serve implements prefix-aware routing using a prefix tree of cached prefixes. The router maintains a lightweight index of which prefixes are cached on which replicas. When a request arrives, it hashes the prefix, looks up which replica(s) have it cached, and routes accordingly.

This transforms routing from “who’s least busy?” to “who already has my context?”

Gateway API EPP:

The Kubernetes ecosystem is developing similar capabilities at the network layer through Gateway API’s Endpoint Picker (EPP) extension. Routing decisions happen in the ingress controller rather than in application code (Ray Serve).

The ingress controller can hash request properties (prompt prefix, user ID, session token) and consistently route matching requests to the same backend. This works without modifying the serving framework, using pure infrastructure-level routing.

The Tradeoff:

Locality-aware routing can cause load imbalance. If one prefix is extremely popular, its designated replica gets hammered while others sit idle. Production systems need to balance cache locality against load distribution, often through techniques like bounded load consistent hashing or spillover policies.

The evolution is clear: routing is becoming inference-aware. The network layer increasingly understands the semantics of the requests it carries, making decisions that would previously require application-level logic.

The Programmable Supercomputer

Step back and consider what this stack achieves. You start with a collection of independent machines, each with its own GPUs, memory, and network interfaces. Through layers of orchestration (Kubernetes managing containers, Ray managing actors, vLLM managing inference) these resources transform into something that behaves like a single, coherent system.

A prompt enters and gets routed to the right place based on cached state. Compute spreads across GPUs that might span multiple machines, synchronized through NCCL collectives that operate faster than the software can observe. Memory fragments across PagedAttention blocks, invisible to the model but critical for efficiency. The response streams back, one token at a time, while the system is already processing the next request.

The orchestration is the product. Without it, you have expensive hardware sitting idle. With it, you have an inference machine that can serve thousands of concurrent users at interactive latencies.

What’s Emerging:

The boundaries between these layers continue to blur. Systems like DistServe push disaggregation further, with prefill and decode pools that scale independently. KV cache transfer technologies (NIXL, LMCache) treat GPU memory across machines as a single addressable space. The trend is toward tighter integration between orchestration and execution, with systems that make placement decisions not just at startup, but continuously during inference.

Key Metrics to Watch:

If you’re operating these systems, the metrics that matter span all three layers:

• Kubernetes: Pod scheduling latency, node resource utilization, network policy drops
• Ray: Placement group creation time, actor restart rate, GCS latency (ray_gcs_* metrics)
• vLLM: vllm:gpu_cache_usage_perc (memory pressure), vllm:num_requests_waiting (queuing), time-to-first-token (prefill latency), inter-token-latency (decode performance)

The system is only as good as its weakest link. A Kubernetes scheduling delay adds latency to every request until the pod is running. A misconfigured NCCL interface tanks throughput. A hot expert without proper load balancing creates tail latencies.

Understanding the choreography (knowing which system is responsible for what, where the handoffs occur, what can go wrong at each boundary) is what separates operators who can debug production issues from those who cannot.

The stack is complex because the problem is complex. Distributed inference across dozens of GPUs, serving thousands of users, with sub-second latency requirements. But the complexity is structured. Each layer has clear responsibilities and well-defined interfaces. Master those interfaces, understand the handoffs, and the system becomes comprehensible.

Eight GPUs thinking as one. Three software systems coordinating invisibly. One simple command that hides a universe of orchestration.

That’s the stack. Now you know what’s underneath.

References

• KubeRay Operator: ray-project/kuberay — Kubernetes operator for Ray
• Ray Placement Groups: Ray docs
• vLLM Distributed Inference: vLLM docs
• NCCL AllReduce Algorithms: NVIDIA NCCL docs
• DistServe (Disaggregated Serving): Zhong et al., 2024
• Multus CNI: k8snetworkplumbingwg/multus-cni

If you’ve followed LLM infrastructure over the past two years, you’ve probably heard the greatest hits: PagedAttention eliminates memory fragmentation, continuous batching keeps GPUs busy, and FlashAttention cuts memory from O(N²) to O(N). These optimizations are real and important. They are not the full story.

Below the waterline sits a stack of specialized libraries that most engineers never encounter directly. CUTLASS generates the fused kernels that make quantization practical. Triton lets researchers write GPU code without drowning in thread indexing. FlashInfer handles the messy reality of serving workloads that FlashAttention wasn’t designed for. And NCCL quietly orchestrates communication when models span multiple GPUs.

This post dives into that hidden layer. We’ll trace the path from silicon to scheduler, examining the libraries that transform NVIDIA’s hardware capabilities into the fast inference you actually experience. If you’re deploying LLMs at scale, or simply curious about what happens beneath vLLM’s Python API, this is the stack worth understanding.

Hardware Contract

Every optimization in this stack exists because of a single physical constraint: the memory wall. Modern GPUs have a dramatic imbalance between compute capability and memory bandwidth.

Consider the H100. Its Tensor Cores can deliver roughly 2,000 TFLOPS of FP8 compute. Its HBM3 memory provides 3.35 TB/s of bandwidth. Simple division gives us a “ridge point” of about 600 ops/byte—if your workload performs fewer than 600 operations per byte loaded from memory, you’re memory-bound. Your expensive Tensor Cores sit idle, waiting for data.

LLM inference during the decode phase operates at roughly 0.5-1 ops/byte. For every token generated, the model loads billions of weight parameters, multiplies them by a single vector, and discards the weights. It’s not even close to compute-bound. This is why a $30,000 GPU often achieves single-digit percentage utilization during autoregressive generation.

To understand why, it helps to see what we’re working with.