Guest Lecture: Dan Fu

Inference is the engine

This lecture belongs to a guest: Dan Fu, who runs a lab at UCSD and works on inference at Together AI. His framing inverts the course’s. CS336 is mostly about how to train these models; today is “what it looks like from the other side” — serving them, doing inference, turning “electricity, into tokens, into intelligence.”

The motivation is scale, told fast. In 2018, at the start of his PhD, the largest models were hundreds of millions of parameters. By 2019, GPT-2 was considered too dangerous to release — a model this class could train if it tried hard enough. Today open-source models sit at a trillion parameters and the frontier is probably 5 to 10 trillion. And the transition is faster than it feels. In 1902 Manhattan had 130,000 working horses, and an 1898 conference on the resulting manure problem concluded there was nothing to be done — “you just have to hold your nose and deal with it.” By 1912, cars outnumbered horses. For language models, Dan argues, “that 1912 moment was probably last year”: it’s when he started writing the majority of his code with them.

If GPUs are the new oil — hundreds of billions of dollars of investment, sovereign wealth funds included — then oil is only useful with an engine. A model is just a DAG of operations, “some mathematical object that exists in the ether”; inference engines and GPU kernels are what map it onto hardware. You’ve seen the training side of that mapping (you wrote FlashAttention); inference has a whole world of complexity of its own.

The lifetime of a token

When a request hits an inference service: tokenize it; schedule it onto GPUs (possibly different machines for different phases); check the KV cache to see whether any of this compute has been done before; execute the model, split across nodes and GPUs; then sample, watch for stop tokens, maybe run a safety check, and return the tokens. The engine runs this scheduling–execution–sampling loop continuously, waiting for requests.

What arrives is shaped by the workload, and production traffic looks like neither training data nor a made-up distribution. A Cursor-style coding workload sends tens of thousands of input tokens and may get back thinking tokens or a short edit. Narrative summarization — pasting entire books into a chat window — looks different again from “explain first-order calculus.” Modern use is turn-based and agentic: the coding agent greps your codebase, does an internet search, feeds results back to the model, over many turns. Sessions differ too — Dan has a ChatGPT conversation about his workouts that he touches once every other week. And each application sets its own targets: first token back in under a second for interactive chat, or a whole 500-token response within a fixed budget. Input length, output length, session stickiness, gaps between turns, latency targets — that’s a workload.

One system, many requests

With many requests alive at once, the engine does continuous batching: requests join and leave the running batch at each step, sharing compute and KV-cache memory; a long request can exhaust GPU memory and force newcomers to queue. The KV cache exploits prefix sharing — everyone starts with “hi, ChatGPT,” and the second turn of a conversation shouldn’t recompute the book pasted in the first — via a traditional tree data structure over token prefixes: look up which tokens have been seen, fetch their activations.

A trillion-parameter model won’t fit on an 80 GB GPU, so you split it: tensor parallelism slices every tensor across GPUs; mixture-of-experts models can place experts on different GPUs. And because prefill and decode have opposite bottlenecks, “pretty much we’ve all started adopting” prefill–decode disaggregation: prefill on one set of workers, decode on another, each specialized. The split now reaches the hardware itself — NVIDIA plans GPUs for prefill and Groq’s LPU chips for decode, OpenAI has a compute partnership with Cerebras, SambaNova is making its own bets. Co-design runs the other way too: size your model to the chip’s memory, train in NVFP4 if you’re serving on NVIDIA (MXFP4 on AMD), compress the KV cache aggressively — DeepSeek’s MLA — if agentic workloads need it hot.

The KV cache wants to be as large as possible, so it overflows GPU memory into CPU DRAM — one reason Jensen Huang got obsessed with CPU performance: “if your $500,000 machine is being bottlenecked by the$1,000 CPU that you purchased to put on top of it, that’s not a great place” — and then onto SSDs, part of why OpenAI is buying up storage. Eviction and prefetching follow: LRU is a decent heuristic, and better still is predicting the future — a user reopening a month-old conversation is a strong signal to prefetch it onto a GPU.

