So I ran it on my own desktop. My first answer was 3.81x, and it was wrong.

The Short Version

One RTX 5070 Ti, 16 GB of VRAM, serving gpt-oss-120b. That’s 117 billion parameters, 61 GB of weights on disk.

FreeToken            42.42 tokens/s
llama.cpp (tuned)    24.20 tokens/s     1.75x slower
Ollama (default)     11.14 tokens/s     3.81x slower

Same MXFP4 weights in all three. Same prompt, 256 output tokens, greedy. Everything separating them is a decision about what lives where.

Here is what that gap looks like. Both engines answering the same question, replayed at the speed they actually ran:

1. The Setup

An ordinary gaming desktop:

• GPU — RTX 5070 Ti, 16 GB VRAM, sm_120 (Blackwell)
• CPU — Ryzen 9 9900X, 12 physical cores
• RAM — 128 GB DDR5
• Link — PCIe 5.0 x16
• OS — Windows 11, everything run inside WSL2 (Ubuntu 24.04)

FreeToken’s CLI is Linux-only, so WSL2 wasn’t optional. It quietly costs about 7%, and I’ll come back to that.

The 128 GB of system RAM is what makes this machine interesting. Expert offloading treats host memory as the model’s real home and lets the GPU borrow pieces of it, so RAM capacity sets the ceiling on what you can serve at all.

Picking the model took longer than expected. The README advertises 290B+ models, but I checked the actual checkpoint sizes first and most of them don’t fit in 128 GB. GLM-5.2 wants about 400 GB. DeepSeek-V4-Flash wants 167 GB. gpt-oss-120b at 65 GB was one of the few that did fit, and it had a property I cared about more than size:

Ollama’s gpt-oss:120b and the Hugging Face checkpoint are the same MXFP4 weights.

Local inference comparisons are usually contaminated. One engine runs a 4-bit GGUF, the other runs something else, and part of the speed difference is really a precision difference nobody mentions. Here I could confirm all three engines load identical numbers, so the benchmark measures engines.

2. FreeToken Measures Your Machine First

Before serving anything, FreeToken runs a calibration pass over the two routes an expert can take. Copy it across PCIe to the GPU, or leave it in RAM and compute on the CPU.

ft bench bw --dtype mxfp4,bf16

PCIe came in at 56.6 GB/s. The CPU’s entire memory read bandwidth is 54.2 GB/s. So shipping an expert to the GPU costs roughly what reading it on the CPU costs, before the CPU does any arithmetic at all. PCIe 5.0 is just that fast next to dual-channel DDR5.

FreeToken only splits work across both paths when the CPU is at least twice as fast. It isn’t close here, so it picked pure PCIe streaming:

Auto-selected MoE backend: offload

Which means the paper’s headline mechanism, the bandwidth-adaptive split, does nothing on my hardware. That’s the policy working. Its job is to look at the machine and decide, and on a PCIe-rich desktop the right decision is to leave the CPU alone. It would earn its keep on a laptop with a narrow PCIe link, or a workstation with eight memory channels feeding a slower bus. I don’t have either, so I can only report that it declined to engage and was right.

3. The First Answer, and Why It Was Wrong

I started with Ollama, since that’s what most people have installed. FreeToken hit 42.42 tok/s against Ollama’s 11.14. A 3.81x gap is larger than I predicted, so I went looking for a bug in my harness instead of believing it.

The harness was fine. The baseline wasn’t. Running ollama ps during the benchmark showed what was actually happening:

gpt-oss:120b   66 GB   78%/22% CPU/GPU   4096 ctx
"offloaded 8/37 layers to GPU"

Eight complete layers on the GPU, twenty-nine running entirely on the CPU. When I went looking for a way to improve that, I found Ollama doesn’t really offer one. Its environment variables cover context length, KV cache type, flash attention. Nothing that changes how the model gets split. The one offload knob is num_gpu, which sets how many whole layers go to the GPU, and Ollama had already filled VRAM to 14.5 of 16.3 GB. Maybe one more layer would fit.

“3.81x faster than Ollama” is a true measurement of a real default configuration. It is not a fair statement about the state of the art. Ollama wraps llama.cpp, and llama.cpp has a much better option that Ollama never exposes.

4. The Real Comparison

The option is --n-cpu-moe, and why it matters is the whole point of this post.

A transformer layer holds two things with completely different shapes. Attention is small but does heavy math. The experts are enormous but only a handful get used per token; gpt-oss-120b has 128 per layer and picks 4. When Ollama sends a layer to the CPU it sends both. The small compute-heavy part gets dragged onto the slow device along with the bulky part, for no reason anyone would choose.

--n-cpu-moe N pulls them apart. Every layer’s attention stays on the GPU, and only the expert weights of the first N layers move to the CPU. Heavy math on the fast device, bulk wherever there’s room.

Getting this measured was more annoying than it should have been. CUDA 13.1’s runtime segfaults inside WSL’s driver stub, so I built llama.cpp against 12.8 instead. Then Ollama’s GGUF refused to load, because it declares its architecture as gptoss while upstream llama.cpp expects gpt-oss. Ollama’s model files aren’t readable by the project it’s built on. I ended up downloading llama.cpp’s own MXFP4 build of the same model.

With that sorted, llama.cpp reaches 24.20 tok/s. That is more than double Ollama’s 11.14, on the same engine, same weights, same GPU. One flag.

5. Why llama.cpp Still Hits a Ceiling

The sweep tells you more than the peak number does. Every time I moved two more layers’ worth of experts onto the GPU, throughput rose about 2%. The climb stops at N=28, with 15.7 of 16.3 GB used. There is nowhere left to go.

At that limit, 78% of the expert work is still on the CPU. Moving layers two at a time picks at the edge of the problem without changing it. The CPU carries the bulk either way, and it manages roughly 16 GB/s effective, under a third of what PCIe delivers here.

FreeToken gets 42.42 tok/s using less VRAM, 13.9 GB. It never makes the layer-level decision. All 36 layers compute on the GPU and individual experts get fetched across PCIe as the router asks for them. You’re moving books instead of shelves.

One measurement shows the payoff cleanly. FreeToken keeps recently-used experts cached in VRAM, 702 of the model’s 4,608, about 15%. If the router picked at random, that cache would hit 15% of the time. It hits 31.5%:

15.2% of experts resident  ->  31.5% of lookups served

Routers have favourites. Some experts are far more popular than others, and a cache working at expert granularity can exploit that. A layer-granularity cache can’t, because it committed to whole layers before the router said anything.

6. What I Take Away

The paper’s framing held up, but not through the mechanism it spends the most pages on. The bandwidth-adaptive split sat inert on my hardware. It measured the machine, concluded the CPU wasn’t worth using, and got out of the way. The gap came from the plainer decision to offload at expert granularity rather than layer granularity, plus a cache that exploits the fact that routers have favourites.

That’s the part worth carrying elsewhere. A layer is a convenient unit when you’re writing the code, because it’s how the model is built. It’s a bad unit when you’re scheduling, because it forces two unrelated decisions to be made together and one of them is usually wrong.

The other thing I’d keep is about method. My first number was measured correctly, reproduced cleanly, and misled anyway. It only became honest after I went and built the thing I had dismissed in a caveat. When a benchmark writeup says “of course, X wasn’t tuned,” that’s usually the experiment that needed running.

Paper: FreeToken: Efficient Edge-Native MoE Serving with Bandwidth-Adaptive Execution
arXiv: 2608.16157
Link: https://arxiv.org/abs/2608.16157

Large Mixture-of-Experts models are interesting for local inference because they are huge in total parameter count while activating only a small fraction of those parameters for each token. On paper, that sounds ideal: a 200B+ parameter MoE model may use only around 10–20B parameters per token, so the actual compute cost can be much smaller than the total model size suggests. But there is an obvious problem:

Even if only a few experts are active at a time, all of the expert weights still need to live somewhere.

For consumer hardware, that usually means keeping most experts in CPU memory while only part of the model fits in GPU VRAM. FreeToken is about making this setup much more efficient, and its main idea is surprisingly simple:

Instead of always moving a missing expert to the GPU, or always computing it on the CPU, FreeToken uses both paths at the same time and decides the split based on the machine’s actual memory bandwidth.

That idea is combined with different strategies for prefill, decode, expert caching, agent workloads, and GPU memory management.

The Core Problem

A typical local MoE setup looks roughly like this:

                GPU
        ┌─────────────────┐
        │ Attention       │
        │ Shared layers   │
        │ Expert cache    │
        └────────┬────────┘
                 │ PCIe
                 │
        ┌────────▼────────┐
        │ CPU DRAM        │
        │ Full expert pool│
        └─────────────────┘

The non-expert parts of the model can stay on the GPU, while the full set of experts remains in host memory and a subset is cached in VRAM. At that point, serving performance depends heavily on how well the runtime manages three resources:

• GPU compute
• PCIe bandwidth
• CPU memory bandwidth

The important part is that prefill and decode behave very differently, and that distinction drives much of FreeToken’s design.

1. Prefill and Decode Are Different Problems

For a single decode token, the router may only select a few experts. For example:

token t
   │
   ▼
 router
   │
   ├── Expert 3
   ├── Expert 7
   ├── Expert 12
   └── Expert 18

This is the normal advantage of MoE: only a small subset of the model is active. But prefill is different. Suppose a prompt contains thousands of tokens. Each token may select different experts:

token 1  → E1, E3
token 2  → E4, E8
token 3  → E2, E9
token 4  → E5, E11
...
token N  → E6, E10

Once we take the union over thousands of tokens, the model may end up touching almost every expert in the layer. So during prefill:

MoE sparsity becomes much less useful.

This leads FreeToken to use two completely different execution strategies.

This separation is one of the cleanest ideas in the paper.

2. Prefill: Stream the Entire Layer

During prefill, trying to predict which experts will be needed is not very useful. Instead, FreeToken assumes that most experts will eventually be used and simply streams an entire MoE layer from CPU memory to the GPU. The key optimization is double buffering. Imagine the GPU is computing layer L.

At the same time, FreeToken transfers the experts for layer L+1.

time ─────────────────────────────────────>

GPU:
[ compute layer L ]
                  [ compute layer L+1 ]
                                      [ compute layer L+2 ]

PCIe:
[ load layer L+1 ]
                  [ load layer L+2 ]
                                      [ load layer L+3 ]

Two buffers alternate:

Buffer A → current layer

Buffer B → next layer being transferred

After layer L finishes, the buffers swap roles. In the ideal case, the execution time becomes roughly:

instead of:

The difference is important because, instead of waiting for transfer and then computing, the runtime tries to hide one behind the other.

Why This Matters

The paper reports an experiment where an expert pool of roughly 64 GB is streamed over a PCIe connection delivering around 52.7 GB/s. The measured prefill time is very close to the time required to simply stream those weights once, which suggests that most of the GPU computation is effectively hidden behind the transfer. That is exactly what a good overlap-based runtime should achieve.

3. Decode: Exploit Expert Locality

Decode has a very different access pattern. A single token activates only a handful of experts, so streaming every expert would obviously be wasteful; FreeToken instead keeps an LRU expert cache inside GPU memory. The intuition is straightforward:

Consecutive tokens often choose similar experts.

For example:

token t-1:
E3 E7 E9 E12 E17 E24

token t:
E3 E7 E9 E12 E17 E24

The selection will not always be identical, but there is enough temporal locality that recently used experts are often useful again, which makes keeping these hot experts in VRAM worthwhile.

Cache Hit

If the router selects an expert that already exists in the GPU cache:

Router
  │
  ▼
Expert cache
  │
  └── HIT → run on GPU

This is the easy case. No transfer is needed.

Cache Miss

The interesting part happens when an expert is not in GPU memory. A normal runtime has two obvious options.

Option A: Move the expert to the GPU

CPU DRAM
   │
   │ PCIe
   ▼
GPU VRAM
   │
   ▼
GPU compute

Option B: Keep it on the CPU

CPU DRAM
   │
   ▼
CPU compute

Many systems strongly favor one of these paths. FreeToken asks a different question:

Why not do both?

4. The Most Important Idea: Bandwidth-Adaptive Execution

Suppose there are m missing experts for the current layer. FreeToken sends some of them to the GPU and computes the rest directly on the CPU.

             Cache misses
                  │
         ┌────────┴─────────┐
         │                  │
         ▼                  ▼
   PCIe → GPU          CPU execution

The two paths run in parallel. The only remaining question is:

How many experts should go to each side?

This is where the paper introduces its most important equation.

5. Deriving ($q^*$)

Let:

• $m$: number of cache misses
• $q$: number of missed experts transferred to the GPU
• $S$: size of one expert
• $B_P$: measured PCIe transfer bandwidth
• $B_H$: measured host-side expert processing bandwidth

If q experts are transferred to the GPU, their transfer time is approximately:

The remaining:

experts stay on the CPU. The important observation is that PCIe DMA is also reading from host DRAM. So CPU execution and PCIe transfer compete for the same memory subsystem. FreeToken approximates the bandwidth left for CPU execution as:

Therefore:

Because both paths execute simultaneously, the best split is approximately the point where both finish at the same time:

So:

After simplifying:

This is probably the single equation worth remembering from the paper.

6. A Simple Example

Suppose:

and:

with four cache misses:

Then:

So FreeToken chooses roughly:

1 expert  → PCIe → GPU
3 experts → CPU execution

Instead of forcing all four experts onto one side, it uses both available execution paths.

7. Why the Policy Depends on the Machine

One subtle but important point is that FreeToken does not use theoretical bandwidth from hardware specifications. It measures the actual system. The paper reports machines with very different ratios between PCIe bandwidth and CPU-side expert bandwidth. For example:

These systems should not use the same policy. Take the RTX 4060 laptop:

If there are four cache misses:

So only around one expert should be transferred to the GPU. The remaining experts are better handled on the CPU. Now look at the RTX 5090 desktop:

In that case, transferring almost every missed expert to the GPU makes much more sense. This is the reason for the name:

Bandwidth-Adaptive Execution

The runtime adapts the execution policy to the actual machine.

8. Better Expert Caching with LRU

Bandwidth-adaptive execution handles cache misses. But FreeToken first tries to reduce the number of misses in the first place. For this, it uses a global LRU expert cache. The paper reports significantly lower miss rates compared to KTransformers and llama.cpp.

For one Qwen model, the reported miss rates are roughly:

FreeToken       16%
KTransformers   41%
llama.cpp       62%

For DeepSeek-V4-Flash:

FreeToken       39%
KTransformers   59%
llama.cpp       89%

The exact numbers depend on the hardware and cache configuration, but the overall trend is clear. Dynamic locality-aware caching works better than more static placement strategies for this workload.

9. Why Static Expert Placement Can Be Limiting

A simple MoE runtime may classify experts into two groups:

GPU experts
CPU experts

Experts assigned to the GPU run on the GPU. Experts assigned to the CPU stay on the CPU. That is easy to implement, but it assumes expert popularity is relatively stable. FreeToken instead makes the system much more dynamic.

router
   │
   ▼
GPU LRU cache
   │
   ├── hit → GPU
   │
   └── miss
          │
          ▼
      q* policy
       /     \
      /       \
 PCIe → GPU   CPU

The execution path can change from token to token. That makes FreeToken less like a simple CPU offloading system and more like a heterogeneous scheduler.

10. Dynamic Scheduling Without Giving Up CUDA Graphs

There is another implementation challenge. Expert routing changes every token. One token may produce:

