# Chapter 2: Models _Inference Engineering_ by Philip Kiely. © 2026 Baseten Labs, Inc. All rights reserved. [Full book index](https://www.baseten.co/inference-engineering/llms.txt) Inference engineering is the practice of making generative AI models faster, less expensive, and more reliable – without sacrificing the quality that makes them so valuable. Both improving performance and preserving quality require a strong intuition for how models work under the hood. Generative AI models are a composition of big, complex neural networks. The history of neural networks stretches back to the 1950s, when the first perceptrons for simple binary classification were implemented in hardware. In the following decades, perceptrons were abandoned but then reinvented from single to multi-layer perceptrons with a new concept, back-propagation, which introduced hidden states between layers and a learning procedure that repeatedly adjusts weights within the network. These neural networks had only a few layers. In the 2000s, research began into deep neural networks with dozens of layers. In 2012, AlexNet became the first deep neural network to show promising real-world capabilities and the effectiveness of GPUs for deep learning, leading to new architectures like word embedding models for text and Generative Adversarial Networks (GANs) for images. But the story truly starts in 2017, when Vaswani and colleagues published the seminal paper “Attention Is All You Need,” introducing the transformer. A transformer is a neural network with an attention mechanism that can learn relationships between various parts of a sequence. Transformers are the foundation of generative AI. Transformers aren’t just for LLMs, they power every modality of model from embedding to voice to image and video generation. Across modalities, there are two important styles of transformer-based models: - **Autoregressive token generation:** Start from a tokenized sequence and predict the most likely next token. - **Iterative denoising:** Start from random noise and refine toward the most likely output via diffusion. This chapter explores the architectural details of LLMs (autoregressive token generation) and image generation models (iterative denoising). ## 2.1 Neural Networks Generations of research into neural networks form the theoretical foundation for generative AI. To be a productive inference engineer, you need a basic intuition for essential concepts in neural networks. This section provides a high-level introduction; Appendix B offers recommendations for further reading. The fundamental unit of a neural network is a node (a.k.a., neuron). A node is a short program that takes an input, multiplies it by some weights, adds some bias, and returns the result. A group of nodes forms a layer. Nodes within a layer are independent of each other – they do their own calculations. The connection between nodes, or the “network” in a neural network, is between layers, where the nodes in a layer receive the output of the previous layer. The neural networks behind LLMs contain dozens to hundreds of layers. There are three types of layers: - **Input layer:** The first layer, which accepts and processes the input to the neural network. - **Hidden layers:** Every layer between the first and last, which iteratively transform the input to arrive at an output. - **Output layer:** The final layer, which returns the prediction from the network. Each layer produces an output that the next layer reads as input. For the hidden layers, these outputs are called hidden states. Hidden states are one kind of internal representation for data within a neural network. A key aspect of internal representation is its dimensionality, or the actual size of the vectors used. ![Figure 2.1: Multi-layer neural networks have one input layer, many hidden layers, and one output layer.](https://www.datocms-assets.com/104802/1788123169-inference-engineering-figure-2-1.png) _Figure 2.1: Multi-layer neural networks have one input layer, many hidden layers, and one output layer._ Internal representations for text input increase the dimensionality, encoding text chunks into vectors of hundreds or thousands of numbers to capture semantic meaning. But internal representations for image models reduce the dimensionality from millions of pixels down to a manageable size. There are neural networks for creating these internal representations, and there are neural networks for using them: - **Encoder:** Takes an input like text or an image and creates an internal representation of the input that includes additional information and semantic meaning. - **Decoder:** Uses the internal representation to generate an output like text or an image. Neural networks are composable. You can combine multiple neural networks together into a single model or use them sequentially to build a pipeline. Modern LLMs are decoder-only, while encoder-only models are somewhat rare today, with old-school text embedding models from the BERT family as a prominent example. Many models in other modalities use an encoder-decoder architecture. Whisper, a popular open model for audio transcription, uses an encoder to process audio input and a decoder to generate text tokens. ### 2.1.1 Linear Layers and Matmul The most essential operation within a neural network is a matrix multiplication, or matmul. A matmul takes an input vector (a list of numbers) and a matrix (a grid of numbers) and multiplies the vector through the matrix to produce an output vector. Within a neural network, a linear layer is the simplest form of matmul. Given an input vector, the linear layer applies a weight matrix and adds a bias vector: ![Figure 2.2: In a matmul, the output vector y is the product of an input vector x and a weights matrix W plus a bias vector b.](https://www.datocms-assets.com/104802/1788123177-inference-engineering-figure-2-2.png) _Figure 2.2: In a matmul, the output vector y is the product of an input vector x and a weights matrix W plus a bias vector b._ The weights of any given linear layer are a small part of a generative AI model’s total weights, and the individual values within the weights matrix are set during training. ### 2.1.2 Activation Functions Matrix multiplication is composable, meaning that multiplying a vector by two matrices is equivalent to multiplying that vector by the product of those matrices. ![Figure 2.3: Two matmul equations representing separate layers collapse due to composition of linearity.](https://www.datocms-assets.com/104802/1788123182-inference-engineering-figure-2-3.png) _Figure 2.3: Two matmul equations representing separate layers collapse due to composition of linearity._ This is a problem for multi-layer neural networks because a series of linear layers, each one a matmul, would collapse into a single layer with all of the matrices multiplied together. Deep multi-layer neural networks are useful because more layers use more parameters effectively and encode more meaning in hidden states. Neural networks separate layers by breaking linearity with an activation function. Activation functions are non-linear to prevent composable matmul from collapsing layers, and are differentiable or mostly-differentiable to support back propagation. One of the most basic activation functions in inference is ReLU, which stands for Rectified Linear Unit. ReLU is a simple function: if X is greater than zero, return X, else return zero. There are dozens of activation functions – including one named “Swish” thanks to its resemblance to the Nike logo – but most follow the same general pattern of mapping negative values to zero or near-zero, while keeping positive values unchanged. ![Figure 2.4: Activation functions like ReLU are used to break linearity in multi-layer neural networks.](https://www.datocms-assets.com/104802/1788123191-inference-engineering-figure-2-4.png) _Figure 2.4: Activation functions like ReLU are used to break linearity in multi-layer neural networks._ Activation functions like ReLU, SiLU, Swish, and SwiGLU are fast to run, easy to train on (as they are mostly differentiable, they have a gradient at least for most values), and break linearity to support multi-layer neural networks. ## 2.2 LLM Inference Mechanics LLMs are autoregressive token generation models. An LLM generates new tokens one at a time based on every previous token. These tokens, the atomic units of language models, are numbers that represent chunks of text. Modern LLMs use subword tokenization, meaning that each token is a word or a fraction of a word. ![Figure 2.5: Subword tokenization uses one token per common word and punctuation mark, but it splits less common words into multiple tokens.](https://www.datocms-assets.com/104802/1788123199-inference-engineering-figure-2-5.png) _Figure 2.5: Subword tokenization uses one token per common word and punctuation mark, but it splits less common words into multiple tokens._ Converting text into tokens and tokens into text does not require any neural networks. Instead, a tokenizer is a simple mapping between strings and their numerical token representation. A language model’s vocabulary is the complete mapping between tokens and strings. Vocabularies and tokenizers vary from model to model, with more recent models employing more efficient tokenization schemes; the fewer tokens required to generate an output, the faster the end-to-end inference. Most models have over 100,000 tokens in their vocabulary. In other modalities, like speech synthesis, the model vocabulary is expanded to let tokens represent other information like audio waveforms. Inference involves two or three sequences of tokens: - **Input sequence:** The prompt, chat, context, functions, and other input passed into the LLM. - **Reasoning sequence:** Optionally, for reasoning models, an intermediate output sequence for thinking. - **Output sequence:** The response generated by the LLM. Combined, these sequences are limited to the model’s context window (the total number of tokens the model can process and generate per request). A request may further limit the output sequence length with a `max_tokens` argument. While the input sequence is a single string, LLMs are trained to accept varied inputs like multi-turn chat sequences with roles, function signatures for tool calls, and, in some cases, multimodal inputs. These inputs need to be combined into a single sequence. This is handled by the chat template, which differs subtly from model to model and must be implemented correctly in the inference engine. Tokenizing the input sequence with the chat template applied is step zero for inference. Then, there are two primary phases of inference: - **Prefill:** Process the input sequence to calculate attention for each input token and store those values in a KV cache. - **Decode:** Perform forward passes through the model to generate tokens autoregressively. Each forward pass in the decode phase must generate a token. This takes a few extra steps, as neural networks output vectors, not tokens. The output layer of the neural network used in LLM decode generates a vector of logits. The length of this vector equals the model’s vocabulary size. After normalization, these logits represent the probability of each potential token in the vocabulary being the correct output. ![Figure 2.6: A decode pass generates a logit for each token in the model’s vocabulary, then normalizes logits to percentages.](https://www.datocms-assets.com/104802/1788123204-inference-engineering-figure-2-6.png) _Figure 2.6: A decode pass generates a logit for each token in the model’s vocabulary, then normalizes logits to percentages._ The output token is selected via a weighted random number generation based on the normalized probability vector. You can nudge that process via inference arguments: - **Temperature:** Adjust the logits themselves before normalization. - **Top-k:** Select the k most likely tokens after normalization, then re-normalize among them. - **Top-p:** Select the smallest set of tokens after normalization whose probabilities add up to p. A lower temperature, top-k, or top-p makes LLM output more predictable as the model is constrained to selecting highly likely tokens. Setting the temperature to 0 or top-k to 1 makes token selection deterministic (always select the highest-probability token). When generating structured output, where the output conforms to a schema like JSON, there are additional tools like logit biasing for further directing the output of an LLM. These apply after each forward pass and support important LLM abilities like tool use; their correct implementation is essential for high-quality inference. This logit generation and token selection process continues until the model decides that the stop token, a special value signifying the end of the output sequence, is generated (unless the context window or max tokens limit is hit first). The two key pieces of this inference loop are generating the KV cache during prefill and generating output tokens (or more specifically the logit vectors that become tokens) during decode. These steps take the overwhelming majority of the time and resources during inference as they rely on large neural networks. ### 2.2.1 LLM Architecture Every LLM on Hugging Face (the largest repository of open models) includes a `config.json` file: a few dozen lines detailing the model’s architecture. The architecture of a model is a collection of decisions made during the training process about the nature and shape of each component of the model. Within a single architecture, there may be: - **Multiple sizes:** Models at different parameter counts, like Llama 8B and 70B. - **Multiple variants:** The “base” and “instruct” variants of a given model share the same architecture. - **Unlimited fine-tunes:** Methods like LoRA (Low-Rank Adaptation) change behavior, not architecture. These architectures matter because they determine runtime and engine support. If you have a highly optimized deployment of a given architecture, you can deploy another variant of the same architecture and enjoy the same performance improvements. Model architecture is one of the first lines in most configuration files. To parse an architecture name like `Qwen3MoeForCausalLM`: - **Qwen:** The model family, or the brand name of the model. - **3:** The major version of the architecture within the family. - **MoE:** Indicates a Mixture of Experts model (see 2.2.4). - **CausalLM**: Indicates a causal language model. A causal language model predicts the next token in a sequence based on previous tokens, as opposed to, for example, a masked language model which fills in the blank based on surrounding tokens to the left and right. All generative LLMs today are causal language models. Beyond the architecture’s name, the `config.json` file contains information about the nature and dimensions of the various layers that form the underlying neural networks of a model and the vectors that pass through them during inference. ### 2.2.2 Transformer Blocks The main body of an LLM is a series of dozens to hundreds of transformer blocks. These blocks form the core of a large neural network with three kinds of layers: - **Embedding layer:** The input layer of the neural network takes tokens and returns embeddings. - **Transformer blocks:** The hidden layers within the network are transformer blocks that generate a prediction. - **Output layer:** Also known as a language modeling head or LMHead, converts the hidden states from the transformer blocks into a vector of logits, one for each token in the model’s vocabulary. Within the transformer blocks, there are sublayers for attention, a feed-forward neural network, and normalization. ![Figure 2.7: Transformer block diagram, adapted from “Attention Is All You Need” (Vaswani et al., 2017).](https://www.datocms-assets.com/104802/1788123212-inference-engineering-figure-2-7.png) _Figure 2.7: Transformer block diagram, adapted from “Attention Is All You Need” (Vaswani et al., 2017)._ The feed-forward neural network is a multi-layer perceptron. These linear sublayers make up the majority of the trainable weights within an LLM, while the attention sublayers are the second-largest component. Other components like normalization and activation functions are a rounding error in the model’s size. While linear sublayers are the largest portion of the weights, the more complex operation for inference is attention. ### 2.2.3 Attention Attention is the mechanism transformers use to relate a given token to other tokens in the sequence. Humans are good at interpreting the relationship between words. Attention brings the same capability to LLMs. Consider the sentence “I decided to write a book because I thought it would be easy, but it was actually hard.” Attention shows that the word “it” in the sentence refers to writing a book. The standard form of attention is scaled dot-product attention, as shown in this equation. ![Figure 2.8: The attention equation, adapted from “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (Dao et al., 2022).](https://www.datocms-assets.com/104802/1788123220-inference-engineering-figure-2-8.png) _Figure 2.8: The attention equation, adapted from “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (Dao et al., 2022)._ Attention takes three inputs: - **Q (queries):** The embedded representation of the token being generated or updated. - **K (keys):** Representations of all prior tokens. - **V (values):** Computed attention values for all prior tokens. Attention sublayers within models are multi-head, where each head is one attention operation. If you visualize the architecture of a neural network, these heads are parallel to each other on the same sublayer. Each head could be responsible for attending to different kinds of relationships between tokens, like subject-verb agreement and co-reference resolution. There are two main types of attention: - **Self-attention:** Q, K, and V are all from the same sequence. - **Cross-attention:** Q is from a different sequence than K and V, conditioning Q on external information. LLMs use self-attention (with a causal mask to prevent looking ahead in the sequence), while image generation and multimodal models also use cross-attention (e.g., between the image being generated and the associated text prompt). Because attention looks for relationships between the current token and every previous token in the sequence, it’s a quadratic-time equation with respect to sequence length. As context expands, attention slows. In practice, attention is linear, not quadratic, thanks to the KV cache. The KV cache is a standard component in attention implementation and stores key-value pairs for each previous token. By looking up this information in the KV cache rather than recomputing it each time, attention runs in linear time. The KV cache is built during LLM prefill, used and updated during decode, and lives on GPU memory by default. Storing, accessing, and re-using the KV cache, a major topic in inference engineering, is covered in detail in section 5.3. ### 2.2.4 Mixture of Experts Models The density of a neural network is determined by the number of connections between layers. Denser networks retain more information, while sparser networks take less compute and memory to run. Mixture of Experts (MoE) is an architecture optimization that adds sparsity to linear layers. Rather than a single giant matrix, an MoE model has hundreds of smaller matrices (the experts) and routes each input to a small selection of experts. This is called activating the experts. For an MoE model like Qwen3-235B-A22B, 22 billion parameters out of the model’s 235 billion total parameters are activated per request. Thanks to their low number of active parameters, MoE models are highly efficient for single-request local inference. However, in batched inference on production servers, different requests activate different experts, and you should expect almost all of the model parameters to be active at any given time unless sparsity is achieved in large-scale Expert Parallelism (section 5.4.2). Expert routing is granular. Each forward pass through a model generates one token by working through every layer of the model. The router, a tiny model within the LLM, picks which experts to activate at each layer of the model. In the Qwen example, with 128 experts, the router picks eight experts at each of the 94 layers for every token that is generated. MoE architectures are especially popular for larger models with 100B+ parameters, though there are MoE models as small as 20 to 30 billion parameters. Mixture of Experts unlocks a new form of inference parallelism called Expert Parallelism, which enables high-throughput inference for large models on multiple GPUs. ![Figure 2.9: Mixture of Experts architecture includes both sharding and replicating to take advantage of multi-GPU inference.](https://www.datocms-assets.com/104802/1788123225-inference-engineering-figure-2-9.png) _Figure 2.9: Mixture of Experts architecture includes both sharding and replicating to take advantage of multi-GPU inference._ Models under 32B parameters, and especially models under 8B parameters, tend to use traditional dense architectures efficiently. Domain-specific models for tasks like tab completion also don’t gain much benefit from MoE as the entire model is effectively one expert. ## 2.3 Image Generation Inference Mechanics Image generation models take a text prompt and create an image based on the prompt. These models slightly predate the public rise of LLMs, with both closed models from Midjourney and open Stable Diffusion models first released in the summer of 2022. Image generation models aren’t monolithic models like LLMs. Instead, they are pipelines of multiple models working together to generate images. Within a foundation model for image generation, there are three essential components: - **Text encoder:** Converts the text prompt into instructions that the image generation model can understand. - **Denoising model:** The heart of the model, iterates from noise to an image based on the prompt. - **Variational Autoencoder (VAE):** Converts the model output from latent space to pixel space. This pipeline mentality extends through the image model ecosystem. Beyond the base model, image inference often includes: - **LoRAs:** Lightweight fine-tunes to change style and enhance quality. - **ControlNets:** Outlines and edges to steer output images to match broad shapes and colors. The vast and rich open-source ecosystem around image generation models includes tools like ComfyUI for building complex pipelines of foundation models and adaptions, swapping components to produce unique outputs. The entire image generation pipeline operates in latent space. An ordinary HD image may be 1024x1024 pixels; that’s well over a million pixels. As the denoising model needs to calculate attention over the entire image in parallel, it would be infeasible to work in pixel space. Latent space is a low-dimensional representation of an image. A latent space matrix for an image may be 128x128, or about one percent of the total values of the pixel space it represents. The latent space is initialized as random values, or noise. The denoising model refines that noise into an image over a series of steps based on the text prompt. Each step updates the entire latent space, unlike LLMs which process tokens one at a time. Most image generation models take 30 to 50 steps to create a high-quality image. ![Figure 2.10: Diffusion-based models iteratively generate an image from noise, generally over 30 to 50 steps.](https://www.datocms-assets.com/104802/1788123234-inference-engineering-figure-2-10.png) _Figure 2.10: Diffusion-based models iteratively generate an image from noise, generally over 30 to 50 steps._ Within each step, the model runs two forward passes: one with conditioning (the text prompt) and one without conditioning. These generations are then combined based on a guidance scale. Because of this two-part step, a 50-step image generation actually takes 100 forward passes. This process, and other essential parts of image generation, are controlled on a request-by-request basis via inference arguments. The most important arguments are: - **Prompt:** Describes what the image should look like. - **Negative prompt:** Separately describes any styles or objects that should not be in the image. - **Number of steps:** Trades off speed and quality with the number of denoising steps, 30 to 50 for most models. - **Guidance scale:** Controls the balance between creativity and prompt adherence, integer value generally around 4. - **Image size:** Selects from a fixed menu of resolutions and aspect ratios for the output image. While these core mechanisms are common across image generation models, their architecture has evolved considerably in the past few years. ### 2.3.1 Image Generation Model Architecture Image generation models are built on transformers, specifically diffusion transformers. A diffusion transformer is very similar to the transformers that LLMs use, but instead of processing embedding representations of discrete tokens, it processes image data. Diffusion transformers look at images in patches. When training a text-to-image model, the images in the training data are fed in via overlapping patches of 2x2 or 4x4 pixels, which are then embedded into latent space. Inference works in the opposite direction, with latent space transformed back to pixels once the image generation is final. Image generation models are pipelines of multiple models, including a text encoder, denoising model, and variational autoencoder. A clean example of this pipeline is Stable Diffusion XL (SDXL). SDXL is an old model, but its architecture remains relevant. ![Figure 2.11: SDXL architecture pipeline, adapted from “SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis” (Podell et al., 2023).](https://www.datocms-assets.com/104802/1788123242-inference-engineering-figure-2-11.png) _Figure 2.11: SDXL architecture pipeline, adapted from “SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis” (Podell et al., 2023)._ SDXL’s pipeline contains two diffusion models for denoising, the base and refiner. These models were trained for separate tasks: the base model goes from pure noise to a coherent image, while the refiner model adds details and ensures prompt adherence. Modern models substantially outperform SDXL with better prompt adherence; accurate faces, hands, and details; legible text rendering; and support for image-to-image inference. These models, like the Qwen Image family, are broadly similar to the SDXL pipeline, but with larger and more capable models at every step. | Component | SDXL (2023) | Qwen Image (2025) | | :----------- | :--------------- | :---------------- | | Text encoder | CLIP-based model | Qwen 2.5 VL (7B) | | Denoiser | <4B parameters | 20B parameters | | VAE | Single encoding | Dual encoding | The substantial growth in capability in modern image generation models like Qwen Image comes from larger component models and more complex pipelines. New abilities like legible text rendering and photorealistic human faces comes from switching from tiny NLP models to full LLMs for text encoding and increasing parameter counts on denoisers by a factor of five. These larger models take more resources to run. Fortunately, as models have grown larger, GPUs have grown more powerful, though inference engineers can’t rely on hardware gains alone to run these models efficiently. The latest research direction in image generation models is blending diffusion transformer architecture with LLM architecture. Anything that can be tokenized can be modeled as an LLM. LLMs have the advantage of baked-in text understanding, and they are already an important component of image generation pipelines. LLMs solve many of the problems inherent to diffusion models. Where diffusion models can only produce a fixed size of output, LLMs are autoregressive and can produce a variable-length output. And where image models need up to 100 forward passes to generate an image, LLMs generate tokens in a single forward pass. Models like HunyuanImage-3.0 use this LLM-style architecture. While this is a new frontier in image generation, there are other existing architectures for accelerating image creation and for extending image models to video generation. ### 2.3.2 Few-Step Image Generation Models The most time-consuming part of image generation is the 30 to 50 denoising steps. Rather than making each step faster, what if there was a way to optimize image models by simply using fewer steps? Few-step image generation models are trained to do just that: create high-resolution images with eight or fewer denoising steps. These models are 80 to 90 percent faster out of the box than traditional image generation models, though their output quality is noticeably lower. There are two primary methods for creating these models: - **Latent consistency:** Train a model to predict the target latent image vector directly and repeat the prediction two to four times to enhance quality. - **Distillation:** Use adversarial distillation and/or progressive distillation to train a small model to emulate a larger one in fewer inference steps. Distillation is more common than latent consistency today. When new image models like FLUX and Qwen Image are released, members of the open image model community create distillations in addition to quality and style-oriented LoRAs. If you have a latency-sensitive use case where quality is less important, like real-time generative filters, consider few-step image generation models. ### 2.3.3 Video Generation Video generation models are architecturally similar to image generation models, just bigger. They have three to five times more parameters and encode ten to one hundred times more information in latent space. The naive approach to video generation is to work frame-by-frame. Early video generation used this framewise approach: first generating a starting frame, then using that frame to generate the next frame, and so forth. The issue with a framewise approach is error accumulation. Small issues early on compound with each frame, and the video goes off the rails quickly. Instead, modern video models hold the entire video in latent space and modify it on each denoising step. Each frame attends to each other frame and is updated on every forward pass. If latent space for an image model represents two physical dimensions, X and Y, then latent space for a video model represents three dimensions: X, Y, and T (time). The main limitation of this approach is that videos are a fixed number of frames, just like image generation models have fixed aspect ratios. Modern video generation models create sequences of a few seconds. The constraint on video length is compute resources. Even with the latest GPUs, attention over massive latent space is extremely expensive, taking several seconds of inference for each second of video. Video generation models are so compute-intensive that they typically run with a batch size of one, meaning a full node of eight GPUs is working on a single request. While attention on each denoising step is expensive, video generation models have the same total number of steps as image models, generally around 50 steps. Video generation is a recent modality compared to LLMs and image generation. Their limitations line up with the limitations of LLMs two years ago. | LLMs (Late 2023) | Video gen models (Late 2025) | | :---------------------- | :---------------------------- | | High TTFT, Low TPS | Slow generation times | | Frequent hallucinations | Unrealistic physics | | Maxed out Ampere GPUs | Max out Blackwell GPUs | | Limited context windows | Short video outputs | Today, these limitations are mostly eliminated from LLMs. Removing them from cinematic video generation, as well as related areas like world models, 3D object generation (XYZ dimensions rather than XYT), and other generative AI rendering models is a highly active area of research. One important direction in research is coming back around to the idea of autoregressive generation, like with blending LLM architecture into image generation models. Rather than pure framewise video generation with its unusable error accumulation, new techniques like Self Forcing combine a global view of quality with an iterative approach to generation. Adding autoregressive components to video models can partially address the bottleneck of attention, though it remains the most important and expensive component of inference. ## 2.4 Calculating Inference Bottlenecks In a perfectly optimized system, every resource is fully utilized at all times. In GPUs, there are two main resources: - **Compute:** The number of floating-point operations per second that the GPU can achieve. - **Memory bandwidth:** The number of bytes that the GPU can move per second. Ideally, compute is never sitting idle waiting for information from memory, and memory bandwidth never goes unused waiting for compute to finish. In the real world, systems have bottlenecks: imbalances where one resource is idle while another is saturated. Discovering these bottlenecks is the first step to improving performance. If a certain operation is bottlenecked on memory bandwidth, no amount of compute optimization will make the system faster, and vice versa. In most cases, inference systems have the following bottlenecks: - LLM prefill (KV cache construction) is compute bound. - LLM decode (token generation) is memory bound. - Image and video generation are compute bound. When optimizing performance on each of these phases, the goal is to make the bottleneck less limiting to system-wide performance. For example, batching multiple requests together makes LLM decode less memory bound because processing a batch of requests uses more compute for the same amount of memory traffic. ### 2.4.1 Ops:Byte Ratio and Arithmetic Intensity Each GPU has a specific compute speed (measured in operations per second) and memory bandwidth (measured in gigabytes or terabytes per second). Compare these to determine the ops:byte ratio of a given GPU. For example, an H100 GPU in FP16 can perform 989 teraFLOPS of dense computation against 3.35 TB/s of memory bandwidth. This yields an ops:byte ratio of about 295. For inference in FP16 to be perfectly balanced (as all things should be) on an H100 GPU, the inference system needs to perform 295 floating point operations for every byte of memory it accesses. To figure out that ratio, calculate the arithmetic intensity of the algorithm. Arithmetic intensity, also known as operational intensity, is the ratio between work and memory traffic for the calculation at hand. ![Figure 2.12: The equation for arithmetic intensity.](https://www.datocms-assets.com/104802/1788123250-inference-engineering-figure-2-12.png) _Figure 2.12: The equation for arithmetic intensity._ Where ops:byte was measured on a per-second scale, arithmetic intensity is measured across the execution of a single function or algorithm. Arithmetic intensity is visualized with a roofline model, which charts performance against the bandwidth ceiling (a diagonal line) and the performance ceiling (a horizontal line). ![Figure 2.13: A roofline chart shows the switch from memory to compute bottleneck based on arithmetic intensity.](https://www.datocms-assets.com/104802/1788123259-inference-engineering-figure-2-13.png) _Figure 2.13: A roofline chart shows the switch from memory to compute bottleneck based on arithmetic intensity._ Plotting against the roofline model reveals if the algorithm is: - **Compute bound:** When the arithmetic intensity is higher than the hardware’s ops:byte ratio and hits the horizontal performance ceiling, it’s compute bound. - **Memory bound:** When the arithmetic intensity is lower than the hardware’s ops:byte ratio and hits the diagonal bandwidth ceiling, it’s memory bound. To find a bottleneck, look at arithmetic intensity for the most expensive calculations in a system. For inference, one such calculation is attention. ### 2.4.2 LLM Inference Bottlenecks LLM inference has two phases: - **Prefill:** Determines the time to first token (TTFT) and is compute-bound. - **Decode:** Determines the tokens per second (TPS) and is memory-bound. For each phase, you can prove the existence of the bottleneck by comparing the arithmetic intensity of the most important operation to the ops:byte ratio of available hardware. In both prefill and decode, the most expensive operation is attention. The exact arithmetic intensity of attention depends on the model architecture (dimensions, heads, etc), the input sequence length, and the implementation of the attention algorithm. The essential difference is that prefill processes the entire input sequence in parallel, while decode generates tokens one at a time. For prefill, the model weights are loaded a single time, then a series of large matrix multiplication between the matrix of inputs and the attention matrices occurs. This is a lot of calculations versus a single read from memory, creating a high arithmetic intensity. On decode, the model weights are loaded for every token, which is generated via much less-expensive vector-matrix multiplication. In this case, relatively few floating-point operations are needed compared to loading the entire model weights, so the arithmetic intensity is low. As an example of calculating exact arithmetic intensity, consider a decode step for a model with a 128-dimensional attention head (d=128) on a sequence of 4096 tokens (N=4096). For this analysis, use the standard algorithm for attention without any optimizations. ![Figure 2.14: Standard attention implementation, adapted from “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (Dao et al., 2022).](https://www.datocms-assets.com/104802/1788123263-inference-engineering-figure-2-14.png) _Figure 2.14: Standard attention implementation, adapted from “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (Dao et al., 2022)._ Based on the parameters of this exercise, establish the size of these matrices: - **N:** The sequence length, established as 4096. - **d:** The dimensionality of the attention head, set to 128. - **Q, K, V:** Given as Nxd, or 4096x128. - **S, P:** Calculated as NxN, or 4096x4096. - **O:** Calculated as Nxd, or 4096x128 Assume FP16 inference, where each value in the matrix is two bytes. For reference, a 4096x4096 matrix is about 32 MiB, or about the same amount of data as a high-resolution RAW DSLR photo. Each of the three lines of the attention algorithm follows the same pattern: load data from memory, perform a calculation, and store the result to memory. ![Figure 2.15: Memory movement (reads and writes) and compute work for the attention implementation in Figure 2.14.](https://www.datocms-assets.com/104802/1788123268-inference-engineering-figure-2-15.png) _Figure 2.15: Memory movement (reads and writes) and compute work for the attention implementation in Figure 2.14._ To calculate the total memory movement, sum the first and third columns, which track reads from and writes to GPU memory: ![Figure 2.16: The total memory movement for a kernel is the sum of all reads and writes across the three steps.](https://www.datocms-assets.com/104802/1788123276-inference-engineering-figure-2-16.png) _Figure 2.16: The total memory movement for a kernel is the sum of all reads and writes across the three steps._ To calculate the total compute, sum the second column: ![Figure 2.17: The total compute for the kernel is the sum of operations across the three steps.](https://www.datocms-assets.com/104802/1788123282-inference-engineering-figure-2-17.png) _Figure 2.17: The total compute for the kernel is the sum of operations across the three steps._ To calculate the arithmetic intensity, compare the work (total compute) to the memory traffic: ![Figure 2.18: The arithmetic intensity of a kernel is the total work (number of compute operations) divided by the memory movement](https://www.datocms-assets.com/104802/1788123290-inference-engineering-figure-2-18.png) _Figure 2.18: The arithmetic intensity of a kernel is the total work (number of compute operations) divided by the memory movement_ For this example, the arithmetic intensity of 62 is much lower than the H100 GPU’s ops:byte ratio of 295. The exact numbers vary by model, sequence length, and hardware, but this example illustrates the general principle that decode is memory bound. Calculating arithmetic intensity like this is an academic exercise, not a routine task for inference engineers. But seeing it once is useful for building intuition. ### 2.4.3 Image Generation Inference Bottlenecks Image and video models are relatively small – they have a tenth as many parameters as frontier language models – but their attention mechanism is just as computationally demanding. Image and video generation models use iterative denoising, not autoregressive token generation. Just like attention for LLM prefill processes the entire input sequence at once, attention for generating media must consider the entire image or video object as represented in latent space. Also like LLM prefill, image and video generation model inference is bottlenecked on compute. Specific techniques for optimizing inference for these modalities are featured in sections 6.5 and 6.6. ## 2.5 Optimizing Attention For LLMs, attention scales quadratically with the length of the input sequence. Each calculation of attention depends on the K and V values of each previous token. In practice, attention is a linear-time operation during decode as the KV cache stores the results of key and value computations for previous tokens. Even a linearly scaling algorithm gets very expensive. Attention is one of the most expensive parts of inference across models and architectures. Naturally, optimizing attention is an important and highly active research area. Attention is a sensitive process because each token depends on every previous token. Small errors in attention can accumulate quickly, making attention optimization a delicate process. Figure 2.14 showed that the attention algorithm itself is straightforward. However, that basic implementation is inefficient. The intermediate matrices `S` and `P` are stored at the end of one step, then immediately loaded in the next step. There are two strategies for optimizing attention: - **Implementation improvements:** Write higher-performance kernels that use memory and compute more efficiently. - **New algorithms:** Create algorithms for attention that scale in better-than-quadratic time with minimal quality loss. Implementation improvements are still limited by attention’s quadratic time complexity, but are lossless (do not affect quality) and make inference feasible for long sequences on today’s hardware. Other algorithmic approaches trade off quality for time and space complexity, though training techniques can minimize the impact. The most famous implementation of attention is the FlashAttention series of papers and kernels. Where the basic algorithm can be implemented in a handful of lines of code, FlashAttention uses tens of thousands of lines to implement attention in hand-fused kernels built for specific GPUs – FlashAttention for H100 uses different code than FlashAttention for B200. FlashAttention works by eliminating excess reads and writes from memory and laying out the attention algorithm to precisely fit the GPU’s capabilities. FlashAttention is especially useful for compute-bound operations like LLM prefill and video generation. Another important implementation is PagedAttention. KV caches quickly grow large, filling GPU memory and taking time to read. PagedAttention partitions the KV cache into blocks (pages) that can be accessed via a lookup table. This means the KV cache can be stored across the GPU with fragmented memory rather than requiring a single contiguous block of memory. While FlashAttention and PagedAttention are valuable optimizations, they don’t change the fact that attention is a quadratic algorithm. New variants of attention improve the underlying time and space complexity: - **Sliding window attention:** Computes attention for a sliding window of the previous `w` tokens, turning attention from `O(N^2)` to `O(Nw)` where `w` is often in the range of 8K to 32K. - **Gated attention:** Various types of layers introduced in training allow for approximating attention for certain chunks of context in linear time with respect to chunk length. - **Linear attention:** Replaces the quadratic softmax equation with a linear-time algorithm that approximates attention. - **Compressed attention:** Periodically compresses context from earlier in the sequence, attention considers both compressed context and uncompressed recent tokens. - **Multi-latent attention:** Approximates attention in low-dimensional latent space. Intuitively, it makes sense that tokens near each other in a sequence affect each other more than tokens from much earlier. The sentence I am writing now follows closely from the previous sentence, but less so from the sentence at the start of this chapter. This intuition can be extended through training. Algorithms like sliding window attention, when applied during training, create models that keep quality high when the same technique is used in inference. Another avenue of research is avoiding attention altogether by using a different architecture than transformers. Mamba is a selective state-space model that replaces self-attention with a recurrent state update, achieving linear scaling on sequence length. Hybrid models sometimes mix Mamba-style state-space model blocks with transformer blocks. Applications of state-space models are still limited, though hybrid models are becoming more popular with open models like NVIDIA Nemotron 3 Nano adopting hybrid architectures.