The first of Dan’s three demonstrations lives here, in the routing layer: cache-aware prefill–decode disaggregation. If an average conversation lasts 10 turns, about 10 percent of requests are brand new — thousands of cold tokens, expensive prefill. You don’t want someone’s pasted book computing next to someone mid-conversation asking why 1 plus 1 equals 2. So route low-cache-hit requests to one pool of prefill nodes and warm requests to another. “It’s two lines of code in the routing layer” — and up to 40 percent faster serving. His assessment of the field: “we’re very early.” In 10 or 20 years this will all look obvious.

Megakernels: decode at the speed of light

Decode’s fundamental problem: generating a single token requires running the whole model, which turns a massively parallel GPU into “basically a glorified memory loader.” Worse, we write kernels one operation at a time — a norm kernel, a matmul kernel, an attention kernel — because kernels are hard enough to write as it is. Plot GPU activity over time (each row one of the H100’s 132 streaming multiprocessors; 148 on a B200) and the gaps dominate: kernel launch and teardown, and tail effects, where the whole batch waits for its longest sequence. Between-kernel gaps add up no matter how well each kernel is written.

A megakernel covers many operations in a single kernel — FlashAttention-style fusion “done more aggressively across a larger number of things” — and treats the one GPU as a distributed system: here is all the work, here are the dependencies, schedule it to maximize utilization. Applied to the attention inference kernel alone, that’s a 30 to 70 percent speedup. Applied to a whole Llama-1B layer, operations overlap in ways separate kernels never could: the KV-cache load for attention starts while QKV-plus-RoPE is still running, and the O-projection’s weights start loading before attention finishes. The implementation is a CUDA framework with an instruction-based abstraction — each sub-kernel in its own file, a virtualized shared-memory system orchestrating them — built on ThunderKittens, “almost like Triton, except more low level.”

The payoff is near-speed-of-light decoding: 72 percent memory-bandwidth utilization on an H100, beating state-of-the-art engines. The cost, from the Q&A, is “people’s blood, sweat, and tears”: a talented kernel engineer working for a year can write megakernels for one hardware platform, two or three models, batch sizes 1 to 16 — “go batch size 17, nope, start over.” Together is building compilers to automate it, and DeepSeek shipped a megakernel for its MoE inference layer that even fuses the NCCL communication calls.

Parcae: recurrence as a third scaling axis

The last demonstration is an architecture from Dan’s UCSD lab. Parcae is their take on looped transformers: instead of tokens flowing through every layer once, a block of layers runs in a loop. Parameters stay constant while a dial increases FLOPs — quality without parameter cost — and older theory says looped models can express things equal-parameter unlooped ones can’t. The driving question: what’s the best intelligence per parameter?

The known problem is instability. Sweep the learning rate on prior looped models and nine times out of ten training blows up — NaNs, loss spikes — with hacks like norms in every layer or “just pick 2e-4” as the fixes. Parcae’s move is analytical. The loop block is a mess of nonlinearities (attention, GELU, RoPE), but empirically the residual changes only a little per loop, so model the loop as a dynamical system: box up the nonlinear part as $R$, leaving a matrix $A$ that transforms the residual each iteration and a matrix $B$ that injects the initial activation. Drop $R$ — the $A$ and $B$ terms dominate the magnitudes anyway — and high-school calculus gives a closed form for the activation at step $T+1$, dominated by powers of $A$.

The result is a stable loss curve at the 6e-4 learning rate that killed prior models. The comparison is instructive: an unconstrained baseline’s activation norms explode to $10^{19}$, and a normed baseline looks fine in its norms yet still loss-spikes — the model tries to expand its activations for representational room while the norm squeezes them back, two pressures fighting. Parcae removes the fight, and quality follows: better perplexity and downstream performance than the prior Recurrent Depth Model and a strong speedrun-tuned transformer baseline of the same architecture.