misses = {3, 8}

The next:

misses = {2, 9, 10}

And the next:

misses = {7}

A naive runtime could return to the CPU after each routing step:

GPU
 ↓
CPU scheduling
 ↓
GPU
 ↓
CPU scheduling
 ↓
GPU

But frequent synchronization and host-side scheduling hurt decode latency. FreeToken instead performs many routing-dependent decisions directly on the GPU. The runtime handles operations such as:

• expert deduplication
• hit/miss classification
• determining (q)
• choosing eviction victims
• rewriting logical expert IDs to physical cache slots

without repeatedly bouncing control back to Python. The CPU-side expert work is managed by persistent native worker threads, and the system is designed to preserve CUDA Graph execution as much as possible. That implementation detail is important because the scheduling policy would not be nearly as useful if the scheduler itself introduced large per-token overhead.

11. Dynamic GPU Memory Management

Another practical problem is VRAM allocation. GPU memory is not only used by expert weights. It also needs space for things like:

• non-expert model weights
• KV cache
• temporary buffers
• expert cache
• other applications using the GPU

And KV cache grows as the conversation gets longer. A static allocation might start like:

Expert cache : 12 GB
KV cache     : 4 GB

But later the same session may need:

Expert cache : 7 GB
KV cache     : 9 GB

FreeToken can resize the expert cache and give more VRAM to the KV cache when necessary. The full expert pool in CPU memory acts as the source of truth, so evicting or rebuilding the GPU expert cache does not affect correctness. This is a small detail compared to the main scheduling idea, but it matters a lot in a real local inference runtime.

12. Agent Workloads Introduce Another Problem

The paper spends quite a bit of attention on agent workloads. That makes sense because modern agents repeatedly alternate between generation and tools. A simplified agent loop looks like:

User request
    │
    ▼
LLM reasoning
    │
    ▼
Tool call
    │
    ▼
Tool result
    │
    ▼
More reasoning
    │
    ▼
Another tool call

The context keeps growing. At the same time, agent frameworks may rewrite or remove parts of older context. For example:

system
reasoning
tool call
tool output
reasoning
tool call
...

Some intermediate reasoning blocks or tool outputs may be removed later. That creates a problem for prefix reuse.

13. Semantic-Aware State Caching

FreeToken does not place reusable checkpoints at arbitrary token positions. Instead, it creates them at semantic boundaries such as:

</think>

</tool_call>

</tool_output>

The reasoning is simple. Agent frameworks usually modify context at meaningful block boundaries. So if the runtime stores state at the same kind of boundaries, there is a better chance that an old checkpoint remains reusable. For example:

system
reasoning
tool call
tool output
--------------------- checkpoint
new suffix

If the prefix above the checkpoint is still unchanged, FreeToken only needs to recompute the new suffix. Without this mechanism:

re-prefill entire context

With semantic state caching:

reuse old state
+
prefill changed suffix

This becomes especially useful for long-running agent sessions.

14. FreeToken vs. Traditional MoE Offloading

A simplified comparison looks like this.

Static-style offloading

Expert placement decided beforehand

GPU experts → GPU

CPU experts → CPU

The placement is relatively fixed.

FreeToken

                     router
                       │
                       ▼
                 GPU LRU cache
                   /       \
                hit         miss
                 │            │
                 ▼            ▼
                GPU        q* policy
                           /        \
                          /          \
                  PCIe → GPU        CPU
                          \          /
                           \        /
                              merge

The important difference is that placement and execution are continuously adapted. An expert that ran on the CPU for one token may later be cached in the GPU. A cache miss may be transferred to the GPU on one machine but executed directly on the CPU on another machine. The hardware determines the policy.

15. Performance Results

The paper reports fairly strong decode performance on a single RTX 5090. For Qwen3.6-35B-A3B, FreeToken reaches roughly:

77–83 tokens/s

with around a:

1.8×–2.3×

speedup over the strongest baseline depending on the workload. For DeepSeek-V4-Flash 284B, it reaches roughly:

22–25 tokens/s

with around:

1.5×–1.9×

speedup over the strongest baseline. These results are interesting because the model is much larger than GPU memory. The system is not trying to make the full model fit in VRAM. It is trying to make not fitting efficient.

16. Tail TTFT Is Especially Important for Agents

One metric I liked in this paper is the emphasis on tail Time to First Token. For a normal chatbot, a slow request is annoying. For an agent, it can be worse. External tool clients may have timeouts.

A very long prefill can therefore turn into an actual system failure rather than just bad latency. The paper reports worst-case TTFT below roughly:

44 seconds

for FreeToken across its evaluated workloads. Some baselines reached much larger values, including cases in the hundreds of seconds. This leads to a useful framing:

For agent systems, tail TTFT can become an availability problem.

I think this is more meaningful than reporting only average tokens per second.

17. Running Very Large Models on Consumer Hardware

One of the more eye-catching results is the RTX 4060 Laptop experiment. The GPU has only 8 GB of VRAM, but FreeToken runs a 35B-class MoE model at around:

39.3 tokens/s

The system obviously relies heavily on host memory, but it shows how far a good heterogeneous runtime can stretch relatively limited GPU hardware. At the other extreme, the paper runs a 753B parameter model on a single RTX PRO 6000. The configuration includes:

GPU VRAM : 96 GB
Host DRAM: 512 GiB
Checkpoint: ~433 GB

FreeToken reports around:

14.9 tokens/s

while llama.cpp reports roughly:

7.3 tokens/s

So “single GPU” here does not mean the entire model fits in GPU memory. It means one GPU is being combined with a very large host-memory pool. That distinction matters.

18. What I Think Are the Strongest Parts

There are a few things I particularly like about the paper.

1. Prefill and Decode Are Treated Separately

This sounds obvious after reading the paper, but it is an important design decision. The access pattern is fundamentally different.

Prefill
→ many tokens
→ almost every expert gets touched
→ streaming makes sense

Decode
→ one or a few tokens
→ sparse expert activation
→ caching makes sense

Trying to force the same optimization strategy onto both phases would leave performance on the table.

2. The $q^*$ Policy Is Very Simple

The policy boils down to:

There is no large predictive model or complicated scheduler. It takes two measured bandwidths and uses them to balance two execution paths. Simple models are attractive in runtime systems because they are easier to understand, profile, and adapt.

3. It Uses Real Hardware Measurements

The paper does not assume that “PCIe 5.0 means X GB/s” or that DRAM bandwidth equals a number from a spec sheet. It profiles the actual expert workload. That matters because effective bandwidth depends on much more than the interface’s theoretical maximum.

4. The Paper Goes Beyond One Isolated Optimization

FreeToken includes:

• prefill streaming
• expert caching
• CPU/GPU scheduling
• CUDA Graph integration
• persistent CPU workers
• agent state caching
• elastic VRAM management

So it feels like a runtime paper rather than a single micro-optimization.

19. Limitations and Things to Keep in Mind

There are also a few places where I would be careful not to over-generalize the results.

The $q^*$ Model Is Still an Approximation

The model is largely bandwidth based:

That makes sense when expert execution is strongly memory bound. But CPU execution can also depend on:

• SIMD efficiency
• quantization and dequantization
• NUMA behavior
• cache effects
• thread scheduling
• expert dimensions
• batch size

FreeToken partly handles this by measuring effective expert bandwidth rather than using theoretical DRAM bandwidth, which is a good choice. Still, the policy may need to be re-profiled when the workload changes significantly.

This Is Mainly a Small-Batch Local Serving System

The target workload is mostly:

single user
interactive inference
small batches
personal machine
agent workload

That is very different from a datacenter inference server running continuous batching across many users. With large batches, expert reuse and scheduling behavior can change significantly. So I would not assume the exact policy transfers directly to something like a large vLLM deployment.

“753B on One GPU” Needs Context

The paper does run a 753B model with one GPU. But that machine also has:

512 GiB of host memory

The full checkpoint is hundreds of gigabytes. So the result demonstrates efficient GPU + CPU heterogeneous execution, not magical compression of a 753B model into 96 GB of VRAM.

20. The Deeper Idea: Do We Need to Move the Data?

I think this is the most interesting systems insight in the paper. A common optimization question is:

How can I move the weights to the GPU faster?

FreeToken asks something slightly different:

Do I even need to move them?

If an expert already lives in CPU memory, there are two possibilities:

CPU memory
   │
   ├── move weights → GPU → compute
   │
   └── compute directly where the weights already are

Moving data is not free. Sometimes the fastest way to execute something on a GPU system is simply not to move it to the GPU. This way of thinking applies far beyond MoE models. It is a general heterogeneous-systems principle:

Compute placement should depend on both compute capability and data movement cost.

21. A Better Way to Think About the Paper

I would not describe FreeToken simply as an “expert offloading system.” A better description is:

FreeToken treats local MoE inference as a heterogeneous resource scheduling problem.

The runtime has to manage:

GPU compute
GPU VRAM
GPU memory bandwidth
PCIe bandwidth
CPU compute
CPU DRAM bandwidth
KV cache
expert locality
agent state

all at the same time. The goal is not:

How do we fit the model on the GPU?

The model often cannot fit. The better question is:

Where should every piece of computation happen, and when should its data move?

That is a much more useful systems perspective.

Now, NVIDIA’s next-generation Vera Rubin platform is entering full production in H2 2026. The Rubin GPU packs 224 SMs (up from 148), HBM4 at 22 TB/s (2.8× Blackwell), and up to 50 PFLOPS of NVFP4 inference. Meanwhile, the Rubin CPX — a separate chip designed for massive-context processing (roughly analogous to the prefill phase in LLM serving) — delivers up to 30 PFLOPS of NVFP4 compute with 128 GB of GDDR7.

The question is: what happens to the attention bottleneck map? And what would a hypothetical “FlashAttention-5” need to look like?

The Asymmetry Gets Worse — And Multiplies

FA4’s insight was clean: one asymmetry (tensor cores 2× faster, everything else flat) → identify the shifted bottlenecks → attack each one. On Rubin, the picture is messier because multiple asymmetries stack on top of each other.

The Numbers We Know

¹ BF16 ~3.6 PFLOPS is a back-of-the-envelope estimate based on SM count scaling (224/148 × 2.25 ≈ 3.4) and likely clock speed improvements, not an official NVIDIA specification. NVIDIA’s public Rubin disclosures emphasize NVFP4 and FP8 performance.

The critical unknowns are SMEM bandwidth per SM and MUFU throughput per SM. NVIDIA hasn’t disclosed these for Rubin. But the pattern from Hopper → Blackwell is instructive: tensor core throughput doubled per SM, SMEM bandwidth stayed at 128 B/cycle, MUFU stayed at 16 ops/cycle. If this trend continues — even partially — the relative gap between matmul and non-matmul only widens further.

BF16: A Modest Shift

For BF16 workloads, Rubin’s tensor core scaling is relatively modest (~1.6×). If SMEM and MUFU scale similarly or slightly less, the bottleneck profile stays roughly similar to Blackwell. FA4’s existing techniques — exp2 emulation, conditional rescaling, 2-CTA SMEM reduction — should port well with minor retuning. The 224 SMs (vs 148) mean more parallelism, which helps with occupancy but doesn’t change the per-SM roofline.

This is the easy case.

FP4: Where Things Get Interesting

FP4 is where Rubin’s design intent is most visible. At 3.5× the dense NVFP4 throughput of Blackwell, and with a new Transformer Engine featuring adaptive compression that can push effective throughput toward 50 PFLOPS, the tensor cores are leaving everything else behind at an unprecedented rate.

Let’s sketch an FA4-style roofline analysis for a hypothetical NVFP4 attention forward pass. Under a mixed-precision regime, portions of the QK^T and PV matmuls could exploit NVFP4 tensor cores (with FP32 accumulators and scaling metadata), but the softmax — max, subtract, exponentiate, sum, normalize — typically relies on higher precision (often FP32) for numerical stability. You can’t meaningfully compute exp2(x) at 4-bit precision. This creates a widening gap:

• MMA time could drop by up to ~3.5× (if NVFP4 tensor cores are fully utilized for QK^T/PV).
• Exponential time stays the same (MUFU is precision-agnostic for transcendentals — it always operates in FP32).
• SMEM traffic changes modestly (NVFP4 operands are smaller, but accumulation is still FP32, and P likely needs FP8 or BF16 precision for the PV matmul).

On Blackwell with BF16, the roofline showed MMA and exponential roughly tied (~1024 cycles each for 128³ tiles). With FP4 on Rubin, if MMA cycles drop by 3.5× but exponential stays flat, the exponential unit would dominate by a factor of ~3.5×. The softmax wouldn’t just be a co-bottleneck — it would become the overwhelming bottleneck.

This means FA4’s current 10-25% FMA emulation split would be wildly insufficient. A Rubin-optimized FP4 attention kernel might need:

1. 50-75% FMA emulation — or whatever ratio matches the new MMA-to-MUFU throughput gap.
2. Higher-degree polynomials — FP4’s quantization error is enormous (~6% relative error), so a degree-3 polynomial is still fine. But if intermediate precision increases (e.g., FP8 softmax outputs), degree-4 or degree-5 might be needed.
3. Alternative approximation strategies — at some point, polynomial emulation on FMA units hits diminishing returns due to register pressure and latency. Lookup table–based approaches, piecewise linear approximations, or even dedicated softmax-aware hardware (which Rubin’s Transformer Engine might provide) could be the next step.

The Adaptive Compression Wild Card

Rubin’s most architecturally novel feature for attention is the 3rd-generation Transformer Engine with adaptive compression. Unlike Blackwell’s structured 2:4 sparsity (which required exactly half of values to be zero, and which almost nobody used in practice), Rubin’s adaptive compression dynamically detects and skips zeros in the data stream without forcing values to zero.

What follows is speculative — NVIDIA has not disclosed how adaptive compression interacts with fused attention kernels at the SM level. But the architectural implications are worth exploring.

One possible implication for attention: the attention matrix P = softmax(QK^T) is naturally sparse in many practical settings — especially with causal masking, local attention windows, or after the first few tokens where most attention weights are near zero. If the Transformer Engine can skip near-zero P values during the PV matmul, the effective FLOPS for PV could drop without any algorithmic change.

But exploiting this — if it is exposed to software at the tensor-input boundary — would likely require the FA kernel to:

1. Produce P in a format the Transformer Engine can compress. This might mean specific memory layouts, quantization schemes, or metadata that the hardware compression can consume.
2. Coordinate compression with tiling. If compression ratios vary across tiles (dense attention regions vs. sparse regions), the pipeline timing becomes data-dependent — potentially disrupting the carefully balanced ping-pong schedules that FA4 uses.
3. Potentially rethink the softmax-to-MMA boundary. Currently, P is computed, stored to TMEM, and consumed by the PV MMA. If the compression engine operates between these stages, there might be an additional pipeline step — or it might be entirely transparent if the hardware handles it on the MMA input path.

These are open design questions, not known hardware contracts. But they illustrate the kind of hardware-algorithm co-design that FA4 demonstrated is essential.

The Backward Pass: SMEM Pressure Intensifies

FA4’s backward pass analysis showed that shared memory traffic exceeds MMA compute by ~30% in the 1-CTA case, reduced to ~5% with 2-CTA mode. On Rubin, if tensor cores get faster while per-SM SMEM bandwidth stays flat (as it did from Hopper to Blackwell — though this is not yet confirmed for Rubin), the backward pass becomes even more SMEM-dominated.

