Deep dive on LLM Inference at Scale — Harshul Jain, Audible & Tanmay Sah, Independent AI Researcher
Read the talk
Deep Dive on LLM Inference at Scale
Harshul Jain and Tanmay Sah build inference capacity from attention and GPU memory, then explain how quantization, caching, batching, and serving engines change the costs of producing tokens.
From a talk by Harshul Jain and Tanmay Sah
At a glance
Ideas worth remembering
KV storage multiplies across tokens, layers, KV heads, precision, and concurrent requests. The workshop’s roughly 4,000-token context across 80 users requires about 42 GB of KV storage alone.
Prefill performs prompt-wide computation before output begins. Decode repeatedly moves weights and preceding context for each new token, making bandwidth and inter-token latency central.
Memory capacity sets an upper bound on concurrency; latency objectives determine how much of that capacity the application can use.
Quantization frees storage, but task quality and execution speed need separate evaluation. The workshop corrects its MLA comparison to 14× savings after restoring the missing layer multiplier.
Paging reduces allocation waste, continuous batching fills available slots, prefix caching saves shared prompt work, and KV quantization reduces retained-state size.
Standard traffic and repeated agent prefixes can produce different engine comparisons. Benchmark the context reuse your application creates, then study eviction, compression, and distributed serving as its bottlenecks demand.
The model fits. The requests still might not.
Loading Mistral 7B onto a GPU consumes roughly 15 GB before requests add their own memory requirements. Longer prompts then increase memory use and delay the first generated token. Send five requests to a plain sequential implementation, and they wait behind one another. Harshul Jain, a senior software engineer at Audible, and Tanmay Sah, an AI researcher and quantitative modeler, begin their workshop with these three symptoms because each points to a different part of inference. 8:18
Inference is recurring work: every user, session, and generated token consumes resources. The workshop frames the response as reducing token usage or making the service use its hardware more efficiently. Its beginner-to-intermediate progression follows the second route, first explaining the bottlenecks and then separating model optimization from serving optimization.
The accompanying project, LLM Inference at Scale, collects slides, a benchmark report, and Jupyter notebooks; its workshop README points to the exercises. The opening demonstration walks through previously run results—Jain does not trust conference Wi-Fi—and establishes three measurements:
- Memory: How much additional storage does a longer context consume?
- Time to first token (TTFT): How long does the user wait before output begins?
- Throughput: How many tokens or requests can the service produce per second?
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Every token adds keys and values across the model
The inference pipeline converts text into tokens, converts tokens into embeddings, and passes those embeddings through transformer layers. The Mistral example has 32 layers. Each generated token extends the sequence used to generate the next one. Within those layers, attention works alongside normalization and feed-forward computation; attention is the part to inspect when explaining growing key-and-value storage. 12:32
Attention projects each token into query, key, and value vectors. Queries and keys determine attention scores against preceding tokens; values supply the information combined through those scores. Ten tokens require ten sets of vectors, and a thousand require a thousand. This happens across layers and KV heads, so storing the vectors costs much more memory than storing the prompt text.
How much does one token cost? The workshop uses 128 dimensions per head, 32 layers, eight KV heads, and a 16-bit representation:
The first factor of two stores both K and V; the final factor uses two bytes per element. That is approximately 131 KB per token. Eight KV heads, rather than 32, already reflect a memory-saving architectural choice: grouped query attention, explained later. 14:44
Follow one request through the capacity change. At roughly 4,000 tokens, its KV storage occupies about half a GB. Increase the context to roughly 16,000 tokens, and the requirement becomes about 2.1 GB. Alternatively, keep the 4,000-token context and serve 80 users concurrently: aggregate KV storage reaches roughly 42 GB. A 24 GB GPU cannot hold that cache, even before loading weights and accounting for overhead. Context length and concurrency multiply the same per-token cost.
GPU memory therefore divides into relatively fixed model weights, working overhead, and the remaining space for KV vectors. For a fixed GPU, model, and representation, adding users leaves less memory for each user’s context. The memory chart illustrates this directly: a target of 160 concurrent users requires shorter contexts than a smaller batch can support.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
The initial wait and the gaps between tokens have different causes
Prefill processes the prompt, constructs its key and value vectors, and performs attention computation across the input sequence. Longer prompts require more matrix-heavy work, explaining the increase in TTFT. Decode then generates output sequentially, computing attention for each new token against the preceding sequence. The gaps between output tokens introduce the workshop’s fourth metric: inter-token latency. 20:33
The GPU memory hierarchy explains why decode can be slow despite ample arithmetic capacity. In the workshop’s simplified picture, high-bandwidth memory, or HBM, holds large amounts of data. Much smaller shared memory supports faster local transfers. Computation moves chunks from HBM into shared memory, performs arithmetic, and writes results back. Generating tokens one after another repeatedly demands this movement.
The roofline model describes the distinction through arithmetic intensity, the number of floating-point operations performed per byte transferred. Prefill processes many input tokens together and does substantial arithmetic with the data it loads. Decode reads weights and preceding KV vectors to calculate one new token, doing less arithmetic per byte. Under this model, prefill tends toward a compute limit and decode toward a memory-bandwidth limit. 25:32
Where does the user wait, and what carries context between steps? The diagram separates prefill from repeated decode. It also looks ahead to the serving section’s KV cache: retaining vectors avoids rebuilding them, while decode still has to read the retained context.
The notebook results match that distinction. Prefill time grows visibly with input length. After excluding cold start, decode time stays nearer an average, but it is not constant: longer context means retrieving more preceding KV data. A comparatively flat decode curve does not remove the context-length cost.
Prompt processing and repeated generation
Read the diagram as text
- Input tokens. initial — Longer prompts supply more initial work.
- Prefill. initial — Construct K/V and perform prompt attention.
- Retained KV state. memory — Cache preceding vectors instead of rebuilding them.
- Decode step. repeated — Read preceding K/V and calculate the new token.
- Generated token. repeated — Extend the sequence before the next step.
- Input tokens → Prefill: Initial processing contributes to TTFT.
- Prefill → Retained KV state: Retain computed vectors.
- Retained KV state → Decode step: Read preceding context.
- Decode step → Generated token: Produce next token.
- Generated token → Decode step: Repeat; step timing sets inter-token latency.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
A batch can fit in memory and still miss its latency target
A first estimate of concurrency divides memory available for KV storage by storage per user:
This estimates a memory ceiling. A service also has a latency objective: larger batches can increase TTFT and inter-token latency, preventing it from using every theoretically available slot. 28:31
The workshop presents a tradeoff triangle of quality, latency, and throughput. Context contributes to quality when it contains information the task needs; shortening it to fit more users can remove that information. Different applications choose different priorities:
- Premium chat: Preserve useful context and fast responses, accepting fewer concurrent users per GPU and a higher cost.
- Asynchronous agents: Preserve task quality and process many concurrent long-running tasks, accepting more latency where the workflow permits it.
The capacity calculator combines GPU VRAM, bandwidth, FLOPS, and hourly cost. Its procedure fixes the most important requirement first—latency for interactive chat or minimum batch size for asynchronous work—then preserves the necessary context and compares feasible hardware. The chat example selects a 10-millisecond latency target as an input; the calculator does not demonstrate that every candidate achieves it. 34:46
A GPU with a higher hourly price may nevertheless produce tokens more cheaply if its useful throughput rises enough. Cost per million tokens therefore matters more than rental price alone. That comparison must use achievable concurrency under the application’s context and latency requirements; throughput the service cannot use is a poor basis for its cost estimate.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Compression changes the representation—and introduces assumptions
Sah introduces two deliberately fake algorithms. The “ostrich algorithm” ignores an inconvenient problem long enough to proceed with an assumption. The “World Cup algorithm” breaks a large problem into smaller pieces and carries useful results forward. These jokes distinguish two recurring optimization moves: assuming compression preserves useful behavior, and arranging work into manageable parts. 38:26
Weight quantization attacks fixed model memory. Sah’s GPT-OSS illustration uses a 120-billion-parameter model, approximately 240 GB of BF16 weights, and one 80 GB H100. Reducing the representation to FP8 gives approximately 120 GB, still too large; MXFP4 is described as bringing it to around 65 GB. These approximate footprints illustrate the speaker’s capacity decision. Compression matters here because it crosses the threshold from several GPUs to one. 39:26
For Mistral 7B, the alternatives include int8, int4, and NF4. Preserving quality is the ostrich assumption that must eventually face a relevant benchmark. Post-training quantization compresses an already trained model; quantization-aware training incorporates quantization into training or fine-tuning. Fewer bits establish a storage benefit, while task quality still needs evaluation.
Attention variants change how much separate K/V information the architecture maintains. Sah illustrates the heads by splitting 4,096 columns into 32 blocks of 128 columns. The important comparison is the degree of sharing:
- Multi-head attention (MHA): Maintains separate heads, with work that can run in parallel.
- Multi-query attention (MQA): Shares a single KV representation across queries.
- Grouped query attention (GQA): Shares KV representations within groups of queries, retaining more separation than MQA while storing less than MHA.
Multi-head latent attention (MLA) compresses key and value information into a latent representation. Position complicates the design: rotary position embeddings, or RoPE, encode token position, so the compressed representation must also accommodate position-related key information. Sparse attention changes a different part of the calculation by selecting important preceding tokens instead of attending to the whole preceding sequence.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
FlashAttention reduces traffic; quantization frees room
FlashAttention returns to the memory hierarchy. The expensive pattern repeatedly loads attention data, calculates intermediate results, and writes those results to HBM. Tiling divides the matrices into pieces that can be processed in fast on-chip memory. Online softmax carries summary information forward across tiles, avoiding repeated movement of full intermediate matrices through HBM. This changes how attention computation moves data. 48:00
The attention scorecard compares storage, throughput, and quality. Reducing 32 KV heads to eight reduces that storage component fourfold—the architectural choice already present in the Mistral memory calculation. GQA is presented as retaining quality close to MHA in many uses, while MQA’s more aggressive sharing can cost quality. Those judgments depend on the task.
The scorecard also briefly includes sliding-window attention, linear attention’s summarization approach, and Mamba state-space models. These alternatives change which preceding information is consulted or how it is retained. Sah also raises diffusion-based generation as a possible future departure from sequential token generation; the workshop’s bottleneck model is built around the autoregressive architecture it actually demonstrates.
The quantization notebook gives the earlier capacity example another lever. FP16 Mistral weights occupy about 15 GB; the int8 result is about 7.5 GB, and the int4 footprint is described as about 4.5 GB. Reducing the weight allocation leaves more of the same GPU for KV storage. That space can support longer contexts or additional concurrent requests, provided latency remains acceptable. 55:55
Memory savings do not establish a speedup. The notebook’s throughput plot uses theoretical values rather than throughput tests, and Jain notes that some benchmarks they studied showed lower throughput with int8. A smaller representation can improve capacity while its execution path still produces fewer tokens per second.
The attention demonstration ends with a consequential correction: Jain revises the earlier MLA savings figure from 56× to 14× relative to MHA because the calculation omitted the number of layers. The corrected figure applies to the workshop’s comparison; MLA’s storage formula varies with the architecture. Every layer’s retained state belongs in the memory calculation. 59:20
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Retain computed vectors, then manage their storage and scheduling
Serving optimization starts with computation repeated inside one request. A simple implementation rebuilds preceding tokens’ keys and values even though earlier steps already produced them. A KV cache retains those vectors for later decode steps. This trades memory for saved computation, making the earlier growing-memory problem the cost of an important optimization. 1:00:36
Paged attention addresses allocation waste. Jain’s deliberately small example reserves 2 KB for a request that needs only 1 KB, wasting half its reservation. Requiring contiguous storage also limits where a growing request can fit. Operating-system paging supplies the alternative: keep a contiguous logical view while mapping it to physical memory blocks at different addresses, allocating blocks as tokens arrive.
How can ordered token state occupy separate memory regions? The diagram separates the request’s logical view from its physical placement. The sequence retains its order, while the allocator can use separate blocks and add storage when the request grows. This removes the requirement to find one large continuous region.
The remaining serving mechanisms address different kinds of waste:
- Continuous batching: A fixed batch waits for all its requests to finish before accepting replacements. Continuous batching admits new work as capacity opens, reducing unused slots while longer requests continue.
- Prefix caching: A KV cache saves repeated computation within one request; prefix caching reuses computation for the same leading tokens across different requests.
- KV quantization: Lower-precision cached keys and values reduce storage per token, leaving room for more context or concurrent requests.
These mechanisms work on allocation, timing, repeated prompt work, and representation size respectively. Jain recommends using vLLM, which packages the serving optimizations, rather than implementing each from scratch. The next benchmark tests what changes when those capabilities replace the plain baseline.
Logical sequence, separate physical blocks
Read the diagram as text
- Logical KV sequence. logical — The request sees token state in order.
- Address mapping. Resolve logical storage to physical addresses.
- Physical memory blocks. physical — Store KV data in separate regions.
- New tokens. growth — Require additional storage as generation continues.
- Logical KV sequence → Address mapping: Preserve logical order.
- Address mapping → Physical memory blocks: Map to separate physical addresses.
- New tokens → Physical memory blocks: Allocate blocks when needed.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
More aggregate tokens can coexist with a longer initial wait
The benchmark holds Mistral 7B and its set of input questions fixed on an H100, comparing a plain Hugging Face baseline with vLLM configurations. The baseline produces about 51 tokens per second. Default vLLM includes paged attention, continuous batching, and KV caching; the workshop reports nearly 15× higher throughput, alongside higher TTFT and lower inter-token latency. Aggregate output and an individual request’s initial wait are separate measurements. 1:07:19
Adding prefix caching increases throughput further and reduces TTFT while leaving inter-token latency approximately unchanged. That fits its purpose: reuse removes prompt-processing work. Adding KV quantization leaves throughput and latency roughly similar in this benchmark while reducing KV usage. Saved computation and saved storage produce different observable changes.
Speculative decoding attacks sequential output generation. A smaller model drafts perhaps four or five tokens; the larger target model becomes the referee that decides how many to accept. The loop repeats. Accepting several drafts advances the output farther per verification cycle, while poor agreement between draft and target reduces the progress obtained from the extra drafting work. 1:12:02
Sah is skeptical of ordinary speculative decoding based on his personal tests and identifies draft-target alignment as the problem. The workshop benchmark also reports approximately similar results after adding it. That limits the conclusion to those tests; Sah suggests domains with more predictable output patterns, such as coding, as possible candidates.
The variants are “different flavors of same kind of soda,” in Sah’s phrasing:
- Self-speculative decoding: An auxiliary head within the target model performs drafting work.
- EAGLE: A smaller predictor uses features from a main-model layer to predict features, changing the relationship between drafting and the target model. This is Sah’s preferred approach.
- Medusa: Parallel token predictions propose output beyond a single next token.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Shared agent context changes the engine comparison
Prefix caching becomes especially useful when agents repeatedly send the same instructions and shared context. Sah contrasts a simple whole-prompt hash lookup with a radix tree. A changed word defeats a whole-prompt hash match. A prefix tree represents common beginnings; a radix tree compresses paths without branches. SGLang uses this structure for prefix caching. Reuse follows shared token prefixes, rather than the semantic similarity of differently worded requests. 1:14:29
TensorRT-LLM introduces another engine choice. TensorRT is an SDK; TensorRT-LLM is an inference engine comparable in role to vLLM and SGLang. Its NVIDIA-specific optimization reaches into layers and hardware execution, making the target hardware part of the engine decision. 1:16:37
The workshop’s vLLM–SGLang comparison uses ShareGPT questions on an H100. The presenters report no statistical difference for that standard workload, with similar request rates, TTFT, and latency. They then change the workload to repeated agentic branching. 1:18:06
The agent example first requests a proposal for solving traffic congestion in a city. The model produces a proposal; a second turn asks it to review that proposal and rate it from one to ten. Repeating those turns retains common instructions and context while introducing changing output. Prefix caching can preserve work for the shared beginning and compute the material that diverges.
For this branching setup, Sah reports SGLang performing three to four times better. The multiplier is setup-dependent and is not tied to a single named metric in the spoken comparison. It provides a reason to benchmark repeated-prefix workflows separately from ordinary API traffic. Jain recommends vLLM as the default for standard serving and trying SGLang when agent workloads are not performing well enough.
The comparison closes with a Clarifai benchmark that also includes TensorRT-LLM. Its value is another set of workload and hardware conditions to compare, rather than a universal engine ranking. Hardware-focused optimization and repeated-context reuse address different opportunities; the application determines which opportunity matters.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
The next layer is engineering the cache itself
The engine landscape extends into session routing and serving multiple models; NVIDIA Dynamo is mentioned in connection with agentic session routing. The deployment sequence remains practical: establish a baseline, choose a model that satisfies the task, compress its weights where appropriate, and use a serving engine for caching and scheduling. Mistral 7B is the workshop’s teaching model, rather than Jain’s recommended production choice.
KV cache engineering is the next layer of study: eviction strategies, cache compression, and hybrid memory. Retaining vectors saves computation, but creates new questions about which state to keep and how much memory to spend on it. These topics extend the workshop’s original capacity calculation into policies for managing scarce memory. 1:24:40
Distributed inference adds another set of problems beyond these single-GPU foundations. Jain treats it as material for a separate hands-on workshop. The useful closing question is whether a technique solves a problem the application actually has: begin with the bottleneck, then evaluate the proposed change through the same memory, latency, quality, and throughput requirements.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Resources
Further reading
Sah and coauthors study how runtime verification affects multi-step agent completion and compute cost. Useful for connecting the workshop’s asynchronous-agent serving discussion to the additional work introduced by verification and recovery.
Related talks
- Introduction to LLM serving with SGLang
Continue exploring the serving engine discussed for shared prefixes and agentic branching.
- From model weights to API endpoint with TensorRT-LLM
Follow the NVIDIA-focused engine choice toward an API-serving deployment.
- KV Cache-Aware Routing and P/D Disaggregation on Kubernetes — Yuchen Fama & Ashish Kamra, Red Hat
Extend the prefill, decode, and cache foundations into routing and distributed serving.
Transcript & slides
Loading transcript and slides…