The scaling laws are the provocative part. Holding parameters fixed and growing data, the iso-param, iso-FLOP curves slope down and to the right in recurrence — meaning as you add data, you should also add recurrence, following classic power laws. At a fixed FLOP budget, the looped model reaches lower validation loss than the fixed-depth one. Yet every model deployed today sits at zero recurrence with maximal data — the far left of these curves — “which suggests that there might be something slightly better that we could be doing.” And the inference angle closes the loop on the whole talk: fewer parameters mean more room for KV cache and less communication, and Dan’s dream is a recurrent block small enough to live in a megakernel on a 250 MB LPU-class chip, weights resident, activations streaming through.

Check yourself

Three exercises in the lecture’s spirit — reason about serving, not training.

Problem 18.1 decode napkin math

A 3B-parameter dense model is served in bf16 on one H100 (memory bandwidth 3.35 TB/s). (a) Roughly how many FLOPs and how many bytes does one decode step cost, and what does that say about its regime? (b) At the megakernel’s 72 percent bandwidth utilization, what decode speed in tokens per second is achievable, ignoring KV-cache traffic?

Show solution

(a) One token through the model costs about $2N = 6\times10^9$ FLOPs, but every parameter must be read: $2N = 6\times10^9$ bytes at 2 bytes each. Arithmetic intensity is about 1 FLOP per byte — orders of magnitude below an H100’s compute/bandwidth ratio — so decode is deeply memory-bandwidth-bound: the GPU is a “glorified memory loader,” and the time per token is set by how fast weights stream in. (b) Effective bandwidth is $0.72 \cdot 3.35\times10^{12} \approx 2.41\times10^{12}$ bytes/s. Each token needs $6\times10^9$ bytes of weights, so about $2.41\times10^{12} / 6\times10^9 \approx 400$ tokens/s. This is why fewer parameters (or looped models, or quantized weights) translate directly into decode speed.

Problem 18.2 loop stability

A looped transformer runs its recurrent block $T = 24$ times, and the residual dynamics are governed by a matrix $A$ with spectral radius 1.08. (a) By roughly what factor do activations grow over the loop, and what training symptom does that produce? (b) State the condition Parcae enforces and how its parameterization of $A$ guarantees it.

Show solution

(a) The closed-form solution is dominated by $A^T$, so activations scale like $\rho(A)^T = 1.08^{24}$. Since $1.08^{24} = e^{24\ln 1.08} \approx e^{1.85} \approx 6.3$, activations grow about 6x per pass through the loop — and this compounds across training until you see loss spikes and NaNs, exactly the nine-out-of-ten blowups seen in learning-rate sweeps of earlier looped models. Even a modest radius above 1 is unstable because it gets powered up every loop. (b) Parcae requires $\rho(A) < 1$, achieved by making $A$ a negative diagonal matrix so that its powers decay to zero rather than explode; $B$, applied only once, just gets a linear norm. Choices like identity $A$ (marginally stable) or fully learnable $A$ (unstable) violate the condition.

Problem 18.3 cache-aware routing

A chat service observes that conversations average 8 turns before the user leaves. (a) What fraction of incoming requests are cache-cold, and why are these disproportionately expensive? (b) Explain what the cache-aware router does with them and why mixing hot and cold requests on the same prefill nodes hurts latency.

Show solution

(a) One turn in 8 starts a fresh conversation, so about 12.5 percent of requests are cold. A cold request has no KV-cache prefix to reuse — its entire prompt (potentially thousands of tokens, or a whole pasted book) must be prefilled from scratch, while a warm request only prefills the small delta since the last turn; the cold minority carries most of the prefill FLOPs. (b) The router checks the expected cache hit rate and sends low-hit-rate requests to a dedicated pool of prefill nodes, warm requests to another. Mixed together, a long cold prefill occupies the GPUs while short warm turns queue behind it — the tail effect at the scheduling level — so the quick conversational turns inherit the latency of someone else’s book. Separating the pools let each specialize, worth up to 40 percent faster serving from two lines of routing code.