FA4’s 2-CTA approach halves operand B traffic by having each CTA stage only half of B. A natural extension would be larger cooperative MMA patterns — e.g., 4-CTA or cluster-wide operand sharing — though whether Rubin’s hardware supports this is unknown at time of writing. Blackwell’s 2-CTA mode was a new capability that didn’t exist on Hopper; Rubin may similarly introduce new cooperative primitives.

Similarly, if TMEM capacity increases on Rubin (not yet disclosed), more intermediate results could stay in TMEM, reducing SMEM round-trips. These are possibilities to watch for in the Rubin architecture documentation, not confirmed features.

The dQ atomic reduction problem also scales: with more SMs (224 vs 148), more CTAs are writing to the same dQ tiles, increasing atomic contention. The deterministic backward pass’s semaphore-based serialization becomes more expensive as the CTA count grows.

Rubin CPX: A Completely Different Attention Problem

The Rubin CPX is perhaps the most interesting challenge for FlashAttention. Designed for massive-context processing (the compute-heavy “context phase” roughly analogous to prefill in LLM serving), it has a radically different compute-to-bandwidth ratio compared to the R200:

² NVIDIA discloses 1.7 PB/s memory bandwidth for the full NVL144 CPX rack (144 GPUs), but per-chip GDDR7 bandwidth has not been officially specified. Third-party estimates place it around 2 TB/s, which would imply ~10,000-15,000 FLOPS/byte — roughly 6-10× more compute-heavy than the R200. The exact ratio matters less than the directional conclusion: CPX is far more compute-dense relative to its memory bandwidth.

This extreme ratio means:

FlashAttention’s IO-Awareness Becomes Even More Critical

The entire point of FlashAttention was to avoid materializing the N×N attention matrix in HBM. On CPX, with relatively limited memory bandwidth compared to its compute, even modest memory traffic is costly. The tiling strategy must be even more aggressive — larger tiles to maximize compute per byte loaded, potentially at the cost of more TMEM/register pressure.

Recomputation Becomes Free(er)

On CPX, compute is abundant relative to bandwidth. FlashAttention already recomputes S = QK^T in the backward pass rather than storing it. On CPX, you could afford to recompute even more aggressively — for example, recomputing partial softmax results rather than storing statistics, or computing QK^T multiple times with different tile decompositions to optimize memory access patterns.

The Kernel Architecture Might Invert

On the R200, FA4 tries to keep the tensor cores busy and hide SMEM/MUFU latency. On CPX, the tensor cores have so much throughput relative to memory bandwidth that the kernel might be memory-bound even with FlashAttention’s tiling. In this regime, the optimization target flips: instead of “keep tensor cores busy,” it becomes “minimize every byte of memory traffic, even if tensor cores idle.” This might mean:

• Longer inner loops (process more KV blocks before writing back) to amortize the cost of loading Q.
• Fused QKV loading — loading Q, K, V in a single TMA operation rather than separate stages.
• KV cache compression in-flight — if the Transformer Engine can decompress KV cache on the fly during the MMA, the effective memory bandwidth for loading K and V increases.

This is essentially a different kernel for the same algorithm — same FlashAttention math, but with the pipeline optimized for a completely different hardware profile.

Million-Token Context: Rack-Scale Attention

Vera Rubin NVL72 connects 72 GPUs with NVLink 6 at 3.6 TB/s per GPU. With 288 GB HBM per GPU, the total memory pool is 20.7 TB — enough for million-token contexts with large KV caches.

But million-token attention means N = 1,000,000, and even with FlashAttention’s linear memory, the compute scales as O(N²). At some point, the KV cache for a single head doesn’t fit on one GPU, and attention becomes a distributed problem.

Current approaches (like Ring Attention and sequence parallelism) tile the sequence across GPUs and overlap communication with computation. But the FA4-style roofline analysis hasn’t been applied to the distributed case. The relevant “feeds and speeds” expand to include:

• NVLink bandwidth (3.6 TB/s) for moving KV blocks between GPUs
• NVLink latency for synchronizing softmax statistics across sequence partitions
• Load balancing across GPUs when causal masking makes some partitions much heavier than others

An FA5 for Rubin might need to co-design the intra-SM pipeline (tile-level, as FA4 does) with the inter-GPU communication schedule (sequence-level). The LPT scheduling that FA4 applies within a single GPU could extend to scheduling across GPUs — placing the heaviest sequence partitions on GPUs that finish their local work first.

The Framework Advantage

Here’s where FA4’s most practical contribution — the CuTe-DSL framework — pays compound dividends. Adapting attention kernels to Rubin’s new hardware requires rapid experimentation:

• Tuning the FMA/MUFU emulation ratio for each precision (FP4, FP8, BF16).
• Exploring new tile sizes and pipeline stages for the different compute/bandwidth ratios (R200 vs CPX).
• Integrating with the new Transformer Engine’s adaptive compression APIs.
• Prototyping cluster-wide MMA patterns (4-CTA or beyond).

In the FA3 era, each such experiment required modifying C++ templates and waiting 55 seconds per compile. With CuTe-DSL’s 2.5-second JIT compilation, a researcher can iterate 20× faster. This isn’t just a convenience — it fundamentally changes what’s practical to explore.

The modular design also helps: the block-sparse, masking, and scheduling primitives are orthogonal, so adapting to Rubin’s new features (adaptive compression, larger clusters) can be done by adding new primitives without rewriting the core softmax or MMA logic.

Summary: The Bottleneck Multiplication Table

FA4 addressed a single clean asymmetry. On Rubin, the attention kernel faces multiple simultaneous asymmetries that vary by precision (FP4 vs BF16), chip variant (R200 vs CPX), and scale (single GPU vs rack). The methodology — roofline analysis per resource, targeted algorithmic mitigation — remains sound. But the number of configurations to optimize explodes.

The good news: FA4’s CuTe-DSL framework, modular primitives, and tuning infrastructure were designed for exactly this kind of rapid adaptation. The question isn’t whether the attention algorithm can keep up with the hardware. It’s whether the iteration speed of the framework can keep up with the proliferation of hardware targets.

If the Hopper → Blackwell transition is any guide, the answer is that the first team to have working, optimized kernels on new silicon wins — and then everyone else (including NVIDIA’s own cuDNN team) converges on the same techniques. FA4 demonstrated this pattern. On Rubin, the race starts again.

FlashAttention-4, from Tri Dao’s group at Princeton (with collaborators at Meta, Colfax Research, and NVIDIA), is the response to this new reality. Rather than porting FlashAttention-3’s Hopper-optimized kernels forward — which would either leave performance on the table or simply not compile — FA4 rebuilds from scratch around Blackwell’s actual bottlenecks. The results: up to 1613 TFLOPS/s on B200 (71% utilization), 1.3× faster than cuDNN 9.13, and 2.7× faster than Triton.

Let’s break down what changed and why.

The Problem: Asymmetric Hardware Scaling

On Hopper (H100), the performance balance between tensor cores, shared memory, and special function units was tight enough that a well-pipelined kernel could keep most resources busy. On Blackwell (B200), that balance is shattered.

Here are the numbers per SM per clock cycle:

The implication is stark. On Hopper, the attention kernel was roughly compute-bound — the bottleneck was tensor core throughput. On Blackwell, the bottleneck has shifted to shared memory traffic and exponential unit throughput. The roofline analysis in the paper confirms this: for a typical 128³ tile configuration, SMEM and exponential each take ~1024 cycles while MMA also takes ~1024 cycles. They’re neck-and-neck. But for larger tiles (256 × 128²), MMA and exponential both take 2048 cycles — meaning the kernel now has two co-dominant bottlenecks that are no longer matmul.

The forward pass bottleneck is the exponential unit. The backward pass bottleneck is shared memory bandwidth, exceeding MMA compute time by ~30%.

What’s New in Blackwell Hardware

Before diving into the algorithmic fixes, it’s worth understanding the new hardware primitives that FA4 exploits.

Tensor Memory (TMEM)

Blackwell introduces 256 KB of tensor memory per SM — a new level in the memory hierarchy specifically coupled to the tensor cores. On Hopper, MMA results were written back to registers, which created enormous register pressure and limited tile sizes. On Blackwell, MMA outputs go directly to TMEM, asynchronously. This is a game-changer for two reasons:

1. Register pressure relief. Intermediate accumulator values no longer consume the limited 256-register-per-thread budget.
2. Decoupled execution. Since MMA writes to TMEM are fully async, the tensor cores don’t block on register writeback, enabling better overlap with other operations.

TMEM is allocated in 32-column (16 KB) granules and requires explicit programmer management — it’s not a transparent cache.

Larger Tile Sizes

Blackwell MMA instructions operate on 128 × 128 tiles (doubled from Hopper’s 64 × 128). This means fewer instructions per output tile, but also that each instruction touches more shared memory — amplifying the SMEM bandwidth bottleneck.

2-CTA MMA Mode

Blackwell supports a cooperative mode where two CTAs within a cluster jointly execute a single MMA. The accumulator is partitioned along the M dimension, and — critically — each CTA only stages half of operand B in its shared memory while the hardware consumes the combined B tile. This effectively halves redundant SMEM traffic for operand B.

Fix #1: Redesigned Pipeline for MMA & Softmax Overlap

The Core Idea

Since softmax computation (exponentials, max, normalization) is now a co-bottleneck with MMA, the pipeline must be designed so that while one tile’s MMA is running, another tile’s softmax is being computed — and vice versa. FA4 uses a ping-pong schedule (similar in spirit to FA3) where two output tiles are computed per thread block, but the implementation is fundamentally different due to Blackwell’s TMEM and larger tiles.

FA3 vs FA4: Thread Assignment

On Hopper (FA3), each row of an accumulator tile was distributed across 4 threads in an interleaved pattern. Computing a row-wise operation like softmax required inter-warp shuffles — threads had to exchange partial results to reconstruct full rows before computing the row max and row sum. This shuffle overhead was a tax on every softmax step.

Here’s how FA3 handles it in C++ — notice the quad_allreduce_ call that performs cross-thread reduction, required because each thread only holds a partial row:

// FA3 (Hopper) — hopper/softmax.h
// Each thread holds a partial row. 4 threads per row in interleaved layout.
// After computing local row_sum, must shuffle across threads to get the full sum.

__forceinline__ __device__ TensorT finalize(float const final_scale=1.f) {
    SumOp<float> sum_op;
    quad_allreduce_(row_sum, row_sum, sum_op);  // ← Cross-thread shuffle required!
    TensorT scores_scale;
    #pragma unroll
    for (int mi = 0; mi < size(row_sum); ++mi) {
        float sum = row_sum(mi);
        float inv_sum = (sum == 0.f || sum != sum) ? 0.f : 1.f / sum;
        scores_scale(mi) = inv_sum * final_scale;
    }
    return scores_scale;
};

On Blackwell (FA4), the 128×128 tile and TMEM layout allow each thread to own an entire row of 128 elements. Two warpgroups of 128 threads each handle softmax, with each thread processing one complete row. This eliminates inter-warp shuffles entirely:

# FA4 (Blackwell) — flash_attn/cute/softmax.py
# Each thread owns an entire row → no shuffle needed.
# SoftmaxSm100 operates on 1 row per thread (num_rows=1).

@staticmethod
def create(
    scale_log2: Float32,
    rescale_threshold: cutlass.Constexpr[float] = 0.0,
    softmax_scale: Float32 | None = None,
):
    num_rows = 1  # ← Each thread owns exactly 1 full row. No shuffle.
    arch = 100
    row_max = cute.make_rmem_tensor(num_rows, Float32)
    row_sum = cute.make_rmem_tensor(num_rows, Float32)
    return SoftmaxSm100(
        scale_log2, num_rows, row_max, row_sum, arch, softmax_scale,
        rescale_threshold=rescale_threshold,
    )

The difference is structural: FA3’s kNRows is typically 4–8 (partial rows distributed across threads), requiring quad_allreduce_ at finalization. FA4’s num_rows = 1 means each thread computes the complete row max and row sum locally — zero cross-thread communication.

Decoupled Rescaling via Correction Warpgroup

In FA3, the online softmax rescaling step — where previous output tiles are multiplied by exp(m_old - m_new) to account for updated row maxima — happened inline within the softmax warpgroup. This put rescaling squarely on the critical path.

// FA3 (Hopper) — hopper/softmax.h
// Rescaling happens inline in the same warpgroup that computes softmax.
// This is on the critical path — must complete before the next MMA.

template<typename Tensor1>
__forceinline__ __device__ void rescale_o(Tensor1 &acc_o, TensorT const &scores_scale) {
    Tensor acc_o_rowcol = make_tensor(acc_o.data(),
        flash::convert_layout_acc_rowcol(acc_o.layout()));
    #pragma unroll
    for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) {
        #pragma unroll
        for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
            acc_o_rowcol(mi, ni) *= scores_scale(mi);  // ← On the critical path
        }
    }
};

In FA4, because intermediate results are communicated via TMEM rather than registers, the rescaling is offloaded to a separate “correction” warpgroup. The softmax warpgroups compute P and hand off rescale statistics through TMEM; the correction warpgroup applies the rescaling independently. This takes rescaling off the critical path of the main softmax + MMA pipeline.

The key enabler is TMEM: on Hopper, accumulators live in registers which are thread-private — there’s no way for a separate warpgroup to access another warpgroup’s accumulator. On Blackwell, accumulators live in TMEM which is accessible by all warpgroups within the same SM. This allows the correction warpgroup to read rescale statistics and apply them to the output accumulator independently.

The overall warpgroup assignment per thread block:

Compare with FA3’s warp specialization, where the division was simpler:

// FA3 (Hopper) — hopper/flash_fwd_kernel_sm90.h
// Warp group 0 = Producer (TMA loads), Warp groups 1+ = Consumer (MMA + softmax)
// Softmax and rescaling happen in the SAME consumer warpgroup.

pipeline_params_k.role = warp_group_idx == 0
    ? MainloopPipelineK::ThreadCategory::Producer   // WG0: data movement
    : MainloopPipelineK::ThreadCategory::Consumer;   // WG1+: MMA + softmax + rescaling

if (warp_group_idx == 0) {  // Producer
    cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
    // ... TMA loads only
} else {  // Consumer — does EVERYTHING: MMA, softmax, rescaling
    cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
    // ... MMA + softmax + rescale_o all interleaved in one warpgroup
}

FA4 splits the consumer into three specialized roles (softmax-H, softmax-L, correction), enabled by TMEM making the accumulator accessible across warpgroups.

The two softmax warpgroups are explicitly synchronized so their exponential-heavy critical sections don’t overlap, ensuring they don’t fight for the same MUFU resource.

TMEM Partitioning Strategy

With 256 KB of TMEM per SM, the allocation must be carefully planned:

Mandatory: Two output tiles (O^H, O^L) for the ping-pong schedule. At head dimension 128 with FP32 accumulators, each output tile is 128 × 128 × 4 bytes = 64 KB, so two tiles consume 128 KB — exactly half.

Remaining 128 KB can store either:

• Option A: 1 S tile (FP32, 64 KB) + 2 P tiles (BF16, 32 KB each) = 128 KB
• Option B: 2 S tiles (FP32, 64 KB each) = 128 KB, with P overwriting S in-place

FA4 chooses Option B. Why?

• Immediate compute start. Two S tiles can be computed back-to-back at pipeline startup, filling the pipeline faster.
• In-place conversion. Once S is consumed by softmax and converted to P (which is BF16, half the size), P overwrites S. This is safe because S is no longer needed after softmax.
• Statistics storage. The leftover space (since P is smaller than S) is used to store rescale statistics for the correction warpgroup.
• And two S tiles can be computed at once, immediately kicking off the pipeline.

In the actual code, you can see the TMEM capacity check and 2-CTA tile configuration:

# FA4 (Blackwell) — flash_attn/cute/flash_fwd_sm100.py
# TMEM capacity validation and tile configuration

self.cta_group_size = 2 if self.use_2cta_instrs else 1
# cta_tiler M includes only 1 CTA, scheduler accounts for cluster shape
self.cta_tiler = (self.q_stage * m_block_size, n_block_size, self.head_dim_padded)
# With 2CTA, the MMA tiler M covers both CTAs
self.mma_tiler_qk = (self.cta_group_size * m_block_size, n_block_size, self.head_dim_padded)
self.mma_tiler_pv = (self.cta_group_size * m_block_size, self.head_dim_v_padded, n_block_size)

# q_stage=2 means two Q tiles (ping-pong), each m_block_size=128 rows
# Total TMEM: 2 output tiles (O^H, O^L) + 2 S tiles (overwritten by P)
assert self.tmem_total <= SM100_TMEM_CAPACITY_COLUMNS  # 256 KB check

Register Pressure Management

Here’s where things get tight. Each SM has 65,536 registers shared across all threads. With 4 warpgroups × 128 threads = 512 threads per thread block, that’s 128 registers per thread.

But each softmax thread needs:

• 128 registers to hold a full row of input (128 BF16 values, packed 2 per register = 64, but with FP32 intermediates during computation, it’s 128)
• ~64 registers for output staging
• Additional registers for temporaries, loop variables, coefficients

This exceeds the budget. FA4’s solution: staged P storage. The first 3/4 of P is computed, stored to TMEM (which triggers the corresponding MMA), and registers are freed. Then the remaining 1/4 is computed separately. This trades a small amount of instruction overhead for staying within register limits.

# FA4 (Blackwell) — flash_attn/cute/flash_fwd_sm100.py
# Staged P storage: write 3/4 of P first to free registers, then the last 1/4.

self.split_P_arrive = n_block_size // 4 * 3           # = 96 for n_block_size=128
self.split_P_arrive = int(self.split_P_arrive / 32) * 32  # align to 32
assert self.split_P_arrive % 32 == 0
assert self.split_P_arrive < self.n_block_size

# Per-config register tuning — different configs need different register budgets
# Key: (is_causal, use_2cta, head_dim, is_sm103)
SM100_TUNING_CONFIGS = {
    (True, False, 128, False): {
        'ex2_emu_freq': 10,        # emit FMA-emulated exp2 every 10th element
        'ex2_emu_start_frg': 1,    # start emulation from fragment 1
        'num_regs_softmax': 176,   # register budget for softmax warpgroups
        'num_regs_correction': 88  # register budget for correction warpgroup
    },
    (False, True, 128, False): {
        'ex2_emu_freq': 16,
        'ex2_emu_start_frg': 1,
        'num_regs_softmax': 192,
        'num_regs_correction': 72
    },
    # ...
}

Fix #2: Software-Emulated Exponential Function

The Bottleneck

The MUFU (Multi-Function Unit) handles transcendental functions like exp2(x). At 16 ops/cycle/SM, it’s 512× slower than the tensor cores (8192 ops/cycle/SM). For attention, every element of the M × N score matrix needs an exponential evaluation during softmax. This makes MUFU a hard bottleneck.

The Trick: Polynomial Approximation on FMA Units

FA4 implements 2^x in software using the classical Cody-Waite range reduction + polynomial approximation:

2^x = 2^⌊x⌋ × 2^(x - ⌊x⌋)

The integer part 2^⌊x⌋ is computed via bit manipulation of the IEEE 754 exponent field — essentially a shift-and-add on integer ALU.

The fractional part 2^frac (where frac ∈ [0, 1)) is approximated by a low-degree polynomial evaluated with FMA (fused multiply-add) instructions.

Here’s the actual implementation — first the polynomial coefficients (computed via the Sollya package to minimize relative error):

# FA4 (Blackwell) — flash_attn/cute/utils.py
# Minimax polynomial coefficients for 2^x on [0, 1)

POLY_EX2 = {
    3: (
        1.0,
        0.695146143436431884765625,     # p1
        0.227564394474029541015625,     # p2
        0.077119089663028717041015625,  # p3
    ),
    # Degree 3 matches hardware MUFU to within 1 BF16 ULP on 99% of inputs
}

And the full emulation algorithm — note the range reduction via bit tricks and polynomial evaluation via Horner’s method with FMA:

# FA4 (Blackwell) — flash_attn/cute/utils.py

@dsl_user_op
def ex2_emulation(x: Float32, *, poly_degree: int = 3) -> Float32:
    fp32_round_int = float(2**23 + 2**22)
    x_clamped = cute.arch.fmax(x, -127.0)  # Step 1: clamp to avoid underflow

    # Step 2: compute ⌊x⌋ via round-down trick
    x_rounded = add_round_down(x_clamped, fp32_round_int)
    x_rounded_back = x_rounded - fp32_round_int
    x_frac = x_clamped - x_rounded_back  # Step 3: fractional part ∈ [0, 1)

    # Step 4: evaluate polynomial (Horner's method using FMA)
    x_frac_ex2 = evaluate_polynomial(x_frac, POLY_EX2[poly_degree])

    # Step 5: combine via bit manipulation — shift ⌊x⌋ into exponent field
    return combine_int_frac_ex2(x_rounded, x_frac_ex2)

The Horner evaluation is pure FMA — each step is a single fused multiply-add:

# FA4 (Blackwell) — flash_attn/cute/utils.py

def evaluate_polynomial(x: Float32, poly: Tuple[Float32, ...]) -> Float32:
    deg = len(poly) - 1
    out = poly[deg]                          # Start from highest coefficient
    for i in cutlass.range_constexpr(deg - 1, -1, -1):
        out = out * x + poly[i]              # Each iteration = 1 FMA instruction
    return out

And the bit manipulation to combine the integer and fractional parts drops to raw PTX:

# FA4 (Blackwell) — flash_attn/cute/utils.py
# Combines 2^⌊x⌋ (integer) and 2^frac (polynomial result) via bit ops

@dsl_user_op
def combine_int_frac_ex2(x_rounded: Float32, frac_ex2: Float32) -> Float32:
    return cutlass.Float32(llvm.inline_asm(
        T.f32(), [x_rounded.ir_value(), frac_ex2.ir_value()],
        "{\n\t"
        ".reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i;\n\t"
        "mov.b32 x_rounded_i, $1;\n\t"        # reinterpret float as int
        "mov.b32 frac_ex_i, $2;\n\t"
        "shl.b32 x_rounded_e, x_rounded_i, 23;\n\t"  # shift ⌊x⌋ into exponent field
        "add.s32 out_i, x_rounded_e, frac_ex_i;\n\t"  # combine: 2^⌊x⌋ × 2^frac
        "mov.b32 $0, out_i;\n\t"
        "}\n",
        "=f,f,f",
    ))

Compare this entire machinery with FA3’s approach — a single hardware instruction:

// FA3 (Hopper) — hopper/softmax.h
// Pure hardware MUFU — simple but bottlenecked at 16 ops/cycle/SM
#pragma unroll
for (int mi = 0; mi < size<0>(tensor); ++mi) {
    #pragma unroll
    for (int ni = 0; ni < size<1>(tensor); ++ni) {
        tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
        //                ^^^^^ Always hardware MUFU. No alternative path.
    }
}

Why Not Emulate Everything?

Full emulation has costs: more registers for polynomial coefficients and intermediates, higher latency per evaluation, and register bandwidth consumption that can cause spills. So FA4 takes a hybrid approach:

# FA4 — flash_attn/cute/softmax.py
# Hybrid exp2: some fragments use hardware MUFU, others use FMA emulation
@cute.jit
def apply_exp2_convert(self, acc_S_row, acc_S_row_converted,
                        ex2_emu_freq=0, ex2_emu_res=4, ex2_emu_start_frg=0):
    frg_tile = 32
    frg_cnt = cute.size(acc_S_row) // frg_tile
    acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile))

    for j in cutlass.range_constexpr(frg_cnt):
        for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2):
            if cutlass.const_expr(ex2_emu_freq == 0):
                # All hardware MUFU (fallback, or SM103/B300 with doubled MUFU)
                acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True)
                acc_S_row_frg[k+1, j] = cute.math.exp2(acc_S_row_frg[k+1, j], fastmath=True)
            else:
                if cutlass.const_expr(
                    k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res
                    or j >= frg_cnt - 1
                    or j < ex2_emu_start_frg
                ):
                    # Hardware MUFU path
                    acc_S_row_frg[k, j] = cute.math.exp2(...)
                    acc_S_row_frg[k+1, j] = cute.math.exp2(...)
                else:
                    # Software FMA emulation path (runs in parallel with MUFU!)
                    acc_S_row_frg[k, j], acc_S_row_frg[k+1, j] = \
                        utils.ex2_emulation_2(acc_S_row_frg[k, j],
                                              acc_S_row_frg[k+1, j])

The ex2_emu_freq and ex2_emu_res parameters control the split. For example, ex2_emu_freq=10, ex2_emu_res=4 means every 10th pair of elements, 4 are emulated and 6 use hardware. The exact ratios are tuned per configuration.

Accuracy

The degree-3 polynomial has a max relative error of ~8.8 × 10⁻⁵ in FP32. Sounds bad? It doesn’t matter. After rounding to BF16, the quantization error (~3.9 × 10⁻³) completely dominates the polynomial approximation error. The degree-3 polynomial matches hardware MUFU to within 1 BF16 ULP on 99% of inputs. For attention, where softmax outputs are consumed at BF16 precision, this is more than sufficient.

Fix #3: Conditional Softmax Rescaling

The Problem

Standard online softmax (as used in all FlashAttention versions) maintains running statistics:

m_j = max(m_{j-1}, rowmax(S_j))
ℓ_j = exp(m_{j-1} - m_j) · ℓ_{j-1} + rowsum(exp(S_j - m_j))
O_j = exp(m_{j-1} - m_j) · O_{j-1} + exp(S_j - m_j) · V_j

The rescaling term exp(m_{j-1} - m_j) · O_{j-1} is a full vector multiplication that happens every iteration, even when the max barely changes.

The Fix

FA4 introduces a threshold τ (set to log₂(256) = 8.0). The implementation is embedded directly in the update_row_max method we saw earlier:

# FA4 — flash_attn/cute/softmax.py
# Conditional rescaling: skip when max hasn't changed significantly
acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
acc_scale = cute.math.exp2(acc_scale_, fastmath=True)

if cutlass.const_expr(self.rescale_threshold > 0.0):
    if acc_scale_ >= -self.rescale_threshold:
        # New max is close to old max → skip rescaling
        row_max_new = row_max_old
        row_max_safe = row_max_old
        acc_scale = 1.0  # No-op rescale (multiply by 1)

Compare with FA3, where rescaling is unconditional every iteration:

// FA3 (Hopper) — hopper/softmax.h
// Unconditional rescaling — no threshold check
#pragma unroll
for (int mi = 0; mi < size(row_max); ++mi) {
    float scores_max_cur = !Check_inf
        ? row_max(mi)
        : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
    scores_scale(mi) = exp2f((scores_max_prev(mi) - scores_max_cur)
                             * softmax_scale_log2);  // ← Always computed
    row_sum(mi) *= scores_scale(mi);                 // ← Always applied
}

When the new max doesn’t exceed the old max by more than τ, FA4 skips rescaling entirely. The key insight: the final normalization step Output = O_final / ℓ_final corrects any accumulated drift. As long as intermediate values don’t overflow (which the threshold of 256× prevents), the final result is exact.

In practice, rescaling is needed only in the first few iterations when the running max is still being established. Once it stabilizes, most iterations skip rescaling entirely — saving a vector multiply per iteration.

Fix #4: 2-CTA Backward Pass

The backward pass is where shared memory pressure is most severe. With five MMA operations (vs. two in the forward pass), SMEM traffic exceeds MMA compute by ~30% in the 1-CTA configuration.

How 2-CTA Helps

In 2-CTA MMA mode, a CTA pair cooperatively executes each MMA with M = 256 (each CTA holds half the M dimension). The critical benefit: each CTA only stages half of operand B in shared memory, while the hardware reads the combined B tile from both CTAs. This roughly halves SMEM traffic for operand B across the five backward GEMMs.

The roofline improvement is significant:

The dQ Problem

There’s a catch. The dQ computation accumulates along the KV sequence dimension (the outer loop), and its reduction axis is N — which is naturally split across the CTA pair. Each CTA still needs the full reduction for its rows.

FA4 solves this using Distributed Shared Memory (DSMEM) to exchange half of the dS tile between the two CTAs in the same cluster. After the exchange, each CTA holds an (M/2 × 2N) slice of dS, enabling a CTA-pair UMMA with doubled reduction dimension. The per-CTA dQ MMA becomes:

dQ tile shape: (M/2, 2N) × (2N, d) → (M/2, d)

This restructuring also halves the number of global atomic reductions for dQ, since each CTA writes only half the dQ tile. Atomics are expensive and introduce nondeterminism, so cutting them in half is a double win.

Deterministic Mode

For reproducible training (critical for RL applications), FA4 provides a deterministic execution mode using semaphore-based serialization of global reductions. The performance overhead is minimized through careful CTA scheduling:

• Batches processed as the outermost dimension
• Heads swizzled within L2 cache capacity
• For causal masking: KV blocks launched in descending order, query blocks in ascending order from the diagonal, dQ reductions ordered by descending query block index (“shortest-processing-time-first”)

This gets the deterministic backward pass to ~75% the speed of the nondeterministic 1-CTA version.

Fix #5: LPT Scheduling

Load imbalance is inherent in attention — causal masking means tiles near the diagonal do more work than tiles far from it; variable-length batches have different sequence lengths.

FA4 applies Longest-Processing-Time-First (LPT) scheduling, a classical result from parallel processing theory. The implementation is clean — just reverse the block order:

# FA4 — flash_attn/cute/tile_scheduler.py
# LPT: reverse block order so heaviest tiles (near diagonal) run first
@cute.jit
def get_current_work(self) -> WorkTileInfo:
    # ... L2-swizzled coordinate mapping ...

    # Longest-processing-time-first: one line does the trick
    if const_expr(params.lpt):
        block = params.num_block - 1 - block  # ← Simply reverse!

    is_valid = self._tile_idx < params.total_blocks
    return WorkTileInfo(
        (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(self._split_idx)),
        is_valid
    )

But the real insight is in how it interacts with L2 cache locality. The scheduler divides heads into sections that fit in L2 cache, and traverses heads-per-section → mblocks (reversed) → sections → batches:

# FA4 — flash_attn/cute/tile_scheduler.py
# L2-aware head swizzling: fit as many KV heads as possible in L2
size_one_kv_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size
size_l2 = 50 * 1024 * 1024  # 40 MB budget for K & V

# swizzle = how many heads fit in L2 (rounded to power of 2 for fast divmod)
log2_floor = lambda n: 31 - clz(n)
swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head))

For GQA/MQA, all query heads per KV head are traversed before varying over mblocks — ensuring maximum KV reuse from L2 cache.

Empirical gains: 4–8% FLOPS improvement for MHA, 7–14% for MQA on H200 (this optimization is architecture-agnostic and also benefits FA3 on Hopper).

The Framework: CuTe-DSL

Perhaps the most practically impactful contribution for the broader ecosystem: FA4 is implemented entirely in CuTe-DSL embedded in Python — no C++ template metaprogramming whatsoever.

C++ Templates vs. Python DSL: A Side-by-Side

The same conceptual operation — computing softmax rescaling — looks fundamentally different in the two frameworks:

// FA3 (C++ CUTLASS templates) — hopper/softmax.h
// 30+ lines of template metaprogramming
template<bool Is_first, bool Check_inf=false, typename Tensor0>
__forceinline__ __device__ TensorT max_get_scale(Tensor0 &acc_s) {
    Tensor scores = make_tensor(acc_s.data(),
        flash::convert_layout_acc_rowcol(acc_s.layout()));
    static_assert(CUTE_STATIC_V(size<0>(scores)) == kNRows);
    TensorT scores_scale;
    if constexpr (Is_first) {
        flash::template reduce_max</*zero_init=*/true>(scores, row_max);
        cute::fill(scores_scale, 1.f);
    } else {
        Tensor scores_max_prev = make_fragment_like(row_max);
        cute::copy(row_max, scores_max_prev);
        flash::template reduce_max</*zero_init=*/false>(scores, row_max);
        #pragma unroll
        for (int mi = 0; mi < size(row_max); ++mi) {
            float scores_max_cur = !Check_inf
                ? row_max(mi)
                : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
            scores_scale(mi) = exp2f((scores_max_prev(mi) - scores_max_cur)
                                     * softmax_scale_log2);
            row_sum(mi) *= scores_scale(mi);
        }
    }
    return scores_scale;
};

# FA4 (Python CuTe-DSL) — flash_attn/cute/softmax.py
# Same algorithm, same low-level control, readable Python
@cute.jit
def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int):
    if cutlass.const_expr(is_first):
        row_max_new = self._compute_row_max(acc_S_row)
        row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
        acc_scale = 0.0
    else:
        row_max_old = self.row_max[0]
        row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old)
        row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0
        acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2
        acc_scale = cute.math.exp2(acc_scale_, fastmath=True)
        if cutlass.const_expr(self.rescale_threshold > 0.0):
            if acc_scale_ >= -self.rescale_threshold:
                row_max_new = row_max_old
                row_max_safe = row_max_old
                acc_scale = 1.0
    self.row_max[0] = row_max_new
    return row_max_safe, acc_scale

Same algorithm, same low-level control (both compile down to PTX), but the Python version is more readable, doesn’t require template<bool Is_first, bool Check_inf=false, typename Tensor0> incantations, and compiles in seconds instead of minutes. When CuTe-DSL’s APIs don’t cover something, raw PTX is available as an escape hatch — as we saw in the combine_int_frac_ex2 function.

Compile Time

And FA3 required precompiling hundreds of kernels for different attention variants. FA4’s JIT compilation means you compile only what you need, when you need it.

Practical Impact

The practical impact is already visible: developers have built FlexAttention and block-sparse attention variants on top of FA4’s framework without modifying core code. The barrier to entry drops from “years of C++ template metaprogramming expertise” to “a few months of GPU programming experience.”

The design philosophy is modular: block-sparse patterns, masking strategies, variable-length handling, and work scheduling are all orthogonal, composable primitives. New attention variants get all existing optimizations for free.

Results

On B200 with BF16, head dimension 128:

• Forward pass: 1.1–1.3× faster than cuDNN 9.13, 2.1–2.7× faster than Triton. Peak: ~1600 TFLOPS (71% of theoretical max).
• Backward pass: Consistent speedups across sequence lengths for both causal and non-causal.
• Deterministic backward: Up to 75% the speed of the non-deterministic 1-CTA backward — a practical option for RL training that demands reproducibility.

Notably, since FA4’s release, the cuDNN team has incorporated many of FA4’s techniques into cuDNN 9.14+, converging to similar performance. This is perhaps the strongest validation of the algorithmic contributions.

Key Takeaway

FlashAttention-4 is a case study in what happens when you take hardware asymmetry seriously. The tensor cores got faster, so the bottleneck moved. Rather than hoping the bottleneck would go away, the authors identified exactly where the cycles were being spent (roofline analysis with concrete cycle counts for every resource), then attacked each bottleneck with a targeted fix:

• SMEM too slow → 2-CTA MMA to halve operand B traffic
• Exponential unit too slow → Polynomial emulation on FMA units
• Rescaling wastes cycles → Conditional rescaling with deferred correction
• Pipeline bubbles → Ping-pong schedule with TMEM-based decoupling
• C++ compile times too slow → CuTe-DSL in Python

The code is open source at github.com/Dao-AILab/flash-attention.

The B300/GB300 will double MUFU throughput to 32 ops/cycle/SM — and FA4 already handles this: the tuning config sets ex2_emu_freq: 0 for SM103, disabling emulation entirely. When the bottleneck map shifts again, the framework is ready.

This release was notable because it marked the first time OpenAI publicly deployed a model powered by Cerebras hardware, following their reported multi-billion-dollar partnership. I had been anticipating this launch, as it seemed like a potential inflection point — one that could challenge the long-standing GPU dominance in AI infrastructure and open the door for alternative chip architectures to play a serious role.

OpenAI described Spark as an initial step toward expanding Cerebras usage across more frontier models as they scale their WSE-based datacenter capacity. Below are benchmark results for GPT-5.3-Codex-Spark.

SWE-Bench Pro Benchmark

Terminal-Bench 2.0 Benchmark

The benchmarks show that the model does not necessarily achieve the highest accuracy across all metrics. However, that is not the primary objective. The defining characteristic is speed. The model delivers ultra-low latency responses, which is precisely where Cerebras hardware is designed to excel.

Importantly, OpenAI did not rely solely on Cerebras hardware to achieve these gains. They also made substantial changes to their inference infrastructure. OpenAI restructured their streaming pipeline to reduce client-server communication overhead, redesigned parts of the inference stack, and improved session initialization so the first token appears more quickly. By introducing persistent WebSocket connections and optimizing the Responses API, they significantly reduced per-request overhead, per-token processing latency, and time-to-first-token. These improvements are now enabled by default for Codex-Spark and are expected to become standard across other models.

This led me to a central question.

“If OpenAI is actively integrating Cerebras hardware and planning broader adoption, does this signal a shift in chip market dominance? Could GPUs be replaced?”

The answer appears to be NO.

OpenAI has clearly stated that GPUs remain foundational to their training and inference pipelines. GPUs continue to provide the most cost-effective token generation at scale and remain the backbone of general-purpose AI workloads. Cerebras hardware, on the other hand, is optimized for specialized demand — particularly scenarios that require extremely low latency. Rather than replacing GPUs, OpenAI’s strategy appears to combine both architectures, leveraging each where it performs best.

After reading the release notes, I became more curious about what makes the Cerebras chip fundamentally different and what kind of software stack supports it.

Cerebras Chip Architecture

OpenAI explicitly stated that Spark runs on the WSE-3 chip.

Unlike traditional accelerators that cut a silicon wafer into many smaller dies, Cerebras keeps the entire 300mm wafer intact and turns it into a single massive processor. Instead of connecting many chips together via high-speed interconnects, WSE-3 eliminates chip-to-chip communication by implementing a large-scale on-wafer mesh network.

WSE-3 Chip

This design raises two natural questions.

• Why did the industry historically avoid wafer-scale processors?
• If we serve models on such hardware, how must the software stack change?

Why Chips Were Traditionally Cut

The answer lies in economics and physics.

Semiconductor manufacturing inevitably produces defects. The yield approximately decreases exponentially as area increases. As chips become larger, the probability that a defect affects the die increases significantly. Smaller dies improve overall manufacturing yield because a defect only invalidates a small portion of the wafer rather than the entire structure. For decades, yield optimization directly translated into economic viability.

Lithography constraints also played a role. Photolithography equipment has a maximum reticle size, meaning only a limited region can be patterned in one exposure. Larger chips require stitching techniques that historically introduced reliability and alignment challenges. These limitations reinforced the industry’s preference for smaller dies.

Packaging and thermal management further favored modular designs. Smaller chips are easier to cool, test, and replace. A wafer-scale processor requires specialized power delivery systems, advanced cooling solutions, and routing mechanisms capable of bypassing defective regions. Cerebras invested heavily in defect-aware routing and custom cooling to make wafer-scale computing feasible.

In short, traditional design philosophy prioritized yield, modularity, and replaceability. Wafer-scale computing only became viable when technological advances and AI workload characteristics justified the trade-off.

What WSE-3 Changes Architecturally

Traditional GPU systems assume a distributed architecture. Multiple accelerators are connected via NVLink, PCIe, or other networking fabrics. Each device has limited memory bandwidth, and large models must be partitioned carefully. Communication primitives such as AllReduce become central to scaling performance.

Wafer-scale computing inverts this assumption. Instead of distributing computation across multiple chips, the system behaves as a single spatially distributed processor. The interconnect exists directly on the wafer, dramatically reducing communication distance and latency. Memory is distributed across on-chip SRAM rather than relying primarily on external HBM.

This is not merely a hardware upgrade. It alters the abstraction model upon which modern ML systems software has been built.

How the Model Serving Pipeline Changes

In a conventional GPU serving stack, models are defined in frameworks such as PyTorch or JAX, compiled through intermediate representations like ONNX, optimized by tools such as TensorRT, and lowered into CUDA kernels. Runtime systems manage tensor parallelism, pipeline parallelism, and KV cache sharding, while NCCL handles collective communication. The serving infrastructure is fundamentally distributed.

With WSE-3, the structure shifts.

If a model fits on a single wafer, tensor parallelism may become less necessary. Rather than partitioning parameters across multiple devices, the optimization problem becomes one of spatial placement across hundreds of thousands of processing elements. The compiler’s focus shifts from kernel fusion and warp scheduling to global graph mapping and routing efficiency.

Memory management also changes. In GPU systems, developers carefully manage HBM, L2 cache, shared memory, and registers. KV cache eviction and sharding policies are designed around distributed memory constraints. In a wafer-scale architecture with distributed SRAM, memory locality becomes a graph-placement problem rather than a multi-device synchronization problem.

Serving strategy shifts as well. GPU clusters often prioritize throughput via dynamic batching and aggressive utilization strategies. Wafer-scale systems, with enormous on-chip bandwidth and minimal inter-device communication overhead, may enable different latency-throughput trade-offs. Some distributed coordination complexity disappears, but placement and routing complexity increases.

What Tooling Exists Today

Cerebras provides a software development kit designed specifically for its hardware. The Cerebras Software Language allows low-level programming of wafer processing elements, and simulator environments enable development without direct hardware access. There are also model repositories adapted for Cerebras systems and Python SDKs for inference through API endpoints.

However, the ecosystem remains more vertically integrated than CUDA-based development. Hardware access is typically restricted, and much of the software stack is tightly coupled to Cerebras infrastructure. While tooling exists, it is not as broadly accessible or mature as mainstream GPU ecosystems.

For those interested in exploring the SDK and CSL in more detail, Cerebras provides documentation and tutorials at the following link: https://sdk.cerebras.net/computing-with-cerebras

Final Thoughts

The release of GPT-5.3-Codex-Spark does not signal the end of GPU dominance. Instead, it represents a diversification of AI hardware strategies.

GPUs remain foundational for large-scale training and cost-efficient inference. Cerebras hardware, particularly WSE-3, appears optimized for extreme low-latency workloads. OpenAI’s approach suggests a hybrid infrastructure model, where different chips are deployed based on workload characteristics rather than ideological preference.

The more interesting question is not whether Cerebras will replace GPUs, but how heterogeneous hardware architectures will reshape compiler design, serving infrastructure, and system-level optimization in the coming years.

LLVM is a compiler infrastructure centered around a single, low-level intermediate representation. It unified frontends and backends by providing a common optimization and code generation layer. But LLVM IR is intentionally minimal and low-level. Control flow is expressed through basic blocks. Structured loops are flattened. High-level semantics disappear quickly.

MLIR exists because that design, while powerful, is insufficient for modern workloads — especially tensor programs, graph-based computations, and domain-specific transformations.

The key insight behind MLIR is not that we need a better IR. It is that we need multiple IRs, each preserving the right level of abstraction at the right time.

The Structural Problem LLVM Couldn’t Solve

LLVM is extremely good at optimizing scalar programs and generating machine code. But imagine you want to perform polyhedral loop transformations, tensor fusion, or structured scheduling. These optimizations rely on structured information — loop boundaries, affine expressions, tensor shapes.

By the time code reaches LLVM IR, that structure is already lost.

Original C File

int sum_prefix(int *a, int n) {
  int s = 0;
  for (int i = 0; i < n; i++) {
    s += a[i];
  }
  return s;
}

Converted LLVM IR

define i32 @sum_prefix(ptr %a, i32 %n) {
entry:
  br label %loop.header

loop.header:
  %i = phi i32 [ 0, %entry ], [ %i.next, %loop.latch ]
  %s = phi i32 [ 0, %entry ], [ %s.next, %loop.latch ]

  ; loop condition: i < n ?
  %cond = icmp slt i32 %i, %n
  br i1 %cond, label %loop.body, label %exit

loop.body:
  ; load a[i]
  %idx = sext i32 %i to i64
  %elem.ptr = getelementptr i32, ptr %a, i64 %idx
  %elem = load i32, ptr %elem.ptr

  ; s += a[i]
  %s.next = add i32 %s, %elem
  br label %loop.latch

loop.latch:
  ; i++
  %i.next = add i32 %i, 1
  br label %loop.header

exit:
  ret i32 %s
}

You can reconstruct it, but it’s painful and brittle.

MLIR changes the timeline. Instead of lowering everything immediately into a flat IR, it allows high-level representations to coexist with lower-level ones. A tensor operation can remain a tensor operation long enough to be optimized meaningfully before being lowered into loops. An affine loop can remain affine until you’ve exhausted algebraic and structural transformations.

This is why MLIR is not just “another IR.” It is a framework for defining IRs at multiple abstraction levels.

Dialects: IR as a Modular Concept

In MLIR, there is no single universal IR. Instead, there are dialects.

A dialect defines its own types, operations, invariants, and semantics. It is effectively a self-contained language embedded within MLIR’s infrastructure.

This modularity is not cosmetic. It changes how optimization works.

In LLVM, every optimization pass must be correct for all IR. In MLIR, passes can target specific dialects. A pass that manipulates affine loops doesn’t need to understand tensor semantics. A pass that rewrites polynomial arithmetic doesn’t need to understand GPU kernels.

General MLIR Diagram

Optimization becomes scoped and composable.

This design dramatically lowers the barrier for building domain-specific compilers. You no longer need to fork LLVM or write a monolithic optimizer. You define a dialect, declare its semantics, and write transformations that operate only within that semantic space.

Progressive Lowering: Abstraction as a Resource

The most important conceptual shift in MLIR is progressive lowering.

Traditional compilation looks like this:

High-level language → LLVM IR → Machine code

MLIR replaces that linear pipeline with a layered one:

High-level dialect → mid-level dialect → ... → lower-level dialect → LLVM dialect → machine code

Lowering is explicit and incremental.

Each layer preserves just enough structure for the optimizations that belong at that layer. Once you lower too far, certain transformations become either impossible or extremely expensive to express.

This idea — that abstraction is a resource to be preserved until no longer needed — is the philosophical core of MLIR.

Two Ways to Transform IR

When you begin writing passes, you quickly encounter two distinct mental models.

The first is explicit traversal. You walk the IR tree and mutate operations directly. This feels familiar if you’ve worked on traditional compilers. It gives you complete control and allows global reasoning, such as common subexpression elimination or whole-function analysis.

#include "lib/Transform/Affine/AffineFullUnroll.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/LoopUtils.h"
#include "mlir/include/mlir/Pass/Pass.h"

using mlir::affine::AffineForOp;
using mlir::affine::loopUnrollFull;

void AffineFullUnroll::runOnOperation() {
	getOperation().walk([&](AffineForOp op) {
		if (failed(loopUnrollFull(op))) {
			op.emitError("unrolling failed");
			signalPassFailure();
		}
	})
}

Explicit AST walking example on the optimization pass

The second model is pattern rewriting. Instead of scanning the IR manually, you declare rewrite rules: whenever a certain shape appears, replace it with another shape. The rewrite engine applies these patterns greedily until no more matches exist.

#include "lib/Transform/Arith/MulToAdd.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/include/mlir/Pass/Pass.h"

using arith::AddIOp;
using arith::ConstantOp;
using arith::MulIOp;

//define pattern
//inherit from OpRewritePattern and anchors to MulIOp op.
struct PowerOfTwoExpand :
	public OpRewritePattern<MulIOp> {
	// inherits from OpRewritePattern with context and benefit
	PowerOfTwoExpand(mlir::MLIRContext *context) 
		: OpRewritePattern<MulIOp>(context, 2){}
		
	LogicalResult matchAndRewrite(MulIOp op, PatternRewriter &rewriter) const override 
	{
		Value lhs = op.getOperand(0); // variable
		Value rhs = op.getOperand(1); // constant
		
		auto rhsDefiningOp = rhs.getDefiningOp<arith::ConstantIntOp>();
		if(!rhsDefiningOp) {
			return failure();
		}
		
		// may also use cast
		int64_t value = rhsDefiningOp.value();
		bool is_power_of_two = (value & (value - 1)) == 0;
		if (!is_power_of_two) {
			return failure();
		}
		
		ConstantOp newConstant = rewriter.create<ConstantOp>(
			rhsDefiningOp.getLoc(), rewriter.getIntegerAttr(rhs.getType(), value / 2)
		);
		MulIOp newMul = rewriter.create<MulIOp>(op.getLoc(), lhs, newConstant);
		AddIOp newAdd = rewriter.create<AddIOp>(op.getLoc(), newMul, newMul);
		
		rewriter.replaceOp(op, {newAdd});
		rewriter.eraseOp(rhsDefiningOp);
		
		return success();
	}
};

struct PeelFromMul :
	public OpRewritePattern<MulIOp> {
	PeelFromMul(mlir::MLIRContext *context)
		: OpRewritePattern<MulIOp>(context, 1){}
	
	LogicalResult matchAndRewrite(MulIOp op, PatternRewriter &rewriter) const override
	{
		Value lhs = op.getOperand(0); // variable
		Value rhs = op.getOperand(1); // constant
		
		auto rhsDefiningOp = rhs.getDefiningOp<arith::ConstantIntOp>();
		if(!rhsDefiningOp) {
			return failure();
		}
		
		// may also use cast
		int64_t value = rhsDefiningOp.value();
		
		ConstantOp newConstant = rewriter.create<ConstantOp>(
			rhsDefiningOp.getLoc(), rewriter.getIntegerAttr(rhs.getType(), value - 1)
		);
		MulIOp newMul = rewriter.create<MulIOp>(op.getLoc(), lhs, newConstant);
		AddIOp newAdd = rewriter.create<AddIOp>(op.getLoc(), newMul, lhs);
		
		rewriter.replaceOp(op, {newAdd});
		rewriter.eraseOp(rhsDefiningOp);
		
		return success();
	}
};

// now override runOnOperation() in order to write full logic of this pass
void MulToAddPass::runOnOperation() {
	//prepare a pattern set
	mlir::RewritePatternSet patterns(&getContext());
	// add PowerOfTwoExpand pattern to pattern set
	patterns.add<PowerOfTwoExpand>(&getContext());
	// add PeelFromMul pattern to pattern set
	patterns.add<PeelFromMul>(&getContext());
	// run rewrite engine greedily (looking at benefit)
	(void)applyPatternsGreedily(getOperation(), std::move(patterns));
}

Pattern rewriting example on the optimization pass

These two approaches reflect different transformation styles. Explicit walking is well-suited for global structural changes. Pattern rewriting is ideal for local algebraic simplifications.

Understanding when to use each is less about API knowledge and more about understanding the nature of your transformation.

TableGen: Power with Opacity

MLIR introduced TableGen as a domain-specific language (DSL) to reduce boilerplate when defining dialects, types, and operations. TableGen can feel magical at first. You define an operation declaratively, and MLIR generates builders, verifiers, type inference logic, and registration code automatically.

include "mlir/Pass/PassBase.td"

def AffineFullUnroll : Pass<"affine-full-unroll"> {
	let summary = "Fully unroll all affine loops";
	let description = [{
		"Fully unroll all affine loops."
	}];
	let dependentDialects = ["mlir::affine::AffineDialect"];
}

LoopUnrolling TableGen Example

#ifdef GEN_PASS_DEF_AFFINEFULLUNROLL

namespace impl {
  std::unique_ptr<::mlir::Pass> createAffineFullUnroll();
} // namespace impl
namespace impl {

template <typename DerivedT>
class AffineFullUnrollBase : public ::mlir::OperationPass<> {
public:
  using Base = AffineFullUnrollBase;

  AffineFullUnrollBase() : ::mlir::OperationPass<>(::mlir::TypeID::get<DerivedT>()) {}
  AffineFullUnrollBase(const AffineFullUnrollBase &other) : ::mlir::OperationPass<>(other) {}
  AffineFullUnrollBase& operator=(const AffineFullUnrollBase &) = delete;
  AffineFullUnrollBase(AffineFullUnrollBase &&) = delete;
  AffineFullUnrollBase& operator=(AffineFullUnrollBase &&) = delete;
  ~AffineFullUnrollBase() = default;

  /// Returns the command-line argument attached to this pass.
  static constexpr ::llvm::StringLiteral getArgumentName() {
    return ::llvm::StringLiteral("affine-full-unroll");
  }
  ::llvm::StringRef getArgument() const override { return "affine-full-unroll"; }

