@AryaTschand: I’ve spent most of my internship at nvidia architecture research thinking about how to frontier llms should inference o…

I’ve spent most of my internship at nvidia architecture research thinking about how to frontier llms should inference on our newest (heterogenous) systems

I wanted to share some of the things I’ve learned about inference from first principles

TLDR + mini experiments in thread!

At its core, EVERYTHING boils down to arithmetic intensity (AI) - the ratio between FLOPs computed and bytes of data moved

Every ML chip has some peak FLOPs and some memory bandwidth, and the AI of a workload tells us whether we will be waiting for MMAs to run or data to move. A roofline model (attached) is a property of any chip, and shows us whether you will be compute bound or memory bound given a workload’s AI

Memory bound = AI is less than the chip’s “ridge point” and we will be waiting for data (tensorcores will be underutilized) Compute bound = AI is greater than the chip’s “ridge point” and we will be waiting for compute (bandwidth will be underutilized) in GPUs, we want this

An LLM inference can be reasonably broken down into the following phases:

1. Prefill - loading in new tokens (tool call outputs), VERY compute-bound which is great for GPUs
2. Decode Attention - autoregressively attending to past tokens in the KV, usually memory-bound but comes close to compute bound with strategies like MLA, sparse attention, etc.
3. Decode FFN - load weights and compute GEMM per token, usually VERY memory-bound especially at low batch, making it terrible for GPUs

Note that we can break this down even further in newer hybrid models

As GPUs move towards bigger tensorcores to support large-scale training (LLM training is VERY compute-bound and thus increasing total FLOPs with lower precision and bigger systolics lets us train GPT 6.3.2 Pro Max Codex faster)

Increasing the FLOPs of a chip is relatively easier than increasing HBM bandwidth, which is why we see the ridge points of GPUs going higher and higher

While this is great for prefill (high AI), this is a really bad fit for decode (low AI) and we usually see low GPU utilization at decode, especially without multi-token speculation tricks to boost decode AI like draft/verify speculative decoding, self-speculation, or diffusion

To address this gap, some folks have started making chips with only SRAM

Without getting to deep into memory systems, all GPUs today use HBM, which is high capacity but actually quite bad latency compared to on-chip memory. This is great when we are compute bound (and usually need lots of memory to store optimizer states or KV caches) but hurts when we are memory bound in decode.

SRAM is on-chip memory (TSMC tapes it out on the wafer instead of buying HBM separately from Micron) and has way better latency but also way lower capacity. If you can offload your KV cache or keep context short, this makes decode WAY faster

Batch size (total number of simultaneous inferences) has very important relationships with arithmetic intensity. And because we know how important AI is to inference efficiency, it also has important relationships with throughput and interactivity

Increasing batch size (meaning we pack more requests together) almost linearly increases the arithmetic intensity of decode FFN, the most memory bound phase of inference

This means that throughput basically comes for free when we increase batch size because we get way better utilization of our tensorcores

So why would we ever run at low batch sizes? If you’ve ever used an LLM api, you’ve probably seen a /fast mode or slower batched inference. Amongst other things, batch size is very likely something they are tuning to make this work

Lower batch sizes means we have to move less KV cache data (we need a unique KV for each batch) and thus reduces the overall token latency. This means that tokens come out faster, but overall system throughput is lower. This is why fast tokens are so much more expensive!

*For anyone who pays attention to system inference characterization, that right graph might be familiar. It shows the pareto between interactivity (how fast your tokens come) and throughput (how many total tokens is your system producing per second). There is no right answer for where you should be on it, but it generally shows that making tokens fast per user make tokens slow across the system

In decode (which I hope I’ve convinced you is really important and also a hard problem to solve on GPUs), we don’t always see the same split between attention and FFN

FFN is pretty straightforward - super memory-bound, AI scales with batch size, context length doesn’t affect anything because the activations are the same size regardless, and we can’t do much about any of this

However, things get much more interesting in attention. Unlike FFN, the compute cost of attention scales quadratically with context length. We love long context because it means our agents can go a ton of turns without compacting away our important info, but this makes attention super expensive (although the silver lining is it also makes attention more compute bound)

To solve this quadratic scaling, many terrific ML researchers have been exploring subquadratic attention, or attention that doesn’t scale quadratically with context length

Simpler variants like GQA and MLA reduce the footprint of the KV cache. More in-depth methods like sparse attention (DeepSeek, GLM), sliding window attention (Gemma), and linear attention (Nemotron) all require less FLOPs at longer context AND reduce the memory footprint of the KV

Because attention scales with context and FFN does not, FFN latency dominates at short context and attention latency dominates at long context

Now that we understand the arithmetic intensity patterns within LLM inference, we can think about how to efficiently execute it on real systems

Default serving serially uses all the GPUs for prefill and then decode. As you can probably imagine, utilization is great during prefill but terrible during decode

At ISCA 2024, some computer architecture researchers that I respect a lot released the “Splitwise” paper. This was the first paper to explore disaggregated serving for inference by serving prefill on a compute-bound system and decode on memory-bound system, now called PD disaggregation. This is still quite popular today (although we’ve had to make a lot of hard engineering improvements to get it working at scale)

Recently, Attn-FFN has gained popularity as a more fine-grained disaggregation scheme. Because attention is moving more compute bound, we can put only decode FFN on the memory-bound system and PD disagg the prefill and decode attention. However, this significantly increases the amount of data moved across the interconnect and frequency of movements.

As we move to hybrid subquadratic attention models with unique AI patterns and more heterogenous systems where the prefill systems remain GPUs and the decode systems can be SRAM-based chips that crush any memory-bound workloads, maybe the way we approach serving will continue to change…?

Thanks for reading and hopefully some pieces can be insightful

Plz reach out with any questions - full blog post coming soon on my website!

Great thought! Boosting arithmetic intensity of FFN is super important but way less explored that attention variants recently

Some sort of fusion to limit off chip data movement or better expert parallelism for MoEs are definitely places to start

learned from the best

thanks!