  ::llvm::StringRef getDescription() const override { return "Fully unroll all affine loops"; }

// More Codes
// ...
// ...

std::unique_ptr<::mlir::Pass> createAffineFullUnroll() {
  return impl::createAffineFullUnroll();
}
#undef GEN_PASS_DEF_AFFINEFULLUNROLL
#endif // GEN_PASS_DEF_AFFINEFULLUNROLL

Autogenerated .inc file

As can be seen here, TableGen automatically generates getArgument() or getDescription() methods for us. But this convenience comes at a cost. TableGen is not an abstraction layer; it is a code generator. When something goes wrong, errors often originate in generated .inc files. Debugging requires understanding both the declarative specification and the generated C++.

More specifically, TableGen does not explicitly indicate which functions must be implemented for the code to compile. Developers often need to inspect the generated macros and .inc files to understand which pieces of auto-generated code must be included and where they should be injected.

To use MLIR effectively, you must eventually become comfortable reading the generated code. The abstraction only works if you understand what it expands into.

Traits: Optimization Contracts

One of MLIR’s most subtle mechanisms is its trait system.

Adding a trait like Pure may seem trivial, but it fundamentally affects which optimization passes can legally transform your operation. For example, loop-invariant code motion will only move operations that declare themselves free of side effects.

class Poly_BinOp<string mnemonic> : Op<Poly_Dialect, mnemonic, [Pure]> {
    let arguments = (ins Polynomial: $lhs, Polynomial: $rhs);
    let results = (outs Polynomial: $output);
    let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($output)";
}

Example of adding Pure trait

class SubOp : public ::mlir::Op<SubOp,
::mlir::OpTrait::ZeroRegions,
::mlir::OpTrait::OneResult,
::mlir::OpTrait::OneTypedResult<::mlir::tutorial::poly::PolynomialType>::Impl,
::mlir::OpTrait::ZeroSuccessors,
::mlir::OpTrait::NOperands<2>::Impl,
::mlir::OpTrait::OpInvariants,
::mlir::ConditionallySpeculatable::Trait,            // <-- new
::mlir::OpTrait::AlwaysSpeculatableImplTrait,   // <-- new
::mlir::MemoryEffectOpInterface::Trait>          // <--- new
{ ... }

Updates on autogenerated .inc file

Optimizations in MLIR are permission-based. Operations must declare their semantic properties explicitly. Without those declarations, the optimizer will conservatively avoid transforming them.

For example, adding the Pure trait allows certain optimization passes to safely operate on an operation. A concrete case is the -cse (Common Subexpression Elimination) pass. If two operations compute the same result with identical operands, -cse may replace the second with the first — but only if the operation has no side effects. By marking an operation as Pure, we explicitly declare that it does not read or write memory and has no observable side effects, making it safe for CSE to eliminate redundant instances.

This design makes optimization safer and more modular. It also means dialect authors must think carefully about semantic contracts.

However, this also reveals one of MLIR’s practical constraints. Because MLIR evolves rapidly, there is no comprehensive or stable documentation clearly specifying which passes depend on which traits or interfaces. In practice, developers often need to inspect pass implementations directly to understand the required semantic contracts and determine why a particular optimization does or does not apply.

Traits are not annotations. They are formal commitments about behavior.

Folding, Canonicalization, and Global Propagation

Folding is operation-local simplification. If an operation’s operands are constant attributes, it can compute its result immediately. Folding happens opportunistically and locally.

Canonicalization is more structural. It rewrites patterns across multiple operations to produce simpler forms.

Sparse conditional constant propagation goes further. It performs global analysis across control flow to deduce constant values. However, it does not eliminate dead code; canonicalization typically follows to clean up.

Understanding how these mechanisms interact is essential when designing optimization pipelines. Folding reduces local redundancy. Canonicalization simplifies structure. Global passes propagate information across control flow.

Each operates at a different scale.

Declarative Rewrite Patterns: When Rewriting Becomes a Language

Earlier we discussed pattern rewriting using C++. But MLIR goes further. It allows you to describe rewrites declaratively.

This is where PDL (Pattern Description Language) comes in.

PDL is not just syntactic sugar. It’s a meta-layer that treats rewrites as data. Instead of embedding rewrite logic in C++ classes, you describe patterns in a declarative form that MLIR can interpret or compile.

Why is this important?

Because rewriting is central to MLIR. Almost every transformation is a pattern-based rewrite.

When rewrites become data:

• They can be reasoned about.
• They can be generated.
• They can be optimized.
• They can be loaded dynamically.

PDL shifts rewriting from “hardcoded compiler logic” to something closer to a transformation DSL.

This is a powerful idea: the compiler’s transformation logic becomes programmable.

Dialect Conversion: The Formal Lowering Framework

Progressive lowering sounds simple conceptually, but in practice it is extremely delicate.

How do you guarantee that after lowering, no illegal operations remain? How do you ensure partial lowering doesn’t leave dangling constructs?

This is where dialect conversion enters.

Dialect conversion is not just pattern rewriting. It is a constrained rewriting system with legality checks.

You define:

• Which dialects are legal
• Which operations are illegal
• How illegal operations must be rewritten

The conversion framework then ensures that all illegal operations are rewritten into legal ones — or it fails.

This is incredibly important.

Without this structure, lowering becomes ad-hoc and unsafe.

Dialect conversion makes lowering declarative and verifiable.

It is the formal mechanism that turns “we should lower this” into “the IR is guaranteed to be in the target dialect.”

In other words, it transforms lowering from convention into contract.

Lowering Through LLVM: The Final Boundary

Eventually, many MLIR pipelines lower into the LLVM dialect.

This is the boundary between structured, semantic-rich IR and low-level code generation.

What’s interesting is that the LLVM dialect is not just LLVM IR pasted into MLIR. It is a structured representation of LLVM concepts within MLIR’s operation system.

Lowering through LLVM typically involves:

• Converting structured control flow into CFG-based control flow
• Converting high-level types into LLVM-compatible types
• Mapping memory semantics carefully

Once in LLVM dialect, MLIR effectively hands control back to LLVM’s optimization and backend infrastructure.

This layered architecture is elegant:

• High-level reasoning happens in structured dialects.
• Low-level instruction scheduling and codegen happen in LLVM.

MLIR does not replace LLVM. It orchestrates the journey toward it.

Why This Matters for ML Systems

Modern ML workloads consist of tensor computations, structured loops, and domain-specific operations. Flattening these into low-level IR too early destroys information that is crucial for optimization.

MLIR allows compilers to reason about:

• Tensor shapes
• Loop bounds
• Affine indexing
• Memory effects
• Operator semantics

And to preserve that reasoning until the right moment to lower.

This is why MLIR underpins many modern ML compilers. It enables aggressive optimization without sacrificing abstraction prematurely.

The Deeper Insight

LLVM unified code generation across languages. MLIR unifies compiler construction across domains.

It provides a way to define new IRs, attach semantics to them, transform them safely, and lower them progressively. It turns IR design itself into a modular and extensible discipline.

It is more complex than traditional compiler infrastructure. But that complexity reflects the complexity of modern workloads.

MLIR is not simply another intermediate representation.

It is a framework for thinking about representations.

Before diving into the vLLM paper itself, it is important to understand KV caching, because it is the core mechanism being rethought and redesigned throughout the paper.

1. KV Cache

1.1 Naive Approach

In naive autoregressive decoding with Transformers, each new token is generated by reprocessing all previously generated tokens. At decoding step t, the model recomputes the key and value vectors for tokens 1 through t, even though those values were already computed in earlier steps. This results in redundant computation that grows quadratically with sequence length, making long-context decoding extremely inefficient.

Image generated by Gemini

1.2 Autoregressive Decoding With KV Cache

KV caching was introduced to eliminate this redundancy. Instead of recomputing keys and values at every decoding step, the model stores the computed K and V vectors in memory and reuses them for subsequent steps. During decoding, only the query corresponding to the newly generated token is computed, and attention is performed between this query and the cached keys and values. Although this changes the attention computation from a GEMM into a matrix–vector operation—which is less favorable for GPU utilization—the overall computation cost is significantly reduced because expensive KV recomputation is avoided. In practice, the benefits of KV reuse far outweigh the loss in arithmetic intensity.

Image generated by Gemini

2. Challenges with KV Cache

2.1 Large KV Cache

However, KV caching introduces a new and severe bottleneck: memory usage. The KV cache grows linearly with sequence length and must be stored for every active request. For example, in the OPT-13B model, each token requires storing both a key and a value vector across all layers. With a hidden size of 5120, 40 layers, and FP16 precision, a single token consumes approximately 800 KB of KV cache memory. Since OPT supports sequences up to 2048 tokens, the KV cache for a single request can occupy as much as 1.6 GB of GPU memory. Even on modern GPUs with tens of gigabytes of VRAM, this severely limits the number of concurrent requests that can be served.

2.2 Unsharable Cache

The problem becomes even more complex when considering modern decoding algorithms. LLM serving is not limited to greedy decoding; users often request multiple sampled outputs from the same prompt or use beam search to retain multiple high-probability continuations. In theory, these decoding strategies could share KV cache for the common prefix of the sequence. In practice, however, once the outputs diverge, their future tokens depend on different contexts and positions. Because most deep learning frameworks require tensors to be stored in contiguous memory, these KV caches must be physically duplicated, even when large portions are logically identical. This prevents effective memory sharing and leads to significant waste.

2.3 Dynamic Length of Input and Output

Another fundamental challenge arises from the fact that both input and output lengths are unknown at runtime. Prompts vary widely in length, and generation can terminate at unpredictable times. Since the KV cache grows dynamically as tokens are generated, the memory management system must be able to accommodate expanding buffers. When GPU memory becomes exhausted mid-generation, the system must either reject new requests or stall existing ones. Pre-allocating KV cache buffers for the maximum possible sequence length avoids reallocation but wastes enormous amounts of memory when sequences are shorter than the maximum.

2.4 Contiguous Memory Allocation

All of these issues stem from a core assumption shared by most deep learning frameworks: tensors must be allocated in contiguous memory. This assumption makes memory management simple but inflexible.

Fragmentation Example ( From the paper vLLM )

This figure illustrates how internal and external fragmentation arise in traditional LLM serving systems. To handle unknown output lengths, systems pre-allocate KV cache slots for the maximum sequence length. As a result, many reserved slots are never used, causing internal fragmentation. At the same time, as requests with different prompt lengths start and finish at different times, free memory becomes scattered into small gaps between active allocations. These gaps cannot be reused effectively, leading to external fragmentation. Even though sufficient total memory exists, it cannot be used to serve new requests efficiently.

Memory Wastes in Different LLM Servings ( From the paper vLLM )

This bar chart quantifies how much KV cache memory is actually useful versus wasted in different serving systems. In Orca-based systems, a large fraction of memory is lost to internal fragmentation, reservation overhead, and external fragmentation, leaving only a small portion for actual token states. In contrast, vLLM achieves near-optimal memory utilization: over 96% of KV cache memory is used for real token states. Let’s see how vLLM achieved this almost-perfect memory management.

3. vLLM : Rethinking Memory for LLM Serving

These limitations are precisely what motivate vLLM. By removing the requirement that KV cache be stored contiguously and instead managing it in fixed-size blocks—much like pages in a virtual memory system—PagedAttention enables efficient allocation, sharing, and reclamation of KV cache memory. This shift in perspective is what allows vLLM to dramatically improve GPU memory utilization and serve many more concurrent requests than previous systems.

vLLM System Overview ( From the paper vLLM )

We will refer back to this diagram later, but the key takeaway is that the scheduler controls all global decisions in vLLM. It determines request scheduling and KV cache management, while the KV Cache Manager maintains block mappings and allocators. GPU workers are intentionally simple: they only execute assigned computations using the provided physical KV block IDs, without making scheduling or memory management decisions themselves.

3.1 PagedAttention

At first glance, PagedAttention sounds like a new attention algorithm. In reality, it is much closer to a memory management system than a change to the attention computation itself. The key idea is simple but powerful: instead of storing the KV cache as one large contiguous tensor, vLLM breaks it into small fixed-size blocks and manages them dynamically, much like pages in a virtual memory system.

During decoding at time step $t$, the Transformer computes attention using:

Here:

• $Q_t$ is the query vector for the newly generated token
• $K_{1:t}$ and $V_{1:t}$ are the cached keys and values from all previous tokens
• $d$ is the hidden dimension

With KV caching, $K_{1:t}$ and $V_{1:t}$ are not recomputed, but they are assumed to be stored contiguously in memory. This assumption is where the real problem begins as mentioned above.

To resolve the problems with traditional KV cache memory management system, PagedAttention divides the KV cache into fixed-size blocks, each containing KV vectors for a small number of tokens. For example, one block might store KV for 16 tokens.

Instead of indexing tokens by position directly, vLLM introduces a logical-to-physical mapping. Each request maintains a table that maps its logical sequence positions to physical blocks in GPU memory. These blocks can be:

• Allocated non-contiguously
• Reused across requests
• Released independently when no longer needed

This is directly analogous to how virtual memory maps virtual pages to physical pages.

With paged KV storage, the attention computation conceptually becomes:

where ${K_{b}}, {V_{b}}$ are KV blocks corresponding to the sequence so far.

Internally, attention is computed block by block:

Crucially, the mathematical result is identical to standard attention. The only difference is that the keys and values are gathered from multiple physical locations instead of a single contiguous tensor.

PagedAttention therefore preserves correctness while fundamentally changing memory behavior.

Simple PagedAttention Example ( From the paper vLLM )

3.2 KV Cache Manager

The KV Cache Manager is the core component of vLLM, and it is what allows vLLM to overcome the severe memory fragmentation problems found in previous LLM serving systems. Because PagedAttention stores the KV cache in small, fixed-size blocks, the KV cache can be managed in a way that closely resembles virtual memory in an operating system.

Just as an OS divides memory into fixed-size pages and maps a process’s virtual pages to physical memory pages, vLLM partitions GPU memory into fixed-size KV cache blocks and maintains a mapping from logical blocks (the token positions seen by the model) to physical blocks (actual locations in GPU DRAM). Through this indirection, logically contiguous KV cache can be backed by non-contiguous physical memory. As a result, vLLM can allocate KV cache space dynamically, eliminating the need to pre-allocate memory for the maximum possible sequence length—a practice that previously caused massive memory waste.

All of this functionality is implemented by the KV Block Manager, which maintains block tables that record the mapping between logical KV blocks and physical KV blocks.

Logical-Physical Mapping Example ( From the paper vLLM )

As illustrated in the figure, logical blocks remain contiguous from the model’s perspective, while their corresponding physical blocks may be scattered across GPU memory. The block manager also tracks how many tokens have been written into each block. When a block becomes full, a new block is allocated and appended to the logical sequence. This enables KV cache to grow incrementally as tokens are generated, without requiring reallocation or data copying.

PagedAttention does introduce a small amount of internal fragmentation. Since blocks are fixed in size, the unused slots in the last block of a sequence are reserved for future tokens. However, this waste is strictly bounded by the block size and is negligible compared to earlier systems, which often reserved memory for thousands of tokens per request regardless of actual usage. In practice, this bounded fragmentation is a small and acceptable trade-off for dramatically improved memory utilization and scalability.

From the discussion so far, it might appear that blocks cannot be shared across requests. Since each KV block reserves unused space for future tokens, one could mistakenly conclude that blocks must remain private to each request, preventing any form of sharing.

Two Simultaneous Requests ( From the paper vLLM )

At first glance, the figure seems to support this interpretation: each request appears to own its own sequence of blocks, with unused slots reserved for upcoming tokens. If this were truly the case, vLLM would fail to address the core challenges discussed earlier—namely excessive memory usage, duplication of KV cache for shared prefixes, and poor scalability under high concurrency.

3.3 Complex Decoding Algorithm

As discussed in Section 2.2, many decoding algorithms share a common prefix. If their KV caches could also be shared, memory redundancy would be dramatically reduced. vLLM addresses this elegantly by introducing reference counting for KV blocks.

Consider parallel sampling, where multiple outputs are generated from a single input prompt. Since all samples start from the same prompt, their KV cache for the prompt tokens is identical. During the prefill stage, vLLM exploits this by mapping all samples to the same physical KV blocks that store the prompt’s keys and values.

Parallel Sampling Example ( From the paper vLLM )

As shown in the figure, two samples initially reference the same physical block, and the block’s reference count is increased to 2. At this stage, the block is read-only shared: both samples can read from it, but neither can modify it.

As decoding continues, the samples begin to generate different tokens. This is where vLLM’s copy-on-write mechanism comes into play. When vLLM attempts to append a new KV value to a block whose reference count is greater than 1, it recognizes that the block is shared. To preserve correctness, vLLM allocates a new physical block, copies the contents of the original block into it, and appends the new KV entry to the newly allocated block. The reference count of the original block is then decremented, allowing the other sample to continue using it without interference.

This process ensures that KV blocks are shared as long as possible and duplicated only when necessary. Sharing is maximized for common prefixes, while divergence is handled safely and efficiently. As a result, vLLM avoids unnecessary KV duplication, significantly reducing memory usage under parallel sampling and other prefix-sharing decoding strategies.

Parallel sampling diverges immediately once decoding begins, since each sample independently generates its own next token. Beam search, however, behaves differently. In beam search, multiple hypotheses are expanded step by step, and divergence happens gradually as new tokens are generated and low-probability beams are pruned.

vLLM is designed to handle this scenario as well. Just like in parallel sampling, beam search hypotheses initially share KV blocks for their common prefix. As long as multiple beams point to the same prefix, they can safely reference the same physical KV blocks in a read-only manner.

Beam Search Example ( From the paper vLLM )

As decoding progresses, beams may diverge at different time steps. When a beam needs to append a new token to a KV block that is shared with other beams, vLLM again relies on the copy-on-write mechanism. If the target block’s reference count is greater than one, vLLM allocates a new block, copies the existing KV contents, and appends the new token’s KV to the new block. This allows the beam to continue decoding independently without affecting other beams that still rely on the shared prefix.

Importantly, beam search also involves frequent pruning: beams with low scores are discarded at each step. When a beam is dropped, the reference counts of its associated KV blocks are decremented. If a block’s reference count reaches zero, it can be immediately reclaimed and reused. This tight coupling between beam lifecycle and reference counting allows vLLM to reclaim memory aggressively, which is especially important in beam search where the set of active hypotheses changes rapidly.

3.4 Beauty of Abstraction

Traditional LLM serving systems struggled to support mixed decoding methods because different decoding strategies imposed incompatible requirements on KV cache allocation and sharing. vLLM overcomes this limitation through a clean abstraction layer: attention kernels operate only on lists of physical KV block IDs and remain completely agnostic to how blocks are shared or which decoding algorithm is being used. All complexity related to memory ownership, sharing, and divergence is handled by the KV Cache Manager through logical-to-physical block mapping. This separation allows vLLM to efficiently process requests with different decoding algorithms—such as greedy decoding, sampling, and beam search—simultaneously, without specialized kernels or decoding-specific memory layouts.

4. Scheduling and Preemption

Even with block-based KV cache management, GPU memory is still a finite resource. When the number of concurrent requests increases or when sequences grow longer than expected, vLLM may eventually reach a point where GPU memory becomes insufficient. Rather than failing requests or severely limiting concurrency, vLLM introduces two complementary mechanisms to handle memory pressure gracefully: swapping and recomputation.

4.1 Swapping

Swapping allows vLLM to temporarily move KV blocks from GPU memory to CPU memory. When a request becomes inactive—for example, when it is waiting for other requests to finish or when it is deprioritized by the scheduler—its KV blocks can be swapped out to host memory. Later, when the request becomes active again, the blocks are swapped back into GPU memory. Because KV cache is managed in fixed-size blocks, swapping can be performed at block granularity, without affecting other requests or requiring large contiguous memory regions. This makes memory eviction and restoration both predictable and efficient.

4.2 Recomputation

Recomputation offers an alternative trade-off. Instead of preserving KV cache across the entire lifetime of a request, vLLM may choose to discard some KV blocks and recompute them later when needed. This is particularly effective for short prefixes or low-cost recomputation scenarios, where recomputing attention is cheaper than storing KV cache in scarce GPU memory. By selectively recomputing KV cache, vLLM trades additional computation for reduced memory usage, allowing the system to sustain higher concurrency under tight memory constraints.

4.3 When to Swap and When to Recompute

The choice between swapping and recomputation depends on several factors, including sequence length, request priority, and system load. Long-running requests with expensive prefixes are more likely to benefit from swapping, since recomputation would be costly. Short or low-priority requests, on the other hand, may be better candidates for recomputation. By combining both mechanisms, vLLM avoids committing to a single rigid strategy and instead adapts dynamically to runtime conditions.

This flexibility is made possible by the same abstraction that underlies PagedAttention. Because attention kernels only operate on physical block IDs, blocks can be swapped out, recomputed, or reclaimed without changing the execution logic. Memory management decisions remain entirely within the KV Cache Manager and scheduler.

5. Performance

The performance results clearly show that vLLM’s advantages in KV cache management translate directly into higher sustainable throughput under realistic workloads. Across all settings, vLLM consistently maintains low and stable latency at much higher request rates than prior serving systems.

Single Sequence Generation Performance Graphs ( From the paper vLLM )

This figure evaluates end-to-end serving performance on the ShareGPT and Alpaca workloads with increasingly large models and GPU counts. As request rates increase, traditional systems such as FasterTransformer and Orca variants exhibit a sharp rise in latency once they approach memory saturation. This behavior reflects the fragmentation and over-reservation issues discussed earlier: even when GPU compute capacity remains available, KV cache memory becomes the limiting factor, causing throughput to collapse.

In contrast, vLLM sustains low normalized latency across a significantly wider range of request rates. This gap becomes more pronounced as model size and GPU count increase, particularly for OPT-66B and OPT-175B. Larger models amplify KV cache pressure, and vLLM’s block-based allocation allows it to utilize GPU memory far more effectively. As a result, vLLM can serve substantially more concurrent requests before hitting the latency cliff that limits other systems.

Parallel Generation and Beam Search Performance Graphs ( From the paper vLLM )

This figure focuses on complex decoding workloads, including parallel generation and beam search with varying parallel sizes and beam widths. These scenarios are especially challenging for traditional serving systems, as they involve frequent branching, prefix sharing, and rapid changes in the set of active sequences. Orca-based systems show steep latency degradation as parallelism increases, indicating that KV cache duplication and poor memory reuse quickly overwhelm available resources.

vLLM, however, scales smoothly with both parallel generation and beam width. Even as the number of parallel hypotheses grows, vLLM maintains low latency at much higher request rates. This behavior directly reflects its use of reference counting and copy-on-write KV blocks: shared prefixes remain shared as long as possible, and memory is duplicated only when divergence actually occurs. Moreover, when beams are pruned, their KV blocks are promptly reclaimed, preventing memory buildup over time.

Taken together, these results highlight a key insight: the performance bottleneck in LLM serving is not raw computation, but memory efficiency under concurrency and dynamic decoding patterns. By eliminating fragmentation and enabling fine-grained sharing and reclamation of KV cache, vLLM shifts the throughput limit much closer to the true hardware capacity of the GPU. This is why its performance advantage grows with model size, decoding complexity, and request concurrency—exactly the regimes that matter most in real-world LLM deployment.

6. Conclusion

The most striking insight is that many long-standing limitations in LLM serving were not inherent to Transformers themselves, but artifacts of how deep learning frameworks treat memory. By questioning the assumption that KV cache must be stored contiguously, vLLM unlocks a surprisingly large design space. Once KV cache is treated as a managed resource rather than a static tensor, features such as prefix sharing, mixed decoding, swapping, and recomputation fall out naturally.

Another takeaway is the power of abstraction. vLLM’s decision to keep attention kernels unaware of sharing and ownership details dramatically simplifies the execution model, while enabling far more complex behaviors at the system level. This separation of concerns is what allows vLLM to scale without becoming fragile or overly specialized.

Ultimately, vLLM changed how I think about LLM inference. Efficient serving is not just about faster kernels or clever batching, but about bringing systems thinking into machine learning. As models grow larger and workloads become more dynamic, approaches like vLLM suggest that the future of LLM systems will be shaped as much by operating systems principles as by advances in model design.

7. What’s Next

In the next post, I plan to dive deeper into vLLM by looking directly at its codebase—how the scheduler, KV Cache Manager, and GPU workers are actually implemented, and how these design ideas show up in real systems code.

One interesting insight I gained while studying vLLM is related to quantization. Intuitively, lowering precision should significantly improve inference performance. However, in practice, running quantized models inside modern serving engines does not always lead to the performance gains one might expect. The main reason is that today’s GPUs are still not well optimized for low-bit arithmetic. As a result, quantized weights often need to be dequantized back to higher precision during execution, which introduces overhead and limits the overall speedup.

This observation made me realize that quantization is not just a model-level optimization, but a systems-level problem that tightly interacts with hardware capabilities and execution engines. To better understand this gap, I plan to read and summarize the QServe paper in a future post, focusing on how it addresses the inefficiencies of quantized inference in real serving systems.

Prerequisites This post assumes that the reader is already familiar with the Transformer architecture, including self-attention, autoregressive decoding, and the role of key–value (KV) caches. If you are new to Transformers, it is recommended to review the basics before reading, as this post focuses on inference-time systems rather than model fundamentals.

Large Language Models (LLMs) are usually discussed from a model-centric perspective: parameter counts, attention mechanisms, training tricks, or dataset scale. However, once a model is trained, a very different problem dominates real-world usage:

How do we serve LLMs efficiently, reliably, and at scale?

LLM Serving has little to do with model accuracy. Instead, it is governed by:

• latency
• throughput (tokens per second)
• GPU memory efficiency
• cost per generated token

This post is a foundational but deep overview of modern LLM serving systems, covering:

• how inference is structured,
• why prefill and decode behave so differently,
• where the real bottlenecks come from,
• and how modern systems mitigate them.

1. What is LLM Serving?

At its core, LLM serving is the inference-time system that sits between users and GPUs. A typical serving pipeline looks like this:

1. User sends a request
2. The request enters a queue
3. A scheduler decides when and how it runs
4. A batcher decides which other requests to group it with
5. The GPU executes inference
6. Tokens are streamed back to the user

What makes this difficult is that:

• prompts are variable-length
• users arrive asynchronously
• GPUs are optimized for large, regular workloads

All of these suggest that serving is fundamentally a systems problem, not a modeling problem.

2. Inference Into Two Stages : Prefill vs. Decode

A critical insight in LLM serving is that inference is split into two different stages.

2.1 Prefill Stage

Prefill processes the entire prompt in one forward pass and builds the KV cache looking at entire input tokens.

Characteristics:

• Input: all prompt tokens (length = $T$)
• Attention: full self-attention ($QK^T \in \mathbb{R}^{T \times T}$)
• Output:
  • KV cache (keys and values for every layer)
  • hidden states for the first token generation

Since prefill can look at entire tokens, it brings a high GPU utilization due to large matrix-matrix mulitplications and high arithmetic intensity. However, it is also where latency explodes for long prompts.

2.2 Decode Stage

Decode stage generates new tokens autoregressively, one token at a time. It means each token generation depends on all previous tokens.

It suggests that this stage forces sequential computing, underutilizing GPUs due to lack of parallelism.

Characteristic:

• Input: a single token per step
• Attention: $\mathbb{R}^{1 \times T}$ (query attends to KV cache)
• Repeated until generation ends (reachs termination token)

Decode is fundamentally different to prefill because:

• sequential dependency across time
• small kernels
• memory-bound execution

This split explains why optimizing inference is not a single problem.

3. Why Prefill Dominates Latency

3.1 Quadratic Attention Cost

Self-attention computes:

For prefill:

• $Q \in \mathbb{R}^{T \times d}$
• $K \in \mathbb{R}^{T \times d}$
• $QK^T \in \mathbb{R}^{T \times T}$

This means doubling prompt length quadraples attention work (= $O(T^2)$) and long-context users dominate wall-clock latency, leading to tail latency by a small fraction of long prompts.

3.2 KV Cache Creation Cost

Prefill must materialize keys and values for every token and for every transformer layer. This results in heavy GPU memory writes, pressure on HBM capacity, and reduced concurrency. This tells that KV cache is not just a decode bottleneck—it is created in prefill.

3.3 Scheduling Side Effects

Prefill kernels are large, long-running, and non-preemptive. Due to these properties, a single long-prompt prefill can monopolize the GPU, block decode requests, and inflate latency. This causes delaying entire service due to few long prompt users. This is why prefill is both a compute bottleneck and a scheduling bottleneck.

4. Why Decode Is Still Slow

Despite much less computation, decode introduces its own challenges.

4.1 Autoregressive Dependency

Each token depends on all previous tokens as mentioned above. This dependency prevents parallelization along the time dimension and enforces strict sequential execution.

4.2 Memory Bandwith Bottleneck

Decode attention reads the entire KV cache from HBM and performs relatively little computation. This incurs a large data exchange between SM and HBM, leading to memory bound. Since FLOPs are cheap and HBM bandwith dominates, GPU utilization drops sharply.

4.3 Kernel Launch Overhead

Decode consists of many small kernels and repeated thousands of times per request. Launch overhead and poor fusion become visible performance costs.

5. Batching: The Central Lever of Serving Performance

Batching is a technique of combining multiple requests and dealing it with a single GPU execution. Without batching, as LLM inference is a large matmul computation, handling single request each time makes most SMs idle, increasing latency and costs per token.

5.1 Why Naive Batching Fails

Prompts vary widely in length.

Request A: 8000 tokens
Request B: 200 tokens
Request C: 100 tokens

Batching forces padding to the maximum length. Because attention is $O(T^2)$, padding causes real computation, not just wasted memory, so large batches can reduce throughput. This is a length variance problem.

5.2 Modern Batching Strategies

Dynamic Batching

• Collect requests within a short time window
• Trade latency for throughput

Length-Aware Batching

• Group requests with similar prompt lengths
• Trade latency for throughput

Token-Level (Continuous) Batching

• Batch tokens, not sequences
• Mix prefill and decode steps
• Prevent long requests from blocking short ones

Token-level batching is the foundation of modern serving engines such as vLLM.

6. Moden Solutions to Core Bottlenecks

6.1 FlashAttention (Prefill Optimization)

FlashAttention algorithm avoids materializing $T \times T$ attention matrix and keeps computation in SRAM, resulting in reducing memory traffic from $O(T^2)$ to $O(T)$. It does not remove quadratic computation, but makes it practical.

6.2 Paged KV Cache

• Store KV cache in page-sized blocks
• Reduce fragmentation
• Increase Concurrency This directly improves decode stability.

6.3 Chunked Prefill

Instead of processing long prompts in one large kernel:

• split prefill into chunks
• interleave with decode steps

Benefits:

• lower peak latency
• better scheduling fairness
• reduced tail latency

6.4 Quantization (Decode-Focused)

Lower-precision formats:

• reduce memory footprint
• benefit decode more than prefill
• improve throughput at scale

6.5 Speculative Decoding

Idea:

• a small model drafts multiple tokens
• a large model verifies them

Effect:

• fewer autoregressive steps
• lower perceived latency

7. Putting All Together

8. Conclusion

LLM serving is not about running a model. It is about orchestrating computation, memory, and time.

This post focused on building a high-level mental model of LLM serving—how inference is structured, where the real bottlenecks come from, and why serving is fundamentally a systems problem.

The individual optimization techniques discussed here, such as FlashAttention, KV cache management, token-level batching, chunked prefill, and decode-aware scheduling, are each deep topics on their own. They will be explored in more detail in future posts.

In practice, these techniques are not applied in isolation. They are already integrated into modern inference engines and serving platforms, such as:

• vLLM for efficient GPU-side execution,
• Ray Serve for flexible and scalable request handling,
• KServe on top of Kubernetes for production deployment and orchestration.

LLM serving is not about a single optimization, but about how many optimizations work together as a system.

This post provides the map; the next posts will dive into the individual components that make modern LLM inference engines work at scale.

What is MXFP4?

• FP4 basics: FP4 is a 4-bit floating point number. Normally 4 bits is too small, you lose a lot of precision.
• Micro-scaling idea: MFXP4 fixes this by grouping numbers intal small blocks (like 32 values, or 16 in NVFP4) and giving each block its own scale. This way, the numbers don’t lose range too quickly.
• Why it matters: Using FP4 reduces memory a lot. Compared to BF16 you save ~4× memory, and compared to FP8 it’s about 2×.

GPT-OSS using MXFP4

OpenAI released GPT-OSS-120B and GPT-OSS-20B with weights already in MXFP4.

• The 120B model can run on about 80 GB GPU memory.
• The 20B version fits in only 16 GB which even can be run on my local desktop (I have 5070TI which has 16GB VRAM). That is crazy small compared to normal large models. GPT-OSS also uses Mixture-of-Experts (MoE) and Grouped MQA, so it is very efficient.

However, GPT-OSS was not trained fully in FP4. It was trained with higher precision, then converted to MXFP4 for release. The research part now is how to really train in FP4 end to end.

FP4 for Training

Before, people only used FP4 for inference. Training was too unstable. But now we see progress:

1. Fine-tuning while keeping FP4
   • NVIDIA shows a recipe: upcast to BF16, fine-tune, then quantization-aware training (QAT) to go back to FP4.
   • LMSYS and Unsloth also made guides to fine-tune GPT-OSS in FP4 on normal GPUs.
2. Full FP4 training (still research)
   • New tricks like stochastic rounding (avoid bias in gradients) help stability.
   • Some papers show FP4 training can reach almost same accuracy as FP8, but with 2× faster GEMM speed.
   • For vision models, methods like Q-EMA and Q-Ramping reduce oscillation problems.
   • NVIDIA Blackwell GPUs (5000 models) even support NVFP4 natively, with smaller block size and better precision.

Why It Solves Memory Bottleneck

Big models are often blocked by memory and bandwidth, not just compute.

• Inference: FP4 makes parameters and KV cache much smaller.
• Training: Optimizer states, activations, gradients take huge memory. If they also move to FP4 (with tricks like error feedback), you can train bigger models or larger batch size on same GPU.

Still Some Problems

• Some tasks (math, code, reasoning) are more sensitive to low precision.
• Not all frameworks support MXFP4/NVFP4 yet. Sometimes only weights are FP4 but activations still higher precision.
• No public model is fully pretrained in FP4 yet, only research results.

My Thoughts

GPT-OSS made FP4 real for the community. Before, FP4 was more like an academic idea. Now people can actually run a 120B parameter model on one 80 GB GPU. That is impressive. The next step is full training in FP4, and we already see good progress.

If FP16 made deep learning practical, and BF16 made LLM training possible, then maybe FP4 will be the format that makes trillion parameter models possible for everyone.

1) CFG Basics

• A context-free grammar G = (V, Σ, R, S)
  • V: nonterminals (variables)
  • Σ: terminals (alphabet, tokens)
  • R: production rules of the form A → α where A ∈ V and α ∈ (V ∪ Σ)*
  • S: start symbol

• Language of G: L(G) = { w ∈ Σ*

  S ⇒* w }

Example

E → E + T | T
T → T * F | F
F → ( E ) | id

• Generates arithmetic expressions with +, *, parentheses, and identifiers.

2) Derivations & Sentential Forms

• Single-step derivation: αAβ ⇒ αγβ if A → γ is a production.
• Leftmost derivation (LMD): always rewrite the leftmost nonterminal first.
• Rightmost derivation (RMD): always rewrite the rightmost nonterminal first.
• Sentential form: any string in (V ∪ Σ)* reachable from S by ⇒*.
• Yield: a sentential form consisting only of terminals (a sentence).

Example (LMD)

S ⇒ E
E ⇒ E + T
E ⇒ T + T
T ⇒ F + T
F ⇒ id + T
T ⇒ F
F ⇒ id
; result: id + id

3) Parse Trees

• A parse tree (derivation tree) visualizes a derivation:
  • Root labeled with the start symbol S.
  • Each internal node A expands to its children according to a production A → α.
  • Leaves read left to right give the derived sentence.
• Parse trees abstract away order of derivation (LMD vs RMD) but reveal structure.

4) Ambiguity

• A grammar is ambiguous if some w ∈ L(G) has two or more distinct parse trees (or two distinct LMD/RMD).
• Classic example:

  E → E + E | E * E | id
  

• For id + id * id there are two parses:
  • (id + id) * id vs id + (id * id).
• Fix with precedence/associativity encoded in grammar:

  E → E + T | T          ; '+' lowest precedence, left-assoc
  T → T * F | F          ; '*' higher precedence, left-assoc
  F → ( E ) | id
  

5) Nullable, FIRST, FOLLOW

• Nullable(A): A ⇒* ε.
• FIRST(α): set of terminals that can begin strings derived from α (and ε if α ⇒* ε).
• FOLLOW(A): set of terminals that can appear immediately to the right of A in some sentential form (and $ for end-of-input when A can finish the input).

FIRST rules (summary)

• If X is terminal: FIRST(X) = {X}
• If X is nonterminal: add FIRST of its alternatives; include ε if any RHS ⇒* ε
• For concatenation αβ: FIRST(αβ) = FIRST(α) minus ε, plus FIRST(β) if α ⇒* ε

FOLLOW rules (summary)

• Put $ in FOLLOW(S) for start symbol S.
• For any A → αBβ: add FIRST(β) \ {ε} to FOLLOW(B).
• If β ⇒* ε: add FOLLOW(A) to FOLLOW(B).

6) LL(1) Conditions (Top-Down Parsing)

A grammar is LL(1) if for every nonterminal A with alternatives A → α | β:

• Disjointness: FIRST(α) and FIRST(β) are disjoint.
• ε-case: If ε ∈ FIRST(α), then FIRST(β) must be disjoint from FOLLOW(A) (and vice versa).

Left Factoring

• If A → aβ1 | aβ2, factor common prefix:

  A → aA'
  A' → β1 | β2
  

Eliminate Immediate Left Recursion

• If A → Aα | β, transform:

  A  → βA'
  A' → αA' | ε
  

7) LR Parsing Glance (Bottom-Up)

• Recognizes handles using items and states (LR(0)/SLR/LALR/LR(1)).
• Shift: push next input symbol.
• Reduce: apply A → α when top-of-stack matches α.
• Goto: transition on nonterminal after reduction.

8) Closure Properties (High-level)

• CFLs (context-free languages) are closed under: union, concatenation, Kleene star, homomorphism, reversal.
• Not closed under: intersection, complement (but are closed with regular languages, e.g., CFL ∩ REG is CFL).

9) Leftmost vs Rightmost Example

Grammar:

S → aS | Sb | ab

• String: aab
  • LMD: S ⇒ aS ⇒ aaS ⇒ aab
  • RMD: S ⇒ Sb ⇒ Sab ⇒ aab

10) Sample Grammar Exercises

1. Balanced parentheses:

   S → SS | (S) | ε
   

2. Binary numbers without leading zeros (except zero itself):

   N → 0 | 1B
   B → 0B | 1B | ε
   

3. Palindromes over {a, b}:

   P → aPa | bPb | a | b | ε
   

Quick Glossary

• Terminal: token from the input alphabet.
• Nonterminal: variable to be expanded.
• Production: rewrite rule.
• ε (epsilon): empty string.
• Handle: substring that matches the RHS of a production and whose reduction is a step in reverse of a rightmost